def lds_sample(As, bs, Qi_sqrts, ms, Ri_sqrts, z=None):
"""
Sample a linear dynamical system
"""
T, D = ms.shape
assert As.shape == (T-1, D, D)
assert bs.shape == (T-1, D)
assert Qi_sqrts.shape == (T-1, D, D)
assert Ri_sqrts.shape == (T, D, D)
# Convert to block form
J_diag, J_lower_diag, h = convert_lds_to_block_tridiag(As, bs, Qi_sqrts, ms, Ri_sqrts)
# Convert blocks to banded form so we can capitalize on Lapack code
J_banded = A_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
L = cholesky_banded(J_banded, lower=True)
U = transpose_banded((2*D-1, 0), L)
# We have (U^T U)^{-1} = U^{-1} U^{-T} = AA^T = Sigma
# where A = U^{-1}. Samples are Az = U^{-1}z = x, or equivalently Ux = z.
z = npr.randn(T*D,) if z is None else np.reshape(z, (T*D,))
samples = np.reshape(solve_banded((0, 2*D-1), U, z), (T, D))
# Get the mean mu = J^{-1} h
mu = np.reshape(solveh_banded(J_banded, np.ravel(h), lower=True), (T, D))
# Add the mean
return samples + mu
def vjp(C_bar):
b_bar = solve_banded((u, l), transpose_banded((l, u), A_banded), C_bar)
A_bar = np.zeros_like(A_banded)
K = b.shape[1] if b.ndim == 2 else 1
_vjp_solve_banded_A(A_bar, b_bar.reshape(-1, K), C_bar.reshape(-1, K),
C.reshape(-1, K), u, A_banded)
return A_bar
def block_tridiagonal_log_probability(x, J_diag, J_lower_diag, h):
T, D = x.shape
assert h.shape == (T, D)
assert J_diag.shape == (T, D, D)
assert J_lower_diag.shape == (T - 1, D, D)
# Convert blocks to banded form so we can capitalize on Lapack code
J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
# -1/2 x^T J x = -1/2 \sum_{t=1}^T x_t.T J_tt x_t
ll = -1 / 2 * np.sum(
np.matmul(x[:, None, :], np.matmul(J_diag, x[:, :, None])))
# -\sum_{t=1}^{T-1} x_t.T J_{t,t+1} x_{t+1}
ll -= np.sum(
np.matmul(x[1:, None, :], np.matmul(J_lower_diag, x[:-1, :, None])))
# h^T x
ll += np.sum(h * x)
# -1/2 h^T J^{-1} h = -1/2 h^T (LL^T)^{-1} h
# = -1/2 h^T L^{-T} L^{-1} h
# = -1/2 (L^{-1}h)^T (L^{-1} h)
# L = cholesky_block_tridiag(J_diag, J_lower_diag, lower=True)
L = cholesky_banded(J_banded, lower=True)
Linv_h = solve_banded((2 * D - 1, 0), L, h.ravel())
ll -= 1 / 2 * np.sum(Linv_h * Linv_h)
# 1/2 log |J| -TD/2 log(2 pi) = log |L| -TD/2 log(2 pi)
L_diag = L[0]
ll += np.sum(np.log(L_diag))
ll -= 1 / 2 * T * D * np.log(2 * np.pi)
return ll
def lds_log_probability(x, As, bs, Qi_sqrts, ms, Ri_sqrts):
"""
Compute the log normalizer of a linear dynamical system.
"""
T, D = x.shape
assert As.shape == (T - 1, D, D)
assert bs.shape == (T - 1, D)
assert Qi_sqrts.shape == (T - 1, D, D)
assert ms.shape == (T, D)
assert Ri_sqrts.shape == (T, D, D)
# Convert to block form
J_diag, J_lower_diag, h = convert_lds_to_block_tridiag(
As, bs, Qi_sqrts, ms, Ri_sqrts)
# Convert blocks to banded form so we can capitalize on Lapack code
J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
# -1/2 x^T J x = -1/2 \sum_{t=1}^T x_t.T J_tt x_t
ll = -1 / 2 * np.sum(
np.matmul(x[:, None, :], np.matmul(J_diag, x[:, :, None])))
# -\sum_{t=1}^{T-1} x_t.T J_{t,t+1} x_{t+1}
ll -= np.sum(
np.matmul(x[1:, None, :], np.matmul(J_lower_diag, x[:-1, :, None])))
# h^T x
ll += np.sum(h * x)
# -1/2 h^T J^{-1} h = -1/2 h^T (LL^T)^{-1} h
# = -1/2 h^T L^{-T} L^{-1} h
# = -1/2 (L^{-1}h)^T (L^{-1} h)
# L = cholesky_block_tridiag(J_diag, J_lower_diag, lower=True)
L = cholesky_banded(J_banded, lower=True)
Linv_h = solve_banded((2 * D - 1, 0), L, h.ravel())
ll -= 1 / 2 * np.sum(Linv_h * Linv_h)
# 1/2 log |J| -TD/2 log(2 pi) = log |L| -TD/2 log(2 pi)
L_diag = L[0]
ll += np.sum(np.log(L_diag))
ll -= 1 / 2 * T * D * np.log(2 * np.pi)
return ll
def block_tridiagonal_sample(J_diag, J_lower_diag, h, z=None):
"""
Sample a Gaussian chain graph represented by a block
tridiagonal precision matrix and a linear potential.
"""
T, D = h.shape
assert J_diag.shape == (T, D, D)
assert J_lower_diag.shape == (T - 1, D, D)
# Convert blocks to banded form so we can capitalize on Lapack code
J_banded = A_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True)
L = cholesky_banded(J_banded, lower=True)
U = transpose_banded((2 * D - 1, 0), L)
# We have (U^T U)^{-1} = U^{-1} U^{-T} = AA^T = Sigma
# where A = U^{-1}. Samples are Az = U^{-1}z = x, or equivalently Ux = z.
z = npr.randn(T * D, ) if z is None else np.reshape(z, (T * D, ))
samples = np.reshape(solve_banded((0, 2 * D - 1), U, z), (T, D))
# Get the mean mu = J^{-1} h
mu = np.reshape(solveh_banded(J_banded, np.ravel(h), lower=True), (T, D))
# Add the mean
return samples + mu
def vjp(C_bar):
return solve_banded((u, l), transpose_banded((l, u), A_banded), C_bar)
