Apply the angle gate to the matrix-free switching sensitivity

docs/src/math/dc-opf.md

@@ -32,7 +32,7 @@ where:
 - ``c_{\text{shed}}`` is the load shedding cost vector
 - ``d_+ = \max(d, 0)`` is the curtailable portion of signed net demand; negative net demand remains in power balance as an injection
 - ``\tau`` is a small regularization parameter for numerical conditioning
-- ``\mathrm{sw}_{\mathrm{eff},e} = \mathrm{sw}_e`` when ``b_e \mathrm{sw}_e \ne 0`` and is zero otherwise, so angle difference limits use the same energized predicate as the island partition while retaining fractional switching on energized branches
+- ``\mathrm{sw}_{\mathrm{eff},e} = \mathrm{sw}_e`` when ``b_e \mathrm{sw}_e \ne 0`` and is zero otherwise, so angle difference limits use the same energized predicate as the island partition. A de-energized branch (``\mathrm{sw}_{\mathrm{eff},e} = 0``) imposes no limit; on an energized branch the gate multiplies both sides, so the factor cancels and the primal constraint is exactly ``\alpha_{\min,e} \leq (A\theta)_e \leq \alpha_{\max,e}`` for every ``\mathrm{sw}_e \in (0, 1]``. Fractional switching therefore rescales the angle limit duals rather than relaxing the limit
 - ``\text{refs}`` contains one reference bus per energized island, including isolated buses. The choice is deterministic: the configured ``\text{ref_bus}`` is the reference for its island, and every other island uses its lowest sequential bus index
 
 The built OPF model stores one reference constraint for each entry of

src/sens/susceptance.jl

@@ -82,7 +82,6 @@ function calc_kkt_jacobian_susceptance(prob::DCOPFProblem, sol::DCOPFSolution)
 
         # dK_flow_def/db_e: only row e is nonzero
         J_b[idx.nu_flow[e], e] = sw[e] * Atheta_e
-
     end
 
     return J_b

src/sens/vjp_jvp.jl

