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

1"""Bordered KKT (saddle-point) solve over a symmetric operator's free block.""" 

2 

3from __future__ import annotations 

4 

5from typing import cast 

6 

7import numpy as np 

8 

9from ..core.types import Matrix, Vector 

10from ..operators import SymmetricOperator 

11 

12 

13def bordered_solve( 

14 operator: SymmetricOperator, 

15 free: object, 

16 c_free: Matrix, 

17 rhs: Vector | Matrix, 

18 d: Vector | Matrix, 

19) -> tuple[Vector | Matrix, Vector | Matrix]: 

20 """Solve the bordered KKT system ``[[H_FF, C.T], [C, 0]] @ [x; nu] = [rhs; d]``. 

21 

22 ``H_FF`` is the principal ``free`` block of *operator*, reached only through 

23 :meth:`~cvx.linalg.SymmetricOperator.solve_free` (so structured backends never 

24 materialise an ``n x n`` matrix). ``C = c_free`` is the ``(mc, len(free))`` 

25 constraint block on the free coordinates. This is the range-space (Schur 

26 complement) solution of an equality-constrained quadratic: the constraint 

27 columns ``C.T`` and the right-hand side are solved against ``H_FF`` together, 

28 so ``H_FF`` is factorised once, then the ``mc x mc`` Schur complement 

29 ``C H_FF^{-1} C.T`` gives the multipliers. 

30 

31 Multiple right-hand sides are supported in one call: pass ``rhs`` and ``d`` as 

32 matrices whose columns are independent systems (they share the ``H_FF`` and 

33 Schur factorisations). ``mc == 0`` (no constraint rows) is the plain free-block 

34 solve, with an empty multiplier block. 

35 

36 Args: 

37 operator: The symmetric operator whose free block is ``H_FF``. 

38 free: Index set of the free coordinates (as accepted by *operator*). 

39 c_free: Constraint block ``C`` of shape ``(mc, len(free))``. 

40 rhs: Right-hand side of shape ``(len(free),)`` or ``(len(free), k)``. 

41 d: Constraint target of shape ``(mc,)`` or ``(mc, k)``. 

42 

43 Returns: 

44 A pair ``(x, nu)``: the free solution (same trailing shape as *rhs*) and the 

45 constraint multipliers (shape ``(mc,)`` or ``(mc, k)``; empty when ``mc == 0``). 

46 

47 Example: 

48 >>> import numpy as np 

49 >>> from cvx.linalg import DenseOperator, bordered_solve 

50 >>> A = np.array([[4.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 2.0]]) 

51 >>> op = DenseOperator(A) 

52 >>> free = np.array([0, 1, 2]) 

53 >>> C = np.array([[1.0, 1.0, 1.0]]) 

54 >>> x, nu = bordered_solve(op, free, C, np.zeros(3), np.array([1.0])) 

55 >>> bool(np.isclose(C @ x, 1.0)) # constraint satisfied 

56 True 

57 >>> np.allclose(A @ x + C.T @ nu, 0.0) # stationarity: A x + C.T nu = rhs (= 0) 

58 True 

59 """ 

60 c_free = np.asarray(c_free, dtype=np.float64) 

61 rhs = np.asarray(rhs, dtype=np.float64) 

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

63 vector_rhs = rhs.ndim == 1 

64 rhs_2d = rhs.reshape(rhs.shape[0], -1) 

65 mc = c_free.shape[0] 

66 

67 if mc == 0: 

68 x_only = operator.solve_free(free, rhs) 

69 empty = np.zeros((0,) if vector_rhs else (0, rhs_2d.shape[1])) 

70 return x_only, empty 

71 

72 # One multi-RHS solve against H_FF covers the constraint columns C.T and the 

73 # right-hand side, so H_FF is factorised a single time. 

74 solved = operator.solve_free(free, np.column_stack([c_free.T, rhs_2d])) 

75 y = solved[:, :mc] # H_FF^{-1} C.T 

76 z = solved[:, mc:] # H_FF^{-1} rhs 

77 

78 # Schur complement C H_FF^{-1} C.T and the multipliers, then back-substitute. 

79 schur = c_free @ y 

80 d_2d = d.reshape(mc, -1) 

81 nu = cast("Matrix", np.linalg.solve(schur, c_free @ z - d_2d)) 

82 x = z - y @ nu 

83 

84 if vector_rhs: 

85 return x[:, 0], nu[:, 0] 

86 return x, nu