Coverage for src/nncg/solver.py: 100%
55 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-01 07:21 +0000
1"""Non-negative conjugate gradients: the active-set / block-principal-pivoting loop.
3Solves the strictly convex non-negative quadratic program
5 min_{x >= 0} 1/2 x^T A x - b^T x, A symmetric positive definite,
7and its equality-augmented variant with a general linear system ``B x = c``,
8by wrapping a matrix-free inner solver in a primal-dual active-set outer loop.
9The working-set toggles are the principal pivots of the linear complementarity
10problem LCP(A, -b); guarding the fast block-pivot path with a least-index Bland
11fallback gives unconditional finite termination at the unique global minimiser
12— no non-degeneracy assumption (Theorem 5.1 of the accompanying paper). See
13https://github.com/Jebel-Quant/mean_variance_solvers.
15:class:`ActiveSetSolver` is the outer loop and the entry point. It knows nothing
16about preconditioning: it asks its :class:`nncg.inner.InnerSolver` for a
17per-free-block solve and drives the pivots around it. The quadratic term enters
18as a :class:`cvx.linalg.SymmetricOperator`, accessed only through block products
19— wrap an explicit SPD array in ``DenseOperator``, or pass ``GramOperator(M,
20ridge)`` for ``A = M^T M + ridge I`` so the ``n x n`` matrix is never formed.
21"""
23from __future__ import annotations
25from dataclasses import dataclass, field, replace
26from typing import Protocol
28import numpy as np
29from cvx.linalg import Matrix, SymmetricOperator, Vector
30from numpy.typing import NDArray
32from ._active_set import ReducedGradient, SubSolve, _drive
33from ._equality import _eq_feasible, _require_eq_shapes, _saddle_solve
34from .certificate import _require_operator
37class InnerSolver(Protocol):
38 """The inner-solver interface the active-set loop depends on (dependency inversion).
40 Structural (a :class:`typing.Protocol`): anything with a matching
41 :meth:`solve` is an inner solver, so implementations need neither import
42 nor subclass this — this module (the high-level loop) owns the interface, and
43 the implementations depend on it, not the other way round. The built-ins live
44 in :mod:`nncg.inner` (:class:`~nncg.inner.CG`, :class:`~nncg.inner.Jacobi`,
45 :class:`~nncg.inner.Nystrom`, :class:`~nncg.inner.Exact`); further ones —
46 Clarabel- or KKT-equation-based — live in Jebel-Quant/mean_variance_solvers.
47 """
49 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
50 """Solve the free-block system ``A[F, F] y = rhs``, warm-started at ``x0``.
52 Returns the free-block solution and the inner iteration count (each
53 direct solve counts as one). Called once per outer step by the
54 bound-constrained loop, and once per ``p + 1`` right-hand side per outer
55 step by the equality-augmented loop.
56 """
57 ...
60@dataclass(frozen=True)
61class ActiveSetConfig:
62 """Configuration of the active-set outer loop (:class:`ActiveSetSolver`).
64 Bundles the outer-loop knobs into one argument; the inner solver and its
65 tolerances live in :class:`nncg.inner.InnerSolver`, and the warm start stays
66 a separate argument.
68 Attributes:
69 tol: Threshold of the primal and dual KKT violator tests.
70 p_max: Patience budget — non-improving batch steps tolerated before a
71 least-index Bland fallback pivot. Any value gives finite termination.
72 track: Record the visited free-set trajectory in ``Result.traj``.
73 max_outer: Optional cap on outer steps; when hit, the current iterate is
74 returned with ``converged=False``.
75 """
77 tol: float = 1e-8
78 p_max: int = 3
79 track: bool = False
80 max_outer: int | None = None
83@dataclass(frozen=True)
84class Result:
85 """Outcome of an active-set solve.
87 Attributes:
88 x: The minimiser (or the final iterate if ``converged`` is False).
89 outer: Number of outer active-set steps taken.
90 inner: Total inner (CG/PCG) iterations across all outer steps; each
91 direct inner solve counts as one.
92 fallback: Number of least-index Bland fallback pivots taken.
93 converged: True when the KKT exit was reached; False when an
94 ``max_outer`` cap stopped the loop first, or (equality-augmented
95 solves) when the final iterate violates ``B x = c``.
96 free: Boolean mask of the final free set.
97 lam: Multipliers of the equality constraints (equality-augmented
98 solves only; None otherwise).
99 traj: The sequence of visited free sets as index tuples when
100 trajectory tracking was requested; None otherwise.
101 """
103 x: Vector
104 outer: int
105 inner: int
106 fallback: int
107 converged: bool
108 free: NDArray[np.bool_]
109 lam: Vector | None = None
110 traj: list[tuple[int, ...]] | None = None
113@dataclass(frozen=True)
114class ActiveSetSolver:
115 """The primal-dual active-set outer loop for the non-negative quadratic program.
117 Holds the outer-loop :class:`ActiveSetConfig` and an
118 :class:`nncg.inner.InnerSolver`, and drives the guarded block-pivot loop
119 around the per-free-block solve the inner solver provides. It never touches a
120 preconditioner — everything about CG/PCG/Nyström lives in ``inner``.
122 Attributes:
123 inner: The inner solver for each free block — e.g. :class:`nncg.inner.CG`
124 (plain CG), :class:`nncg.inner.Jacobi`, :class:`nncg.inner.Nystrom`
125 or :class:`nncg.inner.Exact`.
126 config: Outer-loop configuration (violator tolerance, patience,
127 trajectory tracking, outer-step cap).
129 Examples:
130 ``A`` enters as an operator, never as a bare array:
132 >>> import numpy as np
133 >>> from cvx.linalg import DenseOperator
134 >>> from nncg import ActiveSetSolver, CG, kkt_violation
135 >>> a = DenseOperator(np.array([[2.0, 0.0], [0.0, 2.0]]))
136 >>> b = np.array([2.0, -2.0])
138 The unconstrained minimiser would be ``(1, -1)``, so the bound binds on
139 the second coordinate and the loop returns ``(1, 0)`` with that
140 coordinate active:
142 >>> res = ActiveSetSolver(inner=CG()).solve(a, b)
143 >>> res.converged
144 True
145 >>> bool(np.allclose(res.x, [1.0, 0.0]))
146 True
147 >>> res.free.tolist()
148 [True, False]
150 ``converged`` is the KKT exit, which :func:`nncg.kkt_violation` scores
151 independently — zero certifies the unique global minimiser:
153 >>> round(kkt_violation(a, b, res.x), 12)
154 0.0
156 Generic data never needs the Bland fallback; that it stayed dormant is
157 reported rather than assumed:
159 >>> res.fallback
160 0
162 Swap the inner solver freely — the outer loop is unchanged, and on this
163 problem so is the answer:
165 >>> from nncg import Exact
166 >>> direct = ActiveSetSolver(inner=Exact()).solve(a, b)
167 >>> bool(np.allclose(direct.x, res.x))
168 True
169 """
171 inner: InnerSolver
172 config: ActiveSetConfig = field(default_factory=ActiveSetConfig)
174 def solve(
175 self,
176 a: SymmetricOperator,
177 b: Vector,
178 warm: tuple[NDArray[np.bool_], Vector] | None = None,
179 ) -> Result:
180 """Minimise ``1/2 x^T A x - b^T x`` over ``x >= 0`` by the active-set loop.
182 Each free-block solve is delegated to :attr:`inner`; the reduced matrix
183 is never materialised and ``A`` is never refactorised. The batch
184 block-pivot fast path is guarded by a least-index Bland fallback, so
185 termination at the unique global minimiser is unconditional.
187 Args:
188 a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`) —
189 ``DenseOperator`` for an explicit array, ``GramOperator(M, ridge)``
190 for ``A = M^T M + ridge I`` whose Gram matrix is never formed.
191 b: The linear term ``b``.
192 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
193 Starts the loop from that free set and warm-starts every inner
194 solve from the newest iterate (the :class:`nncg.inner.Exact`
195 inner solver is direct, so it has nothing to seed but still
196 starts from the warm free set) — across a support-stable
197 parameter step the loop then terminates in a single outer step.
199 Returns:
200 A :class:`Result`; ``converged`` is True iff the KKT system was
201 satisfied to ``config.tol``, which certifies the unique global
202 minimiser.
204 Raises:
205 TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
206 ValueError: When the operator dimension does not match ``len(b)``, or
207 on the inner solver's own conditions in
208 :meth:`InnerSolver.solve`.
209 NotImplementedError: When a diagonal-preconditioned inner solver
210 (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
211 (propagated from ``cvx.linalg``).
212 """
213 _require_operator(a, b)
215 def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
216 """Solve the reduced system ``A_F x_F = b_F`` with the chosen inner solver."""
217 xf, k_step = self.inner.solve(a, idx, b[idx], x0)
218 return xf, None, k_step
220 def reduced_gradient(x: Vector, lam: Vector | None) -> Vector: # noqa: ARG001
221 """Return the reduced gradient ``s = A x - b``."""
222 return a.matvec(x) - b
224 return self._run(len(b), sub_solve, reduced_gradient, warm)
226 def solve_eq(
227 self,
228 a: SymmetricOperator,
229 b: Vector,
230 b_eq: Matrix,
231 c_eq: Vector,
232 warm: tuple[NDArray[np.bool_], Vector] | None = None,
233 ) -> Result:
234 """Solve ``min 1/2 x^T A x - b^T x`` subject to ``x >= 0`` and ``B x = c``.
236 On each free set the saddle system is solved by eliminating the
237 multiplier ``lambda`` in R^p through the p-by-p Schur complement
238 ``S = B_F A_F^{-1} B_F^T``: the ``p + 1`` right-hand sides share the
239 operator ``A_F`` and are each one inner solve, then ``S lambda = c - B_F
240 v0`` fixes the multipliers in closed form. The single normalisation
241 ``1^T x = beta`` is the ``p = 1`` case. ``B`` must have full row rank on
242 the visited free sets (automatic for ``p = 1``).
244 Args:
245 a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`).
246 b: The linear term ``b``.
247 b_eq: Equality matrix ``B`` of shape ``(p, n)``, full row rank.
248 c_eq: Equality right-hand side ``c`` of shape ``(p,)``.
249 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
250 Starts the loop from that free set and seeds the ``v0`` solve of
251 every saddle step from the newest iterate; the ``v1`` columns are
252 re-solved cold (their right-hand sides are the rows of ``B_F``,
253 unrelated to ``x_prev``). Across a support-stable parameter step
254 the loop then terminates in a single outer step.
256 Returns:
257 A :class:`Result` with the multipliers in ``lam``. The reduced
258 gradient underlying the dual test is ``s = A x - b - B^T lam``.
259 ``converged`` additionally requires ``max|B x - c| <= tol *
260 max(1, max|c|)``: a rank-deficient ``B`` with an inconsistent ``c``
261 can pass the KKT test while infeasible, and is returned with
262 ``converged=False`` rather than certified.
264 Raises:
265 TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
266 ValueError: When the operator dimension does not match ``len(b)``,
267 when ``b_eq`` is not of shape ``(p, n)`` or ``c_eq`` not of shape
268 ``(p,)``, or on the inner solver's own conditions in
269 :meth:`InnerSolver.solve`.
270 NotImplementedError: When a diagonal-preconditioned inner solver
271 (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
272 (propagated from ``cvx.linalg``).
273 """
274 _require_operator(a, b)
275 _require_eq_shapes(b_eq, c_eq, len(b))
277 def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
278 """Solve the saddle system on the free set via the p-by-p Schur complement."""
279 return _saddle_solve(self.inner, a, b, b_eq, c_eq, idx, x0)
281 def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:
282 """Return the constrained reduced gradient ``s = A x - b - B^T lam``."""
283 correction = b_eq.T @ lam if lam is not None else np.zeros_like(b)
284 return a.matvec(x) - b - correction
286 result = self._run(len(b), sub_solve, reduced_gradient, warm)
287 if result.converged and not _eq_feasible(b_eq, c_eq, result.x, self.config.tol):
288 return replace(result, converged=False)
289 return result
291 def _run(
292 self,
293 n: int,
294 sub_solve: SubSolve,
295 reduced_gradient: ReducedGradient,
296 warm: tuple[NDArray[np.bool_], Vector] | None,
297 ) -> Result:
298 """Run the guarded primal-dual active-set loop.
300 The driver owns everything the termination proof depends on: the primal
301 and dual violator tests, the batch exchange with its patience counter,
302 and the least-index Bland fallback. What is solved on each free set — a
303 single reduced system (:meth:`solve`), or the equality-augmented saddle
304 system (:meth:`solve_eq`) — enters through the ``sub_solve`` callback,
305 with ``reduced_gradient`` supplying the matching dual test quantity. The
306 thresholds (``tol``, ``p_max``, ``track``, ``max_outer``) are read from
307 :attr:`config`.
309 Args:
310 n: Problem dimension.
311 sub_solve: Callback ``(idx, x0) -> (x_F, lam, inner_iters)`` solving
312 the subproblem on the free set ``idx``. ``x0`` is a warm inner
313 guess restricted to ``idx`` (None on a cold start); ``lam`` are
314 the equality multipliers (None for the bound-only problem).
315 reduced_gradient: Callback ``(x, lam) -> s`` computing the reduced
316 gradient that drives the dual violator test.
317 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
318 Starts the loop from that free set and seeds every subproblem
319 solve from the newest iterate.
321 Returns:
322 A :class:`Result`; ``lam`` is whatever the last subproblem returned.
323 """
324 cfg = self.config
325 out = _drive(cfg.tol, cfg.p_max, cfg.track, cfg.max_outer, n, sub_solve, reduced_gradient, warm)
326 return Result(
327 x=out.x,
328 outer=out.outer,
329 inner=out.inner,
330 fallback=out.fallback,
331 converged=out.converged,
332 free=out.free,
333 lam=out.lam,
334 traj=out.traj,
335 )