Add bridge from generator to vector of affine

src/GenOpt.jl

@@ -6,6 +6,7 @@
 module GenOpt
 
 include("MOI_wrapper.jl")
+include("bridge.jl")
 include("JuMP_wrapper.jl")
 include("examodels.jl")
 

src/MOI_wrapper.jl

@@ -13,6 +13,11 @@ end
 Base.copy(array::ContiguousArrayOfVariables) = array
 Base.size(array::ContiguousArrayOfVariables) = array.size
 
+function Base.getindex(A::ContiguousArrayOfVariables, I::Integer...)
+    index = A.offset + Base._to_linear_index(CartesianIndices(A.size), I...)
+    return MOI.VariableIndex(index)
+end
+
 """
     struct Iterator{T}
         values::Vector{T}
@@ -24,11 +29,16 @@ end
 
 Iterator(values::AbstractArray) = Iterator(vec(collect(values)))
 
+Base.length(it::Iterator) = length(it.values)
+
 struct IteratorIndex
     value::Int
 end
 
 Base.copy(i::IteratorIndex) = i
+function Base.isapprox(a::IteratorIndex, b::IteratorIndex; kwargs...)
+    return a.value == b.value
+end
 
 struct FunctionGenerator{F} <: MOI.AbstractVectorFunction
     func::MOI.ScalarNonlinearFunction
@@ -38,14 +48,30 @@ end
 function Base.copy(f::FunctionGenerator{F}) where {F}
     return FunctionGenerator{F}(copy(f.func), f.iterators)
 end
+
+function Base.isapprox(a::FunctionGenerator, b::FunctionGenerator; kwargs...)
+    return isapprox(a.func, b.func; kwargs...) &&
+           length(a.iterators) == length(b.iterators) &&
+           all(
+               isapprox(ai.values, bi.values; kwargs...) for
+               (ai, bi) in zip(a.iterators, b.iterators)
+           )
+end
 function MOI.Utilities.is_canonical(f::FunctionGenerator)
     return MOI.Utilities.is_canonical(f.func)
 end
+
+function MOI.output_dimension(f::FunctionGenerator)
+    return prod(length, f.iterators)
+end
+
 function MOI.Utilities.is_coefficient_type(
-    ::Type{FunctionGenerator{E}},
+    ::Type{<:FunctionGenerator},
     ::Type{T},
-) where {E,T}
-    return MOI.Utilities.is_coefficient_type(E, T)
+) where {T}
+    # Return false so standard MOI bridges (ScalarizeBridge, FlipSignBridge, etc.)
+    # don't try to handle FunctionGenerator. Only FunctionGeneratorBridge should.
+    return false
 end
 
 struct SumGenerator{F} <: MOI.AbstractScalarFunction

src/bridge.jl

@@ -0,0 +1,247 @@
+# Copyright (c) 2024: Benoît Legat and contributors
+#
+# Use of this source code is governed by an MIT-style license that can be found
+# in the LICENSE.md file or at https://opensource.org/licenses/MIT.
+
+"""
+    FunctionGeneratorBridge{T,F,S}
+
+Bridge that expands a `FunctionGenerator{F}` constraint into individual
+scalar `F`-in-`S` constraints.
+
+The `FunctionGenerator` contains a template `MOI.ScalarNonlinearFunction`
+with `IteratorIndex` placeholders. This bridge substitutes concrete values
+from the iterators for each combination and converts each expanded expression
+to the type `F` declared by the `FunctionGenerator{F}` type parameter.
+If the expanded expression cannot be converted to `F`, an error is thrown.
+"""
+struct FunctionGeneratorBridge{T,F,S,VS<:MOI.Utilities.VectorLinearSet} <:
+       MOI.Bridges.Constraint.AbstractBridge
+    constraints::Vector{MOI.ConstraintIndex{F,S}}
+    func::FunctionGenerator{F}
+    set::VS
+end
+
+function MOI.Bridges.Constraint.bridge_constraint(
+    ::Type{FunctionGeneratorBridge{T,F,S,VS}},
+    model::MOI.ModelLike,
+    func::FunctionGenerator{F},
+    set::VS,
+) where {T,F,S,VS<:MOI.Utilities.VectorLinearSet}
+    scalar_set = S(zero(T))
+    constraints = MOI.ConstraintIndex{F,S}[]
+    sizes = Tuple(length.(func.iterators))
+    for idx in CartesianIndices(sizes)
+        values = [
+            func.iterators[k].values[idx[k]] for k in eachindex(func.iterators)
+        ]
+        expanded = _expand(func.func, values)
+        scalar_func = _convert(F, expanded)
+        ci = MOI.Utilities.normalize_and_add_constraint(
+            model,
+            scalar_func,
+            scalar_set,
+        )
+        push!(constraints, ci)
+    end
+    return FunctionGeneratorBridge{T,F,S,VS}(constraints, func, set)
+end
+
+function MOI.supports_constraint(
+    ::Type{FunctionGeneratorBridge{T}},
+    ::Type{<:FunctionGenerator},
+    ::Type{<:MOI.Utilities.VectorLinearSet},
+) where {T}
+    return true
+end
+
+function MOI.Bridges.Constraint.concrete_bridge_type(
+    ::Type{FunctionGeneratorBridge{T}},
+    ::Type{FunctionGenerator{F}},
+    VS::Type{<:MOI.Utilities.VectorLinearSet},
+) where {T,F}
+    S = MOI.Utilities.scalar_set_type(VS, T)
+    return FunctionGeneratorBridge{T,F,S,VS}
+end
+
+function MOI.Bridges.added_constraint_types(
+    ::Type{<:FunctionGeneratorBridge{T,F,S}},
+) where {T,F,S}
+    return Tuple{Type,Type}[(F, S)]
+end
+
+function MOI.Bridges.added_constrained_variable_types(
+    ::Type{<:FunctionGeneratorBridge},
+)
+    return Tuple{Type}[]
+end
+
+function MOI.get(
+    bridge::FunctionGeneratorBridge{T,F,S},
+    ::MOI.NumberOfConstraints{F,S},
+) where {T,F,S}
+    return length(bridge.constraints)
+end
+
+function MOI.get(
+    bridge::FunctionGeneratorBridge{T,F,S},
+    ::MOI.ListOfConstraintIndices{F,S},
+) where {T,F,S}
+    return copy(bridge.constraints)
+end
+
+function MOI.get(
+    ::MOI.ModelLike,
+    ::MOI.ConstraintFunction,
+    bridge::FunctionGeneratorBridge,
+)
+    return copy(bridge.func)
+end
+
+function MOI.get(
+    ::MOI.ModelLike,
+    ::MOI.ConstraintSet,
+    bridge::FunctionGeneratorBridge,
+)
+    return bridge.set
+end
+
+function MOI.delete(model::MOI.ModelLike, bridge::FunctionGeneratorBridge)
+    for ci in bridge.constraints
+        MOI.delete(model, ci)
+    end
+    return
+end
+
+# --- Expression expansion ---
+
+"""
+    _expand(func, values)
+
+Recursively expand an `MOI.ScalarNonlinearFunction` template by substituting
+`IteratorIndex` placeholders with concrete `values`.
+
+Evaluates `:getindex` nodes (for variable and data lookups) and simplifies
+pure-numeric subexpressions.
+"""
+function _expand(func::MOI.ScalarNonlinearFunction, values)
+    args = [_expand(arg, values) for arg in func.args]
+    if func.head == :getindex
+        collection = args[1]
+        indices = [_to_index(a) for a in args[2:end]]
+        return getindex(collection, indices...)
+    end
+    if all(_is_numeric, args)
+        return _eval_op(func.head, args)
+    end
+    return MOI.ScalarNonlinearFunction(func.head, args)
+end
+
+_expand(idx::IteratorIndex, values) = values[idx.value]
+_expand(x, _) = x  # constants, MOI.VariableIndex, etc.
+
+_is_numeric(::Number) = true
+_is_numeric(_) = false
+
+_to_index(x::Integer) = Int(x)
+_to_index(x::AbstractFloat) = Int(x)
+_to_index(x) = x
+
+function _eval_op(head::Symbol, args::Vector)
+    registry = MOI.Nonlinear.OperatorRegistry()
+    float_args = Float64.(args)
+    if length(float_args) == 1
+        return MOI.Nonlinear.eval_univariate_function(
+            registry,
+            head,
+            float_args[1],
+        )
+    else
+        return MOI.Nonlinear.eval_multivariate_function(
+            registry,
+            head,
+            float_args,
+        )
+    end
+end
+
+# --- Conversion from expanded ScalarNonlinearFunction to target type F ---
+
+function _convert(
+    ::Type{MOI.ScalarNonlinearFunction},
+    expr::MOI.ScalarNonlinearFunction,
+)
+    return expr
+end
+_convert(::Type{F}, expr::F) where {F} = expr
+
+function _convert(
+    ::Type{MOI.ScalarAffineFunction{T}},
+    expr::MOI.ScalarNonlinearFunction,
+) where {T}
+    terms, constant = _collect_affine_terms(T, expr)
+    if terms === nothing
+        throw(InexactError(:convert, MOI.ScalarAffineFunction{T}, expr))
+    end
+    return MOI.ScalarAffineFunction(terms, T(constant))
+end
+
+function _collect_affine_terms(
+    ::Type{T},
+    expr::MOI.ScalarNonlinearFunction,
+) where {T}
+    if expr.head == :+ && length(expr.args) == 2
+        t1, c1 = _collect_affine_terms(T, expr.args[1])
+        t2, c2 = _collect_affine_terms(T, expr.args[2])
+        (t1 === nothing || t2 === nothing) && return (nothing, zero(T))
+        return (vcat(t1, t2), c1 + c2)
+    elseif expr.head == :- && length(expr.args) == 2
+        t1, c1 = _collect_affine_terms(T, expr.args[1])
+        t2, c2 = _collect_affine_terms(T, expr.args[2])
+        (t1 === nothing || t2 === nothing) && return (nothing, zero(T))
+        neg_t2 = [MOI.ScalarAffineTerm(-t.coefficient, t.variable) for t in t2]
+        return (vcat(t1, neg_t2), c1 - c2)
+    elseif expr.head == :- && length(expr.args) == 1
+        t1, c1 = _collect_affine_terms(T, expr.args[1])
+        t1 === nothing && return (nothing, zero(T))
+        neg_t1 = [MOI.ScalarAffineTerm(-t.coefficient, t.variable) for t in t1]
+        return (neg_t1, -c1)
+    elseif expr.head == :* && length(expr.args) == 2
+        a1, a2 = expr.args
+        if a1 isa Number && a2 isa MOI.VariableIndex
+            return ([MOI.ScalarAffineTerm(T(a1), a2)], zero(T))
+        elseif a2 isa Number && a1 isa MOI.VariableIndex
+            return ([MOI.ScalarAffineTerm(T(a2), a1)], zero(T))
+        elseif a1 isa Number && a2 isa MOI.ScalarNonlinearFunction
+            t2, c2 = _collect_affine_terms(T, a2)
+            t2 === nothing && return (nothing, zero(T))
+            scaled = [
+                MOI.ScalarAffineTerm(T(a1) * t.coefficient, t.variable) for
+                t in t2
+            ]
+            return (scaled, T(a1) * c2)
+        elseif a2 isa Number && a1 isa MOI.ScalarNonlinearFunction
+            t1, c1 = _collect_affine_terms(T, a1)
+            t1 === nothing && return (nothing, zero(T))
+            scaled = [
+                MOI.ScalarAffineTerm(T(a2) * t.coefficient, t.variable) for
+                t in t1
+            ]
+            return (scaled, T(a2) * c1)
+        elseif a1 isa Number && a2 isa Number
+            return (MOI.ScalarAffineTerm{T}[], T(a1 * a2))
+        else
+            return (nothing, zero(T))
+        end
+    else
+        return (nothing, zero(T))
+    end
+end
+
+function _collect_affine_terms(::Type{T}, x::MOI.VariableIndex) where {T}
+    return ([MOI.ScalarAffineTerm(one(T), x)], zero(T))
+end
+
+function _collect_affine_terms(::Type{T}, x::Number) where {T}
+    return (MOI.ScalarAffineTerm{T}[], T(x))
+end

test/Project.toml

@@ -1,6 +1,8 @@
 [deps]
 GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde"
+HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]