def test_blocks_to_banded(T=5, D=3):
"""
Test blocks_to_banded correctness
"""
Ad = np.zeros((T, D, D))
Aod = np.zeros((T-1, D, D))
M = np.arange(1, D+1)[:, None] * 10 + np.arange(1, D+1)
for t in range(T):
Ad[t, :, :] = 100 * ((t+1)*10 + (t+1)) + M
for t in range(T-1):
Aod[t, :, :] = 100 * ((t+2)*10 + (t+1)) + M
# print("Lower")
# L = blocks_to_bands(Ad, Aod, lower=True)
# print(L)
# print("Upper")
# U = blocks_to_bands(Ad, Aod, lower=False)
# print(U)
# Check inverse with random symmetric matrices
Ad = npr.randn(T, D, D)
Ad = (Ad + np.swapaxes(Ad, -1, -2)) / 2
Aod = npr.randn(T-1, D, D)
Ad2, Aod2 = bands_to_blocks(blocks_to_bands(Ad, Aod, lower=True), lower=True)
assert np.allclose(np.tril(Ad), np.tril(Ad2))
assert np.allclose(Aod, Aod2)
Ad3, Aod3 = bands_to_blocks(blocks_to_bands(Ad, Aod, lower=False), lower=False)
assert np.allclose(np.triu(Ad), np.triu(Ad3))
assert np.allclose(Aod, Aod3)
def test_solveh_banded_grad(T=10, D=4):
"""
Test solveh_banded gradient
"""
J_diag, J_lower_diag, J_full = make_block_tridiag(T, D)
J_diag = np.tile(J_diag[None, :, :], (T, 1, 1))
J_lower_diag = np.tile(J_lower_diag[None, :, :], (T - 1, 1, 1))
b = npr.randn(T * D)
J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
check_grads(solveh_banded, argnum=0, modes=['rev'], order=1)(J_banded,
b,
lower=True)
check_grads(solveh_banded, argnum=1, modes=['rev'], order=1)(J_banded,
b,
lower=True)
J_banded = blocks_to_bands(J_diag,
np.swapaxes(J_lower_diag, -1, -2),
lower=False)
check_grads(solveh_banded, argnum=0, modes=['rev'], order=1)(J_banded,
b,
lower=False)
check_grads(solveh_banded, argnum=1, modes=['rev'], order=1)(J_banded,
b,
lower=False)
def test_transpose_banded_grad(T=25, D=4):
"""
Test transpose_banded gradient
"""
J_diag, J_lower_diag, J_full = make_block_tridiag(T, D)
J_diag = np.tile(J_diag[None, :, :], (T, 1, 1))
J_lower_diag = np.tile(J_lower_diag[None, :, :], (T-1, 1, 1))
J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
check_grads(transpose_banded, argnum=1, modes=['rev'], order=1)((2*D-1, 0), J_banded)
def test_cholesky_banded_grad(T=10, D=4):
"""
Test cholesky_banded gradient
"""
J_diag, J_lower_diag, J_full = make_block_tridiag(T, D)
L = np.linalg.cholesky(J_full)
dJ_bar_true = elementwise_grad(np.linalg.cholesky)(J_full)
# Convert to lower bands
J_diag = np.tile(J_diag[None, :, :], (T, 1, 1))
J_lower_diag = np.tile(J_lower_diag[None, :, :], (T-1, 1, 1))
J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
L_banded = np.vstack([
np.concatenate((np.diag(L, -d), np.zeros(d))) for d in range(2 * D)
])
dJ_banded = grad_cholesky_banded(L_banded, J_banded)(np.ones_like(L_banded))
assert np.allclose(np.diag(dJ_bar_true), dJ_banded[0])
for d in range(1, 2 * D):
assert np.allclose(np.diag(dJ_bar_true, -d), dJ_banded[d, :-d] / 2)
