Coverage for src/cvx/linalg/kkt/projection.py: 100%

1"""Euclidean projection onto an affine set ``{x : C x = d}``.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import cast 

7 

8import numpy as np 

9 

10from ..core.exceptions import DimensionMismatchError, NotAMatrixError 

11from ..core.types import Matrix, Vector 

12from ..decomposition.cholesky import _factored_solver 

13 

14 

15class AffineProjection: 

16 """Euclidean projection onto the affine set ``{x : C x = d}``. 

17 

18 The projection of ``x`` is the nearest point (in the 2-norm) satisfying the 

19 equality constraints, 

20 

21 ``P(x) = x - C.T @ (C C.T)^{-1} @ (C x - d)``. 

22 

23 The Gram matrix ``C C.T`` is formed once at construction (the ``O(mc**2 * n)`` 

24 cost) and factorised on the first projection, so repeated projections -- for 

25 instance an alternating box / affine projection loop -- pay only ``O(mc**2)`` 

26 triangular solves each (with SciPy installed). ``C`` should have full row rank. 

27 

28 Args: 

29 c: Constraint matrix ``C`` of shape ``(mc, n)``. 

30 d: Constraint target ``d`` of shape ``(mc,)``. 

31 

32 Example: 

33 >>> import numpy as np 

34 >>> from cvx.linalg import AffineProjection 

35 >>> proj = AffineProjection(np.ones((1, 3)), np.array([1.0])) 

36 >>> p = proj.project(np.array([1.0, 1.0, 1.0])) 

37 >>> bool(np.isclose(np.ones(3) @ p, 1.0)) # lands on the affine set 

38 True 

39 >>> np.allclose(p, [1 / 3, 1 / 3, 1 / 3]) # nearest point with sum 1 

40 True 

41 """ 

42 

43 def __init__(self, c: Matrix, d: Vector) -> None: 

44 """Store ``C`` and ``d`` and precompute the Gram matrix ``C C.T``.""" 

45 c = np.asarray(c, dtype=np.float64) 

46 d = np.asarray(d, dtype=np.float64) 

47 if c.ndim != 2: 

48 raise NotAMatrixError(c.ndim, func="AffineProjection") 

49 if d.ndim != 1 or d.shape[0] != c.shape[0]: 

50 raise DimensionMismatchError(d.shape[0] if d.ndim == 1 else d.size, c.shape[0]) 

51 self._c = c 

52 self._d = d 

53 self._gram = c @ c.T 

54 self._solve_gram: Callable[[Vector | Matrix], Vector | Matrix] | None = None 

55 

56 @property 

57 def m(self) -> int: 

58 """Number of constraints (rows of ``C``).""" 

59 return int(self._c.shape[0]) 

60 

61 @property 

62 def n(self) -> int: 

63 """Ambient dimension (columns of ``C``).""" 

64 return int(self._c.shape[1]) 

65 

66 def project(self, x: Vector | Matrix) -> Vector | Matrix: 

67 """Return the Euclidean projection of ``x`` onto ``{x : C x = d}``. 

68 

69 Args: 

70 x: Point of shape ``(n,)`` or a stack of points ``(n, k)``. 

71 

72 Returns: 

73 The projected point(s), same shape as ``x``. 

74 """ 

75 x = np.asarray(x, dtype=np.float64) 

76 target = self._d if x.ndim == 1 else self._d[:, None] 

77 residual = self._c @ x - target 

78 correction = self._c.T @ self._gram_solver()(residual) 

79 return x - correction 

80 

81 def _gram_solver(self) -> Callable[[Vector | Matrix], Vector | Matrix]: 

82 """Return the solver for ``C C.T``, factorising it on first use. 

83 

84 A non-finite Gram matrix skips the factorisation and is LU-solved per 

85 call, so NaNs in ``C`` propagate to the projection as before. 

86 """ 

87 if self._solve_gram is None: 

88 if np.all(np.isfinite(self._gram)): 

89 self._solve_gram = _factored_solver(self._gram) 

90 else: 

91 gram = self._gram 

92 self._solve_gram = lambda rhs: cast("Vector | Matrix", np.linalg.solve(gram, rhs)) 

93 return self._solve_gram