Create expression even when there is no JuMP variables

Author: blegat

src/JuMP_wrapper.jl

@@ -10,25 +10,18 @@ include("operators.jl")
 
 """
     struct IteratorInExpr
-        iterator::Iterator
+        iterators::Iterators
         index::IteratorIndex
     end
 
-Iterator `iterator` with values at index `index`.
+Iterator `iterators[index.value]`.
 """
 struct IteratorInExpr
     iterators::Iterators
     index::IteratorIndex
 end
 
-function Base.show(io::IO, i::IteratorInExpr)
-    print(
-        io,
-        values_at(i.iterators[i.index.iterator_index], i.index.value_index),
-    )
-    print(io, "[i]")
-    return
-end
+Base.copy(it::IteratorInExpr) = it
 
 JuMP._is_real(::Union{IteratorInExpr,IteratorIndex}) = true
 JuMP.moi_function(i::Union{IteratorInExpr,IteratorIndex}) = i
@@ -55,20 +48,27 @@ function JuMP.jump_function(model, f::FunctionGenerator{F}) where {F}
     )
 end
 
-_size(expr::ExprGenerator) = getfield.(expr.expr.iterators, :length)
+_size(expr::ExprGenerator) = length.(getfield.(expr.expr.iterators, :values))
 
 index_iterators(func, _) = func
 
 function index_iterators(func::IteratorInExpr, index)
     idx = func.index
-    return values_at(iterators[idx.iterator_index], idx.value_index)[index[idx.iterator_index]]
+    return func.iterators[idx.value].values[index[idx.value]]
 end
 
 function index_iterators(func::JuMP.GenericNonlinearExpr, index)
-    return GenericNonlinearExpr(
-        func.head,
-        map(Base.Fix2(index_iterators, index), func.args),
-    )
+    args = map(Base.Fix2(index_iterators, index), func.args)
+    if any(JuMP._has_variable_ref_type, args)
+        return JuMP.GenericNonlinearExpr(func.head, args)
+    else
+        registry = MOI.Nonlinear.OperatorRegistry()
+        if length(func.args) == 1
+            MOI.Nonlinear.eval_univariate_function(registry, func.head, args[])
+        else
+            MOI.Nonlinear.eval_multivariate_function(registry, func.head, args)
+        end
+    end
 end
 
 function Base.getindex(expr::ExprGenerator, i::Integer)
@@ -104,10 +104,11 @@ function JuMP.build_constraint(
 end
 
 struct IteratedConstraint{
+    E,
     V<:JuMP.GenericVariableRef,
     S<:MOI.AbstractVectorSet,
 } <: JuMP.AbstractConstraint
-    func::ExprGenerator{V}
+    func::ExprGenerator{E,V}
     set::S
 end
 

src/MOI_wrapper.jl

@@ -6,37 +6,23 @@
 import MathOptInterface as MOI
 
 """
-    struct Iterator
-        length::Int
-        values::Vector{Float64}
+    struct Iterator{T}
+        values::Vector{T}
     end
 """
-struct Iterator
-    length::Int
-    values::Vector{Float64}
-    function Iterator(length::Int)
-        return new(length, Float64[])
-    end
-end
-
-function num_values(it::Iterator)
-    return div(length(it.values), it.length)
-end
-
-function values_at(it::Iterator, i)
-    return it.values[(1+it.length*(i-1)):(it.length*i)]
+struct Iterator{T}
+    values::Vector{T}
 end
 
 struct IteratorIndex
-    iterator_index::Int
-    value_index::Int
+    value::Int
 end
 
 Base.copy(i::IteratorIndex) = i
 
 struct FunctionGenerator{F} <: MOI.AbstractVectorFunction
     func::MOI.ScalarNonlinearFunction
-    iterators::Vector{Iterator}
+    iterators::Vector{Iterator} # Slight type instability, we don't have `Iterator{T}`
 end
 
 function Base.copy(f::FunctionGenerator{F}) where {F}

src/operators.jl

@@ -10,6 +10,10 @@ const Iterators = Vector{Iterator}
         expr::JuMP.GenericNonlinearExpr{V}
         iterators::Vector{Iterator}
     end
+
+Represent a `JuMP.GenericNonlinearExpr` containing iterators.
+Thanks to this custom type, we can define a custom `JuMP.build_constraint` method
+to generate constraint of different types.
 """
 struct ExprTemplate{E,V<:JuMP.AbstractVariableRef} <: JuMP.AbstractJuMPScalar
     expr::JuMP.GenericNonlinearExpr{V}
@@ -25,6 +29,8 @@ end
 
 JuMP.variable_ref_type(::Type{ExprTemplate{E,V}}) where {E,V} = V
 
+JuMP.check_belongs_to_model(f::ExprTemplate, model) = JuMP.check_belongs_to_model(f.expr, model)
+
 """
     struct IteratorValues{I}
         iterator::Iterator
@@ -59,7 +65,7 @@ iterator([5, -5])
 """
 struct IteratorValues{I}
     iterators::Iterators
-    index::Int
+    index::IteratorIndex
     values::I
 end
 
@@ -71,50 +77,66 @@ function Base.show(io::IO, i::IteratorValues)
 end
 
 function iterators(axes)
-    iterators = [Iterator(length(axe)) for axe in axes]
-    return IteratorValues.(Ref(iterators), eachindex(axes), axes)
+    iterators = Iterator[Iterator(axe) for axe in axes]
+    return IteratorValues.(Ref(iterators), IteratorIndex.(eachindex(axes)), axes)
 end
 
 iterator(axe) = iterators([axe])[]
 
 function Base.getindex(d::Dict, i::IteratorValues)
-    return IteratorValues(i.iterators, i.index, [d[val] for val in i.values])
+    new_values = [d[val] for val in i.values]
+    i.iterators[i.index.value] = Iterator(new_values)
+    return IteratorValues(i.iterators, i.index, new_values)
 end
 
 # The following is intentionally kept close to JuMP/src/nlp_expr.jl
 const _ScalarWithIterator = Union{ExprTemplate,IteratorValues}
 
+function _univariate(f, op, x)
+    V = something(
+        _variable_ref_type(x),
+        JuMP.VariableRef, # FIXME needed if `x` is an iterator
+    )
+    nl = JuMP.GenericNonlinearExpr{V}(op, _expr(x))
+    E = JuMP._MA.promote_operation(f, _type(x))
+    return ExprTemplate{E}(nl, _iterators(x))
+end
+
 # Univariate operators
 for f in MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS
+    op = Meta.quot(f)
     if isdefined(Base, f)
         @eval function Base.$(f)(x::IteratorValues)
-            return IteratorValues(x.iterators, x.index, $(f).(x.values))
+            return _univariate($f, $op, x)
         end
         @eval function Base.$(f)(x::ExprTemplate)
-            return ExprTemplate($f(x.expr), x.iterators)
+            return _univariate($f, $op, x)
         end
     end
 end
 
 function prepare(it::IteratorValues)
-    append!(it.iterators[it.index].values, it.values)
-    index = IteratorIndex(it.index, num_values(it.iterators[it.index]))
-    return IteratorInExpr(it.iterators, index)
+    @assert it.values == it.iterators[it.index.value].values
+    return IteratorInExpr(it.iterators, it.index)
 end
 
 _expr(f::JuMP.AbstractJuMPScalar) = f
 _expr(it::IteratorValues) = prepare(it)
 _expr(f::ExprTemplate) = f.expr
+_expr(f::Number) = f
 
 _variable_ref_type(f::JuMP.AbstractJuMPScalar) = JuMP.variable_ref_type(f)
 _variable_ref_type(::IteratorValues) = nothing
+_variable_ref_type(::Number) = nothing
 
 _iterators(::JuMP.AbstractJuMPScalar) = nothing
 _iterators(it::_ScalarWithIterator) = it.iterators
+_iterators(::Number) = nothing
 
 _type(it::_ScalarWithIterator) = eltype(it.values)
 _type(::ExprTemplate{E}) where {E} = E
 _type(f::JuMP.AbstractJuMPScalar) = typeof(f)
+_type(f::Number) = typeof(f)
 
 _check_equal(it::Iterators, ::Nothing) = it
 _check_equal(::Nothing, it::Iterators) = it
@@ -139,10 +161,16 @@ for f in [:+, :-, :*, :^, :/, :atan, :min, :max]
     op = Meta.quot(f)
     @eval begin
         function Base.$(f)(x::IteratorValues, y::Number)
-            return IteratorValues(x.iterators, x.index, $(f).(x.values, y))
+            return _multivariate($f, $op, x, y)
         end
         function Base.$(f)(x::Number, y::IteratorValues)
-            return IteratorValues(y.iterators, y.index, $(f).(x, y.values))
+            return _multivariate($f, $op, x, y)
+        end
+        function Base.$(f)(x::ExprTemplate, y::Number)
+            return _multivariate($f, $op, x, y)
+        end
+        function Base.$(f)(x::Number, y::ExprTemplate)
+            return _multivariate($f, $op, x, y)
         end
         function Base.$(f)(x::_ScalarWithIterator, y::JuMP.AbstractJuMPScalar)
             return _multivariate($f, $op, x, y)

src/print.jl

@@ -1,3 +1,9 @@
+function Base.show(io::IO, i::IteratorInExpr)
+    print(io, i.iterators[i.index.value].values)
+    print(io, "[i]")
+    return
+end
+
 function Base.show(io::IO, f::Union{ExprGenerator,ExprTemplate})
     return print(io, JuMP.function_string(MIME("text/plain"), f))
 end

test/JuMP_wrapper.jl

@@ -40,15 +40,15 @@ function test_container()
     @test con isa IteratedConstraint
     con_it_expr = jump_function(con)
     @test con_it_expr isa ExprGenerator
-    con_expr = con_it_expr.expr
+    con_expr = con_it_expr.expr.expr
     @test sprint(show, con_ref) ==
-          "ParametrizedArray(((x + (IteratorIndex(1, 1))) - (IteratorIndex(2, 2))) - 0.0, Iterator(2, [-1.0, 1.0, 3.141592653589793, 0.0]), Iterator(2, [-1.0, 1.0, 3.141592653589793, 0.0]) ∈ MathOptInterface.Nonnegatives(4), (iterator([:a, :b]),))"
+          "GenOpt.ParametrizedArray(((x + (Real[π, 0.0][i])) - (Real[π, 0.0][i])) - 0.0, GenOpt.Iterator{Real}(Real[π, 0.0]) ∈ MathOptInterface.Nonnegatives(2), GenOpt.IteratorValues{Vector{Symbol}}[iterator([:a, :b])])"
     @test sprint(show, MIME"text/latex"(), con_ref) ==
-          "ParametrizedArray(((x + (IteratorIndex(1, 1))) - (IteratorIndex(2, 2))) - 0.0, Iterator(2, [-1.0, 1.0, 3.141592653589793, 0.0]), Iterator(2, [-1.0, 1.0, 3.141592653589793, 0.0]) ∈ MathOptInterface.Nonnegatives(4), (iterator([:a, :b]),))"
+          "GenOpt.ParametrizedArray(((x + (Real[π, 0.0][i])) - (Real[π, 0.0][i])) - 0.0, GenOpt.Iterator{Real}(Real[π, 0.0]) ∈ MathOptInterface.Nonnegatives(2), GenOpt.IteratorValues{Vector{Symbol}}[iterator([:a, :b])])"
     @test sprint(show, con_expr) ==
-          "((x + (IteratorIndex(1, 1))) - (IteratorIndex(2, 2))) - 0.0"
+          "((x + (Real[π, 0.0][i])) - (Real[π, 0.0][i])) - 0.0"
     @test sprint(show, MIME"text/latex"(), con_expr) ==
-          "\$ {\\left({\\left({x} + {\\left(IteratorIndex(1, 1)\\right)}\\right)} - {\\left(IteratorIndex(2, 2)\\right)}\\right)} - {0.0} \$"
+          "\$ {\\left({\\left({x} + {\\left(Real[π, 0.0][i]\\right)}\\right)} - {\\left(Real[π, 0.0][i]\\right)}\\right)} - {0.0} \$"
 
     i = GenOpt.iterator(keys)
     expr = x + d1[i] - d2[i]

test/operators.jl

@@ -22,30 +22,39 @@ end
 
 function _test_iterator(it, values)
     @test it isa IteratorValues
-    @test it.iterators[1].length == length(values)
+    @test it.iterators[1].values == values
     @test it.values == values
 end
 
+function _test_template(et, values)
+    @test et isa ExprTemplate
+    for i in eachindex(values)
+        @test index_iterators(et.expr, (i,)) == values[i]
+    end
+end
+
 function test_getindex()
     d1 = Dict(:a => -1, :b => 1)
     d2 = Dict(:a => π, :b => 0.0)
 
     i = GenOpt.iterator([:a, :b])
 
     _test_iterator(d1[i], [-1, 1])
-    return _test_iterator(d2[i], Real[π, 0.0])
+    _test_iterator(d2[i], Real[π, 0.0])
+    return
 end
 
 function test_univariate()
     i = GenOpt.iterator([2, -3])
-    _test_iterator(+i, [2, -3])
-    return _test_iterator(-i, [-2, 3])
+    _test_template(+i, [2, -3])
+    _test_template(-i, [-2, 3])
+    return
 end
 
 function test_multivariate()
     i, j = GenOpt.iterators(([2, -3], [1, -1]))
-    _test_iterator(i + 1, [3, -2])
-    _test_iterator(2 - i, [0, 5])
+    _test_template(i + 1, [3, -2])
+    _test_template(2 - i, [0, 5])
     ij = i + j
     @test ij isa GenOpt.ExprTemplate{Int,JuMP.VariableRef}
     model = JuMP.Model()