@@ -62,10 +62,8 @@ function _dc_param_vjp!(out::AbstractVector, prob::DCOPFProblem,
         _dc_cost_quadratic_vjp!(out, sol, u, idx, prob.network.k)
     elseif param === :fmax
         _dc_flowlimit_vjp!(out, sol, u, idx, prob.network.m)
-    elseif param === :sw
-        _dc_topo_vjp!(out, prob, sol, u, idx, prob.network.b)
-    else  # :b
-        _dc_topo_vjp!(out, prob, sol, u, idx, prob.network.sw)
+    else  # :sw or :b
+        _dc_topo_vjp!(out, prob, sol, u, idx, param)
     end
     return out
 end
@@ -99,17 +97,33 @@ function _dc_flowlimit_vjp!(out, sol, u, idx, m)
     end
 end
 
-# `:sw`/`:b` — 3 KKT blocks (nu_bal, nu_flow, va), vectorized via sparse matvecs.
-# `coeff` is net.b for :sw, net.sw for :b (the coefficient that multiplies the parameter).
-function _dc_topo_vjp!(out, prob, sol, u, idx, coeff)
-    A = prob.network.A
+# `:sw`/`:b` — nu_bal, nu_flow, and va blocks from B and W, vectorized via sparse
+# matvecs. `coeff` is the other factor of the `b .* sw` product: net.b for :sw and
+# net.sw for :b.
+#
+# `:sw` additionally moves the gated angle difference limits, contributing to the va
+# block and to both gamma blocks. `:b` does not: the effective gate is locally
+# constant in b within a fixed energized regime, and zero by convention at the
+# nonsmooth b == 0 boundary (see `_angle_difference_gate_dsw`). This mirrors the
+# split between `calc_kkt_jacobian_switching` and `calc_kkt_jacobian_susceptance`.
+function _dc_topo_vjp!(out, prob, sol, u, idx, param::Symbol)
+    net = prob.network
+    gated = param === :sw
+    coeff = gated ? net.b : net.sw
+    A = net.A
     Aθ = A * sol.va
     Au_nb = A * @view(u[idx.nu_bal])
     Au_va = A * @view(u[idx.va])
     Aν = A * sol.nu_bal
-    @inbounds for e in 1:prob.network.m
+    @inbounds for e in 1:net.m
         out[e] = -coeff[e] * (Aθ[e] * (Au_nb[e] + u[idx.nu_flow[e]]) -
                                (Aν[e] + sol.nu_flow[e]) * Au_va[e])
+        if gated
+            dgate = _angle_difference_gate_dsw(net, e)
+            out[e] -= dgate * ((sol.gamma_ub[e] - sol.gamma_lb[e]) * Au_va[e] +
+                               sol.gamma_lb[e] * (Aθ[e] - net.angmin[e]) * u[idx.gamma_lb[e]] +
+                               sol.gamma_ub[e] * (net.angmax[e] - Aθ[e]) * u[idx.gamma_ub[e]])
+        end
     end
 end
 
@@ -130,10 +144,8 @@ function _dc_param_jvp!(v::AbstractVector, prob::DCOPFProblem,
         _dc_cost_quadratic_jvp!(v, sol, tang, idx, prob.network.k)
     elseif param === :fmax
         _dc_flowlimit_jvp!(v, sol, tang, idx, prob.network.m)
-    elseif param === :sw
-        _dc_topo_jvp!(v, prob, sol, tang, idx, prob.network.b)
-    else  # :b
-        _dc_topo_jvp!(v, prob, sol, tang, idx, prob.network.sw)
+    else  # :sw or :b
+        _dc_topo_jvp!(v, prob, sol, tang, idx, param)
     end
     return v
 end
@@ -165,9 +177,13 @@ function _dc_flowlimit_jvp!(v, sol, tang, idx, m)
     end
 end
 
-# `:sw`/`:b` — 2 sparse transpose-matvecs + elementwise scatter
-function _dc_topo_jvp!(v, prob, sol, tang, idx, coeff)
-    A = prob.network.A
+# `:sw`/`:b` — 2 sparse transpose-matvecs + elementwise scatter. See
+# `_dc_topo_vjp!` for why only `:sw` carries the gated angle difference terms.
+function _dc_topo_jvp!(v, prob, sol, tang, idx, param::Symbol)
+    net = prob.network
+    gated = param === :sw
+    coeff = gated ? net.b : net.sw
+    A = net.A
     Aθ = A * sol.va
     Aν = A * sol.nu_bal
 
@@ -180,6 +196,19 @@ function _dc_topo_jvp!(v, prob, sol, tang, idx, coeff)
 
     # va block: v[va] += A' * (-coeff .* (Aν .+ nu_flow) .* tang)
     w_va = (-coeff) .* (Aν .+ sol.nu_flow) .* tang
+
+    # `:sw` also moves the gated angle difference limits. The va contribution folds
+    # into w_va so it rides the same transpose-matvec; the gamma blocks scatter
+    # directly.
+    if gated
+        @inbounds for e in 1:net.m
+            dgate = _angle_difference_gate_dsw(net, e)
+            w_va[e] += dgate * (sol.gamma_ub[e] - sol.gamma_lb[e]) * tang[e]
+            v[idx.gamma_lb[e]] += dgate * sol.gamma_lb[e] * (Aθ[e] - net.angmin[e]) * tang[e]
+            v[idx.gamma_ub[e]] += dgate * sol.gamma_ub[e] * (net.angmax[e] - Aθ[e]) * tang[e]
+        end
+    end
+
     mul!(@view(v[idx.va]), A', w_va, 1.0, 1.0)
 end
 

test/test_angle_diff_duals.jl

@@ -124,6 +124,22 @@ import PowerDiff: kkt, kkt_indices, flatten_variables
         @test all(isfinite, Matrix(dpg_dsw))
         @test all(isfinite, Matrix(df_dsw))
 
+        # The matrix-free :sw path must carry the same gated angle difference terms
+        # as the materialized Jacobian. `prob_cold` is never passed to
+        # calc_sensitivity, so prob.cache.dz_dsw stays empty and vjp/jvp take the
+        # inline `_dc_topo_vjp!`/`_dc_topo_jvp!` route rather than the cached
+        # matrix multiply. With gamma_ub[2] binding above, dropping those terms
+        # makes the two paths disagree.
+        prob_cold = DCOPFProblem(net, d)
+        solve!(prob_cold)
+        v_sw = randn(m)
+        w_pg = randn(k)
+        w_f = randn(m)
+        @test jvp(prob_cold, :pg, :sw, v_sw) ≈ Matrix(dpg_dsw) * v_sw atol=1e-8
+        @test jvp(prob_cold, :f, :sw, v_sw) ≈ Matrix(df_dsw) * v_sw atol=1e-8
+        @test vjp(prob_cold, :pg, :sw, w_pg) ≈ Matrix(dpg_dsw)' * w_pg atol=1e-8
+        @test vjp(prob_cold, :f, :sw, w_f) ≈ Matrix(df_dsw)' * w_f atol=1e-8
+
         ε = 1e-5
         for e in 1:m
             sw_pert = copy(net.sw)
@@ -522,6 +538,22 @@ import PowerDiff: kkt, kkt_indices, flatten_variables
         @test Jsw_zero_col == Jsw_zero[:, 3]
         @test Jsw_zero[:, 3] ≈ Jsw_zero_fd atol=1e-10
 
+        # The matrix-free dK/dsw applies the gate through its own
+        # `_angle_difference_gate_dsw` call site, so it can drift from the
+        # materialized column above. Drive the inline scatter/gather directly
+        # rather than through vjp/jvp: those resolve the solution from the cache
+        # and would discard the synthetic gammas, and the solver canonicalizes
+        # gamma to zero on a de-energized branch, which would make the b == 0
+        # gate vacuous.
+        tang_zero = randn(3)
+        u_zero = randn(size(Jsw_zero, 1))
+        jvp_free = zeros(size(Jsw_zero, 1))
+        vjp_free = zeros(3)
+        PowerDiff._dc_param_jvp!(jvp_free, prob_zero, sol_zero_gamma, :sw, tang_zero, idx_zero)
+        PowerDiff._dc_param_vjp!(vjp_free, prob_zero, sol_zero_gamma, :sw, u_zero, idx_zero)
+        @test jvp_free ≈ Jsw_zero * tang_zero atol=1e-12
+        @test vjp_free ≈ -Jsw_zero' * u_zero atol=1e-12
+
         Jb_zero = Matrix(PowerDiff.calc_kkt_jacobian_susceptance(prob_zero, sol_zero_gamma))
         Jb_zero_plain = Matrix(PowerDiff.calc_kkt_jacobian_susceptance(prob_zero, sol_zero))
         @test Jb_zero ≈ Jb_zero_plain atol=1e-12
@@ -549,6 +581,18 @@ import PowerDiff: kkt, kkt_indices, flatten_variables
             @test Jsw[:, 3] ≈ Jsw_fd atol=1e-8 rtol=1e-7
             @test norm(Jsw[:, 3] - Jsw_plain[:, 3]) > 1.0
 
+            # Same matrix-free cross-check as at b == 0, but with dgate == 1 so
+            # the gated terms are non-vacuous on branch 3.
+            idx_side = kkt_indices(prob_side)
+            tang_side = randn(3)
+            u_side = randn(size(Jsw, 1))
+            jvp_side = zeros(size(Jsw, 1))
+            vjp_side = zeros(3)
+            PowerDiff._dc_param_jvp!(jvp_side, prob_side, sol_side_gamma, :sw, tang_side, idx_side)
+            PowerDiff._dc_param_vjp!(vjp_side, prob_side, sol_side_gamma, :sw, u_side, idx_side)
+            @test jvp_side ≈ Jsw * tang_side atol=1e-12
+            @test vjp_side ≈ -Jsw' * u_side atol=1e-12
+
             Jb = Matrix(PowerDiff.calc_kkt_jacobian_susceptance(prob_side, sol_side_gamma))
             Jb_plain = Matrix(PowerDiff.calc_kkt_jacobian_susceptance(prob_side, sol_side))
             Jb_col = PowerDiff.calc_kkt_jacobian_susceptance_column(prob_side, sol_side_gamma, 3)

test/test_kkt_vjp_jvp.jl

@@ -30,14 +30,23 @@
         prob = DCOPFProblem(basic)
         solve!(prob)
 
+        # DC OPF vjp/jvp short-circuit to a cached dz/dp matrix multiply once
+        # calc_sensitivity has materialized one for that parameter. Comparing
+        # against S on the same problem therefore only exercises that fast path.
+        # `prob_cold` is deliberately never passed to calc_sensitivity — the
+        # inline `_dc_param_vjp!`/`_dc_param_jvp!` route is only reachable through
+        # it, so keep the two instances separate.
+        prob_cold = DCOPFProblem(basic)
+        solve!(prob_cold)
+
         @testset "VJP matches S' * w — all (op, param)" begin
             for (op, param) in [(:lmp, :d), (:pg, :d), (:f, :sw), (:va, :cq),
                                 (:psh, :cl), (:lmp, :fmax), (:f, :b)]
                 S = calc_sensitivity(prob, op, param)
                 w = randn(size(S, 1))
                 expected = S.matrix' * w
-                result = vjp(prob, op, param, w)
-                @test result ≈ expected atol=1e-10
+                @test vjp(prob_cold, op, param, w) ≈ expected atol=1e-10  # inline path
+                @test vjp(prob, op, param, w) ≈ expected atol=1e-10       # cached path
             end
         end
 
@@ -47,8 +56,8 @@
                 S = calc_sensitivity(prob, op, param)
                 v = randn(size(S, 2))
                 expected = S.matrix * v
-                result = jvp(prob, op, param, v)
-                @test result ≈ expected atol=1e-10
+                @test jvp(prob_cold, op, param, v) ≈ expected atol=1e-10  # inline path
+                @test jvp(prob, op, param, v) ≈ expected atol=1e-10       # cached path
             end
         end