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

1"""Conjugate-gradient inner solvers. 

2 

3Plain CG and Jacobi-preconditioned CG for symmetric positive definite (SPD) 

4systems, accessed only through a mat-vec callable — the matrix is never 

5required explicitly. These are the inner solvers of the active-set loop in 

6:mod:`nncg.solver`; their convergence is governed by the spectral condition 

7number kappa at the O(sqrt(kappa)) Krylov rate. 

8""" 

9 

10from __future__ import annotations 

11 

12from collections.abc import Callable 

13 

14import numpy as np 

15from cvx.linalg import Vector 

16 

17MatVec = Callable[[Vector], Vector] 

18 

19 

20def cg( 

21 matvec: MatVec, 

22 rhs: Vector, 

23 tol: float = 1e-8, 

24 maxit: int = 100_000, 

25 x0: Vector | None = None, 

26) -> tuple[Vector, int]: 

27 """Solve an SPD system by conjugate gradients. 

28 

29 Args: 

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

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

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

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

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

35 good guess cuts the iteration count by the log of the initial 

36 error. 

37 

38 Returns: 

39 The approximate solution and the number of iterations taken. 

40 """ 

41 if x0 is None: 

42 x = np.zeros_like(rhs) 

43 r = rhs.copy() 

44 else: 

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

46 r = rhs - matvec(x) 

47 p = r.copy() 

48 rs = float(r @ r) 

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

50 if bnorm == 0.0: 

51 return np.zeros_like(rhs), 0 

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

53 ap = matvec(p) 

54 alpha = rs / float(p @ ap) 

55 x += alpha * p 

56 r -= alpha * ap 

57 rs_new = float(r @ r) 

58 if np.sqrt(rs_new) / bnorm <= tol: 

59 return x, it 

60 p = r + (rs_new / rs) * p 

61 rs = rs_new 

62 return x, maxit 

63 

64 

65def pcg( 

66 matvec: MatVec, 

67 rhs: Vector, 

68 dinv: Vector, 

69 tol: float = 1e-8, 

70 maxit: int = 100_000, 

71) -> tuple[Vector, int]: 

72 """Solve an SPD system by Jacobi-preconditioned conjugate gradients. 

73 

74 The preconditioner is ``M^{-1} = diag(dinv)``; for operators that are a 

75 well-conditioned core under a bad diagonal scaling, PCG runs at the core's 

76 condition number regardless of the scaling. 

77 

78 Args: 

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

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

81 dinv: Elementwise inverse of the operator's diagonal. 

82 tol: Relative residual stopping tolerance. 

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

84 

85 Returns: 

86 The approximate solution and the number of iterations taken. 

87 """ 

88 x = np.zeros_like(rhs) 

89 r = rhs.copy() 

90 z = dinv * r 

91 p = z.copy() 

92 rz = float(r @ z) 

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

94 if bnorm == 0.0: 

95 return x, 0 

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

97 ap = matvec(p) 

98 alpha = rz / float(p @ ap) 

99 x += alpha * p 

100 r -= alpha * ap 

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

102 return x, it 

103 z = dinv * r 

104 rz_new = float(r @ z) 

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

106 rz = rz_new 

107 return x, maxit