Expose a unified KKT layout interface

docs/src/advanced.md

@@ -124,9 +124,13 @@ prob = ACOPFProblem(net; silent=true)
 exa_prob = ACOPFProblem(net; backend=:exa, silent=true)
 ```
 
-## KKT System Access (Qualified)
+## KKT System Access
 
-KKT internals are available via qualified access (`PowerDiff.function_name`), not exported:
+`kkt_layout` is the primary public entry point for the flattened KKT system. It
+returns the dimension and named index ranges together so they always describe
+the same layout. `kkt_dims` and `kkt_indices` are conveniences for
+callers that need only one part. Lower-level KKT operators remain available via
+qualified access (`PowerDiff.function_name`).
 
 ```julia
 using PowerDiff
@@ -137,14 +141,16 @@ z = PD.flatten_variables(sol, prob)     # Solution → vector
 vars = PD.unflatten_variables(z, prob)  # Vector → named tuple
 K = PD.kkt(z, prob, d)                  # KKT residuals
 J = PD.calc_kkt_jacobian(prob)          # Sparse Jacobian dK/dz
-dim = PD.kkt_dims(dc_net)              # KKT dimension
-idx = PD.kkt_indices(dc_net)           # Named index ranges
+dim, idx = kkt_layout(dc_net)           # Dimension and named index ranges
 
 # AC OPF — same unified API
 z = PD.flatten_variables(sol, ac_prob)
 J = PD.calc_kkt_jacobian(ac_prob)       # Sparse analytical Jacobian
-dim = PD.kkt_dims(ac_prob)             # KKT dimension
-idx = PD.kkt_indices(ac_prob)          # Named index ranges
+dim, idx = kkt_layout(ac_prob)
+
+# Thin conveniences when only one component is needed
+dim = kkt_dims(ac_prob)
+idx = kkt_indices(ac_prob)
 ```
 
 ## LMP Sign Conventions

docs/src/api.md

@@ -30,11 +30,21 @@ jvp
 jvp!
 vjp
 vjp!
-kkt_dims
 dict_to_vec
 vec_to_dict
 ```
 
