Coverage for src/nncg/inner.py: 100%
64 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-15 07:01 +0000
1"""Inner solvers: one free-block system ``A[F, F] y = rhs`` per active-set step.
3Each concrete inner solver provides ``solve(op, idx, rhs, x0) -> (y, iters)``,
4solving the free-block system ``A[F, F] y = rhs`` (and so satisfies the
5:class:`nncg.solver.InnerSolver` interface). This is the only module that knows
6about preconditioning: the built-in solvers are the identity/Jacobi/Nyström-
7preconditioned CG variants (:class:`CG`, :class:`Jacobi`, :class:`Nystrom`,
8:class:`GlobalNystrom`) and the direct :class:`Exact`. The operator-derived
9builders they run on — the free-block matvec and the diagonal/Nyström
10preconditioners — live in :mod:`nncg.preconditioners`. Further inner solvers —
11e.g. Clarabel- or KKT-equation-based — live in Jebel-Quant/mean_variance_solvers
12and satisfy the same structural interface.
14Examples:
15 The inner solver is the one thing that varies between these runs — the outer
16 loop, and the minimiser it certifies, are the same:
18 >>> import numpy as np
19 >>> from cvx.linalg import DenseOperator
20 >>> from nncg import CG, ActiveSetSolver, Exact, Jacobi
21 >>> a = DenseOperator(np.diag([1.0, 2.0, 4.0, 8.0]))
22 >>> b = np.array([1.0, 2.0, -4.0, -8.0])
23 >>> for inner in (CG(), Jacobi(), Exact()):
24 ... res = ActiveSetSolver(inner=inner).solve(a, b)
25 ... print(type(inner).__name__, res.converged, np.allclose(res.x, [1.0, 1.0, 0.0, 0.0]))
26 CG True True
27 Jacobi True True
28 Exact True True
30 The direct solver counts one inner "iteration" per solve, so it never spends
31 more than the number of outer steps:
33 >>> direct = ActiveSetSolver(inner=Exact()).solve(a, b)
34 >>> direct.inner <= direct.outer
35 True
37 The Nyström solvers sketch the free block at :attr:`NystromConfig.rank`, so
38 they belong on a problem larger than that rank and with a decaying spectrum
39 — here three orders of geometric decay, with a planted optimum on the even
40 coordinates:
42 >>> from nncg import GlobalNystrom, Nystrom
43 >>> d = 10.0 ** -np.linspace(0.0, 3.0, 40)
44 >>> x_star = np.where(np.arange(40) % 2 == 0, 1.0, 0.0)
45 >>> b = d * x_star - (1.0 - x_star)
46 >>> for inner in (Nystrom(), GlobalNystrom()):
47 ... res = ActiveSetSolver(inner=inner).solve(DenseOperator(np.diag(d)), b)
48 ... print(type(inner).__name__, res.converged, np.allclose(res.x, x_star))
49 Nystrom True True
50 GlobalNystrom True True
51"""
53from __future__ import annotations
55from dataclasses import dataclass, field, replace
57import numpy as np
58from cvx.linalg import SymmetricOperator, Vector
59from numpy.typing import NDArray
61from .krylov import KrylovConfig, Preconditioner, pcg
62from .preconditioners import (
63 GlobalNystromSketch,
64 NystromConfig,
65 _free_matvec,
66 _global_nystrom_sketch,
67 _jacobi,
68 _masked_nystrom,
69 _nystrom,
70)
72__all__ = ["CG", "Exact", "GlobalNystrom", "Jacobi", "Nystrom", "NystromConfig"]
74_RCOND_MIN = 1e-12 # matches cvx-linalg's DEFAULT_COND_THRESHOLD of 1e12
77def _default_krylov() -> KrylovConfig:
78 """Inner CG config with ``tol`` two orders below the outer tolerance (Lemma 5.1)."""
79 return KrylovConfig(tol=1e-10)
82def _pcg_block(
83 op: SymmetricOperator,
84 idx: NDArray[np.int_],
85 rhs: Vector,
86 x0: Vector | None,
87 krylov: KrylovConfig,
88 precond: Preconditioner | None,
89) -> tuple[Vector, int]:
90 """Solve the free block ``A[F, F] y = rhs`` by (preconditioned) CG, warm-started at ``x0``."""
91 return pcg(_free_matvec(op, idx), rhs, replace(krylov, precond=precond, x0=x0))
94def _same_free_block(
95 checked_op: SymmetricOperator | None,
96 checked_idx: NDArray[np.int_] | None,
97 op: SymmetricOperator,
98 idx: NDArray[np.int_],
99) -> bool:
100 """Return whether ``(op, idx)`` is the memoised free block already verified.
102 Keyed on operator *identity* (not equality) so the memo can never carry a
103 stale verdict across operators; ``checked_idx is None`` means nothing has
104 been verified yet.
105 """
106 return checked_op is op and checked_idx is not None and np.array_equal(checked_idx, idx)
109def _raise_if_singular(op: SymmetricOperator, idx: NDArray[np.int_]) -> None:
110 """Raise if the free block ``A[F, F]`` is numerically singular (``rcond_free < _RCOND_MIN``)."""
111 rcond = op.rcond_free(idx)
112 if rcond < _RCOND_MIN:
113 msg = f"free block of size {idx.size} is numerically singular (rcond={rcond:.2e})"
114 raise ValueError(msg)
117@dataclass(frozen=True)
118class Exact:
119 """Direct free-block solve via ``op.solve_free`` (one "iteration" per solve).
121 Suits backends whose ``solve_free`` is structured and cheap (e.g.
122 ``FactorOperator``'s Woodbury solve at ``O(|F| r^2)``). It ignores warm starts.
124 The ``rcond_free`` conditioning guard depends only on the free block, not the
125 right-hand side, so it is estimated at most once per free set:
126 :meth:`nncg.solver.ActiveSetSolver.solve_eq` drives ``p + 1`` solves through
127 the *same* free set per outer step, and the (up to ``O(|F|^3)``) estimate must
128 not be paid ``p + 1`` times over. The last verified ``(operator, idx)`` is
129 memoised in a private single slot — keyed on operator identity so the memo can
130 never carry a stale verdict across operators, and excluded from equality/repr
131 so ``Exact`` stays a value.
133 On the plain :meth:`~nncg.solver.ActiveSetSolver.solve` path every outer step
134 visits a *different* free set, so the memo never hits and the guard is paid on
135 every step — where, for a dense free block, the ``O(|F|^3)`` eigendecomposition
136 can cost several times the Cholesky solve it precedes. The guard is also
137 redundant when ``solve_free`` already fails loudly on a rank-deficient block
138 (e.g. ``cvx.linalg.cholesky_solve``'s Cholesky→LU fallback). Set
139 ``check_conditioning=False`` to skip it and let ``solve_free`` surface any
140 singularity itself.
142 Attributes:
143 check_conditioning: Estimate ``rcond_free`` and raise on a numerically
144 singular free block before each (new) solve. Default ``True``; set
145 ``False`` to trade the diagnostic for the raw structured solve.
146 """
148 check_conditioning: bool = True
149 _checked_op: SymmetricOperator | None = field(default=None, compare=False, repr=False)
150 _checked_idx: NDArray[np.int_] | None = field(default=None, compare=False, repr=False)
152 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]: # noqa: ARG002
153 """Solve the free block ``A[F, F] y = rhs`` directly, guarding its conditioning once per free set."""
154 if self.check_conditioning and not _same_free_block(self._checked_op, self._checked_idx, op, idx):
155 _raise_if_singular(op, idx)
156 object.__setattr__(self, "_checked_op", op)
157 object.__setattr__(self, "_checked_idx", idx)
158 return op.solve_free(idx, rhs), 1
161@dataclass(frozen=True)
162class CG:
163 """Plain matrix-free conjugate gradients (the identity preconditioner).
165 Attributes:
166 krylov: Tolerance and iteration cap of the CG solves (``tol`` defaults to ``1e-10``).
167 """
169 krylov: KrylovConfig = field(default_factory=_default_krylov)
171 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
172 """Solve the free block ``A[F, F] y = rhs`` by plain CG."""
173 return _pcg_block(op, idx, rhs, x0, self.krylov, None)
176@dataclass(frozen=True)
177class Jacobi:
178 """Jacobi-preconditioned CG — runs at the operator's condition number, a bad diagonal scaling removed.
180 Attributes:
181 krylov: Tolerance and iteration cap of the CG solves (``tol`` defaults to ``1e-10``).
182 """
184 krylov: KrylovConfig = field(default_factory=_default_krylov)
186 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
187 """Solve the free block ``A[F, F] y = rhs`` by Jacobi-preconditioned CG."""
188 return _pcg_block(op, idx, rhs, x0, self.krylov, _jacobi(op, idx))
191@dataclass(frozen=True)
192class Nystrom:
193 """Randomized Nyström-preconditioned CG — for free blocks with a steeply decaying spectrum.
195 Attributes:
196 krylov: Tolerance and iteration cap of the CG solves (``tol`` defaults to ``1e-10``).
197 nystrom: Sketch rank, oversampling, shift and seed of the low-rank
198 preconditioner (see :class:`nncg.preconditioners.NystromConfig`).
199 """
201 krylov: KrylovConfig = field(default_factory=_default_krylov)
202 nystrom: NystromConfig = field(default_factory=NystromConfig)
204 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
205 """Solve the free block ``A[F, F] y = rhs`` by Nyström-preconditioned CG (plain CG on an empty block)."""
206 precond = _nystrom(op, idx, self.nystrom) if idx.size else None
207 return _pcg_block(op, idx, rhs, x0, self.krylov, precond)
210@dataclass(frozen=True)
211class GlobalNystrom:
212 """Nyström-preconditioned CG sketched once on the full operator, then masked per free block.
214 :class:`Nystrom` resketches ``A[F, F]`` from scratch on every outer step —
215 the `rank + oversample` matrix-free products, a QR, a small Cholesky and an
216 SVD, all paid again each time the free set changes. This class instead
217 sketches the *full* operator ``A`` once: restricting a rank-``rank``
218 factorization to a principal submatrix is exact
219 (``(U diag(lam) U^T)[F, F] = U_F diag(lam) U_F^T`` for ``U_F = U[F, :]``),
220 so masking rows of the one global basis gives a valid free-block
221 preconditioner with no further matrix-free products against ``A`` — only a
222 small ``rank x rank`` factorization per free block (see
223 :func:`nncg.preconditioners._masked_nystrom`). This amortises well when the
224 same operator is solved repeatedly (a parameter sweep, successive warm
225 starts) or the active-set loop takes many outer steps; the trade is a
226 preconditioner not adapted to each free block's own local spectrum, so it
227 can take a few more CG iterations than a freshly-sketched :class:`Nystrom`
228 on a small or spectrally unusual free block.
230 The global sketch is memoised in a private single slot, keyed on operator
231 *identity* so the cache can never carry a stale sketch across operators
232 (mirrors :class:`Exact`'s conditioning memo) — excluded from equality/repr
233 so this class stays a value.
235 Attributes:
236 krylov: Tolerance and iteration cap of the CG solves (``tol`` defaults to ``1e-10``).
237 nystrom: Sketch rank, oversampling, shift and seed of the global sketch
238 (see :class:`nncg.preconditioners.NystromConfig`).
239 """
241 krylov: KrylovConfig = field(default_factory=_default_krylov)
242 nystrom: NystromConfig = field(default_factory=NystromConfig)
243 _checked_op: SymmetricOperator | None = field(default=None, compare=False, repr=False)
244 _sketch: GlobalNystromSketch | None = field(default=None, compare=False, repr=False)
246 def _ensure_sketch(self, op: SymmetricOperator) -> GlobalNystromSketch:
247 """Return the memoised global sketch of ``op``, (re)building it on a new operator."""
248 sketch = self._sketch
249 if self._checked_op is not op or sketch is None:
250 sketch = _global_nystrom_sketch(op, self.nystrom)
251 object.__setattr__(self, "_sketch", sketch)
252 object.__setattr__(self, "_checked_op", op)
253 return sketch
255 def solve(self, op: SymmetricOperator, idx: NDArray[np.int_], rhs: Vector, x0: Vector | None) -> tuple[Vector, int]:
256 """Solve the free block ``A[F, F] y = rhs`` by CG preconditioned from the masked global sketch."""
257 precond = _masked_nystrom(self._ensure_sketch(op), idx) if idx.size else None
258 return _pcg_block(op, idx, rhs, x0, self.krylov, precond)