Fix off-by-one bugs

dynamax/generalized_gaussian_ssm/inference.py

@@ -183,11 +183,11 @@ def _predict(prior_mean: Float[Array, "state_dim"],
         p(x_{t+1}) = N(x_{t+1} | mu_pred, Sigma_pred)
         where
             mu_pred = gaussian_expec(f, m, P)
-                    \approx \int f(x_t, u) N(x_t | m, P) dx_t
-            Sigma_pred = gaussian_crosscov(f(x_t, u_t), f(x_t, u_t); m, P) + Q
-                       \approx \int (f(x_t, u) - mu_pred)(f(x_t, u) - mu_pred)^T N(x_t | m, P)dx_t + Q
-            cross_pred = gaussian_crosscov(x_t, f(x_t, u_t); m, P)
-                       \approx \int (x_t - m)(f(x_t, u_t) - mu_pred)^T N(x_t | m, P)dx_t
+                    \approx \int f(x_t, u_{t+1}) N(x_t | m, P) dx_t
+            Sigma_pred = gaussian_crosscov(f(x_t, u_{t+1}), f(x_t, u_{t+1}); m, P) + Q
+                       \approx \int (f(x_t, u_{t+1}) - mu_pred)(f(x_t, u_{t+1}) - mu_pred)^T N(x_t | m, P)dx_t + Q
+            cross_pred = gaussian_crosscov(x_t, f(x_t, u_{t+1}); m, P)
+                       \approx \int (x_t - m)(f(x_t, u_{t+1}) - mu_pred)^T N(x_t | m, P)dx_t
 
     Args:
         prior_mean (state_dim,): prior mean.
@@ -232,7 +232,7 @@ def _condition_on(prior_mean: Float[Array, "state_dim"],
        Gaussian approximation.
 
        p(x_t | y_t, u_t, y_{1:t-1}, u_{1:t-1})
-         propto p(x_t | y_{1:t-1}, u_{1:t-1}) p(y_t | x_t, u_t)
+         propto p(x_t | y_{1:t-1}, u_{1:t}) p(y_t | x_t, u_t)
          \approx N(x_t | m, P) ArbitraryDist(y_t | emission_mean(x_t, u_t), emission_cov(x_t, u_t))
          \approx N(x_t | mu_cond, Sigma_cond)
 
@@ -305,6 +305,7 @@ def conditional_moments_gaussian_filter(model_params: ParamsGGSSM,
 
     """
     num_timesteps = len(emissions)
+    assert model_params.dynamics_covariance.ndim != 3 or model_params.dynamics_covariance.shape[0] == num_timesteps
 
     # Process dynamics function and conditional emission moments to take in control inputs
     f = model_params.dynamics_function
@@ -319,12 +320,11 @@ def conditional_moments_gaussian_filter(model_params: ParamsGGSSM,
     # Emission distribution
     emission_dist = model_params.emission_dist
 
-    def _step(carry, t):
+    def _step(carry, t, predict=True):
         """One step of the CMGF"""
         ll, pred_mean, pred_cov = carry
 
-        # Get parameters and inputs for time index t
-        Q = _get_params(model_params.dynamics_covariance, 2, t)
+        # Get the input and observation at time t
         u = inputs[t]
         y = emissions[t]
 
@@ -334,13 +334,17 @@ def _step(carry, t):
         ll += log_likelihood
 
         # Predict the next state
-        pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, u, g_ev, g_cov)
+        if predict:
+            Q = _get_params(model_params.dynamics_covariance, 2, t + 1)
+            pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, inputs[t + 1], g_ev, g_cov)
 
         return (ll, pred_mean, pred_cov), (filtered_mean, filtered_cov)
 
     # Run the general linearization filter
     carry = (0.0, model_params.initial_mean, model_params.initial_covariance)
-    (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps))
+    carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1))
+    (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=False)
+    filtered_means, filtered_covs = [jnp.concatenate((a, b[None])) for a, b in zip(outputs, final_outputs)]
     return PosteriorGSSMFiltered(marginal_loglik=ll, 
                                  filtered_means=filtered_means, 
                                  filtered_covariances=filtered_covs)
@@ -366,6 +370,7 @@ def conditional_moments_gaussian_smoother(model_params: ParamsGGSSM,
 
     """
     num_timesteps = len(emissions)
+    assert model_params.dynamics_covariance.ndim != 3 or model_params.dynamics_covariance.shape[0] == num_timesteps
 
     # Get filtered posterior
     if filtered_posterior is None:
@@ -386,9 +391,9 @@ def _step(carry, args):
         smoothed_mean_next, smoothed_cov_next = carry
         t, filtered_mean, filtered_cov = args
 
-        # Get parameters and inputs for time index t
-        Q = _get_params(model_params.dynamics_covariance, 2, t)
-        u = inputs[t]
+        # Get parameters and inputs for time index t + 1
+        Q = _get_params(model_params.dynamics_covariance, 2, t + 1)
+        u = inputs[t + 1]
 
         # Prediction step
         pred_mean, pred_cov, pred_cross = _predict(filtered_mean, filtered_cov, f, Q, u, g_ev, g_cov)

dynamax/generalized_gaussian_ssm/inference_test.py

@@ -1,13 +1,16 @@
 """
 Tests for inference in the generalized Gaussian SSM.
 """
+import numpy as np
+import pytest
 import jax.numpy as jnp
+from jax import jit
 
 from dynamax.generalized_gaussian_ssm.models import ParamsGGSSM
-from dynamax.generalized_gaussian_ssm.inference import conditional_moments_gaussian_smoother, EKFIntegrals, UKFIntegrals
+from dynamax.generalized_gaussian_ssm.inference import conditional_moments_gaussian_filter, conditional_moments_gaussian_smoother, EKFIntegrals, UKFIntegrals, GHKFIntegrals
 from dynamax.nonlinear_gaussian_ssm.inference_ekf import extended_kalman_smoother
 from dynamax.nonlinear_gaussian_ssm.inference_ukf import unscented_kalman_smoother, UKFHyperParams
-from dynamax.nonlinear_gaussian_ssm.inference_test_utils import random_nlgssm_args
+from dynamax.nonlinear_gaussian_ssm.inference_test_utils import random_nlgssm_args, make_affine_timing_case
 from dynamax.utils.utils import has_tpu
 from functools import partial
 
@@ -71,4 +74,36 @@ def test_ukf(key=1, num_timesteps=15):
     assert allclose(ukf_post.filtered_means, ggf_post.filtered_means)
     assert allclose(ukf_post.filtered_covariances, ggf_post.filtered_covariances)
     assert allclose(ukf_post.smoothed_means, ggf_post.smoothed_means)
-    assert allclose(ukf_post.smoothed_covariances, ggf_post.smoothed_covariances)
+    assert allclose(ukf_post.smoothed_covariances, ggf_post.smoothed_covariances)
+
+@pytest.mark.parametrize("num_timesteps", [1, 3])
+@pytest.mark.parametrize("integrals", [EKFIntegrals(), UKFIntegrals(), GHKFIntegrals(order=3)])
+def test_generalized_gaussian_inference_input_indexing(num_timesteps, integrals):
+    """Check GGSSM input and parameter indexing against exact LGSSM filtering and smoothing results.
+
+    Also check that incorrect dynamics covariance lengths are rejected.
+    """
+    params, emissions, inputs, expected = make_affine_timing_case(num_timesteps)
+    generalized_params = ParamsGGSSM(
+        initial_mean=params.initial_mean,
+        initial_covariance=params.initial_covariance,
+        dynamics_function=params.dynamics_function,
+        dynamics_covariance=params.dynamics_covariance,
+        emission_mean_function=params.emission_function,
+        emission_cov_function=lambda state, inpt: params.emission_covariance,
+    )
+    posterior = jit(lambda y, u: conditional_moments_gaussian_smoother(
+        generalized_params, integrals, y, inputs=u
+    ))(emissions, inputs)
+    for field in posterior._fields:
+        value = getattr(posterior, field)
+        if value is not None:
+            np.testing.assert_allclose(value, expected[field], atol=2e-5, rtol=2e-5)
+
+    # Check dynamics covariance lengths in both filtering and smoothing.
+    short_params = generalized_params._replace(dynamics_covariance=params.dynamics_covariance[:-1])
+    with pytest.raises(AssertionError):
+        conditional_moments_gaussian_filter(short_params, integrals, emissions, inputs=inputs)
+    with pytest.raises(AssertionError):
+        conditional_moments_gaussian_smoother(
+            short_params, integrals, emissions, filtered_posterior=posterior, inputs=inputs)

dynamax/hidden_markov_model/inference.py

@@ -127,7 +127,7 @@ def hmm_filter(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -175,7 +175,7 @@ def hmm_backward_filter(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -224,7 +224,7 @@ def hmm_two_filter_smoother(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -280,7 +280,7 @@ def hmm_smoother(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -362,7 +362,7 @@ def hmm_fixed_lag_smoother(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         window_size: the number of future steps to use, $L$
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
@@ -380,7 +380,7 @@ def _step(carry, t):
 
         # Get parameters for time t
         A_fwd = get_trans_mat(transition_matrix, transition_fn, t-1)
-        A_bwd = get_trans_mat(transition_matrix, transition_fn, t)
+        A_bwd = A_fwd
         ll = log_likelihoods[t]
 
         # Shift window forward by 1
@@ -468,7 +468,7 @@ def hmm_posterior_mode(
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -519,7 +519,7 @@ def hmm_posterior_sample(
     Args:
         rng: random number generator
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
         transition_fn: function that takes in an integer time index and returns a $K \times K$ transition matrix.
 
@@ -554,7 +554,7 @@ def _step(carry, args):
     keys = jr.split(key, num_timesteps)
     last_state = jr.choice(keys[-1], a=num_states, p=filtered_probs[-1])
     _, states = lax.scan(
-        _step, last_state, (jnp.arange(1, num_timesteps), keys[:-1], filtered_probs[:-1]),
+        _step, last_state, (jnp.arange(num_timesteps - 1), keys[:-1], filtered_probs[:-1]),
         reverse=True
     )
 

dynamax/hidden_markov_model/inference_test.py

@@ -366,3 +366,25 @@ def test_parallel_posterior_sample(
     # Compare the joint distributions
     assert jnp.allclose(blj_sample, blj, rtol=0, atol=eps)
 
+
+
+def test_hmm_posterior_sample_nonstationary():
+    """Check HMM posterior sampling transition indexing:
+    transitions[t] connects states[t] and states[t + 1].
+    """
+    transitions = jnp.stack([jnp.eye(2), jnp.fliplr(jnp.eye(2))])
+    _, states = core.hmm_posterior_sample(
+        jr.PRNGKey(0), jnp.array([0.5, 0.5]), transitions, jnp.zeros((3, 2)))
+    assert states[0] == states[1]
+    assert states[1] != states[2]
+
+
+def test_hmm_fixed_lag_smoother_nonstationary():
+    """Check HMM fixed-lag smoothing transition indexing:
+    both messages use transition index t - 1 at time t.
+    """
+    transitions = jnp.stack([jnp.eye(2), jnp.fliplr(jnp.eye(2))])
+    log_likelihoods = jnp.log(jnp.array([[1., 1.], [0.9, 0.1], [1., 1.]]))
+    posterior = core.hmm_fixed_lag_smoother(
+        jnp.array([0.5, 0.5]), transitions, log_likelihoods, window_size=2)
+    assert jnp.allclose(posterior.smoothed_probs[1], jnp.array([[0.9, 0.1], [0.9, 0.1]]))

dynamax/hidden_markov_model/models/abstractions.py

@@ -230,7 +230,7 @@ def distribution(self,
         Args:
             params: transition parameters
             state: current latent state
-            inputs: current inputs
+            inputs: next inputs
 
         Returns:
             conditional distribution of next state.
@@ -602,7 +602,7 @@ def _inference_args(self, params: HMMParameterSet,
                         emissions: Array, 
                         inputs: Optional[Array]) -> Tuple:
         """Return the arguments needed for inference."""
-        return (self.initial_component._compute_initial_probs(params.initial, inputs),
+        return (self.initial_component._compute_initial_probs(params.initial, pytree_slice(inputs, 0)),
                 self.transition_component._compute_transition_matrices(params.transitions, inputs),
                 self.emission_component._compute_conditional_logliks(params.emissions, emissions, inputs))
 

dynamax/hidden_markov_model/models/test_models.py

@@ -5,6 +5,9 @@
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+from dynamax.hidden_markov_model.models.abstractions import HMMInitialState
+from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState
+from tensorflow_probability.substrates.jax import distributions as tfd
 
 from jax import vmap
 from jax.tree_util import tree_leaves, tree_map
@@ -369,3 +372,20 @@ def init_kmeans(key, emissions, inputs):
     vmapped_params = vmap(init_kmeans)(batch_keys, batch_emissions, batch_inputs)
     assert vmapped_params.emissions.biases.shape == (batch_size, num_states)
     assert jnp.all(jnp.isfinite(vmapped_params.emissions.biases))
+
+
+def test_hmm_input_dependent_initial_distribution():
+    """Check that HMM inference uses the first input to compute initial probabilities."""
+    class InputInitialState(StandardHMMInitialState):
+        _compute_initial_probs = HMMInitialState._compute_initial_probs
+
+        def distribution(self, params, inputs=None):
+            return tfd.Categorical(logits=inputs)
+
+    hmm = models.GaussianHMM(num_states=2, emission_dim=1)
+    params, _ = hmm.initialize(jr.PRNGKey(0))
+    hmm.initial_component = InputInitialState(2)
+    inputs = jnp.array([[3., 0.], [0., 3.], [2., 1.]])
+    posterior = hmm.filter(params, jnp.zeros((3, 1)), inputs=inputs)
+    expected = hmm.initial_distribution(params, inputs[0]).probs_parameter()
+    assert jnp.allclose(posterior.predicted_probs[0], expected)

dynamax/hidden_markov_model/parallel_inference.py

@@ -49,7 +49,7 @@ def hmm_filter(initial_probs: Float[Array, " num_states"],
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
 
     Returns:
@@ -114,7 +114,7 @@ def hmm_smoother(initial_probs: Float[Array, " num_states"],
 
     Args:
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
 
     Returns:
@@ -185,7 +185,7 @@ def hmm_posterior_sample(key: Array,
     Args:
         key: random number generator
         initial_distribution: $p(z_1 \mid u_1, \theta)$
-        transition_matrix: $p(z_{t+1} \mid z_t, u_t, \theta)$
+        transition_matrix: $p(z_{t+1} \mid z_t, u_{t+1}, \theta)$
         log_likelihoods: $p(y_t \mid z_t, u_t, \theta)$ for $t=1,\ldots, T$.
 
     Returns: