def gen_brsa_data_matnorm_model(U, n_T, n_V, space_cov, time_cov, n_nureg):
n_C = U.shape[0]
beta = rmn(U, space_cov)
X = rmn(np.eye(n_T), np.eye(n_C))
beta_0 = rmn(np.eye(n_nureg), space_cov)
X_0 = rmn(np.eye(n_T), np.eye(n_nureg))
Y_hat = X.dot(beta) + X_0.dot(beta_0)
Y = Y_hat + rmn(time_cov, space_cov)
sizes = {"n_C": n_C, "n_T": n_T, "n_V": n_V}
train = {"beta": beta, "X": X, "Y": Y, "U": U, "X_0": X_0}
return train, sizes
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_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_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_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)