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

1"""Equality-augmented reduction: the Schur-complement saddle solve on a free set. 

2 

3Factored out of :meth:`nncg.solver.ActiveSetSolver.solve_eq` so the outer loop 

4keeps only its orchestration. On a free set the saddle system for ``B x = c`` is 

5solved by eliminating the multiplier ``lambda`` in R^p through the p-by-p Schur 

6complement ``S = B_F A_F^{-1} B_F^T``: the ``p + 1`` right-hand sides share the 

7operator ``A_F`` and are each one inner solve, then ``S lambda = c - B_F v0`` 

8fixes the multipliers in closed form. 

9""" 

10 

11from __future__ import annotations 

12 

13from typing import TYPE_CHECKING 

14 

15import numpy as np 

16from cvx.linalg import Matrix, SymmetricOperator, Vector, cholesky_solve 

17from numpy.typing import NDArray 

18 

19if TYPE_CHECKING: 

20 from .solver import InnerSolver 

21 

22 

23def _saddle_solve( 

24 inner: InnerSolver, 

25 a: SymmetricOperator, 

26 b: Vector, 

27 b_eq: Matrix, 

28 c_eq: Vector, 

29 idx: NDArray[np.int_], 

30 x0: Vector | None, 

31) -> tuple[Vector, Vector, int]: 

32 """Solve the equality-augmented saddle system on the free set ``idx``. 

33 

34 Runs ``p + 1`` inner solves through the shared free-block operator ``A_F`` 

35 (the ``v0`` column warm-started at ``x0``, the ``v1`` columns cold), forms the 

36 SPD Schur complement ``S = B_F A_F^{-1} B_F^T`` and recovers the multipliers 

37 from ``S lambda = c - B_F v0`` before back-substituting ``x_F = v0 + v1 lambda``. 

38 

39 Args: 

40 inner: The inner solver driving each free-block solve. 

41 a: The SPD operator ``A``. 

42 b: The linear term ``b``. 

43 b_eq: Equality matrix ``B`` of shape ``(p, n)``, full row rank on ``idx``. 

44 c_eq: Equality right-hand side ``c`` of shape ``(p,)``. 

45 idx: Integer positions of the free set ``F``. 

46 x0: Warm inner guess for the ``v0`` column restricted to ``idx``, or ``None``. 

47 

48 Returns: 

49 ``(x_F, lam, inner_iters)``: the free-block solution, the equality 

50 multipliers, and the total inner iteration count across all columns. 

51 """ 

52 p = b_eq.shape[0] 

53 b_f = b_eq[:, idx] 

54 v0, k0 = inner.solve(a, idx, b[idx], x0) 

55 v1 = np.zeros((idx.size, p)) 

56 k_cols = 0 

57 for j in range(p): 

58 v1[:, j], kj = inner.solve(a, idx, b_f[j], None) 

59 k_cols += kj 

60 schur = b_f @ v1 # p-by-p Schur complement, SPD 

61 lam = cholesky_solve(schur, c_eq - b_f @ v0) 

62 xf = v0 + v1 @ lam # x_F = A_F^{-1}(b_F + B_F^T lambda) 

63 return xf, lam, k0 + k_cols