def test_conditional_mean_fwd():
a, U, V, P, Y, U_star, V_star, inds = get_matrices(vector=True,
conditional=True)
d, W = driver.factor(U, P, a, np.copy(V))
z = driver.solve(U, P, d, W, Y)
check_basic(
driver.conditional_mean,
ops.conditional_mean,
[U, V, P, z, U_star, V_star, inds],
)
def test_solve(vector):
a, U, V, P, K, Y = get_matrices(vector=vector, include_dense=True)
# First compute the expected value
expect = np.linalg.solve(K, Y)
# Then solve using celerite
d, W = driver.factor(U, P, a, V)
value = driver.solve(U, P, d, W, Y)
# Make sure that no copy is made if possible
assert np.allclose(value, Y)
# Check that the solution is correct
assert np.allclose(value, expect)
def test_conditional_mean():
a, U, V, P, Y, U_star, V_star, inds = get_matrices(vector=True,
conditional=True)
d, W = driver.factor(U, P, a, np.copy(V))
z = driver.solve(U, P, d, W, Y)
mu = driver.conditional_mean(U, V, P, z, U_star, V_star, inds,
np.empty(len(inds), dtype=np.float64))
check_op(
ops.conditional_mean,
[U, V, P, z, U_star, V_star, inds],
[mu],
grad=False,
)
def test_solve_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
X0 = driver.solve(U, P, d, W, np.copy(Y))
X = np.empty_like(Y)
Z = np.empty_like(Y)
if vector:
F = np.empty_like(U)
else:
F = np.empty((U.shape[0], U.shape[1] * Y.shape[1]))
G = np.empty_like(F)
X, Z, F, G = backprop.solve_fwd(U, P, d, W, Y, X, Z, F, G)
assert np.allclose(X0, X)
def test_solve(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
X = driver.solve(U, P, d, W, np.copy(Y))
check_op(ops.solve, [U, P, d, W, Y], [X])
check_op(jit(ops.solve), [U, P, d, W, Y], [X])