def _extract_samples(quantities):
x = quantities.x
samples = quantities.all_samples
return (
B.concat(-x[::-1][:-1], x),
B.concat(samples[::-1, :][:-1, :], samples, axis=0),
)
def project(self, x, y):
"""Project data.
Args:
x (matrix): Locations of data.
y (matrix): Observations of data.
Returns:
tuple: Tuple containing the locations of the projection,
the projection, weights associated with the projection, and
a regularisation term.
"""
n = B.shape(x)[0]
available = ~B.isnan(B.to_numpy(y))
# Optimise the case where all data is available.
if B.all(available):
return self._project_pattern(x, y, (True,) * self.p)
# Extract patterns.
patterns = list(set(map(tuple, list(available))))
if len(patterns) > 30:
warnings.warn(
f"Detected {len(patterns)} patterns, which is more "
f"than 30 and can be slow.",
category=UserWarning,
)
# Per pattern, find data points that belong to it.
patterns_inds = [[] for _ in range(len(patterns))]
for i in range(n):
patterns_inds[patterns.index(tuple(available[i]))].append(i)
# Per pattern, perform the projection.
proj_xs = []
proj_ys = []
proj_ws = []
total_reg = 0
for pattern, pattern_inds in zip(patterns, patterns_inds):
proj_x, proj_y, proj_w, reg = self._project_pattern(
B.take(x, pattern_inds), B.take(y, pattern_inds), pattern
)
proj_xs.append(proj_x)
proj_ys.append(proj_y)
proj_ws.append(proj_w)
total_reg = total_reg + reg
return (
B.concat(*proj_xs, axis=0),
B.concat(*proj_ys, axis=0),
B.concat(*proj_ws, axis=0),
total_reg,
)
def __call__(self, prev, x):
prev_h, _ = prev
# Gate logic:
z = self.f_z(B.concat(prev_h, x, axis=1))
r = self.f_r(B.concat(prev_h, x, axis=1))
h = (1 - z) * prev_h + z * self.f_h(B.concat(r * prev_h, x, axis=1))
y = h
return h, y
def test_concat(dense1, dense2, diag1, diag2):
with AssertDenseWarning("concatenating <dense>, <dense>, <diagonal>..."):
res = B.concat(dense1, dense2, diag1, diag2, axis=1)
dense_res = B.concat(B.dense(dense1),
B.dense(dense2),
B.dense(diag1),
B.dense(diag2),
axis=1)
approx(res, dense_res)
assert isinstance(res, Dense)
def test_combine():
x1 = B.linspace(0, 2, 10)
x2 = B.linspace(2, 4, 10)
m = Measure()
p1 = GP(EQ(), measure=m)
p2 = GP(Matern12(), measure=m)
y1 = p1(x1).sample()
y2 = p2(x2).sample()
# Check the one-argument case.
assert_equal_normals(combine(p1(x1, 1)), p1(x1, 1))
fdd_combined, y_combined = combine((p1(x1, 1), B.squeeze(y1)))
assert_equal_normals(fdd_combined, p1(x1, 1))
approx(y_combined, y1)
# Check the two-argument case.
fdd_combined = combine(p1(x1, 1), p2(x2, 2))
assert_equal_normals(
fdd_combined,
Normal(B.block_diag(p1(x1, 1).var,
p2(x2, 2).var)),
)
fdd_combined, y_combined = combine((p1(x1, 1), B.squeeze(y1)),
(p2(x2, 2), y2))
assert_equal_normals(
fdd_combined,
Normal(B.block_diag(p1(x1, 1).var,
p2(x2, 2).var)),
)
approx(y_combined, B.concat(y1, y2, axis=0))
def test_ess():
# Construct a prior and a likelihood.
prior = Normal(np.array([[0.6, 0.3], [0.3, 0.6]]))
lik = Normal(
np.array([[0.2], [0.3]]),
np.array([[1, 0.2], [0.2, 1]]),
)
# Perform sampling.
sampler = ESS(lik.logpdf, prior.sample)
num_samples = 30_000
samples = B.concat(*sampler.sample(num=num_samples), axis=1)
samples_mean = B.mean(samples, axis=1)[:, None]
samples_cov = (
B.matmul(samples - samples_mean, samples - samples_mean, tr_b=True) /
num_samples)
# Compute posterior statistics.
prec_prior = B.inv(prior.var)
prec_lik = B.inv(lik.var)
cov = B.inv(prec_prior + prec_lik)
mean = cov @ (prec_prior @ prior.mean + prec_lik @ lik.mean)
approx(samples_cov, cov, atol=5e-2)
approx(samples_mean, mean, atol=5e-2)
def test_momean(x):
prior = Measure()
p1 = GP(lambda x: 2 * x, 1 * EQ(), measure=prior)
p2 = GP(1, 2 * EQ().stretch(2), measure=prior)
m = MultiOutputMean(prior, p1, p2)
ms = prior.means
# Check representation.
assert str(m) == "MultiOutputMean(<lambda>, 1)"
# Check computation.
approx(m(x), B.concat(ms[p1](x), ms[p2](x), axis=0))
approx(m(p1(x)), ms[p1](x))
approx(m(p2(x)), ms[p2](x))
approx(m(MultiInput(p2(x), p1(x))), B.concat(ms[p2](x), ms[p1](x), axis=0))
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 _attempt_diagonal(rows):
# Check whether the result is diagonal.
# Check that the blocks form a square.
if not all([len(row) == len(rows) for row in rows]):
return None
# Collect the diagonal blocks.
diagonal_blocks = []
for r in range(len(rows)):
for c in range(len(rows[0])):
block_shape = B.shape(rows[r][c])
if r == c:
# Keep track of all the diagonal blocks.
diagonal_blocks.append(rows[r][c])
# All blocks on the diagonal must be diagonal or zero.
if not isinstance(rows[r][c], (Diagonal, Zero)):
return None
# All blocks on the diagonal must be square.
if not block_shape[0] == block_shape[1]:
return None
else:
# All blocks not on the diagonal must be zero.
if not isinstance(rows[r][c], Zero):
return None
return Diagonal(B.concat(*[B.diag(x) for x in diagonal_blocks]))
def __call__(self, prev, x):
prev_h, prev_y = prev
# Gate logic:
h = self.f_h(B.concat(prev_h, x, axis=1))
y = h
return h, y
def concat(*elements: AbstractMatrix, axis=0):
if structured(*elements):
elements_str = ", ".join(map(str, elements[:3]))
if len(elements) > 3:
elements_str += "..."
warn_upmodule(
f"Concatenating {elements_str}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.concat(*(B.dense(el) for el in elements), axis=axis))
def pack(*objs: B.Numeric):
"""Pack objects.
Args:
*objs (tensor): Objects to pack.
Returns:
tensor: Vector representation of the objects.
"""
return B.concat(*[B.flatten(obj) for obj in objs], axis=0)
def compute_I_uz(model, t):
"""Compute the :math:`I_{uz,t_i}` matrix for :math:`t_i` in `t`.
Args:
model (:class:`.gprv.GPRV`): Model.
t (vector): Time points :math:`t_i` of data.
Returns:
tensor: Value of :math:`I_{uz,t_i}`, with shape
`(len(t), len(model.t_u), len(model.ms))`.
"""
# Compute sorting permutation.
perm = np.argsort(model.ms)
inverse_perm = invert_perm(perm)
# Sort to allow for simple concatenation.
ms = model.ms[perm]
# Part of the factor is absorbed in the integrals.
factor = (
model.alpha_t * model.gamma_t * B.exp(-model.gamma * model.t_u[None, :, None])
)
# Compute I(l, u, 0).
k = ms[ms == 0]
I_uz0 = (
_integral_lu0(
model, t[:, None, None] - model.t_u[None, :, None], t[:, None, None]
)
+ k[None, None, :]
) # This is all zeros, but performs broadcasting.
# Compute I(l, u, k) for 0 < k < M + 1.
k = ms[(0 < ms) * (ms <= model.m_max)]
I_uz_leM = _integral_luk_leq_M(
model,
t[:, None, None] - model.t_u[None, :, None],
t[:, None, None],
k[None, None, :],
)
# Compute I(l, u, k) for k > M.
k = ms[ms > model.m_max]
I_uz_gtM = _integral_luk_g_M(
model,
t[:, None, None] - model.t_u[None, :, None],
t[:, None, None],
k[None, None, :],
)
# Construct result.
result = B.concat(I_uz0, I_uz_leM, I_uz_gtM, axis=2)
# Undo sorting and return.
return factor * B.take(result, inverse_perm, axis=2)
def test_mom():
x = B.linspace(0, 1, 10)
prior = Measure()
p1 = GP(lambda x: 2 * x, 1 * EQ(), measure=prior)
p2 = GP(1, 2 * EQ().stretch(2), measure=prior)
m = MultiOutputMean(prior, p1, p2)
ms = prior.means
# Check dimensionality.
assert dimensionality(m) == 2
# Check representation.
assert str(m) == "MultiOutputMean(<lambda>, 1)"
# Check computation.
approx(m(x), B.concat(ms[p1](x), ms[p2](x), axis=0))
approx(m(p1(x)), ms[p1](x))
approx(m(p2(x)), ms[p2](x))
approx(m((p2(x), p1(x))), B.concat(ms[p2](x), ms[p1](x), axis=0))
def sample(self, x, latent=False):
"""Sample from the model.
Args:
x (matrix): Locations to sample at.
latent (bool, optional): Sample noiseless processes. Defaults
to `False`.
Returns:
matrix: Sample.
"""
fdds = [f(x, *(() if latent else (n,))) for f, n in zip(self.fs, self.noises)]
return B.concat(*[fdd.sample() for fdd in fdds], axis=1)
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))
def test_mokernel(x1, x2, x3):
m = Measure()
p1 = GP(1 * EQ(), measure=m)
p2 = GP(2 * EQ().stretch(2), measure=m)
k = MultiOutputKernel(m, p1, p2)
ks = m.kernels
# Check representation.
assert str(k) == "MultiOutputKernel(EQ(), 2 * (EQ() > 2))"
# Input versus input:
approx(
k(x1, x2),
B.concat2d(
[ks[p1, p1](x1, x2), ks[p1, p2](x1, x2)],
[ks[p2, p1](x1, x2), ks[p2, p2](x1, x2)],
),
)
approx(
k.elwise(x1, x3),
B.concat(ks[p1, p1].elwise(x1, x3), ks[p2, p2].elwise(x1, x3), axis=0),
)
# Input versus `FDD`:
approx(k(p1(x1), x2),
B.concat(ks[p1, p1](x1, x2), ks[p1, p2](x1, x2), axis=1))
approx(k(p2(x1), x2),
B.concat(ks[p2, p1](x1, x2), ks[p2, p2](x1, x2), axis=1))
approx(k(x1, p1(x2)),
B.concat(ks[p1, p1](x1, x2), ks[p2, p1](x1, x2), axis=0))
approx(k(x1, p2(x2)),
B.concat(ks[p1, p2](x1, x2), ks[p2, p2](x1, x2), axis=0))
with pytest.raises(ValueError):
k.elwise(x1, p2(x3))
with pytest.raises(ValueError):
k.elwise(p1(x1), x3)
# `FDD` versus `FDD`:
approx(k(p1(x1), p1(x2)), ks[p1](x1, x2))
approx(k(p1(x1), p2(x2)), ks[p1, p2](x1, x2))
approx(k.elwise(p1(x1), p1(x3)), ks[p1].elwise(x1, x3))
approx(k.elwise(p1(x1), p2(x3)), ks[p1, p2].elwise(x1, x3))
# `MultiInput` versus input:
approx(
k(MultiInput(p2(x1), p1(x2)), x1),
B.concat2d(
[ks[p2, p1](x1, x1), ks[p2, p2](x1, x1)],
[ks[p1, p1](x2, x1), ks[p1, p2](x2, x1)],
),
)
approx(
k(x1, MultiInput(p2(x1), p1(x2))),
B.concat2d(
[ks[p1, p2](x1, x1), ks[p1, p1](x1, x2)],
[ks[p2, p2](x1, x1), ks[p2, p1](x1, x2)],
),
)
with pytest.raises(ValueError):
k.elwise(MultiInput(p2(x1), p1(x3)), p2(x1))
with pytest.raises(ValueError):
k.elwise(p2(x1), MultiInput(p2(x1), p1(x3)))
# `MultiInput` versus `FDD`:
approx(
k(MultiInput(p2(x1), p1(x2)), p2(x1)),
B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0),
)
approx(
k(p2(x1), MultiInput(p2(x1), p1(x2))),
B.concat(ks[p2, p2](x1, x1), ks[p2, p1](x1, x2), axis=1),
)
with pytest.raises(ValueError):
k.elwise(MultiInput(p2(x1), p1(x3)), p2(x1))
with pytest.raises(ValueError):
k.elwise(p2(x1), MultiInput(p2(x1), p1(x3)))
# `MultiInput` versus `MultiInput`:
approx(
k(MultiInput(p2(x1), p1(x2)), MultiInput(p2(x1))),
B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0),
)
with pytest.raises(ValueError):
k.elwise(MultiInput(p2(x1), p1(x3)), MultiInput(p2(x1)))
approx(
k.elwise(MultiInput(p2(x1), p1(x3)), MultiInput(p2(x1), p1(x3))),
B.concat(ks[p2, p2].elwise(x1, x1), ks[p1, p1].elwise(x3, x3), axis=0),
)
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
models = model()
# Perform sampling.
if args.train:
ks, us, fs = sample(model, t, noise_f)
wd.save((ks, us, fs), "samples.pickle")
else:
ks, us, fs = wd.load("samples.pickle")
# Plot.
plt.figure(figsize=(15, 4))
for i, (k, u, f) in enumerate(zip(ks, us, fs)):
plt.subplot(3, 5, 1 + i)
plt.plot(
B.concat(-t[::-1][:-1], t),
B.concat(u[:-1] * 0, u),
lw=1,
)
if hasattr(model, "t_u"):
plt.scatter(model.t_u, model.t_u * 0, s=5, marker="o", c="black")
# plt.xlabel("Time (s)")
if i == 0:
plt.ylabel("$h$")
plt.xlim(-6, 6)
plt.ylim(-0.5, 1)
tweak(legend=False)
plt.subplot(3, 5, 6 + i)
plt.plot(
B.concat(-t[::-1][:-1], t),
starting from the origin.
k (vector): Kernel.
n_zero (int, optional): Zero padding. Defaults to `2_000`.
db (bool, optional): Convert to decibel. Defaults to `False`.
Returns:
vector: PSD, correctly scaled.
"""
# Convert to NumPy for compatibility with frameworks.
t, k = B.to_numpy(t, k)
if t[0] != 0:
raise ValueError("Time points must start at zero.")
# Perform zero padding.
k = B.concat(k, B.zeros(n_zero))
# Symmetrise and Fourier transform.
k_symmetric = B.concat(k, k[1:-1][::-1])
psd = np.fft.fft(k_symmetric)
freqs = np.fft.fftfreq(len(psd)) / (t[1] - t[0])
# Should be real and positive, but the numerics may not be in our favour.
psd = np.abs(np.real(psd))
# Now scale appropriately: the total power should equal `k[0]`.
total_power = np.trapz(y=psd, x=freqs)
psd /= total_power / k[0]
# Convert to dB.
if db:
def diag(a: Diagonal, b: Diagonal):
return Diagonal(B.concat(a.diag, b.diag))
def test_mok():
x1 = B.linspace(0, 1, 10)
x2 = B.linspace(1, 2, 5)
x3 = B.linspace(1, 2, 10)
m = Measure()
p1 = GP(EQ(), measure=m)
p2 = GP(2 * EQ().stretch(2), measure=m)
k = MultiOutputKernel(m, p1, p2)
ks = m.kernels
# Check dimensionality.
assert dimensionality(k) == 2
# Check representation.
assert str(k) == "MultiOutputKernel(EQ(), 2 * (EQ() > 2))"
# Input versus input:
approx(
k(x1, x2),
B.concat2d(
[ks[p1, p1](x1, x2), ks[p1, p2](x1, x2)],
[ks[p2, p1](x1, x2), ks[p2, p2](x1, x2)],
),
)
approx(
k.elwise(x1, x3),
B.concat(ks[p1, p1].elwise(x1, x3), ks[p2, p2].elwise(x1, x3), axis=0),
)
# Input versus `FDD`:
approx(k(p1(x1), x2), B.concat(ks[p1, p1](x1, x2), ks[p1, p2](x1, x2), axis=1))
approx(k(p2(x1), x2), B.concat(ks[p2, p1](x1, x2), ks[p2, p2](x1, x2), axis=1))
approx(k(x1, p1(x2)), B.concat(ks[p1, p1](x1, x2), ks[p2, p1](x1, x2), axis=0))
approx(k(x1, p2(x2)), B.concat(ks[p1, p2](x1, x2), ks[p2, p2](x1, x2), axis=0))
with pytest.raises(ValueError):
k.elwise(x1, p2(x3))
with pytest.raises(ValueError):
k.elwise(p1(x1), x3)
# `FDD` versus `FDD`:
approx(k(p1(x1), p1(x2)), ks[p1](x1, x2))
approx(k(p1(x1), p2(x2)), ks[p1, p2](x1, x2))
approx(k.elwise(p1(x1), p1(x3)), ks[p1].elwise(x1, x3))
approx(k.elwise(p1(x1), p2(x3)), ks[p1, p2].elwise(x1, x3))
# Multiple inputs versus input:
approx(
k((p2(x1), p1(x2)), x1),
B.concat2d(
[ks[p2, p1](x1, x1), ks[p2, p2](x1, x1)],
[ks[p1, p1](x2, x1), ks[p1, p2](x2, x1)],
),
)
approx(
k(x1, (p2(x1), p1(x2))),
B.concat2d(
[ks[p1, p2](x1, x1), ks[p1, p1](x1, x2)],
[ks[p2, p2](x1, x1), ks[p2, p1](x1, x2)],
),
)
with pytest.raises(ValueError):
k.elwise((p2(x1), p1(x3)), p2(x1))
with pytest.raises(ValueError):
k.elwise(p2(x1), (p2(x1), p1(x3)))
# Multiple inputs versus `FDD`:
approx(
k((p2(x1), p1(x2)), p2(x1)),
B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0),
)
approx(
k(p2(x1), (p2(x1), p1(x2))),
B.concat(ks[p2, p2](x1, x1), ks[p2, p1](x1, x2), axis=1),
)
with pytest.raises(ValueError):
k.elwise((p2(x1), p1(x3)), p2(x1))
with pytest.raises(ValueError):
k.elwise(p2(x1), (p2(x1), p1(x3)))
# Multiple inputs versus multiple inputs:
approx(
k((p2(x1), p1(x2)), (p2(x1))),
B.concat(ks[p2, p2](x1, x1), ks[p1, p2](x2, x1), axis=0),
)
with pytest.raises(ValueError):
k.elwise((p2(x1), p1(x3)), (p2(x1),))
approx(
k.elwise((p2(x1), p1(x3)), (p2(x1), p1(x3))),
B.concat(ks[p2, p2].elwise(x1, x1), ks[p1, p1].elwise(x3, x3), axis=0),
)
def elwise(k, x: tuple, y: tuple):
if len(x) != len(y):
raise ValueError('"elwise" must be called with similarly sized tuples.')
return B.concat(*[elwise(k, xi, yi) for xi, yi in zip(x, y)], axis=-2)
def _pad_zero_row(a):
zeros = B.zeros(B.dtype(a), 1, B.shape(a)[1])
return B.concat(a, zeros, axis=0)
def _pad_zero_col(a):
zeros = B.zeros(B.dtype(a), B.shape(a)[0], 1)
return B.concat(a, zeros, axis=1)
