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 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_take_tf(check_lazy_shapes):
# Check that TensorFlow also takes in tensors.
a = Matrix(3, 4, 5)
ref = Tensor(3)
approx(B.take(a.tf(), ref.tf() > 0), B.take(a.np(), ref.np() > 0))
approx(B.take(a.tf(), ref.np() > 0), B.take(a.np(), ref.np() > 0))
approx(B.take(a.tf(), B.range(tf.int64, 2)), B.take(a.np(), B.range(2)))
approx(B.take(a.tf(), B.range(np.int64, 2)), B.take(a.np(), B.range(2)))
def test_take_perm(dtype, check_lazy_shapes):
a = B.range(dtype, 10)
perm = B.randperm(B.dtype_int(dtype), 10)
a2 = B.take(a, perm)
assert B.dtype(perm) == B.dtype_int(dtype)
assert B.shape(a) == B.shape(a2)
assert B.dtype(a) == B.dtype(a2)
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 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_fdd_take():
with Measure():
f1 = GP(1, EQ())
f2 = GP(2, Exp())
f = cross(f1, f2)
x = B.linspace(0, 3, 5)
# Build an FDD with a very complicated input specification.
fdd = f((x, (f2(x), x), f1(x), (f2(x), (f1(x), x))))
n = infer_size(fdd.p.kernel, fdd.x)
fdd = f(fdd.x, matrix.Diagonal(B.rand(n)))
# Flip a coin for every element.
mask = B.randn(n) > 0
taken_fdd = B.take(fdd, mask)
approx(taken_fdd.mean, B.take(fdd.mean, mask))
approx(taken_fdd.var, B.submatrix(fdd.var, mask))
approx(taken_fdd.noise, B.submatrix(fdd.noise, mask))
assert isinstance(taken_fdd.noise, matrix.Diagonal)
# Test that only masks are supported, for now.
with pytest.raises(AssertionError):
B.take(fdd, np.array([1, 2]))
def test_take_consistency(check_lazy_shapes):
# Check consistency between indices and mask.
check_function(
B.take,
(Matrix(3, 3), Value([0, 1], [True, True, False])),
{"axis": Value(0, 1, -1)},
)
# Test PyTorch separately, because it has a separate implementation for framework
# masks or indices.
for indices_or_mask in [
torch.tensor([True, True, False], dtype=torch.bool),
torch.tensor([0, 1], dtype=torch.int32),
torch.tensor([0, 1], dtype=torch.int64),
]:
a = B.randn(torch.float32, 3, 3)
approx(B.take(a, indices_or_mask), a[[0, 1]])
def add(a: LowRank, b: LowRank):
assert_compatible(a, b)
l_a_p, l_b_p, l_a_jp, l_b_jp = align(a.left, b.left)
r_a_p, r_b_p, r_a_jp, r_b_jp = align(a.right, b.right)
# Join left parts.
a_left = B.take(_pad_zero_col(a.left), l_a_jp, axis=-1)
b_left = B.take(_pad_zero_col(b.left), l_b_jp, axis=-1)
join_left = a_left + b_left
# Join right parts.
a_right = B.take(_pad_zero_col(a.right), r_a_jp, axis=-1)
b_right = B.take(_pad_zero_col(b.right), r_b_jp, axis=-1)
join_right = a_right + b_right
# Join middle parts.
a_mid = B.take(B.take(_pad_zero_both(a.middle), l_a_p, axis=-2), r_a_p, axis=-1)
b_mid = B.take(B.take(_pad_zero_both(b.middle), l_b_p, axis=-2), r_b_p, axis=-1)
join_middle = a_mid + b_mid
return LowRank(join_left, join_right, join_middle)
def _check_take(a):
approx(B.take(a, [0, 1]), B.take(B.dense(a), [0, 1]))
def take(a: AbstractMatrix, indices_or_mask, axis=0):
if structured(a):
warn_upmodule(f"Taking from {a}: converting to dense.", category=ToDenseWarning)
return B.take(B.dense(a), indices_or_mask, axis=axis)
def test_take_indices_rank(check_lazy_shapes):
# Check that indices must be rank 1.
for a in Matrix(3, 4).forms():
with pytest.raises(ValueError):
B.take(a, [[0], [1]])
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