Coverage for src/nncg/solver.py: 100%
51 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-15 07:01 +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
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 _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.
95 free: Boolean mask of the final free set.
96 lam: Multipliers of the equality constraints (equality-augmented
97 solves only; None otherwise).
98 traj: The sequence of visited free sets as index tuples when
99 trajectory tracking was requested; None otherwise.
100 """
102 x: Vector
103 outer: int
104 inner: int
105 fallback: int
106 converged: bool
107 free: NDArray[np.bool_]
108 lam: Vector | None = None
109 traj: list[tuple[int, ...]] | None = None
112@dataclass(frozen=True)
113class ActiveSetSolver:
114 """The primal-dual active-set outer loop for the non-negative quadratic program.
116 Holds the outer-loop :class:`ActiveSetConfig` and an
117 :class:`nncg.inner.InnerSolver`, and drives the guarded block-pivot loop
118 around the per-free-block solve the inner solver provides. It never touches a
119 preconditioner — everything about CG/PCG/Nyström lives in ``inner``.
121 Attributes:
122 inner: The inner solver for each free block — e.g. :class:`nncg.inner.CG`
123 (plain CG), :class:`nncg.inner.Jacobi`, :class:`nncg.inner.Nystrom`
124 or :class:`nncg.inner.Exact`.
125 config: Outer-loop configuration (violator tolerance, patience,
126 trajectory tracking, outer-step cap).
128 Examples:
129 ``A`` enters as an operator, never as a bare array:
131 >>> import numpy as np
132 >>> from cvx.linalg import DenseOperator
133 >>> from nncg import ActiveSetSolver, CG, kkt_violation
134 >>> a = DenseOperator(np.array([[2.0, 0.0], [0.0, 2.0]]))
135 >>> b = np.array([2.0, -2.0])
137 The unconstrained minimiser would be ``(1, -1)``, so the bound binds on
138 the second coordinate and the loop returns ``(1, 0)`` with that
139 coordinate active:
141 >>> res = ActiveSetSolver(inner=CG()).solve(a, b)
142 >>> res.converged
143 True
144 >>> bool(np.allclose(res.x, [1.0, 0.0]))
145 True
146 >>> res.free.tolist()
147 [True, False]
149 ``converged`` is the KKT exit, which :func:`nncg.kkt_violation` scores
150 independently — zero certifies the unique global minimiser:
152 >>> round(kkt_violation(a, b, res.x), 12)
153 0.0
155 Generic data never needs the Bland fallback; that it stayed dormant is
156 reported rather than assumed:
158 >>> res.fallback
159 0
161 Swap the inner solver freely — the outer loop is unchanged, and on this
162 problem so is the answer:
164 >>> from nncg import Exact
165 >>> direct = ActiveSetSolver(inner=Exact()).solve(a, b)
166 >>> bool(np.allclose(direct.x, res.x))
167 True
168 """
170 inner: InnerSolver
171 config: ActiveSetConfig = field(default_factory=ActiveSetConfig)
173 def solve(
174 self,
175 a: SymmetricOperator,
176 b: Vector,
177 warm: tuple[NDArray[np.bool_], Vector] | None = None,
178 ) -> Result:
179 """Minimise ``1/2 x^T A x - b^T x`` over ``x >= 0`` by the active-set loop.
181 Each free-block solve is delegated to :attr:`inner`; the reduced matrix
182 is never materialised and ``A`` is never refactorised. The batch
183 block-pivot fast path is guarded by a least-index Bland fallback, so
184 termination at the unique global minimiser is unconditional.
186 Args:
187 a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`) —
188 ``DenseOperator`` for an explicit array, ``GramOperator(M, ridge)``
189 for ``A = M^T M + ridge I`` whose Gram matrix is never formed.
190 b: The linear term ``b``.
191 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
192 Starts the loop from that free set and warm-starts every inner
193 solve from the newest iterate (the :class:`nncg.inner.Exact`
194 inner solver is direct, so it has nothing to seed but still
195 starts from the warm free set) — across a support-stable
196 parameter step the loop then terminates in a single outer step.
198 Returns:
199 A :class:`Result`; ``converged`` is True iff the KKT system was
200 satisfied to ``config.tol``, which certifies the unique global
201 minimiser.
203 Raises:
204 TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
205 ValueError: When the operator dimension does not match ``len(b)``, or
206 on the inner solver's own conditions in
207 :meth:`InnerSolver.solve`.
208 NotImplementedError: When a diagonal-preconditioned inner solver
209 (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
210 (propagated from ``cvx.linalg``).
211 """
212 _require_operator(a, b)
214 def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
215 """Solve the reduced system ``A_F x_F = b_F`` with the chosen inner solver."""
216 xf, k_step = self.inner.solve(a, idx, b[idx], x0)
217 return xf, None, k_step
219 def reduced_gradient(x: Vector, lam: Vector | None) -> Vector: # noqa: ARG001
220 """Return the reduced gradient ``s = A x - b``."""
221 return a.matvec(x) - b
223 return self._run(len(b), sub_solve, reduced_gradient, warm)
225 def solve_eq(
226 self,
227 a: SymmetricOperator,
228 b: Vector,
229 b_eq: Matrix,
230 c_eq: Vector,
231 warm: tuple[NDArray[np.bool_], Vector] | None = None,
232 ) -> Result:
233 """Solve ``min 1/2 x^T A x - b^T x`` subject to ``x >= 0`` and ``B x = c``.
235 On each free set the saddle system is solved by eliminating the
236 multiplier ``lambda`` in R^p through the p-by-p Schur complement
237 ``S = B_F A_F^{-1} B_F^T``: the ``p + 1`` right-hand sides share the
238 operator ``A_F`` and are each one inner solve, then ``S lambda = c - B_F
239 v0`` fixes the multipliers in closed form. The single normalisation
240 ``1^T x = beta`` is the ``p = 1`` case. ``B`` must have full row rank on
241 the visited free sets (automatic for ``p = 1``).
243 Args:
244 a: The SPD operator ``A`` (a :class:`cvx.linalg.SymmetricOperator`).
245 b: The linear term ``b``.
246 b_eq: Equality matrix ``B`` of shape ``(p, n)``, full row rank.
247 c_eq: Equality right-hand side ``c`` of shape ``(p,)``.
248 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
249 Starts the loop from that free set and seeds the ``v0`` solve of
250 every saddle step from the newest iterate; the ``v1`` columns are
251 re-solved cold (their right-hand sides are the rows of ``B_F``,
252 unrelated to ``x_prev``). Across a support-stable parameter step
253 the loop then terminates in a single outer step.
255 Returns:
256 A :class:`Result` with the multipliers in ``lam``. The reduced
257 gradient underlying the dual test is ``s = A x - b - B^T lam``.
259 Raises:
260 TypeError: When ``a`` is not a :class:`cvx.linalg.SymmetricOperator`.
261 ValueError: When the operator dimension does not match ``len(b)``, or
262 on the inner solver's own conditions in
263 :meth:`InnerSolver.solve`.
264 NotImplementedError: When a diagonal-preconditioned inner solver
265 (:class:`nncg.inner.Jacobi`) meets a backend without ``diag``
266 (propagated from ``cvx.linalg``).
267 """
268 _require_operator(a, b)
270 def sub_solve(idx: NDArray[np.int_], x0: Vector | None) -> tuple[Vector, Vector | None, int]:
271 """Solve the saddle system on the free set via the p-by-p Schur complement."""
272 return _saddle_solve(self.inner, a, b, b_eq, c_eq, idx, x0)
274 def reduced_gradient(x: Vector, lam: Vector | None) -> Vector:
275 """Return the constrained reduced gradient ``s = A x - b - B^T lam``."""
276 correction = b_eq.T @ lam if lam is not None else np.zeros_like(b)
277 return a.matvec(x) - b - correction
279 return self._run(len(b), sub_solve, reduced_gradient, warm)
281 def _run(
282 self,
283 n: int,
284 sub_solve: SubSolve,
285 reduced_gradient: ReducedGradient,
286 warm: tuple[NDArray[np.bool_], Vector] | None,
287 ) -> Result:
288 """Run the guarded primal-dual active-set loop.
290 The driver owns everything the termination proof depends on: the primal
291 and dual violator tests, the batch exchange with its patience counter,
292 and the least-index Bland fallback. What is solved on each free set — a
293 single reduced system (:meth:`solve`), or the equality-augmented saddle
294 system (:meth:`solve_eq`) — enters through the ``sub_solve`` callback,
295 with ``reduced_gradient`` supplying the matching dual test quantity. The
296 thresholds (``tol``, ``p_max``, ``track``, ``max_outer``) are read from
297 :attr:`config`.
299 Args:
300 n: Problem dimension.
301 sub_solve: Callback ``(idx, x0) -> (x_F, lam, inner_iters)`` solving
302 the subproblem on the free set ``idx``. ``x0`` is a warm inner
303 guess restricted to ``idx`` (None on a cold start); ``lam`` are
304 the equality multipliers (None for the bound-only problem).
305 reduced_gradient: Callback ``(x, lam) -> s`` computing the reduced
306 gradient that drives the dual violator test.
307 warm: Optional ``(free_mask, x_prev)`` pair from a previous solve.
308 Starts the loop from that free set and seeds every subproblem
309 solve from the newest iterate.
311 Returns:
312 A :class:`Result`; ``lam`` is whatever the last subproblem returned.
313 """
314 cfg = self.config
315 x, outer, inner_total, fallback, converged, free, lam, traj = _drive(
316 cfg.tol, cfg.p_max, cfg.track, cfg.max_outer, n, sub_solve, reduced_gradient, warm
317 )
318 return Result(
319 x=x,
320 outer=outer,
321 inner=inner_total,
322 fallback=fallback,
323 converged=converged,
324 free=free,
325 lam=lam,
326 traj=traj,
327 )