schur_complement.py

importunittest

importnumpyasnp
importpytest


defschur_complement(
mat_a: np.ndarray,
mat_b: np.ndarray,
mat_c: np.ndarray,
pseudo_inv: np.ndarray|None=None,
) ->np.ndarray:
"""
 Schur complement of a symmetric matrix X given as a 2x2 block matrix
 consisting of matrices `A`, `B` and `C`.
 Matrix `A` must be quadratic and non-singular.
 In case `A` is singular, a pseudo-inverse may be provided using
 the `pseudo_inv` argument.

 | Link to Wiki: https://en.wikipedia.org/wiki/Schur_complement
 | See also Convex Optimization - Boyd and Vandenberghe, A.5.5

 >>> import numpy as np
 >>> a = np.array([[1, 2], [2, 1]])
 >>> b = np.array([[0, 3], [3, 0]])
 >>> c = np.array([[2, 1], [6, 3]])
 >>> schur_complement(a, b, c)
 array([[ 5., -5.],
 [ 0., 6.]])
 """
shape_a=np.shape(mat_a)
shape_b=np.shape(mat_b)
shape_c=np.shape(mat_c)

ifshape_a[0] !=shape_b[0]:
msg= (
"Expected the same number of rows for A and B. "
f"Instead found A of size {shape_a} and B of size {shape_b}"
 )
raiseValueError(msg)

ifshape_b[1] !=shape_c[1]:
msg= (
"Expected the same number of columns for B and C. "
f"Instead found B of size {shape_b} and C of size {shape_c}"
 )
raiseValueError(msg)

a_inv=pseudo_inv
ifa_invisNone:
try:
a_inv=np.linalg.inv(mat_a)
exceptnp.linalg.LinAlgError:
raiseValueError(
"Input matrix A is not invertible. Cannot compute Schur complement."
 )

returnmat_c-mat_b.T @ a_inv @ mat_b


classTestSchurComplement(unittest.TestCase):
deftest_schur_complement(self) ->None:
a=np.array([[1, 2, 1], [2, 1, 2], [3, 2, 4]])
b=np.array([[0, 3], [3, 0], [2, 3]])
c=np.array([[2, 1], [6, 3]])

s=schur_complement(a, b, c)

input_matrix=np.block([[a, b], [b.T, c]])

det_x=np.linalg.det(input_matrix)
det_a=np.linalg.det(a)
det_s=np.linalg.det(s)

assertnp.is_close(det_x, det_a*det_s)

deftest_improper_a_b_dimensions(self) ->None:
a=np.array([[1, 2, 1], [2, 1, 2], [3, 2, 4]])
b=np.array([[0, 3], [3, 0], [2, 3]])
c=np.array([[2, 1], [6, 3]])

withpytest.raises(ValueError):
schur_complement(a, b, c)

deftest_improper_b_c_dimensions(self) ->None:
a=np.array([[1, 2, 1], [2, 1, 2], [3, 2, 4]])
b=np.array([[0, 3], [3, 0], [2, 3]])
c=np.array([[2, 1, 3], [6, 3, 5]])

withpytest.raises(ValueError):
schur_complement(a, b, c)


if__name__=="__main__":
importdoctest

doctest.testmod()
unittest.main()