def test_CovIdentity(seeded_rng):
cov = CovIdentity(size=m)
# compute the naive version
cov_np = np.eye(m)
logdet_np, sinv_np, sinvx_np = logdet_sinv_np(X, cov_np)
assert_allclose(logdet_np, cov.logdet, rtol=rtol)
assert_allclose(sinv_np, cov.solve(eye), rtol=rtol)
assert_allclose(sinvx_np, cov.solve(X_tf), rtol=rtol)
def test_against_scipy_mvn_col_conditional(seeded_rng):
# have to be careful for constructing everything as a submatrix of a big
# PSD matrix, else no guarantee that anything's invertible.
cov_np = wishart.rvs(df=m + p + 2, scale=np.eye(m + p))
rowcov = CovIdentity(size=m)
colcov = CovUnconstrainedCholesky(Sigma=cov_np[0:n, 0:n])
A = cov_np[n:, 0:n]
Q = CovUnconstrainedCholesky(Sigma=cov_np[n:, n:])
X = rmn(np.eye(m), np.eye(n))
A_tf = tf.constant(A, "float64")
X_tf = tf.constant(X, "float64")
Q_np = Q._cov
colcov_np = colcov._cov - A.T.dot(np.linalg.inv(Q_np)).dot((A))
scipy_answer = np.sum(
multivariate_normal.logpdf(X, np.zeros([n]), colcov_np))
tf_answer = matnorm_logp_conditional_col(X_tf, rowcov, colcov, A_tf, Q)
assert_allclose(scipy_answer, tf_answer, rtol=rtol)
def test_matnorm_regression_unconstrained(seeded_rng):
# Y = XB + eps
# Y is m x p, B is n x p, eps is m x p
X = norm.rvs(size=(m, n))
B = norm.rvs(size=(n, p))
Y_hat = X.dot(B)
rowcov_true = np.eye(m)
colcov_true = wishart.rvs(p + 2, np.eye(p))
Y = Y_hat + rmn(rowcov_true, colcov_true)
row_cov = CovIdentity(size=m)
col_cov = CovUnconstrainedCholesky(size=p)
model = MatnormalRegression(time_cov=row_cov, space_cov=col_cov)
model.fit(X, Y, naive_init=False)
assert pearsonr(B.flatten(), model.beta_.flatten())[0] >= corrtol
pred_y = model.predict(X)
assert pearsonr(pred_y.flatten(), Y_hat.flatten())[0] >= corrtol
model = MatnormalRegression(time_cov=row_cov, space_cov=col_cov)
model.fit(X, Y, naive_init=True)
assert pearsonr(B.flatten(), model.beta_.flatten())[0] >= corrtol
pred_y = model.predict(X)
assert pearsonr(pred_y.flatten(), Y_hat.flatten())[0] >= corrtol
def test_matnorm_regression_scaledDiag(seeded_rng):
# Y = XB + eps
# Y is m x n, B is n x p, eps is m x p
X = norm.rvs(size=(m, n))
B = norm.rvs(size=(n, p))
Y_hat = X.dot(B)
rowcov_true = np.eye(m)
colcov_true = np.diag(np.abs(norm.rvs(size=p)))
Y = Y_hat + rmn(rowcov_true, colcov_true)
row_cov = CovIdentity(size=m)
col_cov = CovDiagonal(size=p)
model = MatnormalRegression(time_cov=row_cov, space_cov=col_cov)
model.fit(X, Y, naive_init=False)
assert pearsonr(B.flatten(), model.beta_.flatten())[0] >= corrtol
pred_y = model.predict(X)
assert pearsonr(pred_y.flatten(), Y_hat.flatten())[0] >= corrtol
# we only do calibration test on the scaled diag
# model because to hit corrtol on unconstrainedCov
# we'd need a lot more data, which would make the test slow
X_hat = model.calibrate(Y)
assert pearsonr(X_hat.flatten(), X.flatten())[0] >= corrtol
def test_against_scipy_mvn_col(seeded_rng):
rowcov = CovIdentity(size=m)
colcov = CovUnconstrainedCholesky(size=n)
X = rmn(np.eye(m), np.eye(n))
X_tf = tf.constant(X, "float64")
colcov_np = colcov._cov
scipy_answer = np.sum(
multivariate_normal.logpdf(X, np.zeros([n]), colcov_np))
tf_answer = matnorm_logp(X_tf, rowcov, colcov)
assert_allclose(scipy_answer, tf_answer, rtol=rtol)
def test_brsa_rudimentary(seeded_rng):
"""this test is super loose"""
# this is Mingbo's synth example from the paper
U = gen_U_nips2016_example()
n_T = 150
n_V = 250
n_nureg = 5
spacecov_true = np.eye(n_V)
timecov_true = np.diag(np.abs(norm.rvs(size=(n_T))))
tr, sz = gen_brsa_data_matnorm_model(
U,
n_T=n_T,
n_V=n_V,
n_nureg=n_nureg,
space_cov=spacecov_true,
time_cov=timecov_true,
)
spacecov_model = CovIdentity(size=n_V)
timecov_model = CovDiagonal(size=n_T)
model_matnorm = MNRSA(time_cov=timecov_model, space_cov=spacecov_model)
model_matnorm.fit(tr["Y"], tr["X"], naive_init=False)
RMSE = np.mean((model_matnorm.C_ - cov2corr(tr["U"]))**2)**0.5
assert RMSE < 0.1
model_matnorm = MNRSA(time_cov=timecov_model, space_cov=spacecov_model)
model_matnorm.fit(tr["Y"], tr["X"], naive_init=True)
RMSE = np.mean((model_matnorm.C_ - cov2corr(tr["U"]))**2)**0.5
assert RMSE < 0.1
def test_against_scipy_mvn_col_marginal(seeded_rng):
rowcov = CovIdentity(size=m)
colcov = CovUnconstrainedCholesky(size=n)
Q = CovUnconstrainedCholesky(size=p)
X = rmn(np.eye(m), np.eye(n))
A = rmn(np.eye(p), np.eye(n))
A_tf = tf.constant(A, "float64")
X_tf = tf.constant(X, "float64")
Q_np = Q._cov
colcov_np = colcov._cov + A.T.dot(Q_np).dot(A)
scipy_answer = np.sum(multivariate_normal.logpdf(
X, np.zeros([n]), colcov_np))
tf_answer = matnorm_logp_marginal_col(X_tf, rowcov, colcov, A_tf, Q)
assert_allclose(scipy_answer, tf_answer, rtol=rtol)
def test_matnorm_calibration_raises(seeded_rng):
# Y = XB + eps
# Y is m x n, B is n x p, eps is m x p
X = norm.rvs(size=(2, 5))
B = norm.rvs(size=(5, 3))
Y_hat = X.dot(B)
rowcov_true = np.eye(2)
colcov_true = np.diag(np.abs(norm.rvs(size=3)))
Y = Y_hat + rmn(rowcov_true, colcov_true)
row_cov = CovIdentity(size=2)
col_cov = CovDiagonal(size=3)
model = MatnormalRegression(time_cov=row_cov, space_cov=col_cov)
model.fit(X, Y, naive_init=False)
with pytest.raises(RuntimeError):
model.calibrate(Y)