Create ComponentPairKey for service models

src/InfrastructureOptimizationModels.jl

@@ -313,6 +313,7 @@ export ConstraintKey
 export ParameterKey
 export ExpressionKey
 export AuxVarKey
+export ComponentPairKey
 
 # Abstract Key Types (from InfrastructureSystems.Optimization)
 export VariableType

src/common_models/add_variable.jl

@@ -77,45 +77,95 @@ function add_variables!(
 end
 
 """
-Add variables to the OptimizationContainer for a single service and its contributing
+Add variables to the OptimizationContainer for every service of a type and their contributing
 devices.
 
-All services of a given `(VariableType, ServiceType)` share a single sparse container
-keyed by `(service_name, device_name, time)`, rather than one dense container per
-service disambiguated by a `meta = service_name` field. The container is created lazily
-on the first service of a type, and each subsequent call appends that service's slice, so
-separate formulation groups sharing a service type append to the same container.
+Each `(device type, service type)` pair gets its own sparse container keyed on
+`ComponentPairKey{D, U}` and indexed by `(service_name, device_name, time)`, holding every
+service of that type. Keying on the pair keeps devices of different types that share a name,
+and services of different types that share a name, apart.
 """
 function add_service_variables!(
     container::OptimizationContainer,
     ::Type{T},
+    services::Vector{U},
+    model::ServiceModel,
+    ::Type{F},
+) where {
+    T <: VariableType,
+    U <: IS.InfrastructureSystemsComponent,
+    F <: AbstractServiceFormulation,
+}
+    by_device_type = Dict{DataType, Vector{Tuple{U, Vector}}}()
+    for service in services
+        for (device_type, devices) in
+            get_contributing_devices_map(model, IS.get_name(service))
+            isempty(devices) && continue
+            push!(
+                get!(Vector{Tuple{U, Vector}}, by_device_type, device_type),
+                (service, devices),
+            )
+        end
+    end
+    for (device_type, entries) in by_device_type
+        _add_service_variables!(container, T, U, device_type, entries, F)
+    end
+    return
+end
+
+function _add_service_variables!(
+    container::OptimizationContainer,
+    ::Type{T},
+    ::Type{U},
+    ::Type{D},
+    entries::Vector{Tuple{U, Vector}},
+    ::Type{F},
+) where {
+    T <: VariableType,
+    U <: IS.InfrastructureSystemsComponent,
+    D <: IS.InfrastructureSystemsComponent,
+    F <: AbstractServiceFormulation,
+}
+    variable = add_variable_container!(
+        container,
+        T,
+        ComponentPairKey{D, U},
+        String[],
+        String[],
+        Int[];
+        sparse = true,
+    )
+    for (service, devices) in entries
+        _add_service_device_variables!(container, variable, T, service, devices, F)
+    end
+    return
+end
+
+function _add_service_device_variables!(
+    container::OptimizationContainer,
+    variable::SparseAxisArray,
+    ::Type{T},
     service::U,
-    contributing_devices::V,
+    devices::Vector{D},
     ::Type{F},
 ) where {
     T <: VariableType,
     U <: IS.InfrastructureSystemsComponent,
-    V <: Union{Vector{D}, IS.FlattenIteratorWrapper{D}},
+    D <: IS.InfrastructureSystemsComponent,
     F <: AbstractServiceFormulation,
-} where {D <: IS.InfrastructureSystemsComponent}
-    @assert !isempty(contributing_devices)
-    time_steps = get_time_steps(container)
+}
     settings = get_settings(container)
     binary = get_variable_binary(T, U, F)
-    s_name = IS.get_name(service)
-    device_names = [IS.get_name(d) for d in contributing_devices]
-    variable = lazy_container_addition!(
-        container, T, U, [s_name], device_names, time_steps; sparse = true,
-    )
+    service_name = IS.get_name(service)
     jump_model = get_jump_model(container)
-    for t in time_steps, d in contributing_devices
-        name = IS.get_name(d)
+    for d in devices, t in get_time_steps(container)
+        device_name = IS.get_name(d)
         var = JuMP.@variable(
             jump_model,
-            base_name = "$(T)_$(U)_{$(s_name), $(name), $(t)}",
+            base_name = "$(T)_$(D)_$(U)_{$(service_name), $(device_name), $(t)}",
             binary = binary,
         )
-        variable[(s_name, name, t)] = var
+        variable[service_name, device_name, t] = var
         ub = get_variable_upper_bound(T, service, d, F)
         ub !== nothing && JuMP.set_upper_bound(var, ub)
         lb = get_variable_lower_bound(T, service, d, F)

src/core/optimization_container.jl

@@ -1421,7 +1421,7 @@ end
     return :($K(T, U, meta))
 end
 
-# note these 3 lazy_container_addition! definitions have different meta handling and adder 
+# note these 3 lazy_container_addition! definitions have different meta handling and adder
 # functions, else we'd collapse into one generated function.
 function lazy_container_addition!(
     container::OptimizationContainer,

src/core/optimization_container_keys.jl

@@ -51,18 +51,29 @@ get_component_type(
 
 # okay to construct AuxVarKey with abstract component type, but not others.
 maybe_throw_if_abstract(::Type{T}, ::Type{U}) where {T <: OptimizationKeyType, U} =
-    isabstracttype(U) && throw(ArgumentError("Type $U can't be abstract"))
+    _is_abstract_component(U) && throw(ArgumentError("Type $U can't be abstract"))
 
 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}
+"""
+Component type for keying a container on two component types. Used downstream to disambiguate
+service variables where the contributing devices on a service can have the same name, or one device
+can contribute to multiple services with the same name.
+"""
+struct ComponentPairKey{
+    A <: IS.InfrastructureSystemsComponent,
+    B <: IS.InfrastructureSystemsComponent,
+} <: IS.InfrastructureSystemsComponent end
+
+_is_abstract_component(::Type{U}) where {U} = isabstracttype(U)
+_is_abstract_component(::Type{ComponentPairKey{A, B}}) where {A, B} =
+    isabstracttype(A) || isabstracttype(B)
+
+function _canonical_type(::Type{U}) where {U <: InfrastructureSystemsType}
     base = U isa UnionAll ? Base.unwrap_unionall(U) : U
-    base isa DataType || return :($U)
+    base isa DataType || return U
     params = collect(base.parameters)
     n = length(params)
     while n > 0 && (
@@ -71,10 +82,19 @@ const CONTAINER_KEY_EMPTY_META = ""
     )
         n -= 1
     end
-    n == length(params) && return :($U)
+    n == length(params) && return U
     kept = params[1:n]
-    stripped = isempty(kept) ? base.name.wrapper : base.name.wrapper{kept...}
-    return :($stripped)
+    return isempty(kept) ? base.name.wrapper : base.name.wrapper{kept...}
+end
+
+_canonical_type(::Type{ComponentPairKey{A, B}}) where {A, B} =
+    ComponentPairKey{_canonical_type(A), _canonical_type(B)}
+
+# Strip units before making keys for parametric structs
+@generated function canonical_component_type(
+    ::Type{U},
+) where {U <: InfrastructureSystemsType}
+    return :($(_canonical_type(U)))
 end
 
 # see https://discourse.julialang.org/t/parametric-constructor-where-type-being-constructed-is-parameter/129866/3
@@ -105,14 +125,21 @@ end
 
 ### Encoding keys ###
 
+_encode_type_str(::Type{U}) where {U <: InfrastructureSystemsType} =
+    replace(replace(strip_module_name(U), "{" => COMPONENT_NAME_DELIMITER), "}" => "")
+
+_encode_type_str(
+    ::Type{ComponentPairKey{A, B}},
+) where {A <: InfrastructureSystemsType, B <: InfrastructureSystemsType} =
+    _encode_type_str(A) * COMPONENT_NAME_DELIMITER * _encode_type_str(B)
+
 @generated function encode_symbol(
     ::Type{T},
     ::Type{U},
     meta::String = CONTAINER_KEY_EMPTY_META,
 ) where {T <: OptimizationKeyType, U <: InfrastructureSystemsType}
     meta_str = :meta
-    U_str =
-        replace(replace(strip_module_name(U), "{" => COMPONENT_NAME_DELIMITER), "}" => "")
+    U_str = _encode_type_str(U)
     T_str = strip_module_name(T)
 
     :(Symbol(

src/operation/decision_model_store.jl

@@ -125,6 +125,20 @@ function write_output!(
     return
 end
 
+function write_output!(
+    store::DecisionModelStore,
+    name::Symbol,
+    key::OptimizationContainerKey,
+    index::DecisionModelIndexType,
+    update_timestamp::Dates.DateTime,
+    array::DenseAxisArray{T, 3, <:Tuple{Vector{String}, Vector{Int}, UnitRange}},
+) where {T}
+    columns = get_column_names_from_axis_array(array)
+    container = getfield(store, get_store_container_type(key))
+    container[key][index] = DenseAxisArray(array.data, columns..., 1:size(array, 3))
+    return
+end
+
 # Sparse expressions (e.g., post-contingency flows keyed by
 # `(outage_id, branch_name, t)`) are pre-allocated as 2D dense storage with
 # the non-time tuple flattened into encoded `"a__b"` columns. Derive the

test/mocks/mock_services.jl

@@ -2,11 +2,38 @@
 Minimal service mocks.
 """
 
-struct MockReserve
+struct MockUp end
+struct MockDown end
+
+struct MockReserve{D} <: IS.InfrastructureSystemsComponent
     name::String
     requirement::Float64
     contributing_devices::Vector{Any}
 end
 
 get_name(r::MockReserve) = r.name
 get_requirement(r::MockReserve) = r.requirement
+
+struct MockReserveFormulation <: IOM.AbstractServiceFormulation end
+
+IOM.get_default_attributes(::Type{<:MockReserve}, ::Type{MockReserveFormulation}) =
+    Dict{String, Any}()
+IOM.get_default_time_series_names(::Type{<:MockReserve}, ::Type{MockReserveFormulation}) =
+    Dict{Type{<:IOM.ParameterType}, String}()
+IOM.get_variable_binary(
+    ::Type{<:IOM.VariableType},
+    ::Type{<:MockReserve},
+    ::Type{MockReserveFormulation},
+) = false
+IOM.get_variable_upper_bound(
+    ::Type{<:IOM.VariableType},
+    r::MockReserve,
+    ::IS.InfrastructureSystemsComponent,
+    ::Type{MockReserveFormulation},
+) = r.requirement
+IOM.get_variable_lower_bound(
+    ::Type{<:IOM.VariableType},
+    ::MockReserve,
+    ::IS.InfrastructureSystemsComponent,
+    ::Type{MockReserveFormulation},
+) = 0.0

test/test_optimization_container_keys.jl

@@ -99,3 +99,85 @@ IOM.should_write_resulting_value(::Type{MockExpression2}) = false
     )
     @test isa(made_key, VariableKey)
 end
+
+@testset "ComponentPairKey keys" begin
+    up_pair = IOM.ComponentPairKey{MockThermalGen, MockReserve{MockUp}}
+    down_pair = IOM.ComponentPairKey{MockThermalGen, MockReserve{MockDown}}
+
+    up_key = VariableKey(MockVariable, up_pair, "spin")
+    @test IOM.encode_key(up_key) ==
+          Symbol("MockVariable__MockThermalGen__MockReserve__MockUp__spin")
+
+    down_key = VariableKey(MockVariable, down_pair, "spin")
+    @test up_key != down_key
+    @test IOM.encode_key(up_key) != IOM.encode_key(down_key)
+
+    swapped_key = VariableKey(
+        MockVariable,
+        IOM.ComponentPairKey{MockReserve{MockUp}, MockThermalGen},
+        "spin",
+    )
+    @test swapped_key != up_key
+    @test IOM.encode_key(swapped_key) != IOM.encode_key(up_key)
+
+    @test IOM.make_key(VariableKey, MockVariable, up_pair, "spin") == up_key
+end
+
+@testset "add_service_variables! keys on (device type, service type)" begin
+    time_steps = 1:3
+    container = _setup_qa_container(time_steps)
+    bus = MockBus("bus", 1, :PV)
+    thermal = MockThermalGen("g1", true, bus, (min = 0.0, max = 1.0))
+    renewable = MockRenewableGen("g1", true, bus, 1.0)
+    contributors =
+        () -> Dict{DataType, Vector{<:IS.InfrastructureSystemsComponent}}(
+            MockThermalGen => [thermal],
+            MockRenewableGen => [renewable],
+        )
+
+    services = Dict(
+        MockReserve{MockUp} => MockReserve{MockUp}("spin", 0.5, Any[]),
+        MockReserve{MockDown} => MockReserve{MockDown}("spin", 0.25, Any[]),
+    )
+    for (S, service) in services
+        model = ServiceModel(
+            S,
+            MockReserveFormulation;
+            contributing_devices_map = Dict("spin" => contributors()),
+        )
+        IOM.add_service_variables!(
+            container,
+            MockVariable,
+            [service],
+            model,
+            MockReserveFormulation,
+        )
+        @test_logs (:error, r"already stored") match_mode = :any @test_throws IS.InvalidValue IOM.add_service_variables!(
+            container,
+            MockVariable,
+            [service],
+            model,
+            MockReserveFormulation,
+        )
+    end
+
+    pair_keys = [
+        VariableKey(MockVariable, IOM.ComponentPairKey{D, S})
+        for D in (MockThermalGen, MockRenewableGen), S in keys(services)
+    ]
+    @test length(IOM.get_variable_keys(container)) == 4
+    @test Set(IOM.get_variable_keys(container)) == Set(pair_keys)
+    @test length(unique(IOM.encode_key.(pair_keys))) == 4
+
+    refs = JuMP.VariableRef[]
+    for key in pair_keys
+        var = IOM.get_variable(container, key)
+        @test Set(keys(var.data)) == Set(("spin", "g1", t) for t in time_steps)
+        S = IOM.get_component_type(key).parameters[2]
+        for t in time_steps
+            @test JuMP.upper_bound(var["spin", "g1", t]) == services[S].requirement
+            push!(refs, var["spin", "g1", t])
+        end
+    end
+    @test length(unique(refs)) == 4 * length(time_steps)
+end

test/test_quadratic_approximations.jl

@@ -1,4 +1,3 @@
-const MOI = JuMP.MOI
 const TEST_META = "TestVar"
 
 @testset "Quadratic Approximations" begin