Coverage for src/nncg/krylov.py: 100%

1"""Matrix-free (preconditioned) conjugate gradients — the Krylov core. 

2 

3:func:`pcg` solves an SPD system accessed only through a mat-vec callable — the 

4matrix is never required explicitly — and takes its preconditioner as a callable 

5``r -> M^{-1} r`` (``config.precond=None`` recovers plain CG). It is 

6operator-agnostic: the preconditioner *builders* that turn a 

7:class:`cvx.linalg.SymmetricOperator` into such a callable live in 

8:mod:`nncg.inner`, alongside the inner solvers that use them. Convergence is 

9governed by the spectral condition number of ``M^{-1} A`` at the 

10``O(sqrt(kappa))`` Krylov rate. 

11""" 

12 

13from __future__ import annotations 

14 

15from collections.abc import Callable 

16from dataclasses import dataclass 

17 

18import numpy as np 

19from cvx.linalg import Vector 

20 

21MatVec = Callable[[Vector], Vector] 

22 

23#: A preconditioner is the action ``r -> M^{-1} r`` of an SPD preconditioner 

24#: ``M ~ A`` — a general linear map, not necessarily diagonal. PCG only ever 

25#: multiplies by ``M^{-1}``, so the callable is all it needs. Build one from an 

26#: operator with the inner solvers in :mod:`nncg.inner`. 

27Preconditioner = Callable[[Vector], Vector] 

28 

29 

30@dataclass(frozen=True) 

31class KrylovConfig: 

32 """Options for a preconditioned CG solve (:func:`pcg`). 

33 

34 Bundles the solve knobs into one argument so :func:`pcg` keeps a short 

35 signature (matrix and right-hand side, then the config). 

36 

37 Attributes: 

38 precond: The action ``r -> M^{-1} r`` of an SPD preconditioner, applied 

39 once per iteration; ``None`` is the identity, so PCG reduces to 

40 plain CG. 

41 tol: Relative residual stopping tolerance ``||b - A x|| / ||b||``. 

42 maxit: Iteration cap; the current iterate is returned when it is hit. 

43 x0: Optional warm start. The initial residual is ``b - A x0``, so a good 

44 guess cuts the iteration count by the log of the initial error. 

45 """ 

46 

47 precond: Preconditioner | None = None 

48 tol: float = 1e-8 

49 maxit: int = 100_000 

50 x0: Vector | None = None 

51 

52 

53#: Shared default so ``KrylovConfig()`` is not called in argument defaults (ruff B008). 

54_DEFAULT_KRYLOV = KrylovConfig() 

55 

56#: The identity preconditioner: ``M^{-1} = I`` turns PCG back into plain CG. 

57_IDENTITY: Preconditioner = lambda r: r # noqa: E731 

58 

59 

60def _resolve_precond(precond: Preconditioner | None) -> Preconditioner: 

61 """Return the preconditioner action, substituting the identity for ``None``.""" 

62 return _IDENTITY if precond is None else precond 

63 

64 

65def _pcg_start(matvec: MatVec, rhs: Vector, x0: Vector | None) -> tuple[Vector, Vector]: 

66 """Initial iterate and residual: a zero start, or ``(x0, b - A x0)`` for a warm start.""" 

67 if x0 is None: 

68 return np.zeros_like(rhs), rhs.copy() 

69 x = x0.astype(np.float64, copy=True) 

70 return x, rhs - matvec(x) 

71 

72 

73# Why in-house rather than scipy.sparse.linalg.cg? scipy's CG is matrix-free 

74# and preconditionable too, so functionally it could stand in here. We keep our 

75# own for four reasons: (1) the matrix-free PCG is this package's core 

76# contribution — the reference implementation of the paper — and must stay 

77# auditable, not delegated to a black box; (2) the runtime dependency set is 

78# NumPy + cvx-linalg only, and this is ~60 lines NumPy already covers; (3) we 

79# return the iteration count, which the numerical study asserts on, whereas 

80# scipy returns only a convergence flag (recovering the count needs a callback); 

81# (4) scipy's atol/rtol stopping semantics have shifted across releases, so 

82# owning the loop pins the exact criterion and keeps the paper's numbers stable. 

83# Third-party solvers belong in the baseline comparisons (kept in the paper 

84# repo, Jebel-Quant/mean_variance_solvers), not in this inner Krylov core. 

85 

86 

87def pcg( 

88 matvec: MatVec, 

89 rhs: Vector, 

90 config: KrylovConfig = _DEFAULT_KRYLOV, 

91) -> tuple[Vector, int]: 

92 """Solve an SPD system by preconditioned conjugate gradients. 

93 

94 PCG converges at the condition number of ``M^{-1} A`` rather than of ``A``, 

95 where the preconditioner ``M^{-1}`` (``config.precond``) enters only as the 

96 action ``r -> M^{-1} r``; ``config.precond=None`` is the identity, so PCG 

97 reduces to plain CG. The inner solvers in :mod:`nncg.inner` build suitable 

98 preconditioners from an operator (diagonal Jacobi, randomized Nyström). 

99 

100 Args: 

101 matvec: The action ``v -> A v`` of an SPD operator. 

102 rhs: Right-hand side ``b``. 

103 config: Preconditioner, tolerance, iteration cap and warm start of the 

104 solve (see :class:`KrylovConfig`). 

105 

106 Returns: 

107 The approximate solution and the number of iterations taken. 

108 """ 

109 precond = _resolve_precond(config.precond) 

110 tol, maxit = config.tol, config.maxit 

111 bnorm = float(np.linalg.norm(rhs)) 

112 if bnorm == 0.0: 

113 return np.zeros_like(rhs), 0 

114 x, r = _pcg_start(matvec, rhs, config.x0) 

115 z = precond(r) 

116 p = z.copy() 

117 rz = float(r @ z) 

118 # A warm start that already solves the system to tolerance leaves r == 0, 

119 # so p == 0 and the search-direction curvature p @ ap vanishes; returning 

120 # here avoids the 0/0 in the alpha step and reports zero iterations. 

121 if float(np.linalg.norm(r)) / bnorm <= tol: 

122 return x, 0 

123 for it in range(1, maxit + 1): 

124 ap = matvec(p) 

125 alpha = rz / float(p @ ap) 

126 x += alpha * p 

127 r -= alpha * ap 

128 if float(np.linalg.norm(r)) / bnorm <= tol: 

129 return x, it 

130 z = precond(r) 

131 rz_new = float(r @ z) 

132 p = z + (rz_new / rz) * p 

133 rz = rz_new 

134 return x, maxit