Added Marginal Log likelihood to SLDS RBPF

cd_dynamax/dynamax/slds/inference.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 from jax import lax, vmap, jit
+from jax.scipy.special import logsumexp
 from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN
 from functools import partial
 from jaxtyping import Array, Float, Int
@@ -82,15 +83,20 @@ def initialize(self, num_states, state_dim, emission_dim, input_dim=1):
 
 class RBPFiltered(NamedTuple):
     r"""RBPF posterior.
+    :param marginal_loglik: approximate marginal log likelihood from the particle-filter normalizers.
     :param weights: weights of the particles.
     :param means: array of filtered means $\mathbb{E}[z_t \mid y_{1:t}, u_{1:t}]$
     :param covariances: array of filtered covariances $\mathrm{Cov}[z_t \mid y_{1:t}, u_{1:t}]$
-    :param states: array of sampled discrete state sequences (particles) $$.
+    :param states: integer array of sampled discrete states with shape $(\mathrm{ntime}, \mathrm{num\_particles})$.
     """
-    weights: Optional[Float[Array, "num_particles ntime"]] = None
-    states: Optional[Int[Array, "num_particles ntime num_states"]] = None
-    means: Optional[Float[Array, "num_particles ntime state_dim"]] = None
-    covariances: Optional[Float[Array, "num_particles ntime state_dim state_dim"]] = None
+    marginal_loglik: Optional[Float[Array, ""]] = None
+    weights: Optional[Float[Array, "ntime num_particles"]] = None
+    states: Optional[Int[Array, "ntime num_particles"]] = None
+    means: Optional[Float[Array, "ntime num_particles state_dim"]] = None
+    covariances: Optional[Float[Array, "ntime num_particles state_dim state_dim"]] = None
+
+    def __getitem__(self, key):
+        return getattr(self, key) if isinstance(key, str) else tuple.__getitem__(self, key)
 
 
 def resampling(weights, states, means, covariances, key):                                                                  
@@ -208,6 +214,7 @@ def _step(carry, t):
         lls, filtered_means, filtered_covs = vmap(_conditional_kalman_step, in_axes = (0, 0, 0, None, None, None))(new_states, filtered_means, filtered_covs, params.linear_gaussian, u, y)
 
         # Compute weights
+        marginal_loglik = logsumexp(jnp.log(weights) + lls)
         lls -= jnp.max(lls)
         loglik_weights = jnp.exp(lls)
         weights = jnp.multiply(loglik_weights.T, weights)
@@ -219,10 +226,11 @@ def _step(carry, t):
                                                                                 weights, new_states, filtered_means, filtered_covs, next_key)
 
         # Build carry and output states
-        carry = (weights, prev_states, filtered_means, filtered_covs, next_key)
+        carry = (weights, new_states, filtered_means, filtered_covs, next_key)
         outputs = {
+            "marginal_loglik": marginal_loglik,
             "weights": weights,
-            "states": prev_states,
+            "states": new_states,
             "means": filtered_means,
             "covariances": filtered_covs
         }
@@ -247,7 +255,12 @@ def _step(carry, t):
 
     _, out = lax.scan(_step, carry, jnp.arange(num_timesteps))
 
