Merge pull request #332 from calebweinreb/lgssm_with_diagonal_noise

LGSSM with diagonal noise

Author: slinderman

dynamax/linear_gaussian_ssm/inference.py

@@ -1,9 +1,13 @@
 import jax.numpy as jnp
 import jax.random as jr
 from jax import lax
-from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN
 from functools import wraps
 import inspect
+import warnings
+
+from tensorflow_probability.substrates.jax.distributions import (
+    MultivariateNormalDiagPlusLowRankCovariance as MVNLowRank,
+    MultivariateNormalFullCovariance as MVN)
 
 from jax.tree_util import tree_map
 from jaxtyping import Array, Float
@@ -41,10 +45,22 @@ class ParamsLGSSMDynamics(NamedTuple):
     :param cov: dynamics covariance $Q$
 
     """
-    weights: Union[Float[Array, "state_dim state_dim"], Float[Array, "ntime state_dim state_dim"], ParameterProperties]
-    bias: Union[Float[Array, "state_dim"], Float[Array, "ntime state_dim"], ParameterProperties]
-    input_weights: Union[Float[Array, "state_dim input_dim"], Float[Array, "ntime state_dim input_dim"], ParameterProperties]
-    cov: Union[Float[Array, "state_dim state_dim"], Float[Array, "ntime state_dim state_dim"], Float[Array, "state_dim_triu"], ParameterProperties]
+    weights: Union[ParameterProperties, 
+        Float[Array, "state_dim state_dim"], 
+        Float[Array, "ntime state_dim state_dim"]]
+    
+    bias: Union[ParameterProperties,
+        Float[Array, "state_dim"], 
+        Float[Array, "ntime state_dim"]]
+    
+    input_weights: Union[ParameterProperties,
+        Float[Array, "state_dim input_dim"], 
+        Float[Array, "ntime state_dim input_dim"]]
+    
+    cov: Union[ParameterProperties, 
+        Float[Array, "state_dim state_dim"], 
+        Float[Array, "ntime state_dim state_dim"], 
+        Float[Array, "state_dim_triu"]]
 
 
 class ParamsLGSSMEmissions(NamedTuple):
@@ -60,11 +76,24 @@ class ParamsLGSSMEmissions(NamedTuple):
     :param cov: emission covariance $R$
 
     """
-    weights: Union[Float[Array, "emission_dim state_dim"], Float[Array, "ntime emission_dim state_dim"], ParameterProperties]
-    bias: Union[Float[Array, "emission_dim"], Float[Array, "ntime emission_dim"], ParameterProperties]
-    input_weights: Union[Float[Array, "emission_dim input_dim"], Float[Array, "ntime emission_dim input_dim"], ParameterProperties]
-    cov: Union[Float[Array, "emission_dim emission_dim"], Float[Array, "ntime emission_dim emission_dim"], Float[Array, "emission_dim_triu"], ParameterProperties]
-
+    weights: Union[ParameterProperties,
+        Float[Array, "emission_dim state_dim"], 
+        Float[Array, "ntime emission_dim state_dim"]]
+    
+    bias: Union[ParameterProperties,
+        Float[Array, "emission_dim"], 
+        Float[Array, "ntime emission_dim"]]
+    
+    input_weights: Union[ParameterProperties,
+        Float[Array, "emission_dim input_dim"], 
+        Float[Array, "ntime emission_dim input_dim"]]
+    
+    cov: Union[ParameterProperties,
+        Float[Array, "emission_dim emission_dim"], 
+        Float[Array, "ntime emission_dim emission_dim"], 
+        Float[Array, "emission_dim"], 
+        Float[Array, "ntime emission_dim"], 
+        Float[Array, "emission_dim_triu"]]
 
 
 class ParamsLGSSM(NamedTuple):
@@ -115,14 +144,46 @@ class PosteriorGSSMSmoothed(NamedTuple):
 
 
 # Helper functions
-# _get_params = lambda x, dim, t: x[t] if x.ndim == dim + 1 else x
-def _get_params(x, dim, t):
+
+def _get_one_param(x, dim, t):
+    """Helper function to get one parameter at time t."""
     if callable(x):
         return x(t)
     elif x.ndim == dim + 1:
         return x[t]
     else:
         return x
