From 780b85cfe697b3cef73db9487d6f59e7f6366f7c Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Fri, 17 Jul 2026 14:07:25 +1200 Subject: [PATCH 1/2] Change ProductOfSets to remove dependence on implementation in MOI --- src/diff_opt.jl | 3 +- src/product_of_sets.jl | 185 +++++++++++++++++++++++++++++++++-------- 2 files changed, 154 insertions(+), 34 deletions(-) diff --git a/src/diff_opt.jl b/src/diff_opt.jl index 3f3e02a7f..ae36b2d0e 100644 --- a/src/diff_opt.jl +++ b/src/diff_opt.jl @@ -732,7 +732,8 @@ 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 = mapreduce(length, +, cones.rows; init = 0) + return Vector{T}(undef, n) end function _allocate_rows(cones, ::Flattened{T}) where {T} diff --git a/src/product_of_sets.jl b/src/product_of_sets.jl index 09ed0a308..b632d2373 100644 --- a/src/product_of_sets.jl +++ b/src/product_of_sets.jl @@ -4,40 +4,42 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. """ - ProductOfSets{T} <: MOI.Utilities.OrderedProductOfSets{T} + ProductOfSets{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. +This struct is inspired by `MOI.Utilities.@product_of_sets`. + +The difference is that the MOI 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 at +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. +mutable struct ProductOfSets{T} """ - num_rows::Vector{Int} + `rows[i][j]` corresponds to constraint `j` of set type `i`. + The value depends on `final_touch`: + * Before `final_touch`, these are `1:dimension` of the constraint + * After `final_touch`, these are the 1-indexed rows of the full + constraint matrix """ - 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} + rows::Vector{Vector{UnitRange{Int}}} """ A sanity bit to check that we don't call functions out-of-order. """ final_touch::Bool + """ + The set types, and a dictionary mapping S to the integer index. This list is + defined at run-time. + """ set_types::Vector{Type} set_types_dict::Dict{Type,Int} function ProductOfSets{T}() where {T} return new( - Int[], - Dict{Tuple{Int,Int},Int}(), + Vector{UnitRange{Int}}[], false, Type[], Dict{Type,Int}(), @@ -45,30 +47,147 @@ mutable struct ProductOfSets{T} <: MOI.Utilities.OrderedProductOfSets{T} end end -function MOI.Utilities.set_index(set::ProductOfSets, S::Type{<:MOI.AbstractSet}) +function MOI.Utilities.set_index( + set::ProductOfSets, + ::Type{S}, +) where {S<: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 + MOI.empty!(set) + for S in set_types + add_set_types(set, S) 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 +function add_set_types(set::ProductOfSets, ::Type{S}) where {S} + if haskey(set.set_types_dict, S) + return false + end + push!(set.rows, Vector{UnitRange{Int}}[]) + push!(set.set_types, S) + set.set_types_dict[S] = length(set.set_types) + return true +end + +MOI.is_empty(sets::ProductOfSets) = all(isempty, sets.rows) + +function MOI.empty!(sets::ProductOfSets) + map(empty!, sets.rows) + sets.final_touch = false + return +end + +function MOI.dimension(sets::ProductOfSets)::Int + @assert sets.final_touch + for i in reverse(eachindex(sets.rows)) + if !isempty(sets.rows[i]) + return last(sets.rows[i][end]) + end + end + return 0 # All rows were empty. +end + +function MOI.Utilities.rows( + sets::ProductOfSets{T}, + ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},S}, +)::Int where {T,S} + @assert sets.final_touch + i = MOI.Utilities.set_index(sets, S)::Int + return only(sets.rows[i][ci.value]) +end + +function MOI.Utilities.rows( + sets::ProductOfSets{T}, + ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{T},S}, +)::UnitRange{Int} where {T,S} + @assert sets.final_touch + i = MOI.Utilities.set_index(sets, S)::Int + return sets.rows[i][ci.value] +end + +function MOI.Utilities.add_set(sets::ProductOfSets, i::Int, dim::Int = 1)::Int64 + @assert !sets.final_touch + push!(sets.rows[i], 1:dim) + return length(sets.rows[i]) +end + +function MOI.Utilities.final_touch(sets::ProductOfSets)::Nothing + @assert !sets.final_touch + offset = 0 + for (i, rows) in enumerate(sets.rows) + for (j, row) in enumerate(rows) + rows[j] = offset .+ row + offset += length(row) + end + end + sets.final_touch = true + return +end + +function MOI.Utilities.num_rows(sets::ProductOfSets, ::Type{S})::Int where {S} + i = MOI.Utilities.set_index(sets, S)::Int + rows = sets.rows[i] + if isempty(rows) + return 0 + elseif sets.final_touch + return max(0, last(rows[end]) - first(rows[1]) + 1) + else + return mapreduce(length, +, rows) + end +end + +function MOI.get( + sets::ProductOfSets{T}, + ::MOI.ListOfConstraintTypesPresent, +)::Vector{Tuple{Type,Type}} where {T} + ret = Tuple{Type,Type}[] + for (i, S) in enumerate(MOI.Utilities.set_types(sets)) + if isempty(sets.rows[i]) + continue + elseif S <: MOI.AbstractScalarSet + push!(ret, (MOI.ScalarAffineFunction{T}, S)) + else + @assert S <: MOI.AbstractVectorSet + push!(ret, (MOI.VectorAffineFunction{T}, S)) + end + end + return ret +end + +function MOI.get( + sets::ProductOfSets, + ::MOI.NumberOfConstraints{F,S}, +)::Int64 where {F,S} + i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} + if i == nothing + return 0 + end + return length(sets.rows[i]) +end + +function MOI.get( + sets::ProductOfSets, + ::MOI.ListOfConstraintIndices{F,S}, +)::Vector{MOI.ConstraintIndex{F,S}} where {F,S} + i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} + if i == nothing + return MOI.ConstraintIndex{F,S}[] + end + return MOI.ConstraintIndex{F,S}.(1:length(sets.rows[i])) +end + +function MOI.is_valid( + sets::ProductOfSets, + ci::MOI.ConstraintIndex{F,S}, +)::Bool where {F,S} + i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} + if i == nothing + return false end - return false + return 1 <= ci.value <= length(sets.rows[i]) end From b159d031211002b35481cf5def49abfb00a1aa47 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Fri, 24 Jul 2026 15:59:00 +1200 Subject: [PATCH 2/2] Update to use MOI.Utilities.RuntimeProductOfSets --- Project.toml | 3 + src/ConicProgram/ConicProgram.jl | 5 +- src/DiffOpt.jl | 1 - src/diff_opt.jl | 40 +++++-- src/product_of_sets.jl | 193 ------------------------------- test/Project.toml | 3 + test/conic_program.jl | 27 ++--- 7 files changed, 50 insertions(+), 222 deletions(-) delete mode 100644 src/product_of_sets.jl diff --git a/Project.toml b/Project.toml index ebfa1203e..3b281521a 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 80e97d3ce..1079bc255 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 27f16d5ea..f2ccb0834 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 ae36b2d0e..350c2777a 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,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) diff --git a/src/product_of_sets.jl b/src/product_of_sets.jl deleted file mode 100644 index b632d2373..000000000 --- a/src/product_of_sets.jl +++ /dev/null @@ -1,193 +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} - -This struct is inspired by `MOI.Utilities.@product_of_sets`. - -The difference is that the MOI 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 at -run-time. -""" -mutable struct ProductOfSets{T} - """ - `rows[i][j]` corresponds to constraint `j` of set type `i`. - - The value depends on `final_touch`: - * Before `final_touch`, these are `1:dimension` of the constraint - * After `final_touch`, these are the 1-indexed rows of the full - constraint matrix - """ - rows::Vector{Vector{UnitRange{Int}}} - - """ - A sanity bit to check that we don't call functions out-of-order. - """ - final_touch::Bool - - """ - The set types, and a dictionary mapping S to the integer index. This list is - defined at run-time. - """ - set_types::Vector{Type} - set_types_dict::Dict{Type,Int} - - function ProductOfSets{T}() where {T} - return new( - Vector{UnitRange{Int}}[], - false, - Type[], - Dict{Type,Int}(), - ) - end -end - -function MOI.Utilities.set_index( - set::ProductOfSets, - ::Type{S}, -) where {S<: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) - MOI.empty!(set) - for S in set_types - add_set_types(set, S) - end - return -end - -function add_set_types(set::ProductOfSets, ::Type{S}) where {S} - if haskey(set.set_types_dict, S) - return false - end - push!(set.rows, Vector{UnitRange{Int}}[]) - push!(set.set_types, S) - set.set_types_dict[S] = length(set.set_types) - return true -end - -MOI.is_empty(sets::ProductOfSets) = all(isempty, sets.rows) - -function MOI.empty!(sets::ProductOfSets) - map(empty!, sets.rows) - sets.final_touch = false - return -end - -function MOI.dimension(sets::ProductOfSets)::Int - @assert sets.final_touch - for i in reverse(eachindex(sets.rows)) - if !isempty(sets.rows[i]) - return last(sets.rows[i][end]) - end - end - return 0 # All rows were empty. -end - -function MOI.Utilities.rows( - sets::ProductOfSets{T}, - ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},S}, -)::Int where {T,S} - @assert sets.final_touch - i = MOI.Utilities.set_index(sets, S)::Int - return only(sets.rows[i][ci.value]) -end - -function MOI.Utilities.rows( - sets::ProductOfSets{T}, - ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{T},S}, -)::UnitRange{Int} where {T,S} - @assert sets.final_touch - i = MOI.Utilities.set_index(sets, S)::Int - return sets.rows[i][ci.value] -end - -function MOI.Utilities.add_set(sets::ProductOfSets, i::Int, dim::Int = 1)::Int64 - @assert !sets.final_touch - push!(sets.rows[i], 1:dim) - return length(sets.rows[i]) -end - -function MOI.Utilities.final_touch(sets::ProductOfSets)::Nothing - @assert !sets.final_touch - offset = 0 - for (i, rows) in enumerate(sets.rows) - for (j, row) in enumerate(rows) - rows[j] = offset .+ row - offset += length(row) - end - end - sets.final_touch = true - return -end - -function MOI.Utilities.num_rows(sets::ProductOfSets, ::Type{S})::Int where {S} - i = MOI.Utilities.set_index(sets, S)::Int - rows = sets.rows[i] - if isempty(rows) - return 0 - elseif sets.final_touch - return max(0, last(rows[end]) - first(rows[1]) + 1) - else - return mapreduce(length, +, rows) - end -end - -function MOI.get( - sets::ProductOfSets{T}, - ::MOI.ListOfConstraintTypesPresent, -)::Vector{Tuple{Type,Type}} where {T} - ret = Tuple{Type,Type}[] - for (i, S) in enumerate(MOI.Utilities.set_types(sets)) - if isempty(sets.rows[i]) - continue - elseif S <: MOI.AbstractScalarSet - push!(ret, (MOI.ScalarAffineFunction{T}, S)) - else - @assert S <: MOI.AbstractVectorSet - push!(ret, (MOI.VectorAffineFunction{T}, S)) - end - end - return ret -end - -function MOI.get( - sets::ProductOfSets, - ::MOI.NumberOfConstraints{F,S}, -)::Int64 where {F,S} - i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} - if i == nothing - return 0 - end - return length(sets.rows[i]) -end - -function MOI.get( - sets::ProductOfSets, - ::MOI.ListOfConstraintIndices{F,S}, -)::Vector{MOI.ConstraintIndex{F,S}} where {F,S} - i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} - if i == nothing - return MOI.ConstraintIndex{F,S}[] - end - return MOI.ConstraintIndex{F,S}.(1:length(sets.rows[i])) -end - -function MOI.is_valid( - sets::ProductOfSets, - ci::MOI.ConstraintIndex{F,S}, -)::Bool where {F,S} - i = MOI.Utilities.set_index(sets, S)::Union{Nothing,Int} - if i == nothing - return false - end - return 1 <= ci.value <= length(sets.rows[i]) -end diff --git a/test/Project.toml b/test/Project.toml index 2893cbd97..07dfbe33b 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 a45b1a461..6fa280ec7 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)