From bc991f6f754e4b785dc29350be2d3f7ddcfec0fe Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Tue, 9 Jun 2026 17:18:14 +0100 Subject: [PATCH 1/2] Add final factorial state to filtering output --- cuthbert/factorial/README.md | 6 ++++-- cuthbert/factorial/filtering.py | 13 +++++++++---- tests/cuthbert/factorial/test_discrete.py | 5 ++++- tests/cuthbert/factorial/test_kalman.py | 14 ++++++++++++-- tests/cuthbert/factorial/test_smc.py | 4 ++-- tests/cuthbert/factorial/test_taylor.py | 2 +- 6 files changed, 32 insertions(+), 12 deletions(-) diff --git a/cuthbert/factorial/README.md b/cuthbert/factorial/README.md index 9f67a94..df3246e 100644 --- a/cuthbert/factorial/README.md +++ b/cuthbert/factorial/README.md @@ -82,8 +82,10 @@ You can also use `cuthbert.factorial.filter` for convenient offline filtering. Note that associative/parallel filtering is not supported for factorial filtering. ```python -init_factorial_state, local_filter_states = cuthbert.factorial.filter( - kalman_filter, factorializer, model_inputs, output_factorial=False +init_factorial_state, local_filter_states, final_factorial_state = ( + cuthbert.factorial.filter( + kalman_filter, factorializer, model_inputs, output_factorial=False + ) ) ``` diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index b325434..08eac9a 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -16,7 +16,7 @@ def filter( output_factorial: bool = False, key: KeyArray | None = None, ) -> ( - ArrayTree | tuple[ArrayTree, ArrayTree] + ArrayTree | tuple[ArrayTree, ArrayTree, ArrayTree] ): # TODO: Can overload this function so the type checker knows that the output is a ArrayTree if output_factorial is True and a tuple[ArrayTree, ArrayTree] if output_factorial is False """Applies offline factorial filtering for given model inputs. @@ -40,10 +40,15 @@ def filter( with first dimension of length F and temporal dimension. The second being the local states for each time step, i.e. first dimension of length T and no factorial dimension. + The third being the final factorial state. key: The key for the random number generator. Returns: - The filtered states (NamedTuple with leading temporal dimension of length T + 1). + If output_factorial is True, returns a single state with first temporal dimension + of length T + 1 and second factorial dimension of length F. + If output_factorial is False, returns a tuple of + (initial factorial state, local states for each time step, + final factorial state). """ T = tree.leaves(model_inputs)[0].shape[0] - 1 @@ -105,9 +110,9 @@ def body_factorial(prev_factorial_state, prep_inp_and_k): return factorial_states else: - _, local_states = scan( + final_factorial_state, local_states = scan( body_local, init_factorial_state, (prep_model_inputs, prepare_keys[1:]), ) - return init_factorial_state, local_states + return init_factorial_state, local_states, final_factorial_state diff --git a/tests/cuthbert/factorial/test_discrete.py b/tests/cuthbert/factorial/test_discrete.py index c4a02d8..3227cd5 100644 --- a/tests/cuthbert/factorial/test_discrete.py +++ b/tests/cuthbert/factorial/test_discrete.py @@ -160,7 +160,7 @@ def test_factorial_discrete_filter( ) # output_factorial=False - init_state, local_filter_states = factorial.filter( + init_state, local_filter_states, final_state = factorial.filter( filter_obj, factorializer, model_inputs, output_factorial=False ) chex.assert_trees_all_close(init_state.dist, model_params[0], rtol=1e-10, atol=0.0) @@ -174,6 +174,9 @@ def test_factorial_discrete_filter( rtol=1e-10, atol=0.0, ) + chex.assert_trees_all_close( + final_state.dist, factorial_dists_ref[-1], rtol=1e-10, atol=0.0 + ) # output_factorial=True factorial_states = factorial.filter( diff --git a/tests/cuthbert/factorial/test_kalman.py b/tests/cuthbert/factorial/test_kalman.py index bd050d0..428ca92 100644 --- a/tests/cuthbert/factorial/test_kalman.py +++ b/tests/cuthbert/factorial/test_kalman.py @@ -210,7 +210,7 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste fac_covs_t_all = jnp.stack(fac_covs_t_all) # Check output_factorial = False - init_state, local_filter_states = factorial.filter( + init_state, local_filter_states, final_state = factorial.filter( filter_obj, factorializer, model_inputs, output_factorial=False ) local_filter_covs = ( @@ -228,6 +228,16 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste local_filter_states.log_normalizing_constant, ), ) + chex.assert_trees_all_close( + ( + final_state.mean, + final_state.chol_cov @ final_state.chol_cov.transpose(0, 2, 1), + final_state.log_normalizing_constant, + ), + (fac_means_t_all[-1], fac_covs_t_all[-1], ells[-1]), + rtol=1e-10, + atol=0.0, + ) # Check output_factorial = True factorial_filtering_states = factorial.filter( @@ -287,7 +297,7 @@ def test_smoother( smoother_factorial_index=smoother_factorial_index, ) - init_state, local_filter_states = factorial.filter( + init_state, local_filter_states, _ = factorial.filter( filter_obj, factorializer, filter_model_inputs, output_factorial=False ) diff --git a/tests/cuthbert/factorial/test_smc.py b/tests/cuthbert/factorial/test_smc.py index 118fd14..85a8776 100644 --- a/tests/cuthbert/factorial/test_smc.py +++ b/tests/cuthbert/factorial/test_smc.py @@ -139,7 +139,7 @@ def test_factorial_smc_filter( output_factorial=True, ), ) - _, kalman_local_states = factorial.filter( + _, kalman_local_states, _ = factorial.filter( kalman_filter, kalman_factorializer, kalman_model_inputs, output_factorial=False ) kalman_covs = kalman_states.chol_cov @ kalman_states.chol_cov.transpose(0, 1, 3, 2) @@ -157,7 +157,7 @@ def test_factorial_smc_filter( key=random.key(seed + 123), ), ) - _, smc_local_states = factorial.filter( + _, smc_local_states, _ = factorial.filter( smc_filter, smc_factorializer, smc_model_inputs, diff --git a/tests/cuthbert/factorial/test_taylor.py b/tests/cuthbert/factorial/test_taylor.py index d396d0d..9f975fe 100644 --- a/tests/cuthbert/factorial/test_taylor.py +++ b/tests/cuthbert/factorial/test_taylor.py @@ -62,7 +62,7 @@ def observation_log_density(x, y): ) filter_model_inputs = jnp.arange(num_time_steps + 1) - init_state, local_filter_states = jax.jit( + init_state, local_filter_states, _ = jax.jit( factorial.filter, static_argnames=("filter_obj", "factorializer", "output_factorial"), )(filter_obj, factorializer, filter_model_inputs, output_factorial=False) From ba638e777f14c73c62ee87994e8674602828a014 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:37:21 +0000 Subject: [PATCH 2/2] Fix factorial filter review doc clarifications --- cuthbert/factorial/filtering.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index 08eac9a..cc02ae3 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -17,7 +17,7 @@ def filter( key: KeyArray | None = None, ) -> ( ArrayTree | tuple[ArrayTree, ArrayTree, ArrayTree] -): # TODO: Can overload this function so the type checker knows that the output is a ArrayTree if output_factorial is True and a tuple[ArrayTree, ArrayTree] if output_factorial is False +): # TODO: Can overload this function so the type checker knows that the output is an ArrayTree if output_factorial is True and a tuple[ArrayTree, ArrayTree, ArrayTree] if output_factorial is False """Applies offline factorial filtering for given model inputs. `model_inputs` should have leading temporal dimension of length T + 1, @@ -37,10 +37,11 @@ def filter( output_factorial: If True, return a single state with first temporal dimension of length T + 1 and second factorial dimension of length F. If False, return a tuple of states. The first being the initial state - with first dimension of length F and temporal dimension. + with first dimension of length F and no temporal dimension. The second being the local states for each time step, i.e. first dimension of length T and no factorial dimension. - The third being the final factorial state. + The third being the final factorial state with first dimension of + length F and no temporal dimension. key: The key for the random number generator. Returns: