def test_general_matmul_upper():
x, c, a, U, V, Y, t, U2, V2 = get_matrices(conditional=True)
Z = driver.general_matmul_upper(t, x, c, U2, V, Y,
np.zeros((len(t), Y.shape[1])))
check_op(ops.general_matmul_upper, [t, x, c, U2, V, Y], [Z], grad=False)
check_op(jit(ops.general_matmul_upper), [t, x, c, U2, V, Y], [Z],
grad=False)
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_matmul_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
check_basic(
backprop.matmul_fwd,
ops.matmul,
[a, U, V, P, Y],
)
def test_factor_fwd():
x, c, a, U, V, Y = get_matrices()
check_basic(
backprop.factor_fwd,
ops.factor,
[x, c, a, U, V],
)
def test_matmul_rev(vector):
a, U, V, P, Y = get_matrices(vector=vector)
check_grad(
ops.matmul,
[a, U, V, P, Y],
1,
)
def test_matmul_upper_rev():
x, c, a, U, V, Y = get_matrices()
check_grad(
ops.matmul_upper,
[x, c, U, V, Y],
1,
)
def test_general_matmul_upper_fwd():
x, c, a, U, V, Y, t, U2, V2 = get_matrices(conditional=True)
check_basic(
backprop.general_matmul_upper_fwd,
ops.general_matmul_upper,
[t, x, c, U2, V, Y],
)
def test_factor_rev():
x, c, a, U, V, Y = get_matrices()
check_grad(
ops.factor,
[x, c, a, U, V],
2,
)
def test_matmul_upper_fwd():
x, c, a, U, V, Y = get_matrices()
check_basic(
backprop.matmul_upper_fwd,
ops.matmul_upper,
[x, c, U, V, Y],
)
def test_factor_fwd():
a, U, V, P, Y = get_matrices()
check_basic(
backprop.factor_fwd,
ops.factor,
[a, U, V, P],
)
def test_factor_rev():
a, U, V, P, Y = get_matrices()
check_grad(
ops.factor,
[a, U, V, P],
2,
)
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_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_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_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_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_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_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_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_factor_rev():
a, U, V, P, Y = get_matrices()
d = np.empty_like(a)
W = np.empty_like(V)
S = np.empty((len(a), U.shape[1]**2))
d, W, S = backprop.factor_fwd(a, U, V, P, d, W, S)
check_grad(backprop.factor_fwd, backprop.factor_rev, [a, U, V, P], [d, W],
[S])
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_factor_rev():
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]))
d, W, S = backprop.factor_fwd(x, c, a, U, V, d, W, S)
check_grad(backprop.factor_fwd, backprop.factor_rev, [x, c, a, U, V],
[d, W], [S])
def test_matmul_lower_fwd(vector):
x, c, a, U, V, Y = get_matrices(vector=vector)
if vector:
Y = Y[:, None]
Z0 = driver.matmul_lower(x, c, U, V, Y, np.zeros_like(Y))
Z = np.empty_like(Y)
F = np.empty((U.shape[0], U.shape[1], Y.shape[1]))
Z, F = backprop.matmul_lower_fwd(x, c, U, V, Y, Z, F)
assert np.allclose(Z0, Z)
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_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_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_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_matmul_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.dot(np.triu(K, 1), Y)
# Then solve using celerite
value = driver.matmul_upper(x, c, U, V, Y, np.zeros_like(Y))
# Check that the solution is correct
assert np.allclose(value, expect)
def test_conditional_mean():
a, U, V, P, K, Y, U_star, V_star, inds, K_star = get_matrices(
vector=True, conditional=True, include_dense=True)
# First compute the expected value
alpha = np.linalg.solve(K, Y)
expect = np.dot(K_star, alpha)
# Then solve using celerite
mu = np.empty(len(U_star))
value = driver.conditional_mean(U, V, P, alpha, U_star, V_star, inds, mu)
# Check that the solution is correct
assert np.allclose(value, expect)