def test_conditioning():
m = Measure()
p1 = GP(EQ(), measure=m)
p2 = GP(Exp(), measure=m)
p_sum = p1 + p2
# Sample some data to condition on.
x1 = B.linspace(0, 2, 2)
y1 = p1(x1).sample()
x_sum = B.linspace(3, 5, 3)
y_sum = p_sum(x_sum).sample()
# Determine FDDs to check.
x_check = B.linspace(0, 5, 5)
fdds_check = [
cross(p1, p2, p_sum)(x_check),
p1(x_check),
p2(x_check),
p_sum(x_check),
]
assert_equal_measures(
fdds_check,
m.condition(p_sum(x_sum), y_sum),
m.condition((p_sum(x_sum), y_sum)),
m | (p_sum(x_sum), y_sum),
m | ((p_sum(x_sum), y_sum), ),
m | Obs(p_sum(x_sum), y_sum),
m | Obs((p_sum(x_sum), y_sum)),
)
assert_equal_measures(
fdds_check,
m.condition((p1(x1), y1), (p_sum(x_sum), y_sum)),
m | ((p1(x1), y1), (p_sum(x_sum), y_sum)),
m | Obs((p1(x1), y1), (p_sum(x_sum), y_sum)),
)
def test_sparse_conditioning_and_elbo():
m = Measure()
p1 = GP(EQ(), measure=m)
p2 = GP(Exp(), measure=m)
e = GP(Delta(), measure=m)
p_sum = p1 + p2
# Sample some data to condition on.
x1 = B.linspace(0, 2, 2)
y1 = (p1 + e)(x1).sample()
x_sum = B.linspace(3, 5, 3)
y_sum = (p_sum + e)(x_sum).sample()
# Determine FDDs to check.
x_check = B.linspace(0, 5, 5)
fdds_check = [
cross(p1, p2, p_sum)(x_check),
p1(x_check),
p2(x_check),
p_sum(x_check),
]
# Check conditioning and ELBO on one data set.
assert_equal_measures(
fdds_check,
m | ((p_sum + e)(x_sum), y_sum),
m | SparseObs(p_sum(x_sum), e, p_sum(x_sum), y_sum),
m | SparseObs((p_sum(x_sum), ), e, p_sum(x_sum), y_sum),
m | SparseObs((p_sum(x_sum), p1(x1)), e, p_sum(x_sum), y_sum),
m | SparseObs(p_sum(x_sum), (e, p_sum(x_sum), y_sum)),
m | SparseObs((p_sum(x_sum), ), (e, p_sum(x_sum), y_sum)),
m.condition(
SparseObs(
(p_sum(x_sum), p1(x1)),
(e, p_sum(x_sum), y_sum),
)),
)
approx(
m.logpdf(Obs((p_sum + e)(x_sum), y_sum)),
SparseObs(p_sum(x_sum), (e, p_sum(x_sum), y_sum)).elbo(m),
)
# Check conditioning and ELBO on two data sets.
assert_equal_measures(
fdds_check,
m | (((p_sum + e)(x_sum), y_sum), ((p1 + e)(x1), y1)),
m.condition(
SparseObs((p_sum(x_sum), p1(x1)), (e, p_sum(x_sum), y_sum),
(e, p1(x1), y1))),
)
approx(
m.logpdf(Obs(((p_sum + e)(x_sum), y_sum), ((p1 + e)(x1), y1))),
SparseObs((p_sum(x_sum), p1(x1)), (e, p_sum(x_sum), y_sum),
(e, p1(x1), y1)).elbo(m),
)
# The following lose information, so check them separately.
assert_equal_measures(
fdds_check,
m | SparseObs(p_sum(x_sum), (e, p_sum(x_sum), y_sum), (e, p1(x1), y1)),
m | SparseObs((p_sum(x_sum), ), (e, p_sum(x_sum), y_sum),
(e, p1(x1), y1)),
)
# Test lazy computation.
obs = SparseObs(p_sum(x_sum), e, p_sum(x_sum), y_sum)
for name in ["K_z", "elbo", "mu", "A"]:
approx(
getattr(SparseObs(p_sum(x_sum), e, p_sum(x_sum), y_sum), name)(m),
getattr(obs, name)(m),
)
# Test requirement that noise must be diagonal.
with pytest.raises(RuntimeError):
SparseObs(p_sum(x_sum), p_sum, p_sum(x_sum), y_sum).elbo(m)