-    return out
+    return RBPFiltered(
+        marginal_loglik=out["marginal_loglik"].sum(),
+        weights=out["weights"],
+        states=out["states"],
+        means=out["means"],
+        covariances=out["covariances"])
 
 def rbpfilter_optimal(
     num_particles: int,
@@ -289,6 +302,9 @@ def _step(carry, t):
         lls, filtered_means, filtered_covs = vmap(_vec_kalman_step, in_axes=(0, 0))(filtered_means, filtered_covs)
 
         # Compute weights
+        marginal_loglik = logsumexp(
+            jnp.log(weights)[:, None] + jnp.log(params.discrete.transition_matrix[prev_states, :]) + lls
+        )
         lls -= jnp.max(lls)
         loglik_weights = jnp.exp(lls)
         trans_weights = params.discrete.transition_matrix[prev_states, :]
@@ -310,6 +326,7 @@ def _step(carry, t):
         # Build carry and output states
         carry = (res_weights, new_states, filtered_means, filtered_covs, next_key)
         outputs = {
+            "marginal_loglik": marginal_loglik,
             "weights": res_weights,
             "states": new_states,
             "means": filtered_means,
@@ -336,4 +353,9 @@ def _step(carry, t):
 
     _, out = lax.scan(_step, carry, jnp.arange(num_timesteps))
 
-    return out
+    return RBPFiltered(
+        marginal_loglik=out["marginal_loglik"].sum(),
+        weights=out["weights"],
+        states=out["states"],
+        means=out["means"],
+        covariances=out["covariances"])

tests/test_filter_posterior_contracts.py

@@ -41,13 +41,23 @@
 from cd_dynamax.src.continuous_discrete_nonlinear_ssm.models import (
     PosteriorCDNLSSMFiltered,
 )
+from cd_dynamax.dynamax.slds.inference import (
+    DiscreteParamsSLDS,
+    LGParamsSLDS,
+    ParamsSLDS,
+    RBPFiltered,
+    rbpfilter,
+    rbpfilter_optimal,
+)
 
 
 NTIME = 3
 STATE_DIM = 1
 EMISSION_DIM = 1
 ENKF_PARTICLES = 16
 DPF_PARTICLES = 8
+RBPF_PARTICLES = 16
+SLDS_NUM_STATES = 2
 TEST_EMISSIONS = jnp.zeros((NTIME, EMISSION_DIM))
 TEST_T_EMISSIONS = jnp.array([[0.0], [0.1], [0.2]])
 
@@ -153,6 +163,16 @@ def _assert_posterior_cdnlssm_contract(posterior):
     _assert_scalar_total_loglik(posterior.marginal_loglik)
 
 
+def _assert_rbpfiltered_contract(posterior):
+    assert isinstance(posterior, RBPFiltered)
+    _assert_scalar_total_loglik(posterior.marginal_loglik)
+    assert posterior.weights.shape == (NTIME, RBPF_PARTICLES)
+    assert posterior.states.shape == (NTIME, RBPF_PARTICLES)
+    assert posterior.means.shape == (NTIME, RBPF_PARTICLES, STATE_DIM)
+    assert posterior.covariances.shape == (NTIME, RBPF_PARTICLES, STATE_DIM, STATE_DIM)
+    assert posterior["means"] is posterior.means
+
+
 def _expected_gssm_field_shapes(ensemble_size=None):
     expected_shapes = dict(BASE_GSSM_FIELD_SHAPES)
 
@@ -230,6 +250,38 @@ def _assert_posterior_gssm_contract(
         assert posterior.posterior_extras is None
 
 
+def _make_slds_params():
+    transition_matrix = jnp.ones((SLDS_NUM_STATES, SLDS_NUM_STATES)) / SLDS_NUM_STATES
+    return ParamsSLDS(
+        discrete=DiscreteParamsSLDS(
+            initial_distribution=jnp.ones(SLDS_NUM_STATES) / SLDS_NUM_STATES,
+            transition_matrix=transition_matrix,
+            proposal_transition_matrix=transition_matrix,
+        ),
+        linear_gaussian=LGParamsSLDS(
+            initial_mean=jnp.zeros((SLDS_NUM_STATES, STATE_DIM)),
+            initial_cov=jnp.tile(
+                jnp.eye(STATE_DIM)[None, :, :], (SLDS_NUM_STATES, 1, 1)
+            ),
+            dynamics_weights=jnp.tile(
+                (0.9 * jnp.eye(STATE_DIM))[None, :, :], (SLDS_NUM_STATES, 1, 1)
+            ),
+            dynamics_cov=jnp.tile(
+                (0.1 * jnp.eye(STATE_DIM))[None, :, :], (SLDS_NUM_STATES, 1, 1)
+            ),
+            dynamics_bias=jnp.zeros((SLDS_NUM_STATES, STATE_DIM)),
+            dynamics_input_weights=jnp.zeros((SLDS_NUM_STATES, STATE_DIM, 1)),
+            emission_weights=jnp.ones((SLDS_NUM_STATES, EMISSION_DIM, STATE_DIM)),
+            emission_cov=jnp.tile(
+                (0.1 * jnp.eye(EMISSION_DIM))[None, :, :], (SLDS_NUM_STATES, 1, 1)
+            ),
+            emission_bias=jnp.zeros((SLDS_NUM_STATES, EMISSION_DIM)),
+            emission_input_weights=jnp.zeros((SLDS_NUM_STATES, EMISSION_DIM, 1)),
+            initialized=True,
+        ),
+    )
+
+
 def _run_cdlgssm_filter(requested_fields):
     _, params = _initialize_model(ContDiscreteLinearGaussianSSM)
     return cdlgssm_filter(
@@ -442,3 +494,15 @@ def test_filter_output_field_constants_reference_known_posterior_fields():
 def test_particle_filter_posterior_contract(case):
     posterior = case.runner()
     _assert_posterior_cdnlssm_contract(posterior)
+
+
+def test_slds_rbpf_posterior_contract():
+    params = _make_slds_params()
+    posterior = rbpfilter(RBPF_PARTICLES, params, TEST_EMISSIONS, jr.PRNGKey(0))
+    _assert_rbpfiltered_contract(posterior)
+
+
+def test_slds_optimal_rbpf_posterior_contract():
+    params = _make_slds_params()
+    posterior = rbpfilter_optimal(RBPF_PARTICLES, params, TEST_EMISSIONS, jr.PRNGKey(0))
+    _assert_rbpfiltered_contract(posterior)

tests/test_models.py

@@ -12,6 +12,12 @@
 from cd_dynamax.src.continuous_discrete_nonlinear_ssm.models import (
     ContDiscreteNonlinearSSM,
 )
+from cd_dynamax.dynamax.slds import SLDS
+from cd_dynamax.dynamax.slds.inference import (
+    DiscreteParamsSLDS,
+    LGParamsSLDS,
+    ParamsSLDS,
+)
 
 # Useful auxiliary function to sample from model with num_timesteps and t_emissions
 # and check states and emissions match
@@ -183,3 +189,60 @@ def test_cdnlssm_sample_path():
     check_cddynamax_model_sample_match(
         model=model, params=params, key=jr.PRNGKey(0), T=10
     )
+
+
+def make_slds_params(num_states, state_dim, emission_dim):
+    transition_matrix = jnp.ones((num_states, num_states)) / num_states
+    return ParamsSLDS(
+        discrete=DiscreteParamsSLDS(
+            initial_distribution=jnp.ones(num_states) / num_states,
+            transition_matrix=transition_matrix,
+            proposal_transition_matrix=transition_matrix,
+        ),
+        linear_gaussian=LGParamsSLDS(
+            initial_mean=jnp.zeros((num_states, state_dim)),
+            initial_cov=jnp.tile(jnp.eye(state_dim)[None, :, :], (num_states, 1, 1)),
+            dynamics_weights=jnp.tile(
+                (0.9 * jnp.eye(state_dim))[None, :, :], (num_states, 1, 1)
+            ),
+            dynamics_cov=jnp.tile(
+                (0.1 * jnp.eye(state_dim))[None, :, :], (num_states, 1, 1)
+            ),
+            dynamics_bias=jnp.zeros((num_states, state_dim)),
+            dynamics_input_weights=jnp.zeros((num_states, state_dim, 1)),
+            emission_weights=jnp.ones((num_states, emission_dim, state_dim)),
+            emission_cov=jnp.tile(
+                (0.1 * jnp.eye(emission_dim))[None, :, :], (num_states, 1, 1)
+            ),
+            emission_bias=jnp.zeros((num_states, emission_dim)),
+            emission_input_weights=jnp.zeros((num_states, emission_dim, 1)),
+            initialized=True,
+        ),
+    )
+
+
+def test_slds_basic_init():
+    """Test that an SLDS can be created and has correct dimensions."""
+    model = SLDS(num_states=2, state_dim=3, emission_dim=2)
+    assert model.num_states == 2
+    assert model.state_dim == 3
+    assert model.emission_dim == 2
+    assert model.input_dim == 1
+    assert model.inputs_shape == (1,)
+
+
+def test_slds_sample():
+    """Run a short SLDS forward sample and check output shapes."""
+    model = SLDS(num_states=2, state_dim=3, emission_dim=2)
+    params = make_slds_params(
+        num_states=model.num_states,
+        state_dim=model.state_dim,
+        emission_dim=model.emission_dim,
+    )
+    dstates, cstates, emissions = model.sample(
+        params, key=jr.PRNGKey(0), num_timesteps=10, inputs=jnp.zeros((10, 1))
+    )
+
+    assert dstates.shape == (10,)
+    assert cstates.shape == (10, model.state_dim)
+    assert emissions.shape == (10, model.emission_dim)