def cholesky(a: Woodbury):
if a.cholesky is None:
warn_upmodule(
f"Converting {a} to dense to compute its Cholesky decomposition.",
category=ToDenseWarning,
)
a.cholesky = LowerTriangular(B.cholesky(B.reg(B.dense(a))))
return a.cholesky
def ratio(a, b):
"""Compute the ratio between two positive-definite matrices.
Args:
a (matrix): Numerator.
b (matrix): Denominator.
Returns:
matrix: Ratio.
"""
chol = B.cholesky(a)
return B.sum(B.iqf_diag(b, chol))
def sample(a, num=1): # pragma: no cover
"""Sample from covariance matrices.
Args:
a (tensor): Covariance matrix to sample from.
num (int): Number of samples.
Returns:
tensor: Samples as rank 2 column vectors.
"""
chol = B.cholesky(a)
return B.matmul(chol, B.randn(B.dtype_float(a), B.shape(chol)[1], num))
def test_cholesky_solve_ut(dense_pd):
chol = B.cholesky(dense_pd)
with AssertDenseWarning(
[
"solving <upper-triangular> x = <diagonal>",
"matrix-multiplying <upper-triangular> and <lower-triangular>",
]
):
approx(
B.cholesky_solve(B.transpose(chol), B.eye(chol)),
B.inv(B.matmul(chol, chol, tr_a=True)),
)
def pd_inv(a: Union[B.Numeric, AbstractMatrix]):
"""Invert a positive-definite matrix.
Args:
a (matrix): Positive-definite matrix to invert.
Returns:
matrix: Inverse of `a`, which is also positive definite.
"""
a = convert(a, AbstractMatrix)
# The call to `cholesky_solve` will convert the identity matrix to dense, because
# `cholesky(a)` will not have any exploitable structure. We suppress the expected
# warning by converting `B.eye(a)` to dense here already.
return B.cholesky_solve(B.cholesky(a), B.dense(B.eye(a)))
def iqf_diag(a, b, c):
"""Compute the diagonal of `transpose(b) inv(a) c` where `a` is assumed
to be positive definite.
Args:
a (matrix): Matrix `a`.
b (matrix): Matrix `b`.
c (matrix, optional): Matrix `c`. Defaults to `b`.
Returns:
vector: Diagonal of resulting quadratic form.
"""
chol = B.cholesky(a)
chol_b = B.solve(chol, b)
if c is b:
chol_c = chol_b
else:
chol_c = B.solve(chol, c)
return B.matmul_diag(chol_b, chol_c, tr_a=True)
def cholesky(a: LowRank):
_assert_square_cholesky(a)
if a.cholesky is None:
a.cholesky = B.matmul(a.left, B.cholesky(a.middle))
return a.cholesky
def cholesky(a: Dense):
_assert_square_cholesky(a)
if a.cholesky is None:
a.cholesky = LowerTriangular(B.cholesky(B.reg(a.mat)))
return a.cholesky
def test_cholesky_const(const_pd):
chol = B.dense(B.cholesky(const_pd))
approx(B.matmul(chol, chol, tr_b=True), const_pd)
_check_cache(const_pd)
def test_cholesky_square_assertion():
with pytest.raises(AssertionError):
B.cholesky(Dense(B.randn(3, 4)))
def root(a: Diagonal):
_assert_square_root(a)
return B.cholesky(a)
def generate(code):
"""Generate a random tensor of a particular type, specified with a code.
Args:
code (str): Code of the matrix.
Returns:
tensor: Random tensor.
"""
mat_code, shape_code = code.split(":")
# Parse shape.
if shape_code == "":
shape = ()
else:
shape = tuple(int(d) for d in shape_code.split(","))
if mat_code == "randn":
return B.randn(*shape)
elif mat_code == "randn_pd":
mat = B.randn(*shape)
# If it is a scalar or vector, just pointwise square it.
if len(shape) in {0, 1}:
return mat**2 + 1
else:
return B.matmul(mat, mat, tr_b=True) + B.eye(shape[0])
elif mat_code == "zero":
return Zero(B.default_dtype, *shape)
elif mat_code == "const":
return Constant(B.randn(), *shape)
elif mat_code == "const_pd":
return Constant(B.randn()**2 + 1, *shape)
elif mat_code == "lt":
mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1))))
return LowerTriangular(mat)
elif mat_code == "lt_pd":
mat = generate(f"randn_pd:{shape[0]},{shape[0]}")
return LowerTriangular(B.cholesky(B.reg(mat)))
elif mat_code == "ut":
mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1))))
return UpperTriangular(B.transpose(mat))
elif mat_code == "ut_pd":
mat = generate(f"randn_pd:{shape[0]},{shape[0]}")
return UpperTriangular(B.transpose(B.cholesky(B.reg(mat))))
elif mat_code == "dense":
return Dense(generate(f"randn:{shape_code}"))
elif mat_code == "dense_pd":
return Dense(generate(f"randn_pd:{shape_code}"))
elif mat_code == "diag":
return Diagonal(generate(f"randn:{shape_code}"))
elif mat_code == "diag_pd":
return Diagonal(generate(f"randn_pd:{shape_code}"))
else:
raise RuntimeError(f'Cannot parse generation code "{code}".')
def logdet(a: AbstractMatrix):
return 2 * B.logdet(B.cholesky(a))
def test_cholesky_solve_lt(dense_pd):
chol = B.cholesky(dense_pd)
with AssertDenseWarning("solving <lower-triangular> x = <diagonal>"):
approx(B.cholesky_solve(chol, B.eye(chol)), B.inv(dense_pd))
def cholesky(a: Kronecker):
if a.cholesky is None:
a.cholesky = Kronecker(B.cholesky(a.left), B.cholesky(a.right))
return a.cholesky
def _check_cache(a):
chol1 = B.cholesky(a)
chol2 = B.cholesky(a)
assert chol1 is chol2
def root(a: Constant):
_assert_square_root(a)
return B.cholesky(a)
def test_cholesky_zero(zero1):
assert B.cholesky(zero1) is zero1
def assert_positive_definite(x):
"""Assert that a matrix is positive definite."""
# Check that Cholesky decomposition succeeds.
B.cholesky(x)
def test_cholesky_lr(lr_pd):
chol = B.dense(B.cholesky(lr_pd))
approx(B.matmul(chol, chol, tr_b=True), lr_pd)
_check_cache(lr_pd)
def test_cholesky_retry_factor(check_lazy_shapes):
# Try `cholesky_retry_factor = 1`.
B.cholesky_retry_factor = 1
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Try `cholesky_retry_factor = 10`.
B.cholesky_retry_factor = 10
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Try `cholesky_retry_factor = 100`.
B.cholesky_retry_factor = 100
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Reset the factor!
B.cholesky_retry_factor = 1
def f1(x):
dists2 = (x - B.transpose(x))**2
K = B.exp(-0.5 * dists2)
K = K + B.epsilon * B.eye(t, n)
L = B.cholesky(K)
return B.matmul(L, B.ones(t, n, m))
def inverse_transform(x):
chol = B.cholesky(B.reg(x))
return B.concat(B.log(B.diag(chol)), B.tril_to_vec(chol,
offset=-1))