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))