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..3a858603 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,14 +116,23 @@ 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 ) # 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"])