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 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 cholesky_lds(As, bs, Qi_sqrts, ms, Ri_sqrts): J_diag, J_lower_diag, _ = convert_lds_to_block_tridiag(As, bs, Qi_sqrts, ms, Ri_sqrts) J_banded = blocks_to_bands(J_diag, J_lower_diag, lower=True) return cholesky_banded(J_banded, lower=True)