Ignore units in container keys; add constraint dual keys in a consistent order

Project.toml

@@ -39,7 +39,7 @@ JuMP = "^1.28"
 LinearAlgebra = "1"
 Logging = "1"
 MathOptInterface = "1"
-PowerNetworkMatrices = "^0.20"
+PowerNetworkMatrices = "^0.20, ^0.22"
 PrettyTables = "3.1"
 Random = "^1.10"
 Serialization = "1"

src/common_models/add_constraint_dual.jl

@@ -64,6 +64,22 @@ function assign_dual_variable!(
     return
 end
 
+_existing_constraint_keys(
+    container::OptimizationContainer,
+    ::Type{T},
+    ::Type{D},
+) where {T <: ConstraintType, D} = filter(
+    key -> get_entry_type(key) === T && get_component_type(key) === D,
+    get_constraint_keys(container),
+)
+
+_validate_keys(keys) =
+    isempty(keys) && throw(
+        IS.InvalidValue(
+            "No constraint of type $constraint_type for $D is stored; cannot assign a dual variable.",
+        ),
+    )
+
 # device formulation
 function assign_dual_variable!(
     container::OptimizationContainer,
@@ -75,33 +91,32 @@ function assign_dual_variable!(
 } where {D <: IS.InfrastructureSystemsComponent}
     @assert !isempty(devices)
     time_steps = get_time_steps(container)
-    metas = _existing_constraint_metas(container, constraint_type, D)
-    if isempty(metas)
-        device_names = IS.get_name.(devices)
-        add_dual_container!(container, constraint_type, D, device_names, time_steps)
-        return
-    end
-    for meta in metas
-        key = ConstraintKey(constraint_type, D, meta)
+    constraint_keys = _existing_constraint_keys(container, constraint_type, D)
+    _validate_keys(constraint_keys)
+    for key in constraint_keys
         existing = get_constraint(container, key)
         _assign_dual_from_existing!(container, key, existing, D, time_steps)
     end
     return
 end
 
-# Sparse constraints (e.g. post-contingency flow-rate constraints keyed by
-# (outage_id, name, t)) have no `axes`. Mirror the constraint's exact sparse keys
-# into a Float64 dual container so the dual matches the constraint storage one-to-one.
-function _assign_dual_from_existing!(
+# network model with buses
+function assign_dual_variable!(
     container::OptimizationContainer,
-    key::ConstraintKey,
-    existing::SparseAxisArray,
-    ::Type{D},
-    time_steps,
-) where {D}
-    dual_container =
-        SparseAxisArray(Dict(k => zero(Float64) for k in keys(existing.data)))
-    _assign_container!(container.duals, key, dual_container)
+    constraint_type::Type{<:ConstraintType},
+    devices::U,
+    ::NetworkModel{<:AbstractPowerModel},
+) where {
+    U <: Union{Vector{D}, IS.FlattenIteratorWrapper{D}},
+} where {D <: IS.InfrastructureSystemsComponent}
+    IS.@assert_op !isempty(devices)
+    time_steps = get_time_steps(container)
+    constraint_keys = _existing_constraint_keys(container, constraint_type, D)
+    _validate_keys(constraint_keys)
+    for key in constraint_keys
+        existing = get_constraint(container, key)
+        _assign_dual_from_existing!(container, key, existing, D, time_steps)
+    end
     return
 end
 
@@ -129,38 +144,19 @@ function _assign_dual_from_existing!(
     return
 end
 
-function _existing_constraint_metas(
+function _assign_dual_from_existing!(
     container::OptimizationContainer,
-    ::Type{T},
+    key::ConstraintKey,
+    existing::SparseAxisArray,
     ::Type{D},
-) where {T <: ConstraintType, D}
-    metas = String[]
-    for key in get_constraint_keys(container)
-        if get_entry_type(key) === T &&
-           get_component_type(key) === D
-            push!(metas, key.meta)
-        end
-    end
-    return metas
-end
-
-# network model with buses
-function assign_dual_variable!(
-    container::OptimizationContainer,
-    constraint_type::Type{<:ConstraintType},
-    devices::U,
-    ::NetworkModel{<:AbstractPowerModel},
-) where {
-    U <: Union{Vector{D}, IS.FlattenIteratorWrapper{D}},
-} where {D <: IS.InfrastructureSystemsComponent}
-    @assert !isempty(devices)
-    time_steps = get_time_steps(container)
+    time_steps,
+) where {D}
     add_dual_container!(
         container,
-        constraint_type,
+        get_entry_type(key),
         D,
-        IS.get_name.(devices),
-        time_steps,
+        keys(existing.data);
+        sparse = true,
     )
     return
 end

src/core/dual_processing.jl

@@ -2,6 +2,11 @@
 # duals are SparseAxisArray (Dict-backed), where `.data .= …` is undefined, so
 # copy per key instead.
 function _copy_dual_values!(dual::DenseAxisArray, constraint::DenseAxisArray)
+    # The dual container is built by reusing the constraint's axes (see
+    # `assign_dual_variable!` / `_assign_dual_from_existing!`), so a positional
+    # copy is correct only when the axes match. Mismatched axes would write each
+    # value onto the wrong label, so fail loudly instead.
+    IS.@assert_op axes(dual) == axes(constraint)
     dual.data .= jump_value.(constraint).data
     return
 end

src/core/optimization_container_keys.jl

@@ -57,24 +57,46 @@ maybe_throw_if_abstract(::Type{<:ConstraintType}, ::Type{U}) where {U} = nothing
 
 const CONTAINER_KEY_EMPTY_META = ""
 
+# Strip units before making keys for parametric structs
+@generated function canonical_component_type(
+    ::Type{U},
+) where {U <: InfrastructureSystemsType}
+    base = U isa UnionAll ? Base.unwrap_unionall(U) : U
+    base isa DataType || return :($U)
+    params = collect(base.parameters)
+    n = length(params)
+    while n > 0 && (
+        params[n] isa TypeVar ||
+        (params[n] isa Type && params[n] <: IS.AbstractUnitSystem)
+    )
+        n -= 1
+    end
+    n == length(params) && return :($U)
+    kept = params[1:n]
+    stripped = isempty(kept) ? base.name.wrapper : base.name.wrapper{kept...}
+    return :($stripped)
+end
+
 # see https://discourse.julialang.org/t/parametric-constructor-where-type-being-constructed-is-parameter/129866/3
 function (M::Type{S} where {S <: OptimizationContainerKey})(
     ::Type{T},
     ::Type{U},
     meta = CONTAINER_KEY_EMPTY_META,
 ) where {T <: OptimizationKeyType, U <: InfrastructureSystemsType}
     check_meta_chars(meta)
-    maybe_throw_if_abstract(T, U)
-    return M{T, U}(meta)
+    K = canonical_component_type(U)
+    maybe_throw_if_abstract(T, K)
+    return M{T, K}(meta)
 end
 
 function (M::Type{S} where {S <: OptimizationContainerKey})(
     ::Type{T},
     ::Type{U},
     meta::String,
 ) where {T <: OptimizationKeyType, U <: InfrastructureSystemsType}
-    maybe_throw_if_abstract(T, U)
-    return M{T, U}(meta)
+    K = canonical_component_type(U)
+    maybe_throw_if_abstract(T, K)
+    return M{T, K}(meta)
 end
 
 function make_key(
@@ -87,7 +109,7 @@ function make_key(
     T <: OptimizationKeyType,
     U <: InfrastructureSystemsType,
 }
-    return S{T, U}(meta)
+    return S{T, canonical_component_type(U)}(meta)
 end
 
 ### Encoding keys ###

src/operation/decision_model.jl

@@ -59,7 +59,7 @@ function DecisionModel{M}(
         OptimizationContainer(sys, settings, jump_model, ts_type),
     )
 
-    template_ = deepcopy(template)
+    template_ = _deepcopy_template(template)
     finalize_template!(template_, sys)
     model = DecisionModel{M}(
         name,

src/operation/problem_template.jl

@@ -135,6 +135,25 @@ function finalize_template!(template::AbstractProblemTemplate, args...)
     )
 end
 
+# Deep-copy a template while sharing the network model's PNM matrices by reference:
+# their solver caches hold raw factorization handles and deliberately error on deepcopy
+# (PNM #312). The matrices are read-only inputs, so sharing is safe.
+function _deepcopy_template(template::AbstractProblemTemplate)
+    network_model = get_network_model(template)
+    network_model === nothing && return deepcopy(template)
+    ptdf = network_model.PTDF_matrix
+    modf = network_model.MODF_matrix
+    network_model.PTDF_matrix = nothing
+    network_model.MODF_matrix = nothing
+    template_ = deepcopy(template)
+    network_model.PTDF_matrix = ptdf
+    network_model.MODF_matrix = modf
+    copied_network_model = get_network_model(template_)
+    copied_network_model.PTDF_matrix = ptdf
+    copied_network_model.MODF_matrix = modf
+    return template_
+end
+
 """
 Return the set of device types whose formulation is FixedOutput (incompatible with
 service provision).