Coverage for src/nncg/preconditioners.py: 100%
102 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-15 07:01 +0000
1"""Free-block operator and preconditioner builders for the inner solvers.
3Every builder here has the same shape — ``(op, idx, ...) -> callable`` — turning
4a :class:`cvx.linalg.SymmetricOperator` and a free set ``F`` into something the
5matrix-free CG in :mod:`nncg.krylov` can call on the block ``A[F, F]``:
6:func:`_free_matvec` yields the block action ``v -> A[F, F] v``, and
7:func:`_jacobi` / :func:`_nystrom` yield the diagonal / randomized-Nyström
8preconditioners ``r -> M^{-1} r``. The reduced matrix is never materialised.
9:func:`_global_nystrom_sketch` / :func:`_masked_nystrom` are the same Nyström
10idea sketched once on the *full* operator and reused across free blocks by
11masking rows, rather than resketched per block (see :class:`nncg.inner.GlobalNystrom`).
12The inner-solver classes that select among these live in :mod:`nncg.inner`.
13"""
15from __future__ import annotations
17from dataclasses import dataclass
18from typing import cast
20import numpy as np
21from cvx.linalg import SymmetricOperator, Vector
22from numpy.typing import NDArray
24from .krylov import MatVec, Preconditioner
27@dataclass(frozen=True)
28class NystromConfig:
29 """Tuning knobs for the Nyström preconditioner of :class:`nncg.inner.Nystrom`.
31 Attributes:
32 rank: Target sketch rank — the number of leading eigenpairs captured,
33 clamped to the free-block dimension.
34 oversample: Extra sketch columns drawn for accuracy before truncating
35 back to ``rank`` (the standard randomized-SVD oversampling).
36 shift: Explicit scalar tail eigenvalue, or ``None`` for the default
37 (the largest eigenvalue the sketch does not capture).
38 seed: Seed for the Gaussian test matrix; fixed by default so a solve is
39 reproducible. ``None`` draws a fresh one.
41 Raises:
42 ValueError: When ``rank < 1``.
43 """
45 rank: int = 10
46 oversample: int = 10
47 shift: float | None = None
48 seed: int | None = 0
50 def __post_init__(self) -> None:
51 """Validate that the sketch rank is a positive integer."""
52 if self.rank < 1:
53 msg = f"NystromConfig.rank must be a positive integer; got {self.rank}"
54 raise ValueError(msg)
57#: Shared default so ``NystromConfig()`` is not called in argument defaults (ruff B008).
58_DEFAULT_NYSTROM = NystromConfig()
61def _free_matvec(op: SymmetricOperator, idx: NDArray[np.int_]) -> MatVec:
62 """Return the free-block action ``v -> A[F, F] v`` of the operator.
64 The free-set restriction is hoisted out of the inner loop: the pre-sliced
65 free-block operator is built once here and the returned callable is its
66 plain ``matvec``. Restricting per CG iteration instead re-gathers the
67 operator's storage on every call, an order of magnitude more wall clock.
69 Args:
70 op: The symmetric operator ``A``.
71 idx: Integer positions of the free set ``F``.
73 Returns:
74 A callable computing ``A[F, F] @ v``; the reduced matrix is never
75 materialised.
76 """
77 return cast(MatVec, op.restricted(idx).matvec)
80def _jacobi(op: SymmetricOperator, idx: NDArray[np.int_] | None = None) -> Preconditioner:
81 """Return the Jacobi preconditioner ``r -> (1 / diag(A))[F] * r``.
83 With ``idx`` given it is sliced to the free block ``A[F, F]``; ``idx=None``
84 gives the whole operator. The diagonal is read off ``op.diag`` (matrix never
85 materialised). A symmetric positive definite operator has a strictly positive
86 diagonal, so a non-positive or non-finite entry means ``A`` is not SPD there;
87 that is reported eagerly rather than propagated as an ``inf`` into the CG loop.
89 Args:
90 op: The SPD operator ``A``.
91 idx: Integer positions of the free set ``F``, or ``None`` for all of ``A``.
93 Returns:
94 The elementwise preconditioner (sliced to the free set when ``idx`` given).
96 Raises:
97 ValueError: When a diagonal entry is non-positive or non-finite.
98 NotImplementedError: When the backend does not expose ``diag``.
99 """
100 diag = np.asarray(op.diag, dtype=np.float64)
101 bad = np.flatnonzero(~(diag > 0.0) | ~np.isfinite(diag))
102 if bad.size:
103 i = int(bad[0])
104 msg = f"operator diagonal is not strictly positive at index {i} (diag={diag[i]:.2e}); A is not SPD"
105 raise ValueError(msg)
106 dinv = 1.0 / diag
107 if idx is not None:
108 dinv = dinv[idx]
109 return lambda r: dinv * r
112def _nystrom_sketch(
113 matvec: MatVec, n: int, sketch: int, seed: int | None
114) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
115 """Build the randomized Nyström eigendecomposition of the operator.
117 Draws an ``n x sketch`` orthonormal test matrix, forms ``Y = A Omega`` with
118 ``sketch`` matrix-free products, and applies the Frangella-Tropp-Udell
119 stabilised sketch (Alg. 2.1): a nugget shift lifts ``Y`` off the range
120 boundary so the small Cholesky is well conditioned, then a thin SVD yields
121 the orthonormal basis and the nugget-corrected, clipped eigenvalues.
123 Args:
124 matvec: The matrix-free action ``v -> A v`` (already free-block sliced).
125 n: Dimension of the (free-block) operator.
126 sketch: Number of test columns (``rank + oversample``, clamped to ``n``).
127 seed: Seed for the Gaussian test matrix (``None`` draws a fresh one).
129 Returns:
130 ``(u_full, lam_full)``: the orthonormal basis and eigenvalues in
131 descending order, before truncation to the requested rank.
132 """
133 rng = np.random.default_rng(seed)
134 omega = np.linalg.qr(rng.standard_normal((n, sketch)))[0] # n x sketch, orthonormal
135 y = np.column_stack([matvec(omega[:, j]) for j in range(sketch)]) # A_F @ Omega
137 # Stabilising shift (Frangella-Tropp-Udell Alg. 2.1): lift Y off the range
138 # boundary so the small Cholesky is well conditioned, then subtract it back.
139 nu = np.sqrt(n) * np.finfo(np.float64).eps * float(np.linalg.norm(y, ord=2))
140 y_nu = y + nu * omega
141 chol = np.linalg.cholesky(omega.T @ y_nu) # lower, chol @ chol.T = Omega^T Y_nu
142 b = np.linalg.solve(chol, y_nu.T).T # B = Y_nu chol^{-T}, so B B^T = Y_nu (Omega^T Y_nu)^{-1} Y_nu^T
143 u_full, sv, _ = np.linalg.svd(b, full_matrices=False)
144 lam_full = np.maximum(sv**2 - nu, 0.0) # eigenvalues of the Nystrom approximation
145 return u_full, lam_full
148def _check_captured(lam: NDArray[np.float64], rank: int) -> None:
149 """Validate that the rank-``rank`` sketch captured a genuine eigenspace.
151 The smallest captured eigenvalue must be a real positive eigenvalue, not
152 floating-point noise from a rank the block does not possess.
154 Args:
155 lam: The captured eigenvalues in descending order.
156 rank: The requested sketch rank (for the error message).
158 Raises:
159 ValueError: When the smallest captured eigenvalue is non-positive or
160 negligible relative to the largest — ``rank`` exceeds the numerical
161 rank of the block, so reduce it.
162 """
163 if float(lam[0]) <= 0.0 or float(lam[-1]) <= 1e-12 * float(lam[0]):
164 ratio = float(lam[-1]) / float(lam[0]) if float(lam[0]) > 0.0 else 0.0
165 msg = (
166 f"the rank-{rank} Nystrom sketch captured a negligible eigenvalue "
167 f"(lam_min/lam_max={ratio:.2e}); rank exceeds the numerical rank of A — reduce it"
168 )
169 raise ValueError(msg)
172def _nystrom_shift(shift: float | None, lam: NDArray[np.float64], lam_full: NDArray[np.float64], rank: int) -> float:
173 """Choose the scalar deflation shift ``mu`` for the uncaptured spectral tail.
175 Uses an explicit ``shift`` when given; otherwise the largest *uncaptured*
176 eigenvalue (deflation), falling back to the smallest captured one when the
177 sketch spanned the whole spectrum (``oversample=0``).
179 Args:
180 shift: Explicit shift, or ``None`` for the default deflation choice.
181 lam: The captured eigenvalues in descending order.
182 lam_full: All sketched eigenvalues (captured plus tail) in descending order.
183 rank: The requested sketch rank.
185 Returns:
186 The positive scalar shift ``mu``.
188 Raises:
189 ValueError: When the resolved shift is not positive.
190 """
191 if shift is not None:
192 mu = float(shift)
193 elif lam_full.size > rank and float(lam_full[rank]) > 1e-12 * float(lam[0]):
194 mu = float(lam_full[rank]) # largest uncaptured eigenvalue: the deflation shift
195 else:
196 mu = float(lam[-1]) # sketch captured the whole spectrum: fall back to smallest captured
197 if mu <= 0.0:
198 msg = f"shift must be positive; got {mu:.2e}"
199 raise ValueError(msg)
200 return mu
203def _nystrom(
204 op: SymmetricOperator, idx: NDArray[np.int_] | None = None, config: NystromConfig = _DEFAULT_NYSTROM
205) -> Preconditioner:
206 """Return a randomized Nyström preconditioner for ``A`` (or its free block ``A[F, F]``).
208 A rank-``rank`` randomized Nyström sketch approximates the block as
209 ``A_F ~ U diag(lam) U^T`` (Frangella, Tropp & Udell, 2023). Treating the
210 uncaptured tail as a single scalar ``shift`` gives the SPD preconditioner
211 ``M = U diag(lam) U^T + shift (I - U U^T)``, applied by the Woodbury formula
212 ``M^{-1} r = (1/shift) r + U ((1/lam - 1/shift) * (U^T r))`` in ``O(|F| rank)``
213 per call — ``U``, ``lam``, ``shift`` are captured once at build time (the
214 ``rank + oversample`` matrix-free products and a small dense factorisation).
215 The default ``shift`` is the largest *uncaptured* eigenvalue (deflation),
216 falling back to the smallest captured one when ``oversample=0``.
218 Args:
219 op: The SPD operator ``A``.
220 idx: Integer positions of the free set ``F``.
221 config: Sketch rank, oversampling, shift and seed.
223 Returns:
224 A callable applying ``r -> M^{-1} r`` on the free block.
226 Raises:
227 ValueError: When the smallest captured eigenvalue is negligible relative
228 to the largest (``config.rank`` exceeds the numerical rank of the
229 block — reduce it), or when an explicit ``config.shift`` is not
230 positive.
231 """
232 rank, oversample, shift, seed = config.rank, config.oversample, config.shift, config.seed
233 matvec: MatVec
234 if idx is None:
235 matvec, n = op.matvec, op.n
236 else:
237 matvec, n = _free_matvec(op, idx), int(idx.size)
238 rank = min(rank, n)
239 sketch = min(rank + max(oversample, 0), n)
241 u_full, lam_full = _nystrom_sketch(matvec, n, sketch, seed)
242 u = u_full[:, :rank]
243 lam = lam_full[:rank]
244 _check_captured(lam, rank)
245 mu = _nystrom_shift(shift, lam, lam_full, rank)
247 inv_mu = 1.0 / mu
248 coef = 1.0 / lam - inv_mu # low-rank correction weights; <= 0 for the default shift
250 def apply(r: Vector) -> Vector:
251 """Apply ``M^{-1}`` via the captured scalar shift plus low-rank correction."""
252 z: Vector = inv_mu * r + u @ (coef * (u.T @ r))
253 return z
255 return apply
258@dataclass(frozen=True)
259class GlobalNystromSketch:
260 """A rank-``rank`` randomized Nyström sketch of the *full* operator ``A``.
262 Reusable across every free block the active-set loop visits — see
263 :func:`_masked_nystrom`.
265 Attributes:
266 u: The ``n x rank`` orthonormal sketch basis (of the full operator).
267 lam: The ``rank`` captured eigenvalues, descending.
268 mu: The scalar deflation shift for the uncaptured tail.
269 """
271 u: NDArray[np.float64]
272 lam: NDArray[np.float64]
273 mu: float
276def _global_nystrom_sketch(op: SymmetricOperator, config: NystromConfig = _DEFAULT_NYSTROM) -> GlobalNystromSketch:
277 """Sketch the full operator ``A`` once: ``A ~ U diag(lam) U^T + mu (I - U U^T)``.
279 Identical randomized-Nyström machinery to :func:`_nystrom`, but run on the
280 whole operator (``idx=None``) rather than a free block, so the result can be
281 masked down to any free block's rows via :func:`_masked_nystrom` without
282 resketching.
284 Args:
285 op: The SPD operator ``A``.
286 config: Sketch rank, oversampling, shift and seed.
288 Returns:
289 The captured basis, eigenvalues and tail shift.
291 Raises:
292 ValueError: When the smallest captured eigenvalue is negligible relative
293 to the largest (``config.rank`` exceeds the numerical rank of ``A``
294 — reduce it), or when an explicit ``config.shift`` is not positive.
295 """
296 rank, oversample, shift, seed = config.rank, config.oversample, config.shift, config.seed
297 n = op.n
298 rank = min(rank, n)
299 sketch = min(rank + max(oversample, 0), n)
301 u_full, lam_full = _nystrom_sketch(op.matvec, n, sketch, seed)
302 u = u_full[:, :rank]
303 lam = lam_full[:rank]
304 _check_captured(lam, rank)
305 mu = _nystrom_shift(shift, lam, lam_full, rank)
306 return GlobalNystromSketch(u=u, lam=lam, mu=mu)
309def _masked_nystrom(sketch: GlobalNystromSketch, idx: NDArray[np.int_]) -> Preconditioner:
310 """Return the Woodbury preconditioner for the free block ``A[F, F]``, from a global sketch.
312 Restricting a rank-``rank`` factorization to a principal submatrix is exact:
313 ``(U diag(lam) U^T + mu (I - U U^T))[F, F] = U_F diag(lam) U_F^T + mu (I_F -
314 U_F U_F^T) = mu I_F + U_F diag(lam - mu) U_F^T`` where ``U_F = U[F, :]`` — no
315 extra approximation beyond the global sketch's own error, and no matrix-free
316 products against the operator are needed here at all. ``U_F`` is not
317 orthonormal after masking though, so unlike :func:`_nystrom`'s specialised
318 identity-plus-projection inverse, ``M^{-1}`` is applied via the *general*
319 Sherman-Morrison-Woodbury identity ``(mu I + U_F C U_F^T)^{-1} = I/mu - U_F
320 (C^{-1} + U_F^T U_F / mu)^{-1} U_F^T / mu^2`` in ``O(|F| rank + rank^3)`` to
321 build, then ``O(|F| rank)`` per application — the ``rank x rank`` factor is
322 built once per free block, not once per CG iteration.
324 Args:
325 sketch: The global sketch of ``A`` from :func:`_global_nystrom_sketch`.
326 idx: Integer positions of the free set ``F``.
328 Returns:
329 A callable applying ``r -> M^{-1} r`` on the free block.
330 """
331 u_f = sketch.u[idx, :]
332 mu = sketch.mu
333 c_inv = 1.0 / (sketch.lam - mu)
334 gram = u_f.T @ u_f
335 k_inv = np.linalg.inv(np.diag(c_inv) + gram / mu)
337 def apply(r: Vector) -> Vector:
338 """Apply ``M^{-1}`` via the general Woodbury update on the masked basis."""
339 y = k_inv @ (u_f.T @ r)
340 z: Vector = r / mu - (u_f @ y) / mu**2
341 return z
343 return apply