Add initial factorial API

cuthbert/factorial/README.md

@@ -0,0 +1,133 @@
+# Factorial State-Space Models
+
+A factorial state-space model is a state-space model where the dynamics distribution
+factors into a product of independent distributions across factors
+
+$$
+p(x_t \mid x_{t-1}) = \prod_{f=1}^F p(x_t^f \mid x_{t-1}^f),
+$$
+for factorial index $f \in \{1, \ldots, F\}$. We additionally assume that observations
+act locally on some subset of factors $S_t \subseteq \{1, \ldots, F\}$.
+
+$$
+p(y_t \mid x_t) = p(y_t \mid x_t^{S_t}).
+$$
+
+This motivates a factored approximation of filtering and smoothing distributions, e.g.
+
+$$
+p(x_t \mid y_{1:t}) = \prod_{f=1}^F p(x_t^f \mid y_{1:t}).
+$$
+
+A tutorial on factorial state-space models can be found in [Duffield et al](https://doi.org/10.1093/jrsssc/qlae035).
+
+The factorial approximation allows us to exploit significant benefits in terms of
+memory, compute and parallelization.
+
+Note that although the dynamics are factorized, `cuthbert` does not differentiate
+between `predict` and `update` (instead favouring a unified filter operation
+via `filter_prepare` and `filter_combine`). Thus the dynamics and model inputs
+should be specified to act on the joint local state (i.e. block diagonal
+where appropriate).
+
+
+## Factorial filtering with `cuthbert`
+
+Filtering in a factorial state-space model is similar to standard filtering, but with
+an additional step before the filtering operation to extract the relevant 
+factors as well as an additional step after the filtering operation to insert the
+updated factors back into the factorial state.
+
+
+```python
+from jax import tree
+import cuthbert
+
+# Define model_inputs
+model_inputs = ...
+
+# Define function to extract the factorial indices from model inputs
+# Here we assume model_inputs is a NamedTuple with a field `factorial_inds`
+get_factorial_indices = lambda mi: mi.factorial_inds
+
+# Build factorializer for the inference method
+factorializer = cuthbert.factorial.gaussian.build_factorializer(get_factorial_indices)
+
+# Load inference method, with parameter extraction functions defined for factorial inference
+kalman_filter = cuthbert.gaussian.kalman.build_filter(
+    get_init_params=get_init_params,  # Init specified to generate factorial state
+    get_dynamics_params=get_dynamics_params,  # Dynamics specified to act on joint local state
+    get_observation_params=get_observation_params,  # Observation specified to act on joint local state
+)
+
+# Online inference
+factorial_state = kalman_filter.init_prepare(tree.map(lambda x: x[0], model_inputs))
+
+for t in range(1, T):
+    model_inputs_t = tree.map(lambda x: x[t], model_inputs)
+    factorial_inds = get_factorial_indices(model_inputs_t)
+    local_state = factorializer.extract_and_join(factorial_state, factorial_inds)
+    prepare_state = kalman_filter.filter_prepare(model_inputs_t)
+    filtered_local_state = kalman_filter.filter_combine(local_state, prepare_state)
+    factorial_state = factorializer.marginalize_and_insert(
+        filtered_local_state, factorial_state, factorial_inds
+    )
+```
+
+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
+)
+```
+
+## Factorial smoothing with `cuthbert`
+
+Smoothing in factorial state-space models can be performed embarrassingly parallel
+across factors since the dynamics and factorial approximation are independent
+across factors (the observations are fully absorbed in the filtering and
+are not accessed during smoothing).
+
+The model inputs and filter states require some preprocessing to convert from being
+single sequence with each state containing all factors into a sequence or multiple
+sequences with each state corresponding to a single factor. This can be
+fiddly but is left to the user for maximum freedom. Oftentimes, it is easiest to
+specify different parameter functions for smoothing than filtering.
+
+After this preprocessing, smoothing can be performed as usual:
+
+```python
+# Define model_inputs for a single factor
+model_inputs_single_factor = ...
+
+# Similarly, we need to extract the filter states for the single factor we're smoothing.
+filter_states_single_factor = ...
+
+# Load smoother, with parameter extraction functions defined for a single factor
+kalman_smoother = cuthbert.gaussian.kalman.build_smoother(
+    get_dynamics_params=get_dynamics_params,  # Dynamics specified to act on a single factor
+)
+
+smoother_state = kalman_smoother.convert_filter_to_smoother_state(
+    tree.map(lambda x: x[-1], filter_states_single_factor),
+    model_inputs=tree.map(lambda x: x[-1], model_inputs_single_factor),
+)
+
+for t in range(T - 1, -1, -1):
+    model_inputs_single_factor_t = tree.map(lambda x: x[t], model_inputs_single_factor)
+    filter_state_single_factor_t = tree.map(lambda x: x[t], filter_states_single_factor)
+    prepare_state = kalman_smoother.smoother_prepare(
+        filter_state_single_factor_t, model_inputs_single_factor_t
+    )
+    smoother_state = kalman_smoother.smoother_combine(prepare_state, smoother_state)
+```
+
+Or directly using the `cuthbert.smoother`:
+
+```python
+smoother_states = cuthbert.smoother(
+    kalman_smoother, filter_states_single_factor, model_inputs_single_factor
+)
+```

cuthbert/factorial/__init__.py

@@ -0,0 +1,8 @@
+from cuthbert.factorial import gaussian
+from cuthbert.factorial.filtering import filter
+from cuthbert.factorial.types import (
+    ExtractAndJoin,
+    Factorializer,
+    GetFactorialIndices,
+    MarginalizeAndInsert,
+)

cuthbert/factorial/filtering.py

@@ -0,0 +1,110 @@
+"""cuthbert factorial filtering interface."""
+
+from jax import numpy as jnp
+from jax import random, tree
+from jax.lax import scan
+
+from cuthbert.factorial.types import Factorializer
+from cuthbert.inference import Filter
+from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
+
+
+def filter(
+    filter_obj: Filter,
+    factorializer: Factorializer,
+    model_inputs: ArrayTreeLike,
+    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
+    """Applies offline factorial filtering for given model inputs.
+
+    `model_inputs` should have leading temporal dimension of length T + 1,
+    where T is the number of time steps excluding the initial state.
+
+    Parallel associative filtering is not supported for factorial filtering.
+
+    Note that if output_factorial is True, this function will output a factorial state
+    with first temporal dimension of length T + 1 and second factorial dimension of
+    length F. Many of the factors will be unchanged across timesteps where they aren't
+    relevant.
+
+    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).
+        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.
+            The second being the local states for each time step, i.e. first
+            dimension of length T and no factorial dimension.
+        key: The key for the random number generator.
+
+    Returns:
+        The filtered states (NamedTuple with leading temporal dimension of length T + 1).
+    """
+    T = tree.leaves(model_inputs)[0].shape[0] - 1
+
+    if key is None:
+        # This will throw error if used as a key, which is desired behavior
+        # (albeit not a useful error, we could improve this)
+        prepare_keys = jnp.empty(T + 1)
+    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]
+    )
+
+    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
+        factorial_inds = factorializer.get_factorial_indices(prep_inp)
+        local_state = factorializer.extract_and_join(
+            prev_factorial_state, factorial_inds
+        )
+        prep_state = filter_obj.filter_prepare(prep_inp, key=k)
+        filtered_joint_state = filter_obj.filter_combine(local_state, prep_state)
+        factorial_state = factorializer.marginalize_and_insert(
+            filtered_joint_state, prev_factorial_state, factorial_inds
+        )
+
+        def extract(arr):
+            if arr.ndim >= 2:
+                return arr[factorial_inds]
+            else:
+                return arr
+
+        factorial_state_fac_inds_only = tree.map(extract, factorial_state)
+        return factorial_state, factorial_state_fac_inds_only
+
+    if output_factorial:
+
+        def body_factorial(prev_factorial_state, prep_inp_and_k):
+            factorial_state, _ = body_local(prev_factorial_state, prep_inp_and_k)
+            return factorial_state, factorial_state
+
+        _, factorial_states = scan(
+            body_factorial,
+            init_factorial_state,
+            (prep_model_inputs, prepare_keys[1:]),
+        )
+        factorial_states = tree.map(
+            lambda x, y: jnp.concatenate([x[None], y]),
+            init_factorial_state,
+            factorial_states,
+        )
+
+        return factorial_states
+
+    else:
+        _, local_states = scan(
+            body_local,
+            init_factorial_state,
+            (prep_model_inputs, prepare_keys[1:]),
+        )
+        return init_factorial_state, local_states