Coverage for src/nncg/_active_set.py: 100%
65 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-15 07:01 +0000
1"""Primal-dual active-set primitives for the driver loop in :mod:`nncg.solver`.
3The pure working-set algebra the outer loop is built from, factored out of the
4:class:`~nncg.solver.ActiveSetSolver` orchestration: seeding the free set from a
5cold or warm start (:func:`_init_active_set`), splitting the KKT violators into
6primal and dual sets (:func:`_violators`), applying one guarded working-set pivot
7(:func:`_pivot`), and the per-outer-step scaffolding (:func:`_init_run_state`,
8:func:`_solve_free_set`). None of it knows about preconditioning or the inner
9solver — it operates on boolean free masks and index arrays alone.
10"""
12from __future__ import annotations
14import math
15from collections.abc import Callable
17import numpy as np
18from cvx.linalg import Vector
19from numpy.typing import NDArray
21SubSolve = Callable[[NDArray[np.int_], "Vector | None"], "tuple[Vector, Vector | None, int]"]
22"""Subproblem solve on a free set: ``(idx, x0) -> (x_F, lam, inner_iters)``."""
24ReducedGradient = Callable[["Vector", "Vector | None"], "Vector"]
25"""Reduced gradient of the subproblem: ``(x, lam) -> s``."""
28def _init_active_set(n: int, warm: tuple[NDArray[np.bool_], Vector] | None) -> tuple[NDArray[np.bool_], Vector | None]:
29 """Seed the free set and inner warm guess for the outer loop.
31 Cold (``warm is None``): everything free (``F = {1..n}``) with no inner
32 guess. Warm: a copy of the previous free mask and the previous iterate as
33 the seed for every subproblem solve.
34 """
35 if warm is None:
36 return np.ones(n, dtype=bool), None # F = {1..n} initially
37 return warm[0].copy(), warm[1]
40def _violators(
41 free: NDArray[np.bool_], x: Vector, s: Vector, tol: float
42) -> tuple[NDArray[np.int_], NDArray[np.int_], NDArray[np.int_]]:
43 """Split the KKT violators at tolerance ``tol`` into primal and dual sets.
45 Returns ``(prim, dual, viol)``: ``prim`` (set ``D``) are free indices whose
46 primal value went negative, ``dual`` (set ``V``) are bound indices whose
47 reduced gradient went negative, and ``viol`` is their concatenation. An empty
48 ``viol`` certifies the KKT conditions at the unique global minimiser.
49 """
50 prim = np.flatnonzero(free & (x < -tol)) # D: free but negative
51 dual = np.flatnonzero((~free) & (s < -tol)) # V: bound but s < 0
52 return prim, dual, np.concatenate([prim, dual])
55def _pivot(
56 free: NDArray[np.bool_],
57 prim: NDArray[np.int_],
58 dual: NDArray[np.int_],
59 viol: NDArray[np.int_],
60 n_bar: int,
61 patience: int,
62 p_max: int,
63) -> tuple[int, int, int]:
64 """Apply one working-set pivot, mutating ``free`` in place.
66 Takes the fast batch exchange — drop every primal violator ``D``, add every
67 dual violator ``V`` — while the violator count strictly drops below ``n_bar``
68 (patience reset to ``p_max``) or patience remains (decremented). Once patience
69 is exhausted without progress it falls back to a single least-index Bland
70 pivot, the load-bearing anti-cycling guarantee behind finite termination.
72 Returns the updated ``(n_bar, patience, fallback_increment)``; the last is 1
73 when the Bland fallback fired, 0 on the batch fast path.
74 """
75 n_viol = viol.size
76 if n_viol < n_bar or patience > 0: # fast path: progress, or patience remains
77 if n_viol < n_bar:
78 n_bar = n_viol
79 patience = p_max
80 else:
81 patience -= 1
82 free[prim] = False # batch exchange: drop all D, add all V
83 free[dual] = True
84 return n_bar, patience, 0
85 i_star = int(np.min(viol)) # anti-cycling fallback: single Bland least-index pivot
86 free[i_star] = not free[i_star]
87 return n_bar, patience, 1
90def _init_run_state(
91 n: int, track: bool, max_outer: int | None, warm: tuple[NDArray[np.bool_], Vector] | None
92) -> tuple[NDArray[np.bool_], Vector | None, list[tuple[int, ...]] | None, float]:
93 """Seed the free set, inner guess, trajectory log and outer-iteration cap.
95 Delegates the free set and warm inner guess to :func:`_init_active_set`,
96 allocates the trajectory list only when ``track`` is set, and resolves
97 ``max_outer`` (``None`` for uncapped) into a numeric loop bound so the
98 driver's ``while`` condition stays branch-free.
100 Returns:
101 ``(free, x_guess, traj, cap)``: the initial free mask, the inner warm
102 guess (``None`` on a cold start), the trajectory list (or ``None``),
103 and the outer-iteration cap (``math.inf`` when uncapped).
104 """
105 free, x_guess = _init_active_set(n, warm)
106 traj: list[tuple[int, ...]] | None = [] if track else None
107 cap = math.inf if max_outer is None else max_outer
108 return free, x_guess, traj, cap
111def _solve_free_set(
112 free: NDArray[np.bool_],
113 n: int,
114 sub_solve: SubSolve,
115 x_guess: Vector | None,
116 traj: list[tuple[int, ...]] | None,
117) -> tuple[Vector, Vector | None, int]:
118 """Solve the subproblem on the current free set and scatter it into ``R^n``.
120 Records the free set on ``traj`` when tracking, seeds the inner solve
121 from ``x_guess`` restricted to the free set (cold when ``None``), and
122 places the returned free-block solution back into a full zero vector.
124 Args:
125 free: Boolean free-set mask over the ``n`` variables.
126 n: Problem dimension.
127 sub_solve: Subproblem callback ``(idx, x0) -> (x_F, lam, inner_iters)``.
128 x_guess: Inner warm guess over all variables, or ``None`` (cold).
129 traj: Trajectory list to append the free set to, or ``None``.
131 Returns:
132 ``(x, lam, inner_iters)``: the full-length iterate, the equality
133 multipliers (``None`` for the bound-only problem), and the inner
134 iteration count from the callback.
135 """
136 idx = np.flatnonzero(free)
137 if traj is not None:
138 traj.append(tuple(idx.tolist()))
139 x0 = x_guess[idx] if x_guess is not None else None
140 xf, lam, k_step = sub_solve(idx, x0)
141 x: Vector = np.zeros(n)
142 x[idx] = xf
143 return x, lam, k_step
146def _drive(
147 tol: float,
148 p_max: int,
149 track: bool,
150 max_outer: int | None,
151 n: int,
152 sub_solve: SubSolve,
153 reduced_gradient: ReducedGradient,
154 warm: tuple[NDArray[np.bool_], Vector] | None,
155) -> tuple[Vector, int, int, int, bool, NDArray[np.bool_], Vector | None, list[tuple[int, ...]] | None]:
156 """Run the guarded primal-dual active-set loop and return its raw outcome.
158 Owns everything the termination proof depends on: the primal and dual
159 violator tests (:func:`_violators`), the batch exchange with its patience
160 counter and the least-index Bland fallback (:func:`_pivot`). What is solved
161 on each free set enters through ``sub_solve``, with ``reduced_gradient``
162 supplying the matching dual test quantity; the thresholds come straight from
163 :class:`nncg.solver.ActiveSetConfig`. Returns a plain tuple so this module
164 need not depend on :class:`nncg.solver.Result` — the caller wraps it.
166 Args:
167 tol: Violator tolerance of the primal and dual KKT tests.
168 p_max: Patience budget before the least-index Bland fallback pivot.
169 track: Record the visited free-set trajectory.
170 max_outer: Optional cap on outer steps (``None`` for uncapped).
171 n: Problem dimension.
172 sub_solve: Callback ``(idx, x0) -> (x_F, lam, inner_iters)`` solving the
173 subproblem on the free set ``idx``.
174 reduced_gradient: Callback ``(x, lam) -> s`` for the dual violator test.
175 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
177 Returns:
178 ``(x, outer, inner_total, fallback, converged, free, lam, traj)``.
179 """
180 free, x_guess, traj, cap = _init_run_state(n, track, max_outer, warm)
181 x: Vector = np.zeros(n)
182 lam: Vector | None = None
183 n_bar, patience = n + 1, p_max
184 outer = inner_total = fallback = 0
185 converged = True
187 while outer < cap:
188 x, lam, k_step = _solve_free_set(free, n, sub_solve, x_guess, traj)
189 outer += 1
190 inner_total += k_step
191 if x_guess is not None:
192 x_guess = x # warm mode: newest iterate seeds the next reduced solve
193 prim, dual, viol = _violators(free, x, reduced_gradient(x, lam), tol)
194 if viol.size == 0:
195 break # KKT satisfied -> unique global minimiser
196 n_bar, patience, fired = _pivot(free, prim, dual, viol, n_bar, patience, p_max)
197 fallback += fired
198 else:
199 converged = False # outer cap reached without certifying KKT
201 return x, outer, inner_total, fallback, converged, free, lam, traj