Support vector sets in MOI wrapper (#669)

* Support vector sets in MOI wrapper

* Add comment

* Fix

* Update MOI_wrapper.jl

* Update MOI_wrapper.jl

---------

Co-authored-by: Oscar Dowson <odow@users.noreply.github.com>

src/Constraint.jl

@@ -3,17 +3,17 @@
 # Use of this source code is governed by a BSD-style license that can be found
 # in the LICENSE file or at https://opensource.org/license/bsd-2-clause
 
-mutable struct Constraint{S<:MOI.AbstractSet}
+mutable struct Constraint{S<:MOI.AbstractVectorSet}
     child::AbstractExpr
     set::S
     dual::Union{Value,Nothing}
 
-    function Constraint(child::AbstractExpr, set::MOI.AbstractSet)
+    function Constraint(child::AbstractExpr, set::MOI.AbstractVectorSet)
         return new{typeof(set)}(child, set, nothing)
     end
 end
 
-function Constraint{S}(child::AbstractExpr) where {S<:MOI.AbstractSet}
+function Constraint{S}(child::AbstractExpr) where {S<:MOI.AbstractVectorSet}
     return Constraint(child, set_with_size(S, size(child)))
 end
 
@@ -34,7 +34,7 @@ end
 
 head(io::IO, c::Constraint) = head(io, c.set)
 
-function head(io::IO, set::MOI.AbstractSet)
+function head(io::IO, set::MOI.AbstractVectorSet)
     return print(io, replace("$(typeof(set))", "MathOptInterface" => "MOI"))
 end
 
@@ -80,7 +80,7 @@ function vexity(
     return ConvexVexity()
 end
 
-function vexity(::Any, set::MOI.AbstractSet)
+function vexity(::Any, set::MOI.AbstractVectorSet)
     return error(
         "`Convex.vexity(vex, ::$(typeof(set)))`: is not yet implemented. Please open an issue at https://github.com/jump-dev/Convex.jl",
     )

src/MOI_wrapper.jl

@@ -116,10 +116,18 @@ function MOI.add_constraint(
 end
 
 function MOI.supports_constraint(
-    ::Optimizer{T},
-    ::Type{MOI.ScalarNonlinearFunction},
-    ::Type{<:Union{MOI.EqualTo{T},MOI.GreaterThan{T},MOI.LessThan{T}}},
-) where {T}
+    ::Optimizer,
+    ::Type{MOI.VectorNonlinearFunction},
+    ::Type{<:MOI.AbstractVectorSet},
+)
+    # This can cause false positives because:
+    #  1) some sets might not be supported by Convex.jl (e.g., `vexity` might
+    #     be missing)
+    #  2) whether we support the constraint can depend on the vexity of the
+    #     function, which we currently don't know.
+    # Rather than attempt an enumeration of supported sets here, let's just
+    # pass things on and hope that there is a nice error message elsewhere in
+    # the callchain.
     return true
 end
 
@@ -189,28 +197,20 @@ function _expr(model::Optimizer, f::MOI.ScalarNonlinearFunction)
     return throw(MOI.UnsupportedNonlinearOperator(f.head))
 end
 
-function MOI.get(::Optimizer, ::MOI.ListOfSupportedNonlinearOperators)
-    return Symbol[:+, :-, :*, :/, :^, :min, :max, :abs, :sqrt, :exp, :log]
-end
-
-function _constraint(expr::AbstractExpr, set::MOI.EqualTo)
-    return expr == MOI.constant(set)
+function _expr(model::Optimizer, f::MOI.VectorNonlinearFunction)
+    return vcat(_expr.(model, f.rows)...)
 end
 
-function _constraint(expr::AbstractExpr, set::MOI.LessThan)
-    return expr <= MOI.constant(set)
-end
-
-function _constraint(expr::AbstractExpr, set::MOI.GreaterThan)
-    return expr >= MOI.constant(set)
+function MOI.get(::Optimizer, ::MOI.ListOfSupportedNonlinearOperators)
+    return Symbol[:+, :-, :*, :/, :^, :min, :max, :abs, :sqrt, :exp, :log]
 end
 
 function MOI.add_constraint(
     model::Optimizer{T},
-    func::MOI.ScalarNonlinearFunction,
-    set::MOI.AbstractScalarSet,
+    func::MOI.VectorNonlinearFunction,
+    set::MOI.AbstractVectorSet,
 ) where {T}
-    constraint = _constraint(_expr(model, func), set)
+    constraint = Constraint(_expr(model, func), set)
     add_constraint!(model.context, constraint)
     push!(model.constraint_map, model.context.constr_to_moi_inds[constraint])
     return MOI.ConstraintIndex{typeof(func),typeof(set)}(
@@ -317,10 +317,9 @@ end
 function MOI.get(
     model::Optimizer,
     attr::Union{MOI.ConstraintDual,MOI.ConstraintPrimal},
-    ci::MOI.ConstraintIndex{MOI.ScalarNonlinearFunction,S},
-) where {S<:MOI.AbstractScalarSet}
-    ret = MOI.get(model.context.model, attr, model.constraint_map[ci.value])
-    return ret[]
+    ci::MOI.ConstraintIndex{MOI.VectorNonlinearFunction,S},
+) where {S<:MOI.AbstractVectorSet}
+    return MOI.get(model.context.model, attr, model.constraint_map[ci.value])
 end
 
 function MOI.get(model::Optimizer, I::Type{<:MOI.Index}, name::String)