diff --git a/docs/src/submodules/Utilities/reference.md b/docs/src/submodules/Utilities/reference.md index e40184e8a4..832961797f 100644 --- a/docs/src/submodules/Utilities/reference.md +++ b/docs/src/submodules/Utilities/reference.md @@ -58,6 +58,8 @@ Utilities.latex_formulation ```@docs Utilities.automatic_copy_to Utilities.default_copy_to +Utilities.IndexMap +Utilities.identity_index_map ``` ### [Allocate-Load API](@id allocate_load_api_ref) diff --git a/src/Utilities/DoubleDicts.jl b/src/Utilities/DoubleDicts.jl index 51ff63dbdf..326a5ff402 100644 --- a/src/Utilities/DoubleDicts.jl +++ b/src/Utilities/DoubleDicts.jl @@ -417,7 +417,9 @@ end end function initialize_inner!(d::AbstractWithType{F,S,V,D}) where {F,S,V,D} - return d.inner = D() + d.inner = D() + d.dict.dict[(F, S)] = d.inner + return end inner_is_empty(d::AbstractWithType)::Bool = d.inner === nothing diff --git a/src/Utilities/copy.jl b/src/Utilities/copy.jl index d3c3a3826b..c15d406628 100644 --- a/src/Utilities/copy.jl +++ b/src/Utilities/copy.jl @@ -58,6 +58,16 @@ function dense_variable_dict(::Type{V}, n) where {V} ) end +""" + struct IndexMap <: AbstractDict{MOI.Index,MOI.Index} + varmap::DenseVariableDict{MOI.VariableIndex} + conmap::DoubleDicts.MainIndexDoubleDict + end + +Dictionary-like object returned by [`MathOptInterface.copy_to`](@ref) that +contains the mapping between variable indices in `varmap` and between +constraint indices in `conmap`. +""" struct IndexMap <: AbstractDict{MOI.Index,MOI.Index} varmap::DenseVariableDict{MOI.VariableIndex} conmap::DoubleDicts.MainIndexDoubleDict @@ -70,6 +80,33 @@ function IndexMap(n = 0) ) end +function _identity_constraints_map( + model, + index_map::MOIU.DoubleDicts.IndexWithType{F,S}, +) where {F,S} + for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) + index_map[ci] = ci + end +end + +""" + identity_index_map(model::MOI.ModelLike) + +Return an [`IndexMap`](@ref) that maps all variable and constraint indices of +`model` to themselves. +""" +function identity_index_map(model::MOI.ModelLike) + vis = MOI.get(model, MOI.ListOfVariableIndices()) + index_map = IndexMap(length(vis)) + for vi in vis + index_map[vi] = vi + end + for (F, S) in MOI.get(model, MOI.ListOfConstraints()) + _identity_constraints_map(model, index_map.conmap[F, S]) + end + return index_map +end + """ index_map_for_variable_indices(variables) diff --git a/test/Utilities/copy.jl b/test/Utilities/copy.jl index 3014a6868b..8b5b1b69ed 100644 --- a/test/Utilities/copy.jl +++ b/test/Utilities/copy.jl @@ -13,6 +13,22 @@ function compare_without_moi(x::String, y::String) @test remove_moi(x) == remove_moi(y) end +function identity_index_map_test(::Type{T}) where {T} + model = MOI.Utilities.Model{T}() + x, y = MOI.add_variables(model, 2) + fx = MOI.SingleVariable(x) + fy = MOI.SingleVariable(y) + cx = MOI.add_constraint(model, fx, MOI.EqualTo(one(T))) + cy = MOI.add_constraint(model, fy, MOI.EqualTo(zero(T))) + c = MOI.add_constraint(model, one(T) * fx + fy, MOI.LessThan(zero(T))) + index_map = MOI.Utilities.identity_index_map(model) + @test x == index_map[x] + @test y == index_map[y] + @test cx == index_map[cx] + @test cy == index_map[cy] + @test c == index_map[c] +end + @testset "IndexMap" begin map = MOIU.IndexMap() @test length(map) == 0 @@ -695,3 +711,8 @@ end test_pass_copy(Int) test_pass_copy(Float64) end + +@testset "identity_index_map" begin + identity_index_map_test(Int) + identity_index_map_test(Float64) +end