From 780b85cfe697b3cef73db9487d6f59e7f6366f7c Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Fri, 17 Jul 2026 14:07:25 +1200 Subject: [PATCH] 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