def all_different(x, y):
"""Assert that two matrices have all different columns.
Args:
x (matrix): First matrix.
y (matrix): Second matrix.
"""
assert B.all(B.pw_dists(B.transpose(x), B.transpose(y)) > 1e-2)
def m2(self):
"""matrix: Second moment."""
if self._m2 is None:
self._m2 = B.cholsolve(B.chol(self.prec),
self.prec + B.outer(self.lam))
self._m2 = B.cholsolve(B.chol(self.prec), B.transpose(self._m2))
return self._m2
def test_cholesky_solve_lt(lt_pd, dense2):
check_bin_op(B.triangular_solve,
lt_pd,
dense2,
asserted_type=Dense,
check_broadcasting=False)
with pytest.warns(UserWarning):
B.triangular_solve(B.transpose(lt_pd), dense2)
def pinv(a: AbstractMatrix):
"""Compute the left pseudo-inverse.
Args:
a (matrix): Matrix to compute left pseudo-inverse of.
Returns:
matrix: Left pseudo-inverse of `a`.
"""
return B.cholsolve(B.chol(B.matmul(a, a, tr_a=True)), B.transpose(a))
def test_cholesky_solve_ut(ut_pd, dense2):
check_bin_op(
lambda a, b: B.triangular_solve(a, b, lower_a=False),
ut_pd,
dense2,
asserted_type=Dense,
check_broadcasting=False,
)
with pytest.warns(UserWarning):
B.triangular_solve(B.transpose(ut_pd), dense2, lower_a=False)
def test_matmul_multiple(code_a, code_b, code_c):
for tr_a in [True, False]:
for tr_b in [True, False]:
for tr_c in [True, False]:
a = generate(code_a)
b = generate(code_b)
c = generate(code_c)
if tr_a:
a = B.transpose(a)
if tr_b:
b = B.transpose(b)
if tr_c:
c = B.transpose(c)
approx(
B.matmul(a, b, c, tr_a=tr_a, tr_b=tr_b, tr_c=tr_c),
B.matmul(B.matmul(a, b, tr_a=tr_a, tr_b=tr_b),
c,
tr_b=tr_c),
)
def eig(a, compute_eigvecs=True):
if compute_eigvecs:
vals, vecs = B.eig(a, compute_eigvecs=True)
vals = B.flatten(vals)
if B.rank(vecs) == 3:
vecs = B.transpose(vecs, perm=(1, 0, 2))
vecs = B.reshape(vecs, 3, -1)
order = compute_order(vals)
return B.take(vals, order), B.abs(B.take(vecs, order, axis=1))
else:
vals = B.flatten(B.eig(a, compute_eigvecs=False))
return B.take(vals, compute_order(vals))
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 __call__(self, x):
# Put the batch dimension second.
x_rank = B.rank(x)
if x_rank == 2:
x = x[:, None, :]
elif x_rank == 3:
x = B.transpose(x, perm=(1, 0, 2))
else:
raise ValueError(f"Cannot handle inputs of rank {B.rank(x)}.")
# Recurrently apply the cell.
n, batch_size, m = B.shape(x)
y0 = B.zeros(B.dtype(x), batch_size, self.cell.width)
h0 = B.tile(self.h0, batch_size, 1)
res = B.scan(self.cell, x, h0, y0)[1]
# Put the batch dimension first again.
res = B.transpose(res, perm=(1, 0, 2))
# Remove the batch dimension, if that didn't exist before.
if x_rank == 2:
res = res[0, :, :]
return res
def summarise_samples(x, samples, db=False):
"""Summarise samples.
Args:
x (vector): Inputs of samples.
samples (tensor): Samples, with the first dimension corresponding
to different samples.
db (bool, optional): Convert to decibels.
Returns:
:class:`collections.namedtuple`: Named tuple containing various
statistics of the samples.
"""
x, samples = B.to_numpy(x, samples)
random_inds = np.random.permutation(B.shape(samples)[0])[:3]
def transform(x):
if db:
return 10 * np.log10(x)
else:
return x
perm = tuple(reversed(range(B.rank(samples)))) # Reverse all dimensions.
return collect(
x=B.to_numpy(x),
mean=transform(B.mean(samples, axis=0)),
var=transform(B.std(samples, axis=0))**2,
err_68_lower=transform(B.quantile(samples, 0.32, axis=0)),
err_68_upper=transform(B.quantile(samples, 1 - 0.32, axis=0)),
err_95_lower=transform(B.quantile(samples, 0.025, axis=0)),
err_95_upper=transform(B.quantile(samples, 1 - 0.025, axis=0)),
err_99_lower=transform(B.quantile(samples, 0.0015, axis=0)),
err_99_upper=transform(B.quantile(samples, 1 - 0.0015, axis=0)),
samples=transform(B.transpose(samples, perm=perm)[..., random_inds]),
all_samples=transform(B.transpose(samples, perm=perm)),
)
def hessian(f, x):
"""Compute the Hessian of a function at a certain input.
Args:
f (function): Function to compute Hessian of.
x (column vector): Input to compute Hessian at.
differentiable (bool, optional): Make the computation of the Hessian
differentiable. Defaults to `False`.
Returns:
matrix: Hessian.
"""
if B.rank(x) != 2 or B.shape(x)[1] != 1:
raise ValueError("Input must be a column vector.")
# Use RMAD twice to preserve memory.
hess = jax.jacrev(jax.jacrev(lambda x: f(x[:, None])))(x[:, 0])
return (hess + B.transpose(hess)
) / 2 # Symmetrise to counteract numerical errors.
def closest_psd(a, inv=False):
"""Map a matrix to the closest PSD matrix.
Args:
a (tensor): Matrix.
inv (bool, optional): Also invert `a`.
Returns:
tensor: PSD matrix closest to `a` or the inverse of `a`.
"""
a = B.dense(a)
a = (a + B.transpose(a)) / 2
u, s, v = B.svd(a)
signs = B.matmul(u, v, tr_a=True)
s = B.maximum(B.diag(signs) * s, 0)
if inv:
s = B.where(s == 0, 0, 1 / s)
return B.mm(u * B.expand_dims(s, axis=-2), v, tr_b=True)
def cholesky_solve(a: Union[LowerTriangular, UpperTriangular],
b: AbstractMatrix):
return B.solve(B.transpose(a), B.solve(a, b))
def transpose(a: Kronecker):
return Kronecker(B.transpose(a.left), B.transpose(a.right))
def transpose(a: Woodbury):
return Woodbury(B.transpose(a.diag), B.transpose(a.lr))
def transpose(a: LowRank):
return LowRank(a.right, a.left, B.transpose(a.middle))
def matmul(a: AbstractMatrix, b: Kronecker, tr_a=False, tr_b=False):
return B.transpose(B.matmul(b, a, tr_a=not tr_b, tr_b=not tr_a))
def transform(x):
tril = B.vec_to_tril(x, offset=-1)
skew = tril - B.transpose(tril)
eye = B.eye(skew)
return B.solve(eye + skew, eye - skew)
def compute_I_hz(model, t):
"""Compute the :math:`I_{hz,t_i}` matrix for :math:`t_i` in `t`.
Args:
model (:class:`.gprv.GPRV`): Model.
t (vector): Time points of data.
Returns:
tensor: Value of :math:`I_{hz,t_i}`, with shape
`(len(t), len(model.ms), len(model.ms))`.
"""
# Compute sorting permutation.
perm = np.argsort(model.ms)
inverse_perm = invert_perm(perm)
# Sort to allow for simple concatenation.
m_max = model.m_max
ms = model.ms[perm]
# Construct I_hz for m,n <= M.
ns = ms[ms <= m_max]
I_0_cos_1 = _I_hx_0_cos(
model, -ns[None, :, None] + ns[None, None, :], t[:, None, None]
)
I_0_cos_2 = _I_hx_0_cos(
model, ns[None, :, None] + ns[None, None, :], t[:, None, None]
)
I_hz_mnleM = 0.5 * (I_0_cos_1 + I_0_cos_2)
# Construct I_hz for m,n > M.
ns = ms[ms > m_max] - m_max
I_0_cos_1 = _I_hx_0_cos(
model, -ns[None, :, None] + ns[None, None, :], t[:, None, None]
)
I_0_cos_2 = _I_hx_0_cos(
model, ns[None, :, None] + ns[None, None, :], t[:, None, None]
)
I_hz_mngtM = 0.5 * (I_0_cos_1 - I_0_cos_2)
# Construct I_hz for 0 < m <= M and n > M.
ns = ms[(0 < ms) * (ms <= m_max)]
ns2 = ms[ms > m_max] # Do not subtract M!
I_0_sin_1 = _I_hx_0_sin(
model, ns[None, :, None] + ns2[None, None, :], t[:, None, None]
)
I_0_sin_2 = _I_hx_0_sin(
model, -ns[None, :, None] + ns2[None, None, :], t[:, None, None]
)
I_hz_mleM_ngtM = 0.5 * (I_0_sin_1 + I_0_sin_2)
# Construct I_hz for m = 0 and n > M.
ns = ms[ms == 0]
ns2 = ms[ms > m_max] # Do not subtract M!
I_hz_0_gtM = _I_hx_0_sin(
model, ns[None, :, None] + ns2[None, None, :], t[:, None, None]
)
# Concatenate to form I_hz for m <= M and n > M.
I_hz_mleM_ngtM = B.concat(I_hz_0_gtM, I_hz_mleM_ngtM, axis=1)
# Compute the other half by transposing.
I_hz_mgtM_nleM = B.transpose(I_hz_mleM_ngtM, perm=(0, 2, 1))
# Construct result.
result = B.concat(
B.concat(I_hz_mnleM, I_hz_mleM_ngtM, axis=2),
B.concat(I_hz_mgtM_nleM, I_hz_mngtM, axis=2),
axis=1,
)
# Undo sorting.
result = B.take(result, inverse_perm, axis=1)
result = B.take(result, inverse_perm, axis=2)
return result
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 _tr(a, do):
return B.transpose(a) if do else a
def transpose(a: Dense):
return Dense(B.transpose(a.mat))
def _reshape_cols(a, *indices):
return B.transpose(B.reshape(B.transpose(a), *reversed(indices)))
def transpose(a: UpperTriangular):
return LowerTriangular(B.transpose(a.mat))
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 T(self):
return B.transpose(self)
def test_properties(dense1):
approx(dense1.T, B.transpose(dense1))
assert dense1.shape == B.shape(dense1)
assert dense1.dtype == B.dtype(dense1)
def transform(x):
tril = B.vec_to_tril(x, offset=-1)
skew = tril - B.transpose(tril)
return B.expm(skew)
