Add final factorial state to filtering output

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
+    )
 )
 ```
 

cuthbert/factorial/filtering.py

@@ -16,8 +16,8 @@ def filter(
     output_factorial: bool = False,
     key: KeyArray | None = None,
 ) -> (
-    ArrayTree | tuple[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
+    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,
@@ -37,13 +37,19 @@ 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 with first dimension of
+            length F and no temporal dimension.
         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 +111,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

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(

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
     )
 

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,

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)