From 1416ce47131ab5f41866b9b1419b149fc56534ca Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Fri, 26 Jun 2026 15:31:00 +0100 Subject: [PATCH 1/4] Add init_state arg --- cuthbert/factorial/filtering.py | 30 +++++++++++------ cuthbert/filtering.py | 23 ++++++++++--- tests/cuthbert/factorial/test_kalman.py | 43 +++++++++++++++++++++++-- tests/cuthbert/gaussian/test_kalman.py | 40 +++++++++++++++-------- 4 files changed, 105 insertions(+), 31 deletions(-) diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index cc02ae38..3fbfb81e 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -13,6 +13,7 @@ def filter( filter_obj: Filter, factorializer: Factorializer, model_inputs: ArrayTreeLike, + init_state: ArrayTreeLike | None = None, output_factorial: bool = False, key: KeyArray | None = None, ) -> ( @@ -20,7 +21,8 @@ def filter( ): # 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, + `model_inputs` should have leading temporal dimension of + length T if `init_state` is provided, or T + 1 if `init_state` is not provided, where T is the number of time steps excluding the initial state. Parallel associative filtering is not supported for factorial filtering. @@ -34,6 +36,9 @@ def filter( filter_obj: The filter inference object. factorializer: The factorializer object for the inference method. model_inputs: The model inputs (with leading temporal dimension of length T + 1). + init_state: The initial state for the filter. If not provided, + `filter_obj.init_prepare` will be called on the first `model_inputs` to + generate the initial state. 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 @@ -51,7 +56,8 @@ def filter( (initial factorial state, local states for each time step, final factorial state). """ - T = tree.leaves(model_inputs)[0].shape[0] - 1 + len_model_inputs = tree.leaves(model_inputs)[0].shape[0] + T = len_model_inputs if init_state else len_model_inputs - 1 if key is None: # This will throw error if used as a key, which is desired behavior @@ -60,15 +66,19 @@ def filter( else: prepare_keys = random.split(key, T + 1) - init_model_input = tree.map(lambda x: x[0], model_inputs) - init_factorial_state = filter_obj.init_prepare( - init_model_input, key=prepare_keys[0] - ) - init_factorial_state = factorializer.factorialize_init_state( - init_factorial_state, init_model_input - ) + if init_state: + init_factorial_state = init_state + prep_model_inputs = model_inputs + else: + init_model_input = tree.map(lambda x: x[0], model_inputs) + init_factorial_state = filter_obj.init_prepare( + init_model_input, key=prepare_keys[0] + ) + init_factorial_state = factorializer.factorialize_init_state( + init_factorial_state, init_model_input + ) - prep_model_inputs = tree.map(lambda x: x[1:], model_inputs) + prep_model_inputs = tree.map(lambda x: x[1:], model_inputs) def body_local(prev_factorial_state, prep_inp_and_k): prep_inp, k = prep_inp_and_k diff --git a/cuthbert/filtering.py b/cuthbert/filtering.py index ce5c6561..f1896e87 100644 --- a/cuthbert/filtering.py +++ b/cuthbert/filtering.py @@ -13,30 +13,40 @@ def filter( filter_obj: Filter, model_inputs: ArrayTreeLike, + init_state: ArrayTreeLike | None = None, parallel: bool = False, key: KeyArray | None = None, ) -> ArrayTree: """Applies offline filtering given a filter object and model inputs. - `model_inputs` should have leading temporal dimension of length T + 1, + `model_inputs` should have leading temporal dimension of + length T if `init_state` is provided, or T + 1 if `init_state` is not provided, where T is the number of time steps excluding the initial state. + The output will therefore have leading temporal dimension of length T + 1. + + If `init_state` is provided, `filter_obj.init_prepare` will not be called. Args: filter_obj: The filter inference object. model_inputs: The model inputs (with leading temporal dimension of length T + 1). + init_state: The initial state for the filter. If not provided, + `filter_obj.init_prepare` will be called on the first `model_inputs` to + generate the initial state. parallel: Whether to run the filter in parallel. Requires `filter.associative_filter` to be `True`. key: The key for the random number generator. Returns: The filtered states (NamedTuple with leading temporal dimension of length T + 1). + `init_state` is include as the first element of the output. """ if parallel and not filter_obj.associative: warnings.warn( f"Parallel filtering attempted but filter.associative is False for {filter_obj}" ) - T = tree.leaves(model_inputs)[0].shape[0] - 1 + len_model_inputs = tree.leaves(model_inputs)[0].shape[0] + T = len_model_inputs if init_state else len_model_inputs - 1 if key is None: # This will throw error if used as a key, which is desired behavior @@ -45,10 +55,13 @@ def filter( else: prepare_keys = random.split(key, T + 1) - init_model_input = tree.map(lambda x: x[0], model_inputs) - init_state = filter_obj.init_prepare(init_model_input, key=prepare_keys[0]) + if init_state: + prep_model_inputs = model_inputs + else: + init_model_input = tree.map(lambda x: x[0], model_inputs) + init_state = filter_obj.init_prepare(init_model_input, key=prepare_keys[0]) - prep_model_inputs = tree.map(lambda x: x[1:], model_inputs) + prep_model_inputs = tree.map(lambda x: x[1:], model_inputs) if parallel: other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))( diff --git a/tests/cuthbert/factorial/test_kalman.py b/tests/cuthbert/factorial/test_kalman.py index 428ca92a..b6395dbf 100644 --- a/tests/cuthbert/factorial/test_kalman.py +++ b/tests/cuthbert/factorial/test_kalman.py @@ -1,5 +1,4 @@ import itertools -from typing import cast import chex import jax @@ -14,7 +13,6 @@ from cuthbert.gaussian import kalman from cuthbert.inference import Filter, Smoother from cuthbertlib.linalg import block_marginal_sqrt_cov -from cuthbertlib.types import ArrayTree from tests.cuthbert.factorial.gaussian_utils import generate_factorial_kalman_model from tests.cuthbertlib.kalman.test_filtering import std_predict, std_update from tests.cuthbertlib.kalman.test_smoothing import std_kalman_smoother @@ -244,7 +242,6 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste filter_obj, factorializer, model_inputs, output_factorial=True ) - factorial_filtering_states = cast(ArrayTree, factorial_filtering_states) factorial_filtering_covs = ( factorial_filtering_states.chol_cov @ factorial_filtering_states.chol_cov.transpose(0, 1, 3, 2) @@ -257,6 +254,46 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste ells, factorial_filtering_states.log_normalizing_constant[1:] ) + # Check with init_state provided + init_model_inputs = tree.map(lambda x: x[0], model_inputs) + other_model_inputs = tree.map(lambda x: x[1:], model_inputs) + init_state_init_prepare = filter_obj.init_prepare(init_model_inputs) + init_state_init_given, local_filter_states_init_given, final_state_init_given = ( + factorial.filter( + filter_obj, factorializer, other_model_inputs, init_state=init_state + ) + ) + # Check init_state unchanged + chex.assert_trees_all_close( + init_state_init_prepare, init_state_init_given, init_state, rtol=1e-10 + ) + # Check local_filter_states unchanged + chex.assert_trees_all_close( + (local_means, local_covs, ells), + ( + local_filter_states_init_given.mean, + local_filter_states_init_given.chol_cov + @ local_filter_states_init_given.chol_cov.transpose(0, 1, 3, 2), + local_filter_states_init_given.log_normalizing_constant, + ), + rtol=1e-10, + ) + # Check final_state unchanged + 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, + ), + ( + final_state_init_given.mean, + final_state_init_given.chol_cov + @ final_state_init_given.chol_cov.transpose(0, 2, 1), + final_state_init_given.log_normalizing_constant, + ), + rtol=1e-10, + ) + smoother_indices = [0, 1, 5] diff --git a/tests/cuthbert/gaussian/test_kalman.py b/tests/cuthbert/gaussian/test_kalman.py index 8f0b347a..c97758bf 100644 --- a/tests/cuthbert/gaussian/test_kalman.py +++ b/tests/cuthbert/gaussian/test_kalman.py @@ -4,7 +4,7 @@ import jax import jax.numpy as jnp import pytest -from jax import Array, vmap +from jax import Array, tree, vmap from cuthbert import filter, smoother from cuthbert.gaussian import kalman @@ -119,21 +119,20 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps): m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys ) + def states_to_mean_cov_ell(states): + return ( + states.mean, + states.chol_cov @ states.chol_cov.transpose(0, 2, 1), + states.log_normalizing_constant, + ) + # Run sequential sqrt filter seq_states = filter(kalman_filter, model_inputs, parallel=False) - seq_means, seq_chol_covs, seq_ells = ( - seq_states.mean, - seq_states.chol_cov, - seq_states.log_normalizing_constant, - ) + seq_means, seq_covs, seq_ells = states_to_mean_cov_ell(seq_states) # Run parallel sqrt filter par_states = filter(kalman_filter, model_inputs, parallel=True) - par_means, par_chol_covs, par_ells = ( - par_states.mean, - par_states.chol_cov, - par_states.log_normalizing_constant, - ) + par_means, par_covs, par_ells = states_to_mean_cov_ell(par_states) # Run the standard Kalman filter. P0 = chol_P0 @ chol_P0.T @@ -143,8 +142,6 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps): m0, P0, Fs, cs, Qs, Hs, ds, Rs, ys ) - seq_covs = seq_chol_covs @ seq_chol_covs.transpose(0, 2, 1) - par_covs = par_chol_covs @ par_chol_covs.transpose(0, 2, 1) chex.assert_trees_all_close( (seq_means, seq_covs, seq_ells), (par_means, par_covs, par_ells), @@ -152,6 +149,23 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps): rtol=1e-10, ) + # Test filter with init_state provided + init_model_inputs = tree.map(lambda x: x[0], model_inputs) + other_model_inputs = tree.map(lambda x: x[1:], model_inputs) + init_state = kalman_filter.init_prepare(init_model_inputs) + seq_states_init_given = filter( + kalman_filter, other_model_inputs, init_state=init_state, parallel=False + ) + par_states_init_given = filter( + kalman_filter, other_model_inputs, init_state=init_state, parallel=True + ) + chex.assert_trees_all_close( + states_to_mean_cov_ell(seq_states_init_given), + states_to_mean_cov_ell(par_states_init_given), + (des_means, des_covs, des_ells), + rtol=1e-10, + ) + @pytest.mark.parametrize("seed,x_dim,y_dim,num_time_steps", gradient_params) def test_check_gradient(seed, x_dim, y_dim, num_time_steps): From cd50fd11df0fe85f203d16520b290f63098692f2 Mon Sep 17 00:00:00 2001 From: SamDuffield <34280297+SamDuffield@users.noreply.github.com> Date: Fri, 26 Jun 2026 17:16:49 +0100 Subject: [PATCH 2/4] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- cuthbert/factorial/filtering.py | 12 ++++++------ cuthbert/filtering.py | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index 3fbfb81e..3bd7bded 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -35,10 +35,10 @@ def filter( Args: filter_obj: The filter inference object. factorializer: The factorializer object for the inference method. - model_inputs: The model inputs (with leading temporal dimension of length T + 1). - init_state: The initial state for the filter. If not provided, - `filter_obj.init_prepare` will be called on the first `model_inputs` to - generate the initial state. + model_inputs: The model inputs (with leading temporal dimension of length T if init_state is provided, or T + 1 otherwise). + init_state: The initial factorial state for the filter. If not provided, + `filter_obj.init_prepare` will be called on the first `model_inputs` and then + passed through `factorializer.factorialize_init_state` to generate the initial state. 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 @@ -57,7 +57,7 @@ def filter( final factorial state). """ len_model_inputs = tree.leaves(model_inputs)[0].shape[0] - T = len_model_inputs if init_state else len_model_inputs - 1 + T = len_model_inputs if init_state is not None else len_model_inputs - 1 if key is None: # This will throw error if used as a key, which is desired behavior @@ -66,7 +66,7 @@ def filter( else: prepare_keys = random.split(key, T + 1) - if init_state: + if init_state is not None: init_factorial_state = init_state prep_model_inputs = model_inputs else: diff --git a/cuthbert/filtering.py b/cuthbert/filtering.py index f1896e87..adf4550f 100644 --- a/cuthbert/filtering.py +++ b/cuthbert/filtering.py @@ -28,7 +28,7 @@ def filter( Args: filter_obj: The filter inference object. - model_inputs: The model inputs (with leading temporal dimension of length T + 1). + model_inputs: The model inputs (with leading temporal dimension of length T if init_state is provided, or T + 1 otherwise). init_state: The initial state for the filter. If not provided, `filter_obj.init_prepare` will be called on the first `model_inputs` to generate the initial state. @@ -38,7 +38,7 @@ def filter( Returns: The filtered states (NamedTuple with leading temporal dimension of length T + 1). - `init_state` is include as the first element of the output. + `init_state` is included as the first element of the output. """ if parallel and not filter_obj.associative: warnings.warn( @@ -46,7 +46,7 @@ def filter( ) len_model_inputs = tree.leaves(model_inputs)[0].shape[0] - T = len_model_inputs if init_state else len_model_inputs - 1 + T = len_model_inputs if init_state is not None else len_model_inputs - 1 if key is None: # This will throw error if used as a key, which is desired behavior @@ -55,7 +55,7 @@ def filter( else: prepare_keys = random.split(key, T + 1) - if init_state: + if init_state is not None: prep_model_inputs = model_inputs else: init_model_input = tree.map(lambda x: x[0], model_inputs) From 1dfe235f5be5a6352cbc34d16ded66bcd0623447 Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Fri, 26 Jun 2026 17:33:58 +0100 Subject: [PATCH 3/4] Fix tests --- tests/cuthbert/smc/test_backward_sampler.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py index 226af1f6..c500483e 100644 --- a/tests/cuthbert/smc/test_backward_sampler.py +++ b/tests/cuthbert/smc/test_backward_sampler.py @@ -76,7 +76,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method): m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys ) key = random.key(seed + 1) - filtered_states = filter(filter_obj, model_inputs, False, key) + filtered_states = filter(filter_obj, model_inputs, parallel=False, key=key) if method == "tracing": bs_fn = tracing @@ -158,7 +158,9 @@ def log_potential(state_prev, state, model_inputs): filter_key, smoother_key = random.split(key) num_time_steps = 5 model_inputs = jnp.empty(num_time_steps + 1) - filtered_states = filter(filter_obj, model_inputs, False, filter_key) + filtered_states = filter( + filter_obj, model_inputs, parallel=False, key=filter_key + ) n_smoother_particles = 1000 smoother_obj = build_smoother( From d6fc9f8b3a504aa018a8f0a48e0c9335786821e2 Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Mon, 29 Jun 2026 11:13:37 +0100 Subject: [PATCH 4/4] Make some filter args keyword only --- cuthbert/factorial/filtering.py | 1 + cuthbert/filtering.py | 1 + 2 files changed, 2 insertions(+) diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index 3bd7bded..2ce3147c 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -13,6 +13,7 @@ def filter( filter_obj: Filter, factorializer: Factorializer, model_inputs: ArrayTreeLike, + *, init_state: ArrayTreeLike | None = None, output_factorial: bool = False, key: KeyArray | None = None, diff --git a/cuthbert/filtering.py b/cuthbert/filtering.py index adf4550f..bdbd175c 100644 --- a/cuthbert/filtering.py +++ b/cuthbert/filtering.py @@ -13,6 +13,7 @@ def filter( filter_obj: Filter, model_inputs: ArrayTreeLike, + *, init_state: ArrayTreeLike | None = None, parallel: bool = False, key: KeyArray | None = None,