+
+def _get_params(params, num_timesteps, t):
+    """Helper function to get parameters at time t."""
+    assert not callable(params.emissions.cov), "Emission covariance cannot be a callable."
+
+    F = _get_one_param(params.dynamics.weights, 2, t)
+    B = _get_one_param(params.dynamics.input_weights, 2, t)
+    b = _get_one_param(params.dynamics.bias, 1, t)
+    Q = _get_one_param(params.dynamics.cov, 2, t)
+    H = _get_one_param(params.emissions.weights, 2, t)
+    D = _get_one_param(params.emissions.input_weights, 2, t)
+    d = _get_one_param(params.emissions.bias, 1, t)
+
+    if len(params.emissions.cov.shape) == 1: 
+        R = _get_one_param(params.emissions.cov, 1, t)
+    elif len(params.emissions.cov.shape) > 2: 
+        R = _get_one_param(params.emissions.cov, 2, t)
+    elif params.emissions.cov.shape[0] != num_timesteps:
+        R = _get_one_param(params.emissions.cov, 2, t)
+    elif params.emissions.cov.shape[1] != num_timesteps:
+        R = _get_one_param(params.emissions.cov, 1, t)
+    else:
+        R = _get_one_param(params.emissions.cov, 2, t)
+        warnings.warn(
+            "Emission covariance has shape (N,N) where N is the number of timesteps. "
+            "The covariance will be interpreted as static and non-diagonal. To "
+            "specify a dynamic and diagonal covariance, pass it as a 3D array.")
+
+    return F, B, b, Q, H, D, d, R
+
+
 _zeros_if_none = lambda x, shape: x if x is not None else jnp.zeros(shape)
 
 
@@ -199,7 +260,6 @@ def _condition_on(m, P, H, D, d, R, u, y):
          S = (R + H * P * H')
          K = P * H' * S^{-1}
          PP = P - K S K' = Sigma_cond
-     **Note! This can be done more efficiently when R is diagonal.**
 
     Args:
          m (D_hid,): prior mean.
@@ -215,9 +275,25 @@ def _condition_on(m, P, H, D, d, R, u, y):
          mu_pred (D_hid,): predicted mean.
          Sigma_pred (D_hid,D_hid): predicted covariance.
     """
-    # Compute the Kalman gain
-    S = R + H @ P @ H.T
-    K = psd_solve(S, H @ P).T
+    if R.ndim == 2:
+        S = R + H @ P @ H.T
+        K = psd_solve(S, H @ P).T
+    else: 
+        # Optimization using Woodbury identity with A=R, U=H@chol(P), V=U.T, C=I
+        # (see https://en.wikipedia.org/wiki/Woodbury_matrix_identity)
+        I = jnp.eye(P.shape[0])
+        U = H @ jnp.linalg.cholesky(P)
+        X = U / R[:, None]
+        S_inv = jnp.diag(1.0 / R) - X @ psd_solve(I + U.T @ X, X.T) 
+        """
+        # Could alternatively use U=H and C=P
+        R_inv = jnp.diag(1.0 / R)
+        P_inv = psd_solve(P, jnp.eye(P.shape[0]))
+        S_inv = R_inv - R_inv @ H @ psd_solve(P_inv + H.T @ R_inv @ H, H.T @ R_inv)
+        """
+        K = P @ H.T @ S_inv  
+        S = jnp.diag(R) + H @ P @ H.T
+
     Sigma_cond = P - K @ S @ K.T
     mu_cond = m + K @ (y - D @ u - d - H @ m)
     return mu_cond, symmetrize(Sigma_cond)
@@ -285,6 +361,8 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+
+
 def lgssm_joint_sample(
     params: ParamsLGSSM,
     key: PRNGKey,
@@ -302,7 +380,6 @@ def lgssm_joint_sample(
         latent states and emissions
 
     """
-    
     params, inputs = preprocess_params_and_inputs(params, num_timesteps, inputs)
 
     def _sample_transition(key, F, B, b, Q, x_tm1, u):
@@ -311,17 +388,15 @@ def _sample_transition(key, F, B, b, Q, x_tm1, u):
 
     def _sample_emission(key, H, D, d, R, x, u):
         mean = H @ x + D @ u + d
+        R = jnp.diag(R) if R.ndim==1 else R
         return MVN(mean, R).sample(seed=key)
 
     def _sample_initial(key, params, inputs):
         key1, key2 = jr.split(key)
 
         initial_state = MVN(params.initial.mean, params.initial.cov).sample(seed=key1)
 
-        H0 = _get_params(params.emissions.weights, 2, 0)
-        D0 = _get_params(params.emissions.input_weights, 2, 0)
-        d0 = _get_params(params.emissions.bias, 1, 0)
-        R0 = _get_params(params.emissions.cov, 2, 0)
+        H0, D0, d0, R0 = _get_params(params, num_timesteps, 0)[4:]
         u0 = tree_map(lambda x: x[0], inputs)
 
         initial_emission = _sample_emission(key2, H0, D0, d0, R0, initial_state, u0)
@@ -331,15 +406,8 @@ def _step(prev_state, args):
         key, t, inpt = args
         key1, key2 = jr.split(key, 2)
 
-        # Shorthand: get parameters and inputs for time index t
-        F = _get_params(params.dynamics.weights, 2, t)
-        B = _get_params(params.dynamics.input_weights, 2, t)
-        b = _get_params(params.dynamics.bias, 1, t)
-        Q = _get_params(params.dynamics.cov, 2, t)
-        H = _get_params(params.emissions.weights, 2, t)
-        D = _get_params(params.emissions.input_weights, 2, t)
-        d = _get_params(params.emissions.bias, 1, t)
-        R = _get_params(params.emissions.cov, 2, t)
+        # Get parameters and inputs for time index t
+        F, B, b, Q, H, D, d, R = _get_params(params, num_timesteps, t)
 
         # Sample from transition and emission distributions
         state = _sample_transition(key1, F, B, b, Q, prev_state, inpt)
@@ -386,23 +454,26 @@ def lgssm_filter(
     num_timesteps = len(emissions)
     inputs = jnp.zeros((num_timesteps, 0)) if inputs is None else inputs
 
+    def _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y):
+        m = H @ pred_mean + D @ u + d
+        if R.ndim==2:
+            S = R + H @ pred_cov @ H.T
+            return MVN(m, S).log_prob(y)
+        else:
+            L = H @ jnp.linalg.cholesky(pred_cov)
+            return MVNLowRank(m, R, L).log_prob(y)
+            
+
     def _step(carry, t):
         ll, pred_mean, pred_cov = carry
 
         # Shorthand: get parameters and inputs for time index t
-        F = _get_params(params.dynamics.weights, 2, t)
-        B = _get_params(params.dynamics.input_weights, 2, t)
-        b = _get_params(params.dynamics.bias, 1, t)
-        Q = _get_params(params.dynamics.cov, 2, t)
-        H = _get_params(params.emissions.weights, 2, t)
-        D = _get_params(params.emissions.input_weights, 2, t)
-        d = _get_params(params.emissions.bias, 1, t)
-        R = _get_params(params.emissions.cov, 2, t)
+        F, B, b, Q, H, D, d, R = _get_params(params, num_timesteps, t)
         u = inputs[t]
         y = emissions[t]
 
         # Update the log likelihood
-        ll += MVN(H @ pred_mean + D @ u + d, H @ pred_cov @ H.T + R).log_prob(y)
+        ll += _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y)
 
         # Condition on this emission
         filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y)
@@ -450,11 +521,8 @@ def _step(carry, args):
         smoothed_mean_next, smoothed_cov_next = carry
         t, filtered_mean, filtered_cov = args
 
-        # Shorthand: get parameters and inputs for time index t
-        F = _get_params(params.dynamics.weights, 2, t)
-        B = _get_params(params.dynamics.input_weights, 2, t)
-        b = _get_params(params.dynamics.bias, 1, t)
-        Q = _get_params(params.dynamics.cov, 2, t)
+        # Get parameters and inputs for time index t
+        F, B, b, Q = _get_params(params, num_timesteps, t)[:4]
         u = inputs[t]
 
         # This is like the Kalman gain but in reverse
@@ -522,10 +590,7 @@ def _step(carry, args):
         key, filtered_mean, filtered_cov, t = args
 
         # Shorthand: get parameters and inputs for time index t
-        F = _get_params(params.dynamics.weights, 2, t)
-        B = _get_params(params.dynamics.input_weights, 2, t)
-        b = _get_params(params.dynamics.bias, 1, t)
-        Q = _get_params(params.dynamics.cov, 2, t)
+        F, B, b, Q = _get_params(params, num_timesteps, t)[:4]
         u = inputs[t]
 
         # Condition on next state