Prepare integration with DiffOpt (#161)

* POI + DiffOpt = S2 * fix POI (pp), add reverse mode, add tests * fix format * add more tests * fix test for correct quadratic constraint interpretation * fix test * add format * fix test * format diff tests * add tests * remove GLPK * fix format * concentrate changes * remove diffopt from poi * remove from tests * format * add tests for new attributes * Update Project.toml * Document attributes --------- Co-authored-by: Joaquim Garcia <[email protected]>
jump-dev · Jan 7, 2025 · 455541d · 455541d
1 parent 4e3757e
commit 455541d
Show file tree

Hide file tree

Showing 4 changed files with 450 additions and 26 deletions.
diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl
@@ -99,7 +99,8 @@ function MOI.is_empty(model::Optimizer)
            #
            isempty(model.multiplicative_parameters) &&
            isempty(model.dual_value_of_parameters) &&
-           model.number_of_parameters_in_model == 0
+           model.number_of_parameters_in_model == 0 &&
+           isempty(model.ext)
 end
 
 function MOI.empty!(model::Optimizer{T}) where {T}
@@ -133,6 +134,7 @@ function MOI.empty!(model::Optimizer{T}) where {T}
     empty!(model.dual_value_of_parameters)
     #
     model.number_of_parameters_in_model = 0
+    empty!(model.ext)
     return
 end
 
@@ -1124,24 +1126,161 @@ end
 # Special Attributes
 #
 
+"""
+    ParametricObjectiveType <: MOI.AbstractModelAttribute
+
+A model attribute for the type `P` of the ParametricOptInterface's parametric
+function type in the objective function. The value os `P` can be `Nothing` if
+the objective function is not parametric. The parametric function type can be
+queried using the [`ParametricObjectiveFunction{P}`](@ref) attribute. The type
+`P` can be `ParametricAffineFunction{T}` or `ParametricQuadraticFunction{T}`.
+"""
+struct ParametricObjectiveType <: MOI.AbstractModelAttribute end
+
+function MOI.get(model::Optimizer{T}, ::ParametricObjectiveType) where {T}
+    if model.quadratic_objective_cache !== nothing
+        return ParametricQuadraticFunction{T}
+    elseif model.affine_objective_cache !== nothing
+        return ParametricAffineFunction{T}
+    end
+    return Nothing
+end
+
+"""
+    ParametricObjectiveFunction{P} <: MOI.AbstractModelAttribute
+
+A model attribute for the parametric objective function of type `P`. The type
+`P` can be `ParametricAffineFunction{T}` or `ParametricQuadraticFunction{T}`.
+"""
+struct ParametricObjectiveFunction{T} <: MOI.AbstractModelAttribute end
+
+function MOI.get(
+    model::Optimizer{T},
+    ::ParametricObjectiveFunction{ParametricQuadraticFunction{T}},
+) where {T}
+    if model.quadratic_objective_cache === nothing
+        error("
+            There is no parametric quadratic objective function in the model.
+        ")
+    end
+    return model.quadratic_objective_cache
+end
+
+function MOI.get(
+    model::Optimizer{T},
+    ::ParametricObjectiveFunction{ParametricAffineFunction{T}},
+) where {T}
+    if model.affine_objective_cache === nothing
+        error("
+            There is no parametric affine objective function in the model.
+        ")
+    end
+    return model.affine_objective_cache
+end
+
+"""
+    ListOfParametricConstraintTypesPresent()
+
+A model attribute for the list of tuples of the form `(F,S,P)`, where `F` is a
+MOI function type, `S` is a set type and `P` is a ParametricOptInterface
+parametric function type indicating that the attribute
+[`DictOfParametricConstraintIndicesAndFunctions{F,S,P}`](@ref) returns a
+non-empty dictionary.
+"""
+struct ListOfParametricConstraintTypesPresent <: MOI.AbstractModelAttribute end
+
+function MOI.get(
+    model::Optimizer{T},
+    ::ListOfParametricConstraintTypesPresent,
+) where {T}
+    output = Set{Tuple{DataType,DataType,DataType}}()
+    for (F, S) in MOI.Utilities.DoubleDicts.nonempty_outer_keys(
+        model.affine_constraint_cache,
+    )
+        push!(output, (F, S, ParametricAffineFunction{T}))
+    end
+    for (F, S) in MOI.Utilities.DoubleDicts.nonempty_outer_keys(
+        model.vector_affine_constraint_cache,
+    )
+        push!(output, (F, S, ParametricVectorAffineFunction{T}))
+    end
+    for (F, S) in MOI.Utilities.DoubleDicts.nonempty_outer_keys(
+        model.quadratic_constraint_cache,
+    )
+        push!(output, (F, S, ParametricQuadraticFunction{T}))
+    end
+    return collect(output)
+end
+
+"""
+    DictOfParametricConstraintIndicesAndFunctions{F,S,P}
+
+A model attribute for a dictionary mapping constraint indices to parametric
+functions. The key is a constraint index with scalar function type `F`
+and set type `S` and the value is a parametric function of type `P`.
+"""
+struct DictOfParametricConstraintIndicesAndFunctions{F,S,P} <:
+       MOI.AbstractModelAttribute end
+
+function MOI.get(
+    model::Optimizer,
+    ::DictOfParametricConstraintIndicesAndFunctions{F,S,P},
+) where {F,S,P<:ParametricAffineFunction}
+    return model.affine_constraint_cache[F, S]
+end
+
+function MOI.get(
+    model::Optimizer,
+    ::DictOfParametricConstraintIndicesAndFunctions{F,S,P},
+) where {F,S,P<:ParametricVectorAffineFunction}
+    return model.vector_affine_constraint_cache[F, S]
+end
+
+function MOI.get(
+    model::Optimizer,
+    ::DictOfParametricConstraintIndicesAndFunctions{F,S,P},
+) where {F,S,P<:ParametricQuadraticFunction}
+    return model.quadratic_constraint_cache[F, S]
+end
+
+"""
+    NumberOfPureVariables
+
+A model attribute for the number of pure variables in the model.
+"""
 struct NumberOfPureVariables <: MOI.AbstractModelAttribute end
 
 function MOI.get(model::Optimizer, ::NumberOfPureVariables)
     return length(model.variables)
 end
 
+"""
+    ListOfPureVariableIndices
+
+A model attribute for the list of pure variable indices in the model.
+"""
 struct ListOfPureVariableIndices <: MOI.AbstractModelAttribute end
 
 function MOI.get(model::Optimizer, ::ListOfPureVariableIndices)
     return collect(keys(model.variables))::Vector{MOI.VariableIndex}
 end
 
+"""
+    NumberOfParameters
+
+A model attribute for the number of parameters in the model.
+"""
 struct NumberOfParameters <: MOI.AbstractModelAttribute end
 
 function MOI.get(model::Optimizer, ::NumberOfParameters)
     return length(model.parameters)
 end
 
+"""
+    ListOfParameterIndices
+
+A model attribute for the list of parameter indices in the model.
+"""
 struct ListOfParameterIndices <: MOI.AbstractModelAttribute end
 
 function MOI.get(model::Optimizer, ::ListOfParameterIndices)

diff --git a/src/ParametricOptInterface.jl b/src/ParametricOptInterface.jl
@@ -166,6 +166,9 @@ mutable struct Optimizer{T,OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
     number_of_parameters_in_model::Int64
     constraints_interpretation::ConstraintsInterpretationCode
     save_original_objective_and_constraints::Bool
+
+    # extension data
+    ext::Dict{Symbol,Any}
     function Optimizer(
         optimizer::OT;
         evaluate_duals::Bool = true,
@@ -223,6 +226,7 @@ mutable struct Optimizer{T,OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
             0,
             ONLY_CONSTRAINTS,
             save_original_objective_and_constraints,
+            Dict{Symbol,Any}(),
         )
     end
 end