+## KKT Layout
+
+`kkt_layout` is the primary interface for inspecting a DC or AC KKT system.
+`kkt_dims` and `kkt_indices` return the individual parts of the same layout.
+
+```@docs
+kkt_layout
+kkt_dims
+kkt_indices
+```
+
 ## Introspection
 
 ```@docs

src/PowerDiff.jl

@@ -101,7 +101,8 @@ export IDMapping
 export calc_sensitivity, calc_sensitivity_column
 export Sensitivity, silence
 export operand_symbols, parameter_symbols
-export jvp, vjp, jvp!, vjp!, dict_to_vec, vec_to_dict, kkt_dims
+export jvp, vjp, jvp!, vjp!, dict_to_vec, vec_to_dict
+export kkt_layout, kkt_dims, kkt_indices
 export parse_file, parse_matpower, parse_matpower_struct, get_path
 
 # DC Power Flow Types

src/prob/kkt_ac_opf.jl

@@ -32,51 +32,12 @@
 _ref_bus_indices(prob::ACOPFProblem) = prob.data.ref_bus_keys
 
 # =============================================================================
-# Dimension Calculations
+# KKT Layout
 # =============================================================================
 
-"""
-    kkt_dims(prob::ACOPFProblem)
-
-Compute the dimension of the flattened KKT variable vector for AC OPF.
+_ac_kkt_dims(n::Int, m::Int, k::Int, n_ref::Int) = 6n + 12m + 6k + n_ref
 
-The KKT system includes:
-- Primal: va (n), vm (n), pg (k), qg (k)
-- Dual (equality): ν_p_bal (n), ν_q_bal (n), ν_ref_bus (n_ref)
-- Dual (inequality): λ_thermal_fr (m), λ_thermal_to (m),
-                     λ_angle_lb (m), λ_angle_ub (m),
-                     μ_vm_lb (n), μ_vm_ub (n),
-                     ρ_pg_lb (k), ρ_pg_ub (k), ρ_qg_lb (k), ρ_qg_ub (k),
-                     σ_p_fr_lb (m), σ_p_fr_ub (m), σ_q_fr_lb (m), σ_q_fr_ub (m),
-                     σ_p_to_lb (m), σ_p_to_ub (m), σ_q_to_lb (m), σ_q_to_ub (m)
-
-Total: 6n + 12m + 6k + n_ref
-"""
-function kkt_dims(prob::ACOPFProblem)
-    n, m, k = prob.network.n, prob.network.m, prob.n_gen
-    n_ref = length(prob.data.ref_bus_keys)
-    return 6n + 12m + 6k + n_ref
-end
-
-"""
-    kkt_indices(n, m, k, n_ref) → NamedTuple
-
-Compute all KKT variable indices from problem dimensions.
-Single source of truth for index calculations.
-
-# Variable ordering
-[va(n), vm(n), pg(k), qg(k),
- ν_p_bal(n), ν_q_bal(n), ν_ref_bus(n_ref),
- λ_thermal_fr(m), λ_thermal_to(m), λ_angle_lb(m), λ_angle_ub(m),
- μ_vm_lb(n), μ_vm_ub(n),
- ρ_pg_lb(k), ρ_pg_ub(k), ρ_qg_lb(k), ρ_qg_ub(k),
- σ_p_fr_lb(m), σ_p_fr_ub(m), σ_q_fr_lb(m), σ_q_fr_ub(m),
- σ_p_to_lb(m), σ_p_to_ub(m), σ_q_to_lb(m), σ_q_to_ub(m)]
-
-# Returns
-NamedTuple with index ranges for each variable block.
-"""
-function kkt_indices(n::Int, m::Int, k::Int, n_ref::Int)
+function _ac_kkt_indices(n::Int, m::Int, k::Int, n_ref::Int)
     i = 0
     # Primal
     idx_va = (i+1):(i+n); i += n
@@ -130,9 +91,34 @@ function kkt_indices(n::Int, m::Int, k::Int, n_ref::Int)
     )
 end
 
-function kkt_indices(prob::ACOPFProblem)
+@inline function _ac_kkt_layout(n::Int, m::Int, k::Int, n_ref::Int)
+    return _ac_kkt_dims(n, m, k, n_ref), _ac_kkt_indices(n, m, k, n_ref)
+end
+
+"""Return the AC KKT index ranges for the supplied bus, branch, generator, and reference-bus counts."""
+kkt_indices(n::Int, m::Int, k::Int, n_ref::Int) =
+    last(_ac_kkt_layout(n, m, k, n_ref))
+
+"""
+    kkt_layout(network::ACNetwork) → (dim, indices)
+    kkt_layout(prob::ACOPFProblem) → (dim, indices)
+
+Return the dimension and named index ranges of the flattened AC KKT variable
+vector.
+
+The layout contains the `va`, `vm`, `pg`, and `qg` primal blocks followed by
+the power-balance, reference-bus, thermal-limit, angle-limit, voltage-limit,
+generation-limit, and reduced-space flow-bound dual blocks. With `n` buses,
+`m` branches, `k` generators, and `n_ref` reference buses, the total dimension
+is `6n + 12m + 6k + n_ref`.
+"""
+function kkt_layout(net::ACNetwork)
+    return _ac_kkt_layout(net.n, net.m, length(net.gen_bus), length(net.ref_bus_keys))
+end
+
+function kkt_layout(prob::ACOPFProblem)
     n_ref = length(prob.data.ref_bus_keys)
-    kkt_indices(prob.network.n, prob.network.m, prob.n_gen, n_ref)
+    return _ac_kkt_layout(prob.network.n, prob.network.m, prob.n_gen, n_ref)
 end
 
 # =============================================================================
@@ -143,7 +129,7 @@ end
     flatten_variables(sol::ACOPFSolution, prob::ACOPFProblem)
 
 Flatten solution primal and dual variables into a single vector for KKT evaluation.
-Ordering matches `kkt_indices`.
+Ordering matches the indices returned by `kkt_layout`.
 """
 function flatten_variables(sol::ACOPFSolution, prob::ACOPFProblem)
     return vcat(
@@ -633,7 +619,7 @@ function calc_kkt_jacobian(prob::ACOPFProblem; sol::Union{ACOPFSolution,Nothing}
         sol = _ensure_ac_solved!(prob)
     end
 
-    idx = kkt_indices(prob)
+    dim, idx = kkt_layout(prob)
     constants = _extract_kkt_constants(prob)
     prob.cache.kkt_constants = constants
     cq = _extract_gen_cq(prob)
@@ -642,7 +628,6 @@ function calc_kkt_jacobian(prob::ACOPFProblem; sol::Union{ACOPFSolution,Nothing}
     vars = unflatten_variables(flatten_variables(sol, prob), idx)
     n, m, k = prob.network.n, prob.network.m, prob.n_gen
 
-    dim = kkt_dims(prob)
     row_idxs = Int[]
     col_idxs = Int[]
     vals = Float64[]
@@ -1203,9 +1188,9 @@ function calc_kkt_jacobian_param(prob::ACOPFProblem, sol::ACOPFSolution, param::
     haskey(_AC_PARAM_EXTRACT, param) || throw(ArgumentError(
         "Unknown AC OPF parameter: $param. Valid: $(keys(_AC_PARAM_EXTRACT))"))
 
-    idx = kkt_indices(prob)
+    dim, idx = kkt_layout(prob)
     p0 = _AC_PARAM_EXTRACT[param](prob)
-    Jp = zeros(Float64, kkt_dims(prob), length(p0))
+    Jp = zeros(Float64, dim, length(p0))
 
     constants = _require_kkt_constants(prob)
 
@@ -1359,10 +1344,10 @@ Compute a single analytical column of ∂K/∂param without materializing the fu
 parameter Jacobian.
 """
 function _calc_ac_kkt_param_column(prob::ACOPFProblem, sol::ACOPFSolution, param::Symbol, col_idx::Int)
-    idx = kkt_indices(prob)
+    dim, idx = kkt_layout(prob)
     vars = sol
     p0 = _AC_PARAM_EXTRACT[param](prob)
-    Kcol = zeros(Float64, kkt_dims(prob))
+    Kcol = zeros(Float64, dim)
 
     constants = _require_kkt_constants(prob)
 

src/prob/kkt_dc_opf.jl

@@ -268,54 +268,11 @@ function _extract_dz_column(prob::DCOPFProblem, dz_col::Vector{Float64}, op::Sym
 end
 
 # =============================================================================
-# Dimension Calculations
+# KKT Layout
 # =============================================================================
 
-"""
-    kkt_dims(prob::DCOPFProblem)
-    kkt_dims(network::DCNetwork)
-
-Compute the dimension of the flattened KKT variable vector.
-
-The KKT system includes:
-- Primal: va (n), pg (k), f (m), psh (n)
-- Dual (inequality): lam_lb (m), lam_ub (m), gamma_lb (m), gamma_ub (m), rho_lb (k), rho_ub (k), mu_lb (n), mu_ub (n)
-- Dual (equality): nu_bal (n), nu_flow (m)
-- Reference bus constraints: n_ref
-
-Total: 5n + 6m + 3k + n_ref
-"""
-function kkt_dims(prob::DCOPFProblem)
-    net = getfield(prob, :network)
-    return _dc_kkt_dims(getfield(net, :n), getfield(net, :m), getfield(net, :k),
-                        getfield(prob, :_n_ref))
-end
-kkt_dims(n::Int, m::Int, k::Int) = _dc_kkt_dims(n, m, k, 1)
 _dc_kkt_dims(n::Int, m::Int, k::Int, n_ref::Int) = 5n + 6m + 3k + n_ref
 
-function kkt_dims(net::DCNetwork)
-    n = getfield(net, :n)
-    m = getfield(net, :m)
-    k = getfield(net, :k)
-    n_ref = length(_reference_buses(net))
-    # va(n) + pg(k) + f(m) + psh(n) + lam_lb(m) + lam_ub(m) + gamma_lb(m) + gamma_ub(m) + rho_lb(k) + rho_ub(k) + mu_lb(n) + mu_ub(n) + nu_bal(n) + nu_flow(m) + ref(n_ref)
-    return _dc_kkt_dims(n, m, k, n_ref)
-end
-
-"""
-    kkt_indices(n, m, k) → NamedTuple
-
-Compute all KKT variable indices from network dimensions.
-Single source of truth for index calculations.
-
-# Variable ordering
-[va(n), pg(k), f(m), psh(n), lam_lb(m), lam_ub(m), gamma_lb(m), gamma_ub(m), rho_lb(k), rho_ub(k), mu_lb(n), mu_ub(n), nu_bal(n), nu_flow(m), eta(n_ref)]
-
-# Returns
-NamedTuple with index ranges for each variable block.
-"""
-kkt_indices(n::Int, m::Int, k::Int) = _dc_kkt_indices(n, m, k, 1)
-
 function _dc_kkt_indices(n::Int, m::Int, k::Int, n_ref::Int)
     i = 0
     idx_θ = (i+1):(i+n); i += n
@@ -344,29 +301,43 @@ function _dc_kkt_indices(n::Int, m::Int, k::Int, n_ref::Int)
     )
 end
 
-kkt_indices(net::DCNetwork) = _dc_kkt_indices(net.n, net.m, net.k, length(_reference_buses(net)))
-
-function kkt_indices(prob::DCOPFProblem)
-    net = getfield(prob, :network)
-    return _dc_kkt_indices(getfield(net, :n), getfield(net, :m), getfield(net, :k),
-                           getfield(prob, :_n_ref))
+@inline function _dc_kkt_layout(n::Int, m::Int, k::Int, n_ref::Int)
+    return _dc_kkt_dims(n, m, k, n_ref), _dc_kkt_indices(n, m, k, n_ref)
 end
 
-function _dc_kkt_layout(net::DCNetwork)
-    n_ref = length(_reference_buses(net))
+"""Return the DC KKT dimension for `n` buses, `m` branches, and `k` generators, assuming one reference bus."""
+kkt_dims(n::Int, m::Int, k::Int) = first(_dc_kkt_layout(n, m, k, 1))
+
+"""Return the DC KKT index ranges for `n` buses, `m` branches, and `k` generators, assuming one reference bus."""
+kkt_indices(n::Int, m::Int, k::Int) = last(_dc_kkt_layout(n, m, k, 1))
+
+"""
+    kkt_layout(network::DCNetwork) → (dim, indices)
+    kkt_layout(prob::DCOPFProblem) → (dim, indices)
+
+Return the dimension and named index ranges of the flattened DC KKT variable
+vector.
+
+The variable ordering is `[va, pg, f, psh, lam_lb, lam_ub, gamma_lb,
+gamma_ub, rho_lb, rho_ub, mu_lb, mu_ub, nu_bal, nu_flow, η]`. With `n`
+buses, `m` branches, `k` generators, and `n_ref` energized islands, the total
+dimension is `5n + 6m + 3k + n_ref`.
+"""
+function kkt_layout(net::DCNetwork)
     n = getfield(net, :n)
     m = getfield(net, :m)
     k = getfield(net, :k)
-    return _dc_kkt_dims(n, m, k, n_ref), _dc_kkt_indices(n, m, k, n_ref)
+    n_ref = length(_reference_buses(net))
+    return _dc_kkt_layout(n, m, k, n_ref)
 end
 
-function _dc_kkt_layout(prob::DCOPFProblem)
+function kkt_layout(prob::DCOPFProblem)
     net = getfield(prob, :network)
-    n_ref = getfield(prob, :_n_ref)
     n = getfield(net, :n)
     m = getfield(net, :m)
     k = getfield(net, :k)
-    return _dc_kkt_dims(n, m, k, n_ref), _dc_kkt_indices(n, m, k, n_ref)
+    n_ref = getfield(prob, :_n_ref)
+    return _dc_kkt_layout(n, m, k, n_ref)
 end
 
 # =============================================================================
@@ -615,7 +586,7 @@ function calc_kkt_jacobian(net::DCNetwork, d::AbstractVector, prob::DCOPFProblem
     k = getfield(net, :k)
     refs = _reference_buses(net)
     n_ref = length(refs)
-    dim = _dc_kkt_dims(n, m, k, n_ref)
+    dim, idx = kkt_layout(prob)
     A = getfield(net, :A)
     G_inc = getfield(net, :G_inc)
     b = getfield(net, :b)
@@ -644,8 +615,7 @@ function calc_kkt_jacobian(net::DCNetwork, d::AbstractVector, prob::DCOPFProblem
     W = Diagonal(-b .* sw)
     B_mat = sparse(A' * W * A)
 
-    # Build Jacobian blocks using centralized index calculation
-    idx = _dc_kkt_indices(n, m, k, n_ref)
+    # Build Jacobian blocks using the centralized layout.
     A_rowptr, A_rowcols, A_rowvals = _sparse_row_storage(A)
     G_rowptr, G_rowcols, G_rowvals = _sparse_row_storage(G_inc)
     rowval = Int[]
@@ -859,7 +829,7 @@ function calc_kkt_jacobian_demand(net::DCNetwork, d::AbstractVector, sol::DCOPFS
     n = getfield(net, :n)
     m = getfield(net, :m)
     k = getfield(net, :k)
-    dim, idx = _dc_kkt_layout(net)
+    dim, idx = kkt_layout(net)
     mu_ub = getfield(sol, :mu_ub)
 
     colptr = Vector{Int}(undef, n + 1)
@@ -887,7 +857,7 @@ Compute column `j` of the KKT parameter Jacobian ∂K/∂d.
 Only 1-2 nonzeros: always `nu_bal[j]`, plus `mu_ub[j]` if `d[j] > 0`.
 """
 function calc_kkt_jacobian_demand_column(net::DCNetwork, d::AbstractVector, sol::DCOPFSolution, j::Int)
-    dim, idx = _dc_kkt_layout(net)
+    dim, idx = kkt_layout(net)
     col = zeros(dim)
     col[idx.nu_bal[j]] = -1.0
     col[idx.mu_ub[j]] = sol.mu_ub[j] * _shed_capacity_derivative(d[j])
@@ -922,7 +892,7 @@ enabling gradient-based optimization for topology control.
 function calc_kkt_jacobian_switching(prob::DCOPFProblem, sol::DCOPFSolution)
     net = prob.network
     n, m, k = net.n, net.m, net.k
-    dim, idx = _dc_kkt_layout(prob)
+    dim, idx = kkt_layout(prob)
 
     θ = sol.va
 
@@ -986,7 +956,7 @@ Compute column `e` of the KKT parameter Jacobian ∂K/∂sw.
 """
 function calc_kkt_jacobian_switching_column(prob::DCOPFProblem, sol::DCOPFSolution, e::Int)
     net = prob.network
-    dim, idx = _dc_kkt_layout(prob)
+    dim, idx = kkt_layout(prob)
     col = zeros(dim)
     A = net.A; b = net.b; θ = sol.va
     Aθ_e = dot(A[e, :], θ)