def test_compatible():
assert compatible(B.ones(5, 5), B.ones(5, 5))
assert compatible(B.ones(5, 5), B.ones(5, 1))
assert not compatible(B.ones(5, 5), B.ones(5))
assert not compatible(B.ones(5, 5), B.ones(5, 4))
with pytest.raises(AssertionError):
assert_compatible(B.ones(5), B.ones(4))
def test_lowertriangular_formatting():
assert (str(LowerTriangular(B.ones(
3, 3))) == "<lower-triangular matrix: shape=3x3, dtype=float64>")
assert (repr(LowerTriangular(B.ones(
3, 3))) == "<lower-triangular matrix: shape=3x3, dtype=float64\n"
" mat=[[1. 1. 1.]\n"
" [1. 1. 1.]\n"
" [1. 1. 1.]]>")
def test_normal_logpdf(normal1):
normal1_sp = multivariate_normal(normal1.mean[:, 0], B.dense(normal1.var))
x = B.randn(3, 10)
approx(normal1.logpdf(x), normal1_sp.logpdf(x.T))
# Test the the output of `logpdf` is flattened appropriately.
assert B.shape(normal1.logpdf(B.ones(3, 1))) == ()
assert B.shape(normal1.logpdf(B.ones(3, 2))) == (2, )
def test_dense_formatting():
assert str(Dense(B.ones(3, 3))) == "<dense matrix: shape=3x3, dtype=float64>"
assert (
repr(Dense(B.ones(3, 3))) == "<dense matrix: shape=3x3, dtype=float64\n"
" mat=[[1. 1. 1.]\n"
" [1. 1. 1.]\n"
" [1. 1. 1.]]>"
)
def sum(a: Constant, axis=None):
if axis is None:
return a.const * a.rows * a.cols
elif axis == 0:
return a.const * a.rows * B.ones(B.dtype(a.const), a.cols)
elif axis == 1:
return a.const * a.cols * B.ones(B.dtype(a.const), a.rows)
else:
_raise(axis)
def test_downrank(rank, preserve, shape, expected_shape, check_lazy_shapes):
kw_args = {}
if rank is not None:
kw_args["rank"] = rank
if preserve is not None:
kw_args["preserve"] = preserve
approx(
B.downrank(B.ones(*shape), **kw_args),
B.ones(*expected_shape),
)
def test_normal_lazy_nonzero_mean():
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3))
# Nothing should be populated yet.
assert dist._mean is None
assert dist._var is None
# But they should be populated upon request.
approx(dist.mean, B.ones(3, 1))
assert dist._var is None
approx(dist.var, B.eye(3))
def constant_to_lowrank(a):
dtype = B.dtype(a)
rows, cols = B.shape(a)
middle = B.fill_diag(a.const, 1)
if rows == cols:
return LowRank(B.ones(dtype, rows, 1), middle=middle)
else:
return LowRank(B.ones(dtype, rows, 1),
B.ones(dtype, cols, 1),
middle=middle)
def test_normal_lazy_nonzero_mean():
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3))
assert not dist.mean_is_zero
approx(dist._mean, B.ones(3, 1))
assert dist._var is None
approx(dist.mean, B.ones(3, 1))
assert dist._var is None
approx(dist.var, B.eye(3))
def test_kronecker_formatting():
left = Diagonal(B.ones(2))
right = Diagonal(B.ones(3))
assert (str(Kronecker(
left, right)) == "<Kronecker product: shape=6x6, dtype=float64>")
assert (repr(Kronecker(
left, right)) == "<Kronecker product: shape=6x6, dtype=float64\n"
" left=<diagonal matrix: shape=2x2, dtype=float64\n"
" diag=[1. 1.]>\n"
" right=<diagonal matrix: shape=3x3, dtype=float64\n"
" diag=[1. 1. 1.]>>")
def test_noise_as_matrix():
def check(noise, dtype, n, asserted_type):
noise = _noise_as_matrix(noise, dtype, n)
assert isinstance(noise, asserted_type)
assert B.dtype(noise) == dtype
assert B.shape(noise) == (n, n)
check(None, int, 5, matrix.Zero)
check(None, float, 5, matrix.Zero)
check(1, np.int64, 5, matrix.Diagonal)
check(1.0, np.float64, 5, matrix.Diagonal)
check(B.ones(int, 5), np.int64, 5, matrix.Diagonal)
check(B.ones(float, 5), np.float64, 5, matrix.Diagonal)
check(matrix.Dense(B.ones(int, 5, 5)), np.int64, 5, matrix.Dense)
check(matrix.Dense(B.randn(float, 5, 5)), np.float64, 5, matrix.Dense)
def test_woodbury_formatting():
diag = Diagonal(B.ones(3))
lr = LowRank(B.ones(3, 1), 2 * B.ones(3, 1))
assert str(Woodbury(diag,
lr)) == "<Woodbury matrix: shape=3x3, dtype=float64>"
assert (repr(Woodbury(
diag, lr)) == "<Woodbury matrix: shape=3x3, dtype=float64\n"
" diag=<diagonal matrix: shape=3x3, dtype=float64\n"
" diag=[1. 1. 1.]>\n"
" lr=<low-rank matrix: shape=3x3, dtype=float64, rank=1\n"
" left=[[1.]\n"
" [1.]\n"
" [1.]]\n"
" right=[[2.]\n"
" [2.]\n"
" [2.]]>>")
def test_logpdf_missing_data():
# Setup model.
m = 3
noise = 1e-2
latent_noises = 2e-2 * B.ones(m)
kernels = [0.5 * EQ().stretch(0.75) for _ in range(m)]
x = B.linspace(0, 10, 20)
# Concatenate two orthogonal matrices, to make the missing data
# approximation exact.
u1 = B.svd(B.randn(m, m))[0]
u2 = B.svd(B.randn(m, m))[0]
u = Dense(B.concat(u1, u2, axis=0) / B.sqrt(2))
s_sqrt = Diagonal(B.rand(m))
# Construct a reference model.
oilmm_pp = ILMMPP(kernels, u @ s_sqrt, noise, latent_noises)
# Sample to generate test data.
y = oilmm_pp.sample(x, latent=False)
# Throw away data, but retain orthogonality.
y[5:10, 3:] = np.nan
y[10:, :3] = np.nan
# Construct OILMM to test.
oilmm = OILMM(kernels, u, s_sqrt, noise, latent_noises)
# Check that evidence is still exact.
approx(oilmm_pp.logpdf(x, y), oilmm.logpdf(x, y), atol=1e-7)
def fill_diag(a: B.Numeric, diag_len: B.Int):
"""Fill the diagonal of a diagonal matrix with a particular scalar.
Args:
a (scalar): Scalar to fill diagonal with.
diag_len (int): Length of diagonal.
"""
return Diagonal(a * B.ones(B.dtype(a), diag_len))
def test_noise_as_matrix():
def check(noise, dtype, n, asserted_type):
noise = _noise_as_matrix(noise, dtype, n)
assert isinstance(noise, asserted_type)
assert B.dtype(noise) == dtype
assert B.shape(noise) == (n, n)
# Check that the type for `n` is appropriate.
for d in [5, Dimension(5)]:
check(None, int, d, matrix.Zero)
check(None, float, d, matrix.Zero)
check(1, np.int64, d, matrix.Diagonal)
check(1.0, np.float64, d, matrix.Diagonal)
check(B.ones(int, 5), np.int64, d, matrix.Diagonal)
check(B.ones(float, 5), np.float64, d, matrix.Diagonal)
check(matrix.Dense(B.ones(int, 5, 5)), np.int64, d, matrix.Dense)
check(matrix.Dense(B.randn(float, 5, 5)), np.float64, d, matrix.Dense)
def f(x, option=None):
# Check that keyword arguments are passed on correctly. Only the compilation
# with PyTorch does not pass on its keyword arguments.
if not isinstance(x, B.TorchNumeric):
assert option is not None
# This requires lazy shapes:
x = B.ones(B.dtype(x), *B.shape(x))
# This requires control flow cache:
return B.cond(x[0, 0] > 0, lambda: x[1], lambda: x[2])
def matmul(a: AbstractMatrix, b: Constant, tr_a=False, tr_b=False):
_assert_composable(a, b, tr_a=tr_a, tr_b=tr_b)
a = _tr(a, tr_a)
b = _tr(b, tr_b)
return LowRank(
B.expand_dims(B.sum(a, axis=1), axis=1),
B.ones(B.dtype(b), b.cols, 1),
B.fill_diag(b.const, 1),
)
def test_normal_lazy_var_diag():
# If `var_diag` isn't set, the variance will be constructed to get the diagonal.
dist = Normal(lambda: B.eye(3))
approx(dist.var_diag, B.ones(3))
approx(dist._var, B.eye(3))
# If `var_diag` is set, the variance will _not_ be constructed to get the diagonal.
dist = Normal(lambda: B.eye(3), var_diag=lambda: 9)
approx(dist.var_diag, 9)
assert dist._var is None
def _integrate_half2(self, name1, name2, **var_map):
if self.poly.highest_power(name1) != 2 or self.poly.highest_power(
name2) != 2:
raise RuntimeError(f'Dependency on "{name1}" and {name2}" must '
f"be quadratic.")
a11 = self.poly.collect_for(Factor(name1, 2))
a22 = self.poly.collect_for(Factor(name2, 2))
a12 = self.poly.collect_for(Factor(name1,
1)).collect_for(Factor(name2, 1))
b1 = self.poly.collect_for(Factor(name1, 1)).reject(name2)
b2 = self.poly.collect_for(Factor(name2, 1)).reject(name1)
c = self.poly.reject(name1).reject(name2)
# Evaluate and scale A.
a11 = -2 * a11.eval(**var_map)
a22 = -2 * a22.eval(**var_map)
a12 = -1 * a12.eval(**var_map)
b1 = b1.eval(**var_map)
b2 = b2.eval(**var_map)
c = c.eval(**var_map)
# Determinant of A:
a_det = a11 * a22 - a12**2
# Inverse of A, which corresponds to variance of distribution after
# completing the square:
ia11 = a22 / a_det
ia12 = -a12 / a_det
ia22 = a11 / a_det
# Mean of distribution after completing the square:
mu1 = ia11 * b1 + ia12 * b2
mu2 = ia12 * b1 + ia22 * b2
# Normalise and compute CDF part.
x1 = -mu1 / safe_sqrt(ia11)
x2 = -mu2 / safe_sqrt(ia22)
rho = ia12 / safe_sqrt(ia11 * ia22)
# Evaluate CDF for all `x1` and `x2`.
orig_shape = B.shape(mu1)
num = reduce(operator.mul, orig_shape, 1)
x1 = B.reshape(x1, num)
x2 = B.reshape(x2, num)
rho = rho * B.ones(x1)
cdf_part = B.reshape(B.bvn_cdf(x1, x2, rho), *orig_shape)
# Compute exponentiated part.
quad_form = 0.5 * (ia11 * b1**2 + ia22 * b2**2 + 2 * ia12 * b1 * b2)
det_part = 2 * B.pi / safe_sqrt(a_det)
exp_part = det_part * B.exp(quad_form + c)
return self.const * cdf_part * exp_part
def sample(model, t, noise_f):
"""Sample from a model.
Args:
model (:class:`gpcm.model.AbstractGPCM`): Model to sample from.
t (vector): Time points to sample at.
noise_f (vector): Noise for the sample of the function. Should have the
same size as `t`.
Returns:
tuple[vector, ...]: Tuple containing kernel samples, filter samples, and
function samples.
"""
ks, us, fs = [], [], []
# In the below, we look at the third inducing point, because that is the one
# determining the value of the filter at zero: the CGPCM adds two extra inducing
# points to the left.
# Get a smooth sample.
u1 = B.ones(model.n_u)
while B.abs(u1[2]) > 1e-2:
u1 = B.sample(model.compute_K_u())[:, 0]
u = GP(model.k_h())
u = u | (u(model.t_u), u1)
u1_full = u(t).mean.flatten()
# Get a rough sample.
u2 = B.zeros(model.n_u)
while u2[2] < 0.5:
u2 = B.sample(model.compute_K_u())[:, 0]
u = GP(model.k_h())
u = u | (u(model.t_u), u2)
u2_full = u(t).mean.flatten()
with wbml.out.Progress(name="Sampling", total=5) as progress:
for c in [0, 0.1, 0.23, 0.33, 0.5]:
# Sample kernel.
K = model.kernel_approx(t, t, c * u2 + (1 - c) * u1)
wbml.out.kv("Sampled variance", K[0, 0])
K = K / K[0, 0]
ks.append(K[0, :])
# Store filter.
us.append(c * u2_full + (1 - c) * u1_full)
# Sample function.
f = B.matmul(B.chol(closest_psd(K)), noise_f)
fs.append(f)
progress()
return ks, us, fs
def test_dtype_kron_promotion():
kron = Kronecker(B.ones(int, 5, 5), B.ones(int, 5, 5))
assert B.dtype(kron) == np.int64
kron = Kronecker(B.ones(float, 5, 5), B.ones(int, 5, 5))
assert B.dtype(kron) == np.float64
kron = Kronecker(B.ones(int, 5, 5), B.ones(float, 5, 5))
assert B.dtype(kron) == np.float64
def test_dtype_wb_promotion():
wb = LowRank(B.ones(int, 5, 5)) + Diagonal(B.ones(int, 5))
assert B.dtype(wb) == np.int64
wb = LowRank(B.ones(float, 5, 5)) + Diagonal(B.ones(int, 5))
assert B.dtype(wb) == np.float64
wb = LowRank(B.ones(int, 5, 5)) + Diagonal(B.ones(float, 5))
assert B.dtype(wb) == np.float64
def test_broadcast():
s = Shape(Dimension(5), Dimension(4))
assert broadcast(B.ones(5, 4), B.ones(5, 4)) == s
assert broadcast(B.ones(5, 4), B.ones(5, 1)) == s
with pytest.raises(AssertionError):
broadcast(B.ones(5, 5), B.ones(5, 4))
with pytest.raises(RuntimeError):
broadcast(Dimension(5), Dimension(4))
def test_normal_lazy_mean_var():
# The lazy `mean_var` should only be called when neither the mean nor the variance
# exists. Otherwise, it's more efficient to just construct the other one. We
# go over all branches in the `if`-statement.
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.mean_var, (8, 9))
approx(dist.mean, 8)
approx(dist.var, 9)
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.mean, B.ones(3, 1))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
approx(dist.var, B.eye(3))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.var, B.eye(3))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
approx(dist.mean, B.ones(3, 1))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.var, B.eye(3))
approx(dist.mean, B.ones(3, 1))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
def test_normalise():
layer = Normalise(epsilon=0)
x = B.randn(10, 5, 3)
# Check number of weights and width.
assert layer.num_weights(10) == 0
assert layer.width == 10
# Check initialisation and width.
layer.initialise(3, None)
assert layer.width == 3
# Check correctness
out = layer(x)
approx(B.std(out, axis=2), B.ones(10, 5), rtol=1e-4)
approx(B.mean(out, axis=2), B.zeros(10, 5), atol=1e-4)
def _project_pattern(self, x, y, pattern):
# Check whether all data is available.
no_missing = all(pattern)
if no_missing:
# All data is available. Nothing to be done.
u = self.u
else:
# Data is missing. Pick the available entries.
y = B.take(y, pattern, axis=1)
# Ensure that `u` remains a structured matrix.
u = Dense(B.take(self.u, pattern))
# Get number of data points and outputs in this part of the data.
n = B.shape(x)[0]
p = sum(pattern)
# Perform projection.
proj_y_partial = B.matmul(y, B.pinv(u), tr_b=True)
proj_y = B.matmul(proj_y_partial, B.inv(self.s_sqrt), tr_b=True)
# Compute projected noise.
u_square = B.matmul(u, u, tr_a=True)
proj_noise = (
self.noise_obs / B.diag(self.s_sqrt) ** 2 * B.diag(B.pd_inv(u_square))
)
# Convert projected noise to weights.
noises = self.model.noises
weights = noises / (noises + proj_noise)
proj_w = B.ones(B.dtype(weights), n, self.m) * weights[None, :]
# Compute Frobenius norm.
frob = B.sum(y ** 2)
frob = frob - B.sum(proj_y_partial * B.matmul(proj_y_partial, u_square))
# Compute regularising term.
reg = 0.5 * (
n * (p - self.m) * B.log(2 * B.pi * self.noise_obs)
+ frob / self.noise_obs
+ n * B.logdet(B.matmul(u, u, tr_a=True))
+ n * 2 * B.logdet(self.s_sqrt)
)
return x, proj_y, proj_w, reg
def test_expand_dims(check_lazy_shapes):
check_function(B.expand_dims, (Tensor(3, 4, 5),), {"axis": Value(0, 1)})
# Test keyword `times`.
assert B.shape(B.expand_dims(B.ones(2), axis=-1, times=1)) == (2, 1)
assert B.shape(B.expand_dims(B.ones(2), axis=-1, times=2)) == (2, 1, 1)
assert B.shape(B.expand_dims(B.ones(2), axis=-1, times=3)) == (2, 1, 1, 1)
assert B.shape(B.expand_dims(B.ones(2), axis=0, times=1)) == (1, 2)
assert B.shape(B.expand_dims(B.ones(2), axis=0, times=2)) == (1, 1, 2)
assert B.shape(B.expand_dims(B.ones(2), axis=0, times=3)) == (1, 1, 1, 2)
# Test keyword `ignore_scalar`.
assert B.expand_dims(1, axis=-1, ignore_scalar=False) is not 1
assert B.expand_dims(1, axis=-1, ignore_scalar=True) is 1
def test_lowrank_formatting():
assert (str(LowRank(B.ones(3, 1), 2 * B.ones(
3, 1))) == "<low-rank matrix: shape=3x3, dtype=float64, rank=1>")
assert (repr(LowRank(B.ones(3, 2), 2 * B.ones(
3, 2))) == "<low-rank matrix: shape=3x3, dtype=float64, rank=2\n"
" left=[[1. 1.]\n"
" [1. 1.]\n"
" [1. 1.]]\n"
" right=[[2. 2.]\n"
" [2. 2.]\n"
" [2. 2.]]>")
assert (repr(LowRank(B.ones(
3, 2))) == "<low-rank matrix: shape=3x3, dtype=float64, rank=2\n"
" left=[[1. 1.]\n"
" [1. 1.]\n"
" [1. 1.]]>")
assert (repr(LowRank(B.ones(3, 2), middle=B.ones(
2, 2))) == "<low-rank matrix: shape=3x3, dtype=float64, rank=2\n"
" left=[[1. 1.]\n"
" [1. 1.]\n"
" [1. 1.]]\n"
" middle=[[1. 1.]\n"
" [1. 1.]]>")
def test_lowrank_shape_checks():
# Check that matrices must be given.
with pytest.raises(AssertionError):
LowRank(B.ones(3))
with pytest.raises(AssertionError):
LowRank(B.ones(3, 1), B.ones(3))
with pytest.raises(AssertionError):
LowRank(B.ones(3, 1), B.ones(3, 1), B.ones(1))
# Check that needs need to be compatible
with pytest.raises(AssertionError):
LowRank(B.ones(3, 1), B.ones(3, 2))
with pytest.raises(AssertionError):
LowRank(B.ones(3, 1), B.ones(3, 2), B.ones(1, 1))
with pytest.raises(AssertionError):
LowRank(B.ones(3, 1), middle=B.ones(2, 1))
def test_lowrank_attributes():
# Check default right and middle factor.
left = B.ones(3, 2)
lr = LowRank(left)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.eye(2))
assert lr.rank == 2
# Check given identical right.
lr = LowRank(left, left)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.eye(2))
assert lr.rank == 2
# Check given identical right and middle.
middle = B.ones(2, 2)
lr = LowRank(left, left, middle=middle)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.ones(2, 2))
assert lr.rank == 2
# Check given other right and middle factor.
right = 2 * B.ones(3, 2)
lr = LowRank(left, right, middle=middle)
assert lr.left is left
assert lr.right is right
approx(lr.middle, B.ones(2, 2))
assert lr.rank == 2
# Check rank in non-square case.
assert LowRank(B.ones(3, 2), B.ones(3, 2), B.ones(2, 2)).rank == 2
assert LowRank(B.ones(3, 2), B.ones(3, 1), B.ones(2, 1)).rank == 1
