Add init_state arg to filtering API

cuthbert/factorial/filtering.py

@@ -13,14 +13,17 @@ def filter(
     filter_obj: Filter,
     factorializer: Factorializer,
     model_inputs: ArrayTreeLike,
+    *,
+    init_state: ArrayTreeLike | None = None,
     output_factorial: bool = False,
     key: KeyArray | None = None,
 ) -> (
     ArrayTree | tuple[ArrayTree, ArrayTree, ArrayTree]
 ):  # 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

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))(

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]
 

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,15 +142,30 @@ 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),
         (des_means, des_covs, des_ells),
         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):

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(