diff --git a/cd_dynamax/dynamax/slds/inference.py b/cd_dynamax/dynamax/slds/inference.py index 66aa5af6..3945a0ef 100644 --- a/cd_dynamax/dynamax/slds/inference.py +++ b/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 \ No newline at end of file + return RBPFiltered( + marginal_loglik=out["marginal_loglik"].sum(), + weights=out["weights"], + states=out["states"], + means=out["means"], + covariances=out["covariances"]) \ No newline at end of file diff --git a/tests/test_filter_posterior_contracts.py b/tests/test_filter_posterior_contracts.py index df499ff1..d5ffe1e1 100644 --- a/tests/test_filter_posterior_contracts.py +++ b/tests/test_filter_posterior_contracts.py @@ -41,6 +41,14 @@ 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 @@ -48,6 +56,8 @@ 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) diff --git a/tests/test_models.py b/tests/test_models.py index 5dd2c27e..ef7fd75a 100644 --- a/tests/test_models.py +++ b/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) diff --git a/tests/test_slds_rbpf_mll_match.py b/tests/test_slds_rbpf_mll_match.py new file mode 100644 index 00000000..041b067d --- /dev/null +++ b/tests/test_slds_rbpf_mll_match.py @@ -0,0 +1,151 @@ +import jax.numpy as jnp +import jax.random as jr +import pytest + +from cd_dynamax.dynamax.linear_gaussian_ssm.inference import lgssm_filter +from cd_dynamax.dynamax.linear_gaussian_ssm.models import LinearGaussianSSM +from cd_dynamax.dynamax.slds.inference import ( + DiscreteParamsSLDS, + LGParamsSLDS, + ParamsSLDS, + rbpfilter, + rbpfilter_optimal, +) + + +NUM_STATES = 3 +STATE_DIM = 2 +EMISSION_DIM = 2 +NUM_TIMESTEPS = 8 +NUM_PARTICLES = 1000 +MLL_ATOL = 0.05 +MEAN_ATOL = 0.1 + + +@pytest.fixture +def seed(): + return 0 + + +def init_lgssm_model(): + """Initialize a small stable LGSSM for the degenerate SLDS comparison.""" + model = LinearGaussianSSM(state_dim=STATE_DIM, emission_dim=EMISSION_DIM) + params, _ = model.initialize( + initial_mean=jnp.array([0.2, -0.1]), + initial_covariance=0.5 * jnp.eye(STATE_DIM), + dynamics_weights=jnp.array([[0.75, 0.1], [-0.05, 0.8]]), + dynamics_bias=jnp.zeros(STATE_DIM), + dynamics_covariance=0.05 * jnp.eye(STATE_DIM), + emission_weights=jnp.array([[1.0, 0.3], [-0.2, 0.7]]), + emission_bias=jnp.zeros(EMISSION_DIM), + emission_covariance=0.1 * jnp.eye(EMISSION_DIM), + ) + return model, params + + +def make_degenerate_slds_params(lgssm_params): + """Duplicate one LGSSM across modes so the discrete state has no effect.""" + 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.tile(lgssm_params.initial.mean[None, :], (NUM_STATES, 1)), + initial_cov=jnp.tile( + lgssm_params.initial.cov[None, :, :], (NUM_STATES, 1, 1) + ), + dynamics_weights=jnp.tile( + lgssm_params.dynamics.weights[None, :, :], (NUM_STATES, 1, 1) + ), + dynamics_cov=jnp.tile( + lgssm_params.dynamics.cov[None, :, :], (NUM_STATES, 1, 1) + ), + dynamics_bias=jnp.tile( + lgssm_params.dynamics.bias[None, :], (NUM_STATES, 1) + ), + dynamics_input_weights=jnp.zeros((NUM_STATES, STATE_DIM, 1)), + emission_weights=jnp.tile( + lgssm_params.emissions.weights[None, :, :], (NUM_STATES, 1, 1) + ), + emission_cov=jnp.tile( + lgssm_params.emissions.cov[None, :, :], (NUM_STATES, 1, 1) + ), + emission_bias=jnp.tile( + lgssm_params.emissions.bias[None, :], (NUM_STATES, 1) + ), + emission_input_weights=jnp.zeros((NUM_STATES, EMISSION_DIM, 1)), + initialized=True, + ), + ) + + +def make_rbpf_implied_lgssm_params(lgssm_model, lgssm_params): + """Return the LGSSM prior implied by the current RBPF initialization. + + The RBPF initializes each particle by sampling an initial mean from + N(m, S) while retaining covariance S. Its first Kalman step then predicts + through the dynamics before conditioning on y[0].""" + initial_mean = ( + lgssm_params.dynamics.weights @ lgssm_params.initial.mean + + lgssm_params.dynamics.bias + ) + initial_cov = ( + 2 + * lgssm_params.dynamics.weights + @ lgssm_params.initial.cov + @ lgssm_params.dynamics.weights.T + + lgssm_params.dynamics.cov + ) + + params, _ = lgssm_model.initialize( + initial_mean=initial_mean, + initial_covariance=initial_cov, + dynamics_weights=lgssm_params.dynamics.weights, + dynamics_bias=lgssm_params.dynamics.bias, + dynamics_covariance=lgssm_params.dynamics.cov, + emission_weights=lgssm_params.emissions.weights, + emission_bias=lgssm_params.emissions.bias, + emission_covariance=lgssm_params.emissions.cov, + ) + return params + + +def test_degenerate_slds_rbpf_mll_matches_lgssm(seed): + lgssm_model, lgssm_params = init_lgssm_model() + _, emissions = lgssm_model.sample(lgssm_params, jr.PRNGKey(seed), NUM_TIMESTEPS) + + slds_params = make_degenerate_slds_params(lgssm_params) + kf_params = make_rbpf_implied_lgssm_params(lgssm_model, lgssm_params) + kf_post = lgssm_filter(kf_params, emissions) + + rbpf_post = rbpfilter(NUM_PARTICLES, slds_params, emissions, jr.PRNGKey(seed + 1)) + optimal_post = rbpfilter_optimal( + NUM_PARTICLES, slds_params, emissions, jr.PRNGKey(seed + 2) + ) + + rbpf_mean = jnp.einsum("tp,tpm->tm", rbpf_post.weights, rbpf_post.means) + optimal_mean = jnp.einsum("tp,tpm->tm", optimal_post.weights, optimal_post.means) + + print( + "KF MLL:", + kf_post.marginal_loglik, + "RBPF MLL:", + rbpf_post.marginal_loglik, + "optimal RBPF MLL:", + optimal_post.marginal_loglik, + ) + + assert jnp.isfinite(rbpf_post.marginal_loglik) + assert jnp.isfinite(optimal_post.marginal_loglik) + assert jnp.allclose( + rbpf_post.marginal_loglik, kf_post.marginal_loglik, atol=MLL_ATOL + ) + assert jnp.allclose( + optimal_post.marginal_loglik, kf_post.marginal_loglik, atol=MLL_ATOL + ) + assert jnp.allclose(rbpf_mean, kf_post.filtered_means, atol=MEAN_ATOL) + assert jnp.allclose(optimal_mean, kf_post.filtered_means, atol=MEAN_ATOL)