From 466f9a8c66cc69547a328e801a71c9c815448417 Mon Sep 17 00:00:00 2001 From: nkbranigan <81596699+nkbranigan@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:46:42 -0700 Subject: [PATCH] Filter with a single scan and a clamped final predict index --- dynamax/generalized_gaussian_ssm/inference.py | 12 +++++------- dynamax/linear_gaussian_ssm/inference.py | 13 +++++-------- dynamax/linear_gaussian_ssm/info_inference.py | 19 ++++++++----------- .../nonlinear_gaussian_ssm/inference_ekf.py | 13 +++++-------- .../nonlinear_gaussian_ssm/inference_ukf.py | 13 +++++-------- 5 files changed, 28 insertions(+), 42 deletions(-) diff --git a/dynamax/generalized_gaussian_ssm/inference.py b/dynamax/generalized_gaussian_ssm/inference.py index 8f1daf36e..d35bbbdfb 100644 --- a/dynamax/generalized_gaussian_ssm/inference.py +++ b/dynamax/generalized_gaussian_ssm/inference.py @@ -320,7 +320,7 @@ def conditional_moments_gaussian_filter(model_params: ParamsGGSSM, # Emission distribution emission_dist = model_params.emission_dist - def _step(carry, t, predict=True): + def _step(carry, t): """One step of the CMGF""" ll, pred_mean, pred_cov = carry @@ -334,17 +334,15 @@ def _step(carry, t, predict=True): ll += log_likelihood # Predict the next state - 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) + t_next = jnp.minimum(t + 1, num_timesteps - 1) + Q = _get_params(model_params.dynamics_covariance, 2, t_next) + pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, inputs[t_next], 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) - 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)] + (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps)) return PosteriorGSSMFiltered(marginal_loglik=ll, filtered_means=filtered_means, filtered_covariances=filtered_covs) diff --git a/dynamax/linear_gaussian_ssm/inference.py b/dynamax/linear_gaussian_ssm/inference.py index 6c98235ad..d2d7a12ba 100644 --- a/dynamax/linear_gaussian_ssm/inference.py +++ b/dynamax/linear_gaussian_ssm/inference.py @@ -496,7 +496,7 @@ def _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y): return MVNLowRank(m, R, L).log_prob(y) - def _step(carry, t, predict=True): + def _step(carry, t): """Run one step of the Kalman filter.""" ll, pred_mean, pred_cov = carry @@ -512,18 +512,15 @@ def _step(carry, t, predict=True): filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y) # Predict the next state - if predict: - F, B, b, Q = _get_params(params, num_timesteps, t + 1)[:4] - pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, F, B, b, Q, inputs[t + 1]) + t_next = jnp.minimum(t + 1, num_timesteps - 1) + F, B, b, Q = _get_params(params, num_timesteps, t_next)[:4] + pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, F, B, b, Q, inputs[t_next]) return (ll, pred_mean, pred_cov), (filtered_mean, filtered_cov) # Run the Kalman filter carry = (0.0, params.initial.mean, params.initial.cov) - carry, (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps - 1)) - (ll, _, _), (filtered_mean, filtered_cov) = _step(carry, num_timesteps - 1, predict=False) - filtered_means = jnp.concatenate((filtered_means, filtered_mean[None])) - filtered_covs = jnp.concatenate((filtered_covs, filtered_cov[None])) + (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps)) return PosteriorGSSMFiltered(marginal_loglik=ll, filtered_means=filtered_means, filtered_covariances=filtered_covs) diff --git a/dynamax/linear_gaussian_ssm/info_inference.py b/dynamax/linear_gaussian_ssm/info_inference.py index 7d31d7328..0ef554fb0 100644 --- a/dynamax/linear_gaussian_ssm/info_inference.py +++ b/dynamax/linear_gaussian_ssm/info_inference.py @@ -192,7 +192,7 @@ def lgssm_info_filter( if param is not None and not callable(param) and param.ndim == ndim + 1: assert param.shape[0] == num_timesteps inputs = jnp.zeros((num_timesteps, 0)) if inputs is None else inputs - def _filter_step(carry, t, predict=True): + def _filter_step(carry, t): """Run a single step of the Kalman filter.""" ll, pred_eta, pred_prec = carry @@ -212,22 +212,19 @@ def _filter_step(carry, t, predict=True): filtered_eta, filtered_prec = _info_condition_on(pred_eta, pred_prec, H, R_prec, D, u, d, y) # Predict the next state - if predict: - F = _get_params(params.dynamics_weights, 2, t + 1) - Q_prec = _get_params(params.dynamics_precision, 2, t + 1) - B = _get_params(params.dynamics_input_weights, 2, t + 1) - b = _get_params(params.dynamics_bias, 1, t + 1) - pred_eta, pred_prec = _info_predict(filtered_eta, filtered_prec, F, Q_prec, B, inputs[t + 1], b) + t_next = jnp.minimum(t + 1, num_timesteps - 1) + F = _get_params(params.dynamics_weights, 2, t_next) + Q_prec = _get_params(params.dynamics_precision, 2, t_next) + B = _get_params(params.dynamics_input_weights, 2, t_next) + b = _get_params(params.dynamics_bias, 1, t_next) + pred_eta, pred_prec = _info_predict(filtered_eta, filtered_prec, F, Q_prec, B, inputs[t_next], b) return (ll, pred_eta, pred_prec), (filtered_eta, filtered_prec) # Run the Kalman filter initial_eta = params.initial_precision @ params.initial_mean carry = (0.0, initial_eta, params.initial_precision) - carry, (filtered_etas, filtered_precisions) = lax.scan(_filter_step, carry, jnp.arange(num_timesteps - 1)) - (ll, _, _), (filtered_eta, filtered_prec) = _filter_step(carry, num_timesteps - 1, predict=False) - filtered_etas = jnp.concatenate((filtered_etas, filtered_eta[None])) - filtered_precisions = jnp.concatenate((filtered_precisions, filtered_prec[None])) + (ll, _, _), (filtered_etas, filtered_precisions) = lax.scan(_filter_step, carry, jnp.arange(num_timesteps)) return PosteriorGSSMInfoFiltered(marginal_loglik=ll, filtered_etas=filtered_etas, filtered_precisions=filtered_precisions) diff --git a/dynamax/nonlinear_gaussian_ssm/inference_ekf.py b/dynamax/nonlinear_gaussian_ssm/inference_ekf.py index 71a553060..275502255 100644 --- a/dynamax/nonlinear_gaussian_ssm/inference_ekf.py +++ b/dynamax/nonlinear_gaussian_ssm/inference_ekf.py @@ -121,7 +121,7 @@ def extended_kalman_filter(params: ParamsNLGSSM, f, h, F, H = (_process_fn(fn, inputs) for fn in (f, h, F, H)) inputs = _process_input(inputs, num_timesteps) - def _step(carry, t, predict=True): + def _step(carry, t): """Iteratively update the state estimate and log likelihood.""" ll, pred_mean, pred_cov = carry @@ -138,10 +138,9 @@ def _step(carry, t, predict=True): filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, h, H, R, u, y, num_iter) # Predict the next state - if predict: - Q = _get_params(params.dynamics_covariance, 2, t + 1) - u_next = inputs[t + 1] if not predict_final else u - pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, f, F, Q, u_next) + t_next = jnp.minimum(t + 1, num_timesteps - 1) + Q = _get_params(params.dynamics_covariance, 2, t_next) + pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, f, F, Q, inputs[t_next]) # Build carry and output states carry = (ll, pred_mean, pred_cov) @@ -158,9 +157,7 @@ def _step(carry, t, predict=True): # Run the extended Kalman filter carry = (0.0, params.initial_mean, params.initial_covariance) - carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1)) - (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=predict_final) - outputs = {key: jnp.concatenate((val, final_outputs[key][None])) for key, val in outputs.items()} + (ll, _, _), outputs = lax.scan(_step, carry, jnp.arange(num_timesteps)) if not predict_final: outputs = {key: val[:-1] if key.startswith("predicted") else val for key, val in outputs.items()} outputs = {"marginal_loglik": ll, **outputs} diff --git a/dynamax/nonlinear_gaussian_ssm/inference_ukf.py b/dynamax/nonlinear_gaussian_ssm/inference_ukf.py index 11371fa38..7e5437714 100644 --- a/dynamax/nonlinear_gaussian_ssm/inference_ukf.py +++ b/dynamax/nonlinear_gaussian_ssm/inference_ukf.py @@ -199,7 +199,7 @@ def unscented_kalman_filter(params: ParamsNLGSSM, f, h = (_process_fn(fn, inputs) for fn in (f, h)) inputs = _process_input(inputs, num_timesteps) - def _step(carry, t, predict=True): + def _step(carry, t): """One step of the UKF""" ll, pred_mean, pred_cov = carry @@ -217,10 +217,9 @@ def _step(carry, t, predict=True): ll += log_likelihood # Predict the next state - if predict: - Q = _get_params(params.dynamics_covariance, 2, t + 1) - u_next = inputs[t + 1] if not predict_final else u - pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, lamb, w_mean, w_cov, u_next) + t_next = jnp.minimum(t + 1, num_timesteps - 1) + Q = _get_params(params.dynamics_covariance, 2, t_next) + pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, lamb, w_mean, w_cov, inputs[t_next]) # Build carry and output states carry = (ll, pred_mean, pred_cov) @@ -237,9 +236,7 @@ def _step(carry, t, predict=True): # Run the Unscented Kalman Filter carry = (0.0, params.initial_mean, params.initial_covariance) - carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1)) - (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=predict_final) - outputs = {key: jnp.concatenate((val, final_outputs[key][None])) for key, val in outputs.items()} + (ll, _, _), outputs = lax.scan(_step, carry, jnp.arange(num_timesteps)) if not predict_final: outputs = {key: val[:-1] if key.startswith("predicted") else val for key, val in outputs.items()} outputs = {"marginal_loglik": ll, **outputs}