def test_factor():
x, c, a, U, V, K, Y = get_matrices(include_dense=True)
d, W = driver.factor(x, c, a, U, V, a, V)
# Make sure that no copy is made if possible
assert np.allclose(a, d)
assert np.allclose(V, W)
def test_norm_rev():
a, U, V, P, Y = get_matrices(vector=True)
d, W = driver.factor(U, P, a, V)
check_grad(
ops.norm,
[U, P, d, W, Y],
1,
)
def test_solve_upper_fwd():
x, c, a, U, V, Y = get_matrices()
d, W = driver.factor(x, c, a, U, V, a, V)
check_basic(
backprop.solve_upper_fwd,
ops.solve_upper,
[x, c, U, W, Y],
)
def test_solve_rev(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
check_grad(
ops.solve,
[U, P, d, W, Y],
1,
)
def test_solve_upper_rev():
x, c, a, U, V, Y = get_matrices()
d, W = driver.factor(x, c, a, U, V, a, V)
check_grad(
ops.solve_upper,
[x, c, U, W, Y],
1,
)
def test_norm_fwd():
a, U, V, P, Y = get_matrices(vector=True)
d, W = driver.factor(U, P, a, V)
check_basic(
backprop.norm_fwd,
ops.norm,
[U, P, d, W, Y],
)
def test_dot_tril_rev(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
check_grad(
ops.dot_tril,
[U, P, d, W, Y],
1,
)
def test_solve_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
check_basic(
backprop.solve_fwd,
ops.solve,
[U, P, d, W, Y],
)
def test_dot_tril_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
check_basic(
backprop.dot_tril_fwd,
ops.dot_tril,
[U, P, d, W, Y],
)
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_norm():
a, U, V, P, K, Y = get_matrices(vector=True, include_dense=True)
# First compute the expected value
expect = np.dot(Y, np.linalg.solve(K, Y))
# Then solve using celerite
d, W = driver.factor(U, P, a, V)
value = driver.norm(U, P, d, W, Y)
# Check that the solution is correct
assert np.allclose(value, expect)
def test_norm_rev():
a, U, V, P, Y = get_matrices(vector=True)
d, W = driver.factor(U, P, a, V)
X = np.empty((1, 1))
Z = np.empty_like(Y)
F = np.empty_like(U)
X, Z, F = backprop.norm_fwd(U, P, d, W, Y, X, Z, F)
check_grad(backprop.norm_fwd, backprop.norm_rev, [U, P, d, W, Y], [X],
[Z, F])
def test_solve_lower_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_lower(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_lower_fwd(x, c, U, W, Y, Z, F)
assert np.allclose(Z0, Z)
def test_norm_fwd():
a, U, V, P, Y = get_matrices(vector=True)
d, W = driver.factor(U, P, a, V)
X0 = driver.norm(U, P, d, W, np.copy(Y))
X = np.empty((1, 1))
Z = np.empty_like(Y)
F = np.empty_like(U)
X, Z, F = backprop.norm_fwd(U, P, d, W, Y, X, Z, F)
assert np.allclose(X0, X)
def test_factor_fwd():
x, c, a, U, V, Y = get_matrices()
d = np.empty_like(a)
W = np.empty_like(V)
S = np.empty((len(a), U.shape[1], U.shape[1]))
d0, W0 = driver.factor(x, c, a, U, V, np.copy(a), np.copy(V))
d, W, S = backprop.factor_fwd(x, c, a, U, V, d, W, S)
assert np.allclose(d, d0)
assert np.allclose(W, W0)
def test_dot_tril_rev(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
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]))
Z, F = backprop.dot_tril_fwd(U, P, d, W, Y, Z, F)
check_grad(backprop.dot_tril_fwd, backprop.dot_tril_rev, [U, P, d, W, Y],
[Z], [F])
def test_dot_tril_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
Z0 = driver.dot_tril(U, P, d, W, np.copy(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]))
Z, F = backprop.dot_tril_fwd(U, P, d, W, Y, Z, F)
assert np.allclose(Z0, Z)
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_dot_tril(vector):
a, U, V, P, K, Y = get_matrices(vector=vector, include_dense=True)
# First compute the expected value
expect = np.dot(np.linalg.cholesky(K), Y)
# Then solve using celerite
d, W = driver.factor(U, P, a, V)
value = driver.dot_tril(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_solve_rev(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
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)
check_grad(backprop.solve_fwd, backprop.solve_rev, [U, P, d, W, Y], [X],
[Z, F, G])
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_lower_rev(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)
Z = np.empty_like(Y)
F = np.empty((U.shape[0], U.shape[1], Y.shape[1]))
Z, F = backprop.solve_lower_fwd(x, c, U, W, Y, Z, F)
check_grad(
backprop.solve_lower_fwd,
backprop.solve_lower_rev,
[x, c, U, W, Y],
[Z],
[F],
)
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])
def test_factor():
x, c, a, U, V, Y = get_matrices()
d, W = driver.factor(x, c, a, U, V, np.copy(a), np.copy(V))
check_op(ops.factor, [x, c, a, U, V], [d, W])
check_op(jit(ops.factor), [x, c, a, U, V], [d, W])
def test_dot_tril(vector):
a, U, V, P, Y = get_matrices(vector=vector)
d, W = driver.factor(U, P, a, V)
X = driver.dot_tril(U, P, d, W, np.copy(Y))
check_op(ops.dot_tril, [U, P, d, W, Y], [X])
check_op(jit(ops.dot_tril), [U, P, d, W, Y], [X])
def test_norm():
a, U, V, P, Y = get_matrices(vector=True)
d, W = driver.factor(U, P, a, V)
X = driver.norm(U, P, d, W, np.copy(Y))
check_op(ops.norm, [U, P, d, W, Y], [X])
check_op(jit(ops.norm), [U, P, d, W, Y], [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])
def test_factor():
a, U, V, P, Y = get_matrices()
d, W = driver.factor(U, P, np.copy(a), np.copy(V))
check_op(ops.factor, [a, U, V, P], [d, W])
check_op(jit(ops.factor), [a, U, V, P], [d, W])