diff --git a/Project.toml b/Project.toml index ebfa1203..3b281521 100644 --- a/Project.toml +++ b/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"} diff --git a/src/ConicProgram/ConicProgram.jl b/src/ConicProgram/ConicProgram.jl index 80e97d3c..1079bc25 100644 --- a/src/ConicProgram/ConicProgram.jl +++ b/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 diff --git a/src/DiffOpt.jl b/src/DiffOpt.jl index 27f16d5e..f2ccb083 100644 --- a/src/DiffOpt.jl +++ b/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") diff --git a/src/diff_opt.jl b/src/diff_opt.jl index 3f3e02a7..350c2777 100644 --- a/src/diff_opt.jl +++ b/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,11 @@ end # Allocate a vector for storing the output of `map_rows`. function _allocate_rows(cones, ::Nested{T}) where {T} - return Vector{T}(undef, length(cones.dimension)) + 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 function _allocate_rows(cones, ::Flattened{T}) where {T} @@ -740,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. @@ -752,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) diff --git a/src/product_of_sets.jl b/src/product_of_sets.jl deleted file mode 100644 index 09ed0a30..00000000 --- a/src/product_of_sets.jl +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) 2020: Akshay Sharma 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. - -""" - ProductOfSets{T} <: MOI.Utilities.OrderedProductOfSets{T} - -The `MOI.Utilities.@product_of_sets` macro requires to know the list of sets -at compile time. In DiffOpt however, the list depends on what the user is going -to use as set as DiffOpt supports any set as long as it implements the -required function of MathOptSetDistances. -For this type, the list of sets can be given a run-time. -""" -mutable struct ProductOfSets{T} <: MOI.Utilities.OrderedProductOfSets{T} - """ - During the copy, this counts the number of rows corresponding to - each set. At the end of copy, `final_touch` is called, which - converts this list into a cumulative ordering. - """ - num_rows::Vector{Int} - - """ - A dictionary which maps the `set_index` and `offset` of a set to the - dimension, i.e., `dimension[(set_index,offset)] → dim`. - """ - dimension::Dict{Tuple{Int,Int},Int} - - """ - A sanity bit to check that we don't call functions out-of-order. - """ - final_touch::Bool - - set_types::Vector{Type} - set_types_dict::Dict{Type,Int} - - function ProductOfSets{T}() where {T} - return new( - Int[], - Dict{Tuple{Int,Int},Int}(), - false, - Type[], - Dict{Type,Int}(), - ) - end -end - -function MOI.Utilities.set_index(set::ProductOfSets, S::Type{<:MOI.AbstractSet}) - return get(set.set_types_dict, S, nothing) -end - -MOI.Utilities.set_types(set::ProductOfSets) = set.set_types - -function set_set_types(set::ProductOfSets, set_types) - resize!(set.num_rows, length(set_types)) - fill!(set.num_rows, 0) - resize!(set.set_types, length(set_types)) - copy!(set.set_types, set_types) - empty!(set.set_types_dict) - for i in eachindex(set_types) - set.set_types_dict[set_types[i]] = i - end - return -end - -function add_set_types(set::ProductOfSets, S::Type) - if !haskey(set.set_types_dict, S) - push!(set.num_rows, 0) - push!(set.set_types, S) - set.set_types_dict[S] = length(set.set_types) - return true - end - return false -end diff --git a/test/Project.toml b/test/Project.toml index 2893cbd9..07dfbe33 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -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"} diff --git a/test/conic_program.jl b/test/conic_program.jl index a45b1a46..6fa280ec 100644 --- a/test/conic_program.jl +++ b/test/conic_program.jl @@ -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)