Unify the trait-vs-list pattern in power flow aux var provisioning

ext/PowerFlowsExt/pf_input_mapping.jl

@@ -267,42 +267,71 @@ function _add_two_terminal_elements_map!(
     return
 end
 
-# Trait that determines what branch aux vars we can get from each PowerFlowContainer
-branch_aux_vars(::PFS.ACPowerFlowData) =
-    [
-        POM.PowerFlowBranchReactivePowerFromTo,
-        POM.PowerFlowBranchReactivePowerToFrom,
-        POM.PowerFlowBranchActivePowerFromTo,
-        POM.PowerFlowBranchActivePowerToFrom,
-        POM.PowerFlowBranchActivePowerLoss,
-    ]
-branch_aux_vars(::PFS.ABAPowerFlowData) =
-    [POM.PowerFlowBranchActivePowerFromTo, POM.PowerFlowBranchActivePowerToFrom]
-branch_aux_vars(::PFS.PTDFPowerFlowData) =
-    [POM.PowerFlowBranchActivePowerFromTo, POM.PowerFlowBranchActivePowerToFrom]
-branch_aux_vars(::PFS.vPTDFPowerFlowData) =
-    [POM.PowerFlowBranchActivePowerFromTo, POM.PowerFlowBranchActivePowerToFrom]
-branch_aux_vars(::PFS.PSSEExporter) = DataType[]
-
-# Same for bus aux vars. Loss/voltage-stability factors are registered ONLY when their
-# `get_calculate_*` flag is set — the same flag under which `_get_pf_result` returns a
-# non-`nothing` matrix (`PFS.get_loss_factors` / `get_voltage_stability_factors` are `nothing`
-# otherwise). Keep these two conditions in lockstep so the read-back never indexes a `nothing`.
-function bus_aux_vars(data::PFS.ACPowerFlowData)
-    vars = [POM.PowerFlowVoltageAngle, POM.PowerFlowVoltageMagnitude]
-    if PFS.get_calculate_loss_factors(data)
-        push!(vars, POM.PowerFlowLossFactors)
-    end
-    if PFS.get_calculate_voltage_stability_factors(data)
-        push!(vars, POM.PowerFlowVoltageStabilityFactors)
-    end
-    return vars
+# ─── Which auxiliary variables can we read back from a given power flow? ──────────────
+# `_pf_provides_aux_var(T, pf_data)` answers "does this power flow
+# container compute a value for aux variable `T`". Both the registration path
+# (`add_power_flow_data!`, via `branch_aux_vars`/`bus_aux_vars` below) and the read-back
+# guard (`calculate_aux_variable_value!`) go through it.
+#
+# The fallback is `false`: a container with no specialization registers nothing and no-ops on
+# read-back.
+_pf_provides_aux_var(::Type{<:POM.PowerFlowAuxVariableType}, ::PFS.PowerFlowContainer) =
+    false
+
+# AC power flow solves the full branch flows and the full bus voltage state.
+_pf_provides_aux_var(::Type{<:POM.BranchFlowAuxVariableType}, ::PFS.ACPowerFlowData) = true
+_pf_provides_aux_var(::Type{POM.PowerFlowBranchActivePowerLoss}, ::PFS.ACPowerFlowData) =
+    true
+_pf_provides_aux_var(::Type{POM.PowerFlowVoltageAngle}, ::PFS.ACPowerFlowData) = true
+_pf_provides_aux_var(::Type{POM.PowerFlowVoltageMagnitude}, ::PFS.ACPowerFlowData) = true
+
+# Loss/voltage-stability factors are provided ONLY when their `get_calculate_*` flag is
+# set — the same flag under which `_get_pf_result` returns a non-`nothing` matrix
+# (`PFS.get_loss_factors` / `get_voltage_stability_factors` are `nothing` otherwise). Keep
+# these two conditions in lockstep so the read-back never indexes a `nothing`. Runtime
+# state, but the same trait shape: only the right-hand side differs.
+_pf_provides_aux_var(::Type{POM.PowerFlowLossFactors}, data::PFS.ACPowerFlowData) =
+    PFS.get_calculate_loss_factors(data)
+_pf_provides_aux_var(
+    ::Type{POM.PowerFlowVoltageStabilityFactors},
+    data::PFS.ACPowerFlowData,
+) = PFS.get_calculate_voltage_stability_factors(data)
+
+# Every DC power flow (ABA, PTDF, vPTDF) gives active branch flows, and only those.
+_pf_provides_aux_var(
+    ::Type{POM.PowerFlowBranchActivePowerFromTo},
+    ::PFS.PowerFlowData{<:PFS.AbstractDCPowerFlow},
+) = true
+_pf_provides_aux_var(
+    ::Type{POM.PowerFlowBranchActivePowerToFrom},
+    ::PFS.PowerFlowData{<:PFS.AbstractDCPowerFlow},
+) = true
+
+# ABA additionally solves for the bus angles; the PTDF formulations never form them.
+_pf_provides_aux_var(::Type{POM.PowerFlowVoltageAngle}, ::PFS.ABAPowerFlowData) = true
+
+"""
+The aux variable types over components of type `C` that `pf_data` provides, derived from
+the `_pf_provides_aux_var` trait so there is no per-container list to keep in sync.
+
+Written so that all type logic inside happens at compile time:
+1. `POM.pf_aux_var_types(C)` returns a tuple, so each element's type is known statically.
+2. `map` ensures `_pf_provides_aux_var` resolves statically.
+3. `provides` is a named function with a type parameter, not a lambda, so Julia specializes.
+"""
+function _provided_aux_vars(
+    pf_data::PFS.PowerFlowContainer,
+    ::Type{C},
+) where {C <: PSY.Component}
+    candidates = POM.pf_aux_var_types(C)
+    provides(::Type{T}) where {T} = _pf_provides_aux_var(T, pf_data)
+    provided = map(provides, candidates)
+    return DataType[candidates[i] for i in eachindex(candidates) if provided[i]]
 end
 
-bus_aux_vars(::PFS.ABAPowerFlowData) = [POM.PowerFlowVoltageAngle]
-bus_aux_vars(::PFS.PTDFPowerFlowData) = DataType[]
-bus_aux_vars(::PFS.vPTDFPowerFlowData) = DataType[]
-bus_aux_vars(::PFS.PSSEExporter) = DataType[]
+branch_aux_vars(pf_data::PFS.PowerFlowContainer) =
+    _provided_aux_vars(pf_data, PSY.ACBranch)
+bus_aux_vars(pf_data::PFS.PowerFlowContainer) = _provided_aux_vars(pf_data, PSY.ACBus)
 
 # TODO: Needs update for MultiTerminal HVDC
 _get_branch_component_tuples(sys::PSY.System) = [

ext/PowerFlowsExt/pf_solve_and_aux.jl

@@ -149,9 +149,7 @@ function IOM.calculate_aux_variable_value!(
     # Skip the aux vars that the current power flow isn't meant to update
     pf_e_data = latest_solved_power_flow_evaluation_data(container)
     pf_data = IOM.get_inner_data(pf_e_data)
-    key_type = IOM.get_entry_type(key)
-    (key_type in branch_aux_vars(pf_data) || key_type in bus_aux_vars(pf_data)) ||
-        return
+    _pf_provides_aux_var(IOM.get_entry_type(key), pf_data) || return
     IOM.calculate_aux_variable_value!(container, key, system, pf_e_data)
     return
 end

src/core/auxiliary_variables.jl

@@ -88,6 +88,36 @@ convert_output_to_natural_units(
     },
 ) = true
 
+"""
+The `PowerFlowAuxVariableType`s that are defined over components of type `C` — i.e. whose
+values are indexed by branch or by bus. This is the complete universe of power flow
+auxiliary variables; which *subset* of it a particular power flow evaluator actually
+provides is the `_pf_provides_aux_var` trait in the `PowerFlowsExt`.
+
+Returns a tuple rather than a `Vector` deliberately: callers iterate it with `map`, so each
+element keeps its concrete `Type{T}` and the `_pf_provides_aux_var` calls resolve at
+compile time. Adding a `PowerFlowAuxVariableType` means adding it here and giving it
+`_pf_provides_aux_var` methods; `test_power_flow_in_the_loop.jl` asserts by reflection that
+no concrete subtype is missing from these tuples, so a forgotten entry fails in CI rather
+than silently never registering.
+"""
+function pf_aux_var_types end
+
+pf_aux_var_types(::Type{PSY.ACBranch}) = (
+    PowerFlowBranchReactivePowerFromTo,
+    PowerFlowBranchReactivePowerToFrom,
+    PowerFlowBranchActivePowerFromTo,
+    PowerFlowBranchActivePowerToFrom,
+    PowerFlowBranchActivePowerLoss,
+)
+
+pf_aux_var_types(::Type{PSY.ACBus}) = (
+    PowerFlowVoltageAngle,
+    PowerFlowVoltageMagnitude,
+    PowerFlowLossFactors,
+    PowerFlowVoltageStabilityFactors,
+)
+
 "Whether the auxiliary variable is calculated using a `PowerFlowEvaluationModel`"
 # Default is_from_evaluator(::Type{<:AuxVariableType}) = false is in IOM interfaces.jl
 is_from_evaluator(::Type{<:PowerFlowAuxVariableType}) = true

test/test_power_flow_in_the_loop.jl

@@ -1,3 +1,23 @@
+@testset "pf_aux_var_types enumerates every PowerFlowAuxVariableType" begin
+    # `_provided_aux_vars` walks the `pf_aux_var_types` tuples rather than reflecting over
+    # the type tree at runtime, so that the `_pf_provides_aux_var` trait calls resolve and
+    # fold at compile time. This test guards the gap that trade opens: a new
+    # `PowerFlowAuxVariableType` left out of the tuples would silently never be registered
+    # by any power flow evaluator.
+    branch_types = POM.pf_aux_var_types(PSY.ACBranch)
+    bus_types = POM.pf_aux_var_types(PSY.ACBus)
+    declared = union(Set(branch_types), Set(bus_types))
+    defined = Set(IS.get_all_concrete_subtypes(POM.PowerFlowAuxVariableType))
+
+    @test setdiff(defined, declared) == Set{DataType}()   # a type nothing would register
+    @test setdiff(declared, defined) == Set{DataType}()   # a stale/removed entry
+    # No type may be claimed as both branch- and bus-indexed.
+    @test isempty(intersect(Set(branch_types), Set(bus_types)))
+    # Tuples, not vectors: the compile-time folding depends on it.
+    @test branch_types isa Tuple
+    @test bus_types isa Tuple
+end
+
 @testset "AC Power Flow in the loop with headroom-proportional slack" begin
     system = build_system(PSITestSystems, "c_sys5_uc")