Update to use MOI.Utilities.RuntimeProductOfSets

Author: odow

Project.toml

@@ -25,3 +25,6 @@ MathOptInterface = "1.18"
 MathOptSetDistances = "0.2.9"
 ParametricOptInterface = "0.15.3"
 julia = "1.10"
+
+[sources]
+MathOptInterface = {url = "https://github.com/jump-dev/MathOptInterface.jl", rev = "od/runtime-product-of-sets"}

src/ConicProgram/ConicProgram.jl

@@ -47,7 +47,7 @@ const Form{T} = MOI.Utilities.GenericModel{
             MOI.Utilities.OneBasedIndexing,
         },
         Vector{T},
-        DiffOpt.ProductOfSets{T},
+        MOI.Utilities.RuntimeProductOfSets{T},
     },
 }
 
@@ -133,7 +133,8 @@ function MOI.supports_constraint(
     F::Type{MOI.VectorAffineFunction{Float64}},
     ::Type{S},
 ) where {S<:MOI.AbstractVectorSet}
-    if DiffOpt.add_set_types(model.model.constraints.sets, S)
+    if MOI.Utilities.set_index(model.model.constraints.sets, S) === nothing
+        MOI.Utilities.add_set_type(model.model.constraints.sets, S)
         push!(model.model.constraints.caches, Tuple{F,S}[])
         push!(model.model.constraints.are_indices_mapped, BitSet())
     end

src/DiffOpt.jl

@@ -16,7 +16,6 @@ import ParametricOptInterface as POI
 import SparseArrays
 
 include("utils.jl")
-include("product_of_sets.jl")
 include("diff_opt.jl")
 include("moi_wrapper.jl")
 include("parameters.jl")

src/diff_opt.jl

@@ -659,14 +659,22 @@ function MOI.get(
 end
 
 """
-    π(v::Vector{Float64}, model::MOI.ModelLike, cones::ProductOfSets)
+    π(
+        v::Vector{T},
+        model::MOI.ModelLike,
+        cones::MOI.Utilities.RuntimeProductOfSets,
+    ) where {T}
 
 Given a `model`, its `cones`, find the projection of the vectors `v`
 of length equal to the number of rows in the conic form onto the cartesian
 product of the cones corresponding to these rows.
 For more info, refer to https://github.com/matbesancon/MathOptSetDistances.jl
 """
-function π(v::Vector{T}, model::MOI.ModelLike, cones::ProductOfSets) where {T}
+function π(
+    v::Vector{T},
+    model::MOI.ModelLike,
+    cones::MOI.Utilities.RuntimeProductOfSets,
+) where {T}
     return map_rows(model, cones, Flattened{T}()) do ci, r
         return MOSD.projection_on_set(
             MOSD.DefaultDistance(),
@@ -677,14 +685,22 @@ function π(v::Vector{T}, model::MOI.ModelLike, cones::ProductOfSets) where {T}
 end
 
 """
-    Dπ(v::Vector{Float64}, model, cones::ProductOfSets)
+    Dπ(
+        v::Vector{T},
+        model::MOI.ModelLike,
+        cones::MOI.Utilities.RuntimeProductOfSets,
+    ) where {T}
 
 Given a `model`, its `cones`, find the gradient of the projection of
 the vectors `v` of length equal to the number of rows in the conic form onto the
 cartesian product of the cones corresponding to these rows.
 For more info, refer to https://github.com/matbesancon/MathOptSetDistances.jl
 """
-function Dπ(v::Vector{T}, model::MOI.ModelLike, cones::ProductOfSets) where {T}
+function Dπ(
+    v::Vector{T},
+    model::MOI.ModelLike,
+    cones::MOI.Utilities.RuntimeProductOfSets,
+) where {T}
     return BlockDiagonals.BlockDiagonal(
         map_rows(model, cones, Nested{Matrix{T}}()) do ci, r
             return MOSD.projection_gradient_on_set(
@@ -716,7 +732,7 @@ function _map_rows!(
     f::Function,
     x::Vector,
     model,
-    cones::ProductOfSets,
+    cones::MOI.Utilities.RuntimeProductOfSets,
     ::Type{F},
     ::Type{S},
     map_mode,
@@ -732,7 +748,10 @@ end
 
 # Allocate a vector for storing the output of `map_rows`.
 function _allocate_rows(cones, ::Nested{T}) where {T}
-    n = mapreduce(length, +, cones.rows; init = 0)
+    n = 0
+    for (F, S) in MOI.get(cones, MOI.ListOfConstraintTypesPresent())
+        n += MOI.get(cones, MOI.NumberOfConstraints{F,S}())
+    end
     return Vector{T}(undef, n)
 end
 
@@ -741,7 +760,12 @@ function _allocate_rows(cones, ::Flattened{T}) where {T}
 end
 
 """
-    map_rows(f::Function, model, cones::ProductOfSets, map_mode::Union{Nested{T}, Flattened{T}})
+    map_rows(
+        f::Function,
+        model,
+        cones::MOI.Utilities.RuntimeProductOfSets,
+        map_mode::Union{Nested{T},Flattened{T}},
+    )
 
 Given a `model`, its `cones` and `map_mode` of type `Nested` (resp.
 `Flattened`), return a `Vector{T}` of length equal to the number of cones (resp.
@@ -753,7 +777,7 @@ form.
 function map_rows(
     f::Function,
     model,
-    cones::ProductOfSets,
+    cones::MOI.Utilities.RuntimeProductOfSets,
     map_mode::Union{Nested,Flattened},
 )
     x = _allocate_rows(cones, map_mode)

src/product_of_sets.jl

@@ -25,3 +25,6 @@ HiGHS = "1"
 Ipopt = "1.0.2"
 SCS = "1"
 MLDatasets = "0.7.18"
+
+[sources]
+MathOptInterface = {url = "https://github.com/jump-dev/MathOptInterface.jl", rev = "od/runtime-product-of-sets"}

test/Project.toml

@@ -181,13 +181,10 @@ function test_simple_psd()
     MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
     MOI.optimize!(model)
     x = MOI.get(model, MOI.VariablePrimal(), X)
-    cone_types = unique([
-        S for (F, S) in
-        MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
-    ])
     conic_form = DiffOpt.ConicProgram.Form{Float64}()
-    cones = conic_form.constraints.sets
-    DiffOpt.set_set_types(cones, cone_types)
+    for (F, S) in MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
+        MOI.Utilities.add_set_type(conic_form.constraints.sets, S)
+    end
     index_map = MOI.copy_to(conic_form, model)
     # s = DiffOpt.map_rows((ci, r) -> MOI.get(model.optimizer, MOI.ConstraintPrimal(), ci), model.optimizer, cones, index_map, DiffOpt.Flattened{Float64}())
     # y = DiffOpt.map_rows((ci, r) -> MOI.get(model.optimizer, MOI.ConstraintDual(), ci), model.optimizer, cones, index_map, DiffOpt.Flattened{Float64}())
@@ -337,13 +334,10 @@ function test_differentiating_conic_with_PSD_and_SOC_constraints()
     MOI.optimize!(model)
     _x = MOI.get(model, MOI.VariablePrimal(), x)
     _X = MOI.get(model, MOI.VariablePrimal(), X)
-    cone_types = unique([
-        S for (F, S) in
-        MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
-    ])
     conic_form = DiffOpt.ConicProgram.Form{Float64}()
-    cones = conic_form.constraints.sets
-    DiffOpt.set_set_types(cones, cone_types)
+    for (F, S) in MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
+        MOI.Utilities.add_set_type(conic_form.constraints.sets, S)
+    end
     index_map = MOI.copy_to(conic_form, model)
     # s = DiffOpt.map_rows((ci, r) -> MOI.get(model.optimizer, MOI.ConstraintPrimal(), ci), model.optimizer, cones, index_map, DiffOpt.Flattened{Float64}())
     # y = DiffOpt.map_rows((ci, r) -> MOI.get(model.optimizer, MOI.ConstraintDual(), ci), model.optimizer, cones, index_map, DiffOpt.Flattened{Float64}())
@@ -725,13 +719,10 @@ function test_differentiating_simple_PSD_back()
     MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
     MOI.optimize!(model)
     x = MOI.get(model, MOI.VariablePrimal(), X)
-    cone_types = unique([
-        S for (F, S) in
-        MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
-    ])
     conic_form = DiffOpt.ConicProgram.Form{Float64}()
-    cones = conic_form.constraints.sets
-    DiffOpt.set_set_types(cones, cone_types)
+    for (F, S) in MOI.get(model.optimizer, MOI.ListOfConstraintTypesPresent())
+        MOI.Utilities.add_set_type(conic_form.constraints.sets, S)
+    end
     index_map = MOI.copy_to(conic_form, model)
     @test x ≈ ones(3) atol = ATOL rtol = RTOL
     MOI.set(model, DiffOpt.ReverseVariablePrimal(), X[1], 1.0)