From bc9ce951cc731c9f46769c008fecdca0561bd60b Mon Sep 17 00:00:00 2001 From: Dan Waxman Date: Thu, 1 Oct 2026 15:13:12 -0400 Subject: [PATCH 1/2] Fix indexing errors in particle smoothers Resolves #281. This fixes an indexing bug in the particle smoothers, which results in the incorrect smoothing distribution. Concretely, `state_2.model_inputs` was previously used as part of the `log_potential`: https://github.com/state-space-models/cuthbert/blob/ec53db22e17b6ab29975e41af8a0758a2a6ff665/cuthbert/smc/backward_sampler.py#L181-L188 But the `model_inputs` convention for particle smoothers differ from the rest of `cuthbert`, and `state_1.model_inputs` should be used accordingly: https://github.com/state-space-models/cuthbert/blob/ec53db22e17b6ab29975e41af8a0758a2a6ff665/cuthbert/smc/backward_sampler.py#L128-L137 The proper `state_2.model_inputs` is now properly used. This issue wasn't found by the tests because the model definition was wrong in tests: in particular, backward sampling smoothers require the joint log potential, but only the observation log potential was given. We now use the proper model definition, which would fail before but passes now. Finally, there was a key reshape that caused issues with typed keys. We now split with the correct shape directly, which makes this compatible with both JAX key types. --- cuthbert/smc/backward_sampler.py | 2 +- cuthbertlib/smc/smoothing/mcmc.py | 2 +- tests/cuthbert/smc/test_backward_sampler.py | 16 +++++++++++----- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/cuthbert/smc/backward_sampler.py b/cuthbert/smc/backward_sampler.py index c16c64c0..96931c12 100644 --- a/cuthbert/smc/backward_sampler.py +++ b/cuthbert/smc/backward_sampler.py @@ -183,7 +183,7 @@ def smoother_combine( x0_all=state_1.particles, x1_all=state_2.particles, log_weight_x0_all=state_1.log_weights, - log_density=lambda s1, s2: log_potential(s1, s2, state_2.model_inputs), + log_density=lambda s1, s2: log_potential(s1, s2, state_1.model_inputs), x1_ancestor_indices=state_2.ancestor_indices, ) diff --git a/cuthbertlib/smc/smoothing/mcmc.py b/cuthbertlib/smc/smoothing/mcmc.py index 86ef08ad..d55c4c00 100644 --- a/cuthbertlib/smc/smoothing/mcmc.py +++ b/cuthbertlib/smc/smoothing/mcmc.py @@ -49,7 +49,7 @@ def simulate( ) n_samples = x1_ancestor_indices.shape[0] - keys = random.split(key, (n_steps * 2)).reshape((n_steps, 2)) + keys = random.split(key, (n_steps, 2)) def body(carry, keys_t): # IMH proposal diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py index efdaadd3..2526e5ba 100644 --- a/tests/cuthbert/smc/test_backward_sampler.py +++ b/tests/cuthbert/smc/test_backward_sampler.py @@ -42,6 +42,12 @@ def log_potential(state_prev, state, model_inputs: int): Hs[idx] @ state + ds[idx], ys[idx], chol_Rs[idx], nan_support=False ) + def log_joint(state_prev, state, model_inputs: int): + idx = model_inputs - 1 + return log_potential(state_prev, state, model_inputs) + logpdf( + Fs[idx] @ state_prev + cs[idx], state, chol_Qs[idx], nan_support=False + ) + n_filter_particles = 5000 resampling_fn = systematic.resampling ess_threshold = 0.7 @@ -54,7 +60,7 @@ def log_potential(state_prev, state, model_inputs: int): adaptive_resampler, ) model_inputs = jnp.arange(len(ys) + 1) - return filter_obj, model_inputs, log_potential + return filter_obj, model_inputs, log_joint class Test(chex.TestCase): @@ -72,7 +78,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method): ) # Run the particle filter. - filter_obj, model_inputs, log_potential = load_inference( + filter_obj, model_inputs, log_joint = load_inference( m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys ) init_key, filter_key, smoother_key = random.split(random.key(seed + 1), 3) @@ -95,9 +101,9 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method): raise ValueError(f"{method} is not a valid backward sampling method.") # Run the particle smoother. - n_smoother_particles = 1000 + n_smoother_particles = 5000 smoother_obj = build_smoother( - log_potential, bs_fn, systematic.resampling, n_smoother_particles + log_joint, bs_fn, systematic.resampling, n_smoother_particles ) smoothed_states = self.variant( smoother, static_argnames=("smoother_obj", "parallel") @@ -110,7 +116,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method): filtered_means, filtered_covs, _ = std_kalman_filter( m0, P0, Fs, cs, Qs, Hs, ds, Rs, ys ) - (smoothed_means, _), _ = std_kalman_smoother( + (smoothed_means, smoothed_covs), _ = std_kalman_smoother( filtered_means, filtered_covs, Fs, cs, Qs ) From 28a5505b1f33b3cdf50edb50964df42f3f26d62b Mon Sep 17 00:00:00 2001 From: Dan Waxman Date: Fri, 2 Oct 2026 08:59:26 -0400 Subject: [PATCH 2/2] Tighten tests Tightens the MSE threshold, and tests covs. --- tests/cuthbert/smc/test_backward_sampler.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py index 2526e5ba..3a858603 100644 --- a/tests/cuthbert/smc/test_backward_sampler.py +++ b/tests/cuthbert/smc/test_backward_sampler.py @@ -123,7 +123,16 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method): # Compare smoothed particles with Kalman smoother results particle_means = jnp.mean(smoothed_states.particles, axis=1) mse = jnp.mean(jnp.square(particle_means - smoothed_means)) - assert mse <= 0.1, "Mean squared error is too high" + assert mse <= 0.05, "Mean squared error is too high" + particle_covs = jax.vmap(lambda particles: jnp.cov(particles.T))( + smoothed_states.particles + ) + chex.assert_trees_all_close( + particle_covs, + smoothed_covs, + atol=0.9 if method == "tracing" else 0.25, + rtol=0.0, + ) @chex.variants(with_jit=True, without_jit=True) @parameterized.product(method=["tracing", "exact", "mcmc"])