def test_solve_upper_fwd(vector):
x, c, a, U, V, Y = get_matrices(vector=vector)
if vector:
Y = Y[:, None]
d, W = driver.factor(x, c, a, U, V, a, V)
Z0 = driver.solve_upper(x, c, U, W, Y, np.copy(Y))
Z = np.empty_like(Y)
F = np.empty((U.shape[0], U.shape[1], Y.shape[1]))
Z, F = backprop.solve_upper_fwd(x, c, U, W, Y, Z, F)
assert np.allclose(Z0, Z)
def test_solve_upper(vector):
x, c, a, U, V, K, Y = get_matrices(vector=vector, include_dense=True)
if vector:
Y = Y[:, None]
# First compute the expected value
expect = np.linalg.solve(np.linalg.cholesky(K).T, Y)
# Then solve using celerite
d, W = driver.factor(x, c, a, U, V, a, V)
Y /= np.sqrt(d)[:, None]
value = driver.solve_upper(x, c, U, W, Y, 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_solve_upper():
x, c, a, U, V, Y = get_matrices()
d, W = driver.factor(x, c, a, U, V, np.copy(a), np.copy(V))
Z = driver.solve_upper(x, c, U, W, Y, np.zeros_like(Y))
check_op(ops.solve_upper, [x, c, U, W, Y], [Z])
check_op(jit(ops.solve_upper), [x, c, U, W, Y], [Z])