diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py index cc02ae3..2ce3147 100644 --- a/cuthbert/factorial/filtering.py +++ b/cuthbert/factorial/filtering.py @@ -13,6 +13,8 @@ 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 +22,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. @@ -33,7 +36,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). + 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 @@ -51,7 +57,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 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 @@ -60,15 +67,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 is not None: + 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 ce5c656..bdbd175 100644 --- a/cuthbert/filtering.py +++ b/cuthbert/filtering.py @@ -13,30 +13,41 @@ 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). + 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. 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 included 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 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 @@ -45,10 +56,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 is not None: + 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 428ca92..b6395db 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 8f0b347..c97758b 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): diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py index 226af1f..c500483 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(