Refactor PrimalDualMap (#183)

* Cleanup renames

* Restructure PDMap

* Complete restructure of PrimalDualMap

* add docs, fix names and add errors

* format

Author: joaquimg

docs/src/reference.md

@@ -4,4 +4,7 @@
 Dualization.supported_constraints
 Dualization.supported_objective
 Dualization.DualNames
+Dualization.VariableData
+Dualization.ConstraintData
+Dualization.PrimalDualMap
 ```

src/MOI_wrapper.jl

@@ -126,7 +126,7 @@ function _change_constant(
     model,
     ci::MOI.ConstraintIndex{<:MOI.ScalarAffineFunction,S},
     constant,
-    idx,
+    ::Int,
 ) where {S}
     MOI.set(model, MOI.ConstraintSet(), ci, S(constant))
     return
@@ -136,7 +136,7 @@ function _change_constant(
     model,
     ci::MOI.ConstraintIndex{<:MOI.VectorAffineFunction},
     constant,
-    idx,
+    idx::Int,
 )
     func = MOI.get(model, MOI.ConstraintFunction(), ci)
     constants = copy(func.constant)
@@ -159,23 +159,19 @@ function MOI.modify(
         constant = -constant
     end
     vi = obj_change.variable
-    if haskey(primal_dual_map.primal_convar_to_primal_convarcon_and_index, vi)
-        ci_primal, index =
-            primal_dual_map.primal_convar_to_primal_convarcon_and_index[vi]
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci_primal]
-        if ci_dual === NO_CONSTRAINT
-            return
-        end
+    data = get(primal_dual_map.primal_variable_data, vi, nothing)
+    if data === nothing
+        # error
+    elseif data.dual_constraint === nothing
+        return
+    elseif data.primal_constrained_variable_constraint !== nothing
         constant = -constant
-    else
-        ci_dual = primal_dual_map.primal_var_to_dual_con[vi]
-        index = 1
     end
     _change_constant(
         optimizer.dual_problem.dual_model,
-        ci_dual,
+        data.dual_constraint,
         constant,
-        index,
+        data.primal_constrained_variable_index,
     )
     return
 end
@@ -246,32 +242,31 @@ function MOI.get(
     optimizer::DualOptimizer{T},
     ::MOI.VariablePrimal,
     vi::MOI.VariableIndex,
-) where {T}
+)::T where {T}
     primal_dual_map = optimizer.dual_problem.primal_dual_map
-    if haskey(primal_dual_map.primal_convar_to_primal_convarcon_and_index, vi)
-        ci_primal, idx =
-            primal_dual_map.primal_convar_to_primal_convarcon_and_index[vi]
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci_primal]
-        if ci_dual === NO_CONSTRAINT
-            return zero(T)
-        elseif ci_dual isa MOI.ConstraintIndex{<:MOI.AbstractVectorFunction}
-            return MOI.get(
-                optimizer.dual_problem.dual_model,
-                MOI.ConstraintDual(),
-                ci_dual,
-            )[idx]
-        else
-            return MOI.get(
-                optimizer.dual_problem.dual_model,
-                MOI.ConstraintDual(),
-                ci_dual,
-            )
-        end
-    else
+    data = get(primal_dual_map.primal_variable_data, vi, nothing)
+    if data === nothing
+        # error
+    elseif data.dual_constraint === nothing
+        return zero(T)
+    elseif data.primal_constrained_variable_constraint === nothing
         return -MOI.get(
             optimizer.dual_problem.dual_model,
             MOI.ConstraintDual(),
-            primal_dual_map.primal_var_to_dual_con[vi],
+            data.dual_constraint,
+        )
+    elseif data.dual_constraint isa
+           MOI.ConstraintIndex{<:MOI.AbstractVectorFunction}
+        return MOI.get(
+            optimizer.dual_problem.dual_model,
+            MOI.ConstraintDual(),
+            data.dual_constraint,
+        )[data.primal_constrained_variable_index]
+    else
+        return MOI.get(
+            optimizer.dual_problem.dual_model,
+            MOI.ConstraintDual(),
+            data.dual_constraint,
         )
     end
 end
@@ -282,16 +277,17 @@ function MOI.get(
     ci::MOI.ConstraintIndex{F,S},
 ) where {F<:MOI.AbstractScalarFunction,S<:MOI.AbstractScalarSet}
     primal_dual_map = optimizer.dual_problem.primal_dual_map
-    if haskey(primal_dual_map.primal_convarcon_to_dual_con, ci)
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci]
-        if ci_dual === NO_CONSTRAINT
+    if haskey(primal_dual_map.primal_constrained_variables, ci)
+        vi = primal_dual_map.primal_constrained_variables[ci][]
+        ci_dual = primal_dual_map.primal_variable_data[vi].dual_constraint
+        if ci_dual === nothing
             return MOI.Utilities.eval_variables(
-                primal_dual_map.primal_convarcon_to_dual_function[ci],
-            ) do vi
+                primal_dual_map.primal_variable_data[vi].dual_function,
+            ) do inner_vi
                 return MOI.get(
                     optimizer.dual_problem.dual_model,
                     MOI.VariablePrimal(),
-                    vi,
+                    inner_vi,
                 )
             end
         end
@@ -309,7 +305,7 @@ function MOI.get(
         return MOI.get(
             optimizer.dual_problem.dual_model,
             MOI.VariablePrimal(),
-            primal_dual_map.primal_con_to_dual_var_vec[ci][],
+            primal_dual_map.primal_constraint_data[ci].dual_variables[],
         )
     end
 end
@@ -320,29 +316,32 @@ function MOI.get(
     ci::MOI.ConstraintIndex{F,S},
 ) where {F<:MOI.AbstractVectorFunction,S<:MOI.AbstractVectorSet}
     primal_dual_map = optimizer.dual_problem.primal_dual_map
-    if haskey(primal_dual_map.primal_convarcon_to_dual_con, ci)
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci]
-        if ci_dual === NO_CONSTRAINT
-            return MOI.Utilities.eval_variables(
-                primal_dual_map.primal_convarcon_to_dual_function[ci],
-            ) do vi
-                return MOI.get(
-                    optimizer.dual_problem.dual_model,
-                    MOI.VariablePrimal(),
-                    vi,
-                )
-            end
+    if !haskey(primal_dual_map.primal_constraint_data, ci)
+        vis = primal_dual_map.primal_constrained_variables[ci]
+        ci_dual = primal_dual_map.primal_variable_data[vis[1]].dual_constraint
+        if ci_dual === nothing
+            return [
+                MOI.Utilities.eval_variables(
+                    primal_dual_map.primal_variable_data[vi].dual_function,
+                ) do inner_vi
+                    return MOI.get(
+                        optimizer.dual_problem.dual_model,
+                        MOI.VariablePrimal(),
+                        inner_vi,
+                    )
+                end for vi in vis
+            ]
         end
         return MOI.get(
             optimizer.dual_problem.dual_model,
             MOI.ConstraintPrimal(),
-            primal_dual_map.primal_convarcon_to_dual_con[ci],
+            ci_dual,
         )
     else
         return MOI.get.(
             optimizer.dual_problem.dual_model,
             MOI.VariablePrimal(),
-            primal_dual_map.primal_con_to_dual_var_vec[ci],
+            primal_dual_map.primal_constraint_data[ci].dual_variables,
         )
     end
 end
@@ -353,9 +352,11 @@ function MOI.get(
     ci::MOI.ConstraintIndex{F,S},
 ) where {T,F<:MOI.AbstractScalarFunction,S<:MOI.AbstractScalarSet}
     primal_dual_map = optimizer.dual_problem.primal_dual_map
-    if haskey(primal_dual_map.primal_convarcon_to_dual_con, ci)
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci]
-        if ci_dual === NO_CONSTRAINT
+    data = get(primal_dual_map.primal_constraint_data, ci, nothing)
+    if data === nothing
+        first_vi = primal_dual_map.primal_constrained_variables[ci][1]
+        ci_dual = primal_dual_map.primal_variable_data[first_vi].dual_constraint
+        if ci_dual === nothing
             return zero(T)
         else
             return MOI.get(
@@ -365,17 +366,16 @@ function MOI.get(
             )
         end
     else
-        primal_ci_constant =
-            primal_dual_map.primal_con_to_primal_constants_vec[ci][1]
+        primal_ci_constant = data.primal_set_constants[1]
         # If it has no key then there is no dual constraint
-        if !haskey(primal_dual_map.primal_con_to_dual_convarcon, ci)
+        ci_dual = data.dual_constrained_variable_constraint
+        if ci_dual === nothing
             return -primal_ci_constant
         end
-        ci_dual_problem = primal_dual_map.primal_con_to_dual_convarcon[ci]
         return MOI.get(
             optimizer.dual_problem.dual_model,
             MOI.ConstraintDual(),
-            ci_dual_problem,
+            ci_dual,
         ) - primal_ci_constant
     end
 end
@@ -386,13 +386,12 @@ function MOI.get(
     ci::MOI.ConstraintIndex{F,S},
 ) where {T,F<:MOI.AbstractVectorFunction,S<:MOI.AbstractVectorSet}
     primal_dual_map = optimizer.dual_problem.primal_dual_map
-    if haskey(primal_dual_map.primal_convarcon_to_dual_con, ci)
-        ci_dual = primal_dual_map.primal_convarcon_to_dual_con[ci]
-        if ci_dual === NO_CONSTRAINT
-            n = MOI.output_dimension(
-                primal_dual_map.primal_convarcon_to_dual_function[ci],
-            )
-            return zeros(T, n)
+    data = get(primal_dual_map.primal_constraint_data, ci, nothing)
+    if data === nothing
+        vis = primal_dual_map.primal_constrained_variables[ci]
+        ci_dual = primal_dual_map.primal_variable_data[vis[1]].dual_constraint
+        if ci_dual === nothing
+            return zeros(T, length(vis))
         else
             return MOI.get(
                 optimizer.dual_problem.dual_model,
@@ -401,18 +400,17 @@ function MOI.get(
             )
         end
     else
+        ci_dual = data.dual_constrained_variable_constraint
         # If it has no key then there is no dual constraint
-        if !haskey(primal_dual_map.primal_con_to_dual_convarcon, ci)
+        if ci_dual === nothing
             # The number of dual variable associated with the primal constraint is the ci dimension
-            ci_dimension =
-                length(primal_dual_map.primal_con_to_dual_var_vec[ci])
+            ci_dimension = length(data.dual_variables)
             return zeros(T, ci_dimension)
         end
-        ci_dual_problem = primal_dual_map.primal_con_to_dual_convarcon[ci]
         return MOI.get(
             optimizer.dual_problem.dual_model,
             MOI.ConstraintDual(),
-            ci_dual_problem,
+            ci_dual,
         )
     end
 end

src/constrained_variables.jl

@@ -14,26 +14,27 @@ function _select_constrained_variables(
             primal_model,
         )
     params = Set(variable_parameters)
+    selected_variables = Set{MOI.VariableIndex}()
     for S in single_or_vector_variables_types
         _select_constrained_variables(
             dual_problem.primal_dual_map,
             primal_model,
             S,
             params,
+            selected_variables,
         )
     end
     return
 end
 
-const NO_CONSTRAINT = MOI.ConstraintIndex{Nothing,Nothing}(0)
-
 # Function barrier for the type instability of `F` and `S`.
 function _select_constrained_variables(
-    m::PrimalDualMap,
+    m::PrimalDualMap{T},
     primal_model,
     ::Type{S},
     params,
-) where {S<:MOI.AbstractVectorSet}
+    selected_variables,
+) where {T,S<:MOI.AbstractVectorSet}
     cis = MOI.get(
         primal_model,
         MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}(),
@@ -44,46 +45,40 @@ function _select_constrained_variables(
         if all(
             # no element of the VectorOfVariables is a constrained variable
             # and not a parameter
-            vi ->
-                !haskey(m.primal_convar_to_primal_convarcon_and_index, vi) &&
-                    !(vi in params),
+            vi -> !(vi in selected_variables) && !(vi in params),
             f.variables,
         )
-            for (i, vi) in enumerate(f.variables)
-                m.primal_convar_to_primal_convarcon_and_index[vi] = (ci, i)
+            m.primal_constrained_variables[ci] = f.variables
+            for vi in f.variables
+                push!(selected_variables, vi)
             end
-            # Placeholder to indicate this constraint is part of constrained variables,
-            # it will be replaced later with a dual constraints
-            m.primal_convarcon_to_dual_con[ci] = NO_CONSTRAINT
         end
     end
     return
 end
 
 function _select_constrained_variables(
-    m::PrimalDualMap,
+    m::PrimalDualMap{T},
     primal_model,
     ::Type{S},
     params,
-) where {S<:MOI.AbstractScalarSet}
+    selected_variables,
+) where {T,S<:MOI.AbstractScalarSet}
     cis = MOI.get(
         primal_model,
         MOI.ListOfConstraintIndices{MOI.VariableIndex,S}(),
     )
     for ci in cis
-        f = MOI.get(primal_model, MOI.ConstraintFunction(), ci)
+        vi = MOI.get(primal_model, MOI.ConstraintFunction(), ci)
         # no element of the VectorOfVariables is a constrained variable
         # and not a parameter
-        if !haskey(m.primal_convar_to_primal_convarcon_and_index, f) &&
-           !(f in params)
+        if !(vi in selected_variables) && !(vi in params)
             set = MOI.get(primal_model, MOI.ConstraintSet(), ci)
             if !iszero(MOI.constant(set))
                 continue
             end
-            m.primal_convar_to_primal_convarcon_and_index[f] = (ci, 1)
-            # Placeholder to indicate this constraint is part of constrained variables,
-            # it will be replaced later with a dual constraints
-            m.primal_convarcon_to_dual_con[ci] = NO_CONSTRAINT
+            m.primal_constrained_variables[ci] = [vi]
+            push!(selected_variables, vi)
         end
     end
     return