diff --git a/src/model/expression_manipulation.jl b/src/model/expression_manipulation.jl index 16cfc0cc65..fe4fe6e7be 100644 --- a/src/model/expression_manipulation.jl +++ b/src/model/expression_manipulation.jl @@ -2,11 +2,13 @@ # Create dense arrays of expressions filled with zeros to be added to later ###### ###### ###### ###### ###### ###### +# Single element version function create_empty_expression!(EP::Model, exprname::Symbol) EP[exprname] = AffExpr(0.0) return nothing end +# Vector version, to avoid needing to wrap the dimension in a tuple or array function create_empty_expression!(EP::Model, exprname::Symbol, dim1::Int64) temp = Array{AffExpr}(undef, dim1) fill_with_zeros!(temp) @@ -14,6 +16,15 @@ function create_empty_expression!(EP::Model, exprname::Symbol, dim1::Int64) return nothing end +@doc raw""" + create_empty_expression!(EP::Model, exprname::Symbol, dims::NTuple{N, Int64}) where N + +Create an dense array filled with zeros which can be altered later. +Other approaches to creating zero-filled arrays will often return an array of floats, not expressions. +This can lead to errors later if a method can only operate on expressions. + +We don't currently have a method to do this with non-contiguous indexing. +""" function create_empty_expression!(EP::Model, exprname::Symbol, dims::NTuple{N, Int64}) where N temp = Array{AffExpr}(undef, dims) fill_with_zeros!(temp) @@ -21,8 +32,8 @@ function create_empty_expression!(EP::Model, exprname::Symbol, dims::NTuple{N, I return nothing end +# Version with the dimensions wrapped in an array. This requires slightly more memory than using tuples function create_empty_expression!(EP::Model, exprname::Symbol, dims::Vector{Int64}) - # The version using tuples is slightly less memory temp = Array{AffExpr}(undef, dims...) fill_with_zeros!(temp) EP[exprname] = temp @@ -33,15 +44,28 @@ end # Fill dense arrays of expressions with zeros or a constant ###### ###### ###### ###### ###### ###### -function fill_with_zeros!(arr::Array{GenericAffExpr{C,T}, dims}) where {C,T,dims} +@doc raw""" + fill_with_zeros!(arr::AbstractArray{GenericAffExpr{C,T}, dims}) where {C,T,dims} + +Fill an array of expressions with zeros in-place. +""" +function fill_with_zeros!(arr::AbstractArray{GenericAffExpr{C,T}, dims}) where {C,T,dims} for i::Int64 in eachindex(IndexLinear(), arr)::Base.OneTo{Int64} arr[i] = AffExpr(0.0) end return nothing end -function fill_with_const!(arr::Array{GenericAffExpr{C,T}, dims}, con::Real) where {C,T,dims} - for i::Int64 in eachindex(IndexLinear(), arr)::Base.OneTo{Int64} +@doc raw""" + fill_with_const!(arr::AbstractArray{GenericAffExpr{C,T}, dims}, con::Real) where {C,T,dims} + +Fill an array of expressions with the specified constant, in-place. + +In the future we could expand this to non AffExpr, using GenericAffExpr +e.g. if we wanted to use Float32 instead of Float64 +""" +function fill_with_const!(arr::AbstractArray{GenericAffExpr{C,T}, dims}, con::Real) where {C,T,dims} + for i in eachindex(arr) arr[i] = AffExpr(con) end return nothing @@ -79,23 +103,19 @@ end # Both arrays must have the same dimensions ###### ###### ###### ###### ###### ###### +# Version for single element function add_similar_to_expression!(expr1::GenericAffExpr{C,T}, expr2::V) where {C,T,V} add_to_expression!(expr1, expr2) return nothing end -function add_similar_to_expression!(expr1::Array{GenericAffExpr{C,T}, dim1}, expr2::Array{V, dim2}) where {C,T,V,dim1,dim2} - # This is defined for Arrays of different dimensions - # despite the fact it will definitely throw an error - # because the error will tell the user / developer - # the dimensions of both arrays - check_sizes_match(expr1, expr2) - for i in eachindex(IndexLinear(), expr1) - add_to_expression!(expr1[i], expr2[i]) - end - return nothing -end +@doc raw""" + add_similar_to_expression!(expr1::AbstractArray{GenericAffExpr{C,T}, dim1}, expr2::AbstractArray{V, dim2}) where {C,T,V,dim1,dim2} +Add an array of some type `V` to an array of expressions, in-place. +This will work on JuMP DenseContainers which do not have linear indexing from 1:length(arr). +However, the accessed parts of both arrays must have the same dimensions. +""" function add_similar_to_expression!(expr1::AbstractArray{GenericAffExpr{C,T}, dim1}, expr2::AbstractArray{V, dim2}) where {C,T,V,dim1,dim2} # This is defined for Arrays of different dimensions # despite the fact it will definitely throw an error @@ -113,18 +133,18 @@ end # Both arrays must have the same dimensions ###### ###### ###### ###### ###### ###### +# Version for single element function add_term_to_expression!(expr1::GenericAffExpr{C,T}, expr2::V) where {C,T,V} add_to_expression!(expr1, expr2) return nothing end -function add_term_to_expression!(expr1::Array{GenericAffExpr{C,T}, dims}, expr2::V) where {C,T,V,dims} - for i in eachindex(IndexLinear(), expr1) - add_to_expression!(expr1[i], expr2) - end - return nothing -end +@doc raw""" + add_term_to_expression!(expr1::AbstractArray{GenericAffExpr{C,T}, dims}, expr2::V) where {C,T,V,dims} +Add an entry of type `V` to an array of expressions, in-place. +This will work on JuMP DenseContainers which do not have linear indexing from 1:length(arr). +""" function add_term_to_expression!(expr1::AbstractArray{GenericAffExpr{C,T}, dims}, expr2::V) where {C,T,V,dims} for i in eachindex(expr1) add_to_expression!(expr1[i], expr2) @@ -136,6 +156,12 @@ end # Check that two arrays have the same dimensions ###### ###### ###### ###### ###### ###### +@doc raw""" + check_sizes_match(expr1::AbstractArray{C, dim1}, expr2::AbstractArray{T, dim2}) where {C,T,dim1, dim2} + +Check that two arrays have the same dimensions. +If not, return an error message which includes the dimensions of both arrays. +""" function check_sizes_match(expr1::AbstractArray{C, dim1}, expr2::AbstractArray{T, dim2}) where {C,T,dim1, dim2} # After testing, this appears to be just as fast as a method for Array{GenericAffExpr{C,T}, dims} or Array{AffExpr, dims} if size(expr1) != size(expr2) @@ -145,27 +171,36 @@ function check_sizes_match(expr1::AbstractArray{C, dim1}, expr2::AbstractArray{T end end +@doc raw""" + check_addable_to_expr(C::DataType, T::DataType) + +Check that two datatype can be added using add_to_expression!(). Raises an error if not. + +This needs some work to make it more flexible. Right now it's challenging to use with GenericAffExpr{C,T} +as the method only works on the constituent types making up the GenericAffExpr, not the resulting expression type. +Also, the default MethodError from add_to_expression! is sometime more informative than the error message here. +""" +function check_addable_to_expr(C::DataType, T::DataType) + if !(hasmethod(add_to_expression!, (C,T))) + error("No method found for add_to_expression! with types $(C) and $(T)") + end +end + ###### ###### ###### ###### ###### ###### # Sum an array of expressions into a single expression ###### ###### ###### ###### ###### ###### -function sum_expression(expr::AbstractArray{C, dims}) where {C,dims} - # This appears to work just as well as a separate method for Array{AffExpr, dims} +@doc raw""" + sum_expression(expr::AbstractArray{C, dims}) where {C,dims} :: C + +Sum an array of expressions into a single expression and return the result. +We're using errors from add_to_expression!() to check that the types are compatible. +""" +function sum_expression(expr::AbstractArray{C, dims}) :: AffExpr where {C,dims} + # check_addable_to_expr(C,C) total = AffExpr(0.0) for i in eachindex(expr) add_to_expression!(total, expr[i]) end return total -end - -function sum_expression(expr::AbstractArray{GenericAffExpr{C,T}, dims}) where {C,T,dims} - return _sum_expression(expr) -end - -function sum_expression(expr::AbstractArray{GenericVariableRef{C}, dims}) where {C,dims} - return _sum_expression(expr) -end - -function sum_expression(expr::AbstractArray{AbstractJuMPScalar, dims}) where {dims} - return _sum_expression(expr) -end +end \ No newline at end of file diff --git a/test/expression_manipulation_test.jl b/test/expression_manipulation_test.jl index a83c234a7b..aae5d442ec 100644 --- a/test/expression_manipulation_test.jl +++ b/test/expression_manipulation_test.jl @@ -1,6 +1,66 @@ using JuMP using HiGHS +function setup_sum_model() + EP = Model(HiGHS.Optimizer) + @variable(EP, x[i=1:100,j=1:4:200]>=0) + @variable(EP, y[i=1:100,j=1:50]>=0) + @expression(EP, eX[i=1:100,j=1:4:200], 2.0*x[i,j]+i+10.0*j) + @expression(EP, eY[i=1:100,j=1:50], 3.0*y[i,j]+2*i+j) + @expression(EP, eZ[i=1:100,j=1:50], 2.0*x[i,(j-1)*4+1] + 4.0*y[i,j]) + return EP +end + +function sum_disjoint_var() + EP = setup_sum_model() + try + GenX.sum_expression(EP[:x]) + catch + return false + end + return true +end + +function sum_dense_var() + EP = setup_sum_model() + try + GenX.sum_expression(EP[:y]) + catch + return false + end + return true +end + +function sum_disjoint_expr() + EP = setup_sum_model() + try + GenX.sum_expression(EP[:eX]) + catch + return false + end + return true +end + +function sum_dense_expr() + EP = setup_sum_model() + try + GenX.sum_expression(EP[:eY]) + catch + return false + end + return true +end + +function sum_combo_expr() + EP = setup_sum_model() + try + GenX.sum_expression(EP[:eZ]) + catch + return false + end + return true +end + let EP = Model(HiGHS.Optimizer) @@ -49,6 +109,13 @@ let GenX.add_term_to_expression!(EP[:large_expr], AffExpr(3.0)) @test EP[:large_expr][100] == test_var[100] + 22.0 + # Test sum_expression + @test sum_dense_var() == true + @test sum_disjoint_var() == true + @test sum_dense_expr() == true + @test sum_disjoint_expr() == true + @test sum_combo_expr() == true + # Test add_term_to_expression! for variable @variable(EP, single_var >= 0) GenX.add_term_to_expression!(EP[:large_expr], single_var)