Skip to content

Filtering

cuthbert.filtering

Unified cuthbert filtering interface.

filter(filter_obj, model_inputs, init_state, *, parallel=False, key=None)

Applies offline filtering given a filter object and model inputs.

model_inputs should have leading temporal dimension of length T, where T is the number of time steps excluding the initial state.

The output will have leading temporal dimension of length T + 1, and include the initial state as the first element.

Parameters:

Name Type Description Default
filter_obj Filter

The filter inference object.

 required
model_inputs ArrayTreeLike

The model inputs for filtering (with leading temporal dimension of length T).

 required
init_state ArrayTree

The initial state (with no temporal dimension). Generated by filter_obj.init_prepare.

 required
parallel bool

Whether to run the filter in parallel. Requires filter.associative_filter to be True.

 False
key KeyArray | None

The key for the random number generator.

 None

Returns:

Type Description
ArrayTree

The filtered states (NamedTuple with leading temporal dimension of length T + 1).
Source code in cuthbert/filtering.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def filter(
    filter_obj: Filter,
    model_inputs: ArrayTreeLike,
    init_state: ArrayTree,
    *,
    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,
    where T is the number of time steps excluding the initial state.

    The output will have leading temporal dimension of length T + 1, and
    include the initial state as the first element.

    Args:
        filter_obj: The filter inference object.
        model_inputs: The model inputs for filtering
            (with leading temporal dimension of length T).
        init_state: The initial state (with no temporal dimension).
            Generated by `filter_obj.init_prepare`.
        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).
    """
    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]

    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)
    else:
        prepare_keys = random.split(key, T)

    init_state = init_state._replace(model_inputs=dummy_leading_element(model_inputs))

    if parallel:
        other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
            model_inputs, prepare_keys
        )
        prep_states = tree.map(
            lambda x, y: jnp.concatenate([x[None], y]), init_state, other_prep_states
        )
        states = associative_scan(
            vmap(filter_obj.filter_combine),
            prep_states,
        )
    else:

        def body(prev_state, prep_inp_and_k):
            prep_inp, k = prep_inp_and_k
            prep_state = filter_obj.filter_prepare(prep_inp, key=k)
            state = filter_obj.filter_combine(prev_state, prep_state)
            return state, state

        _, states = scan(
            body,
            init_state,
            (model_inputs, prepare_keys),
        )
        states = tree.map(
            lambda x, y: jnp.concatenate([x[None], y]), init_state, states
        )

    return states

cuthbert.inference

Provides protocols and types for representing unified inference objects.

Filter

Bases: NamedTuple

Filter object.

Typically passed to cuthbert.filtering.filter.

Attributes:

Name Type Description
init_prepare InitPrepare

Function to prepare the initial state for the filter.
filter_prepare FilterPrepare

Function to prepare intermediate states for the filter.
filter_combine FilterCombine

Function that combines two filter states to produce another.
associative bool

Whether filter_combine is an associative operator. Temporally parallelized filters are guaranteed to produce correct results only if associative=True.

init_prepare instance-attribute

filter_prepare instance-attribute

filter_combine instance-attribute

associative = False class-attribute instance-attribute

InitPrepare

Bases: Protocol

Protocol for preparing the initial state for the inference.

__call__(*, key=None)

Prepare the initial state for the inference.

The state at the first time point, prior to any observations. Initial parameters are supplied when building the filter. The returned state has model_inputs=None; offline filtering pads this field to match the step input tree.

Parameters:

Name Type Description Default
key KeyArray | None

The key for the random number generator. Optional, as only used for stochastic inference methods

 None

Returns:

Type Description
ArrayTree

The initial state, a NamedTuple with inference-specific fields.
Source code in cuthbert/inference.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def __call__(self, *, key: KeyArray | None = None) -> ArrayTree:
    """Prepare the initial state for the inference.

    The state at the first time point, prior to any observations. Initial
    parameters are supplied when building the filter. The returned state has
    `model_inputs=None`; offline filtering pads this field to match the step
    input tree.

    Args:
        key: The key for the random number generator.
            Optional, as only used for stochastic inference methods

    Returns:
        The initial state, a NamedTuple with inference-specific fields.
    """
    ...

FilterPrepare

Bases: Protocol

Protocol for preparing the state for the filter at the next time point.

__call__(model_inputs, *, key=None)

Prepare the state for the filter at the next time point.

Converts the model inputs (and any stochasticity) into a unified state object which can be combined with a state (of the same form) from the previous time point with FilterCombine.

state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))

Parameters:

Name Type Description Default
model_inputs ArrayTreeLike

The model inputs at the next time point.

 required
key KeyArray | None

The key for the random number generator. Optional, as only used for stochastic inference methods

 None

Returns:

Type Description
ArrayTree

The state prepared for FilterCombine, a NamedTuple with inference-specific fields.
Source code in cuthbert/inference.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __call__(
    self,
    model_inputs: ArrayTreeLike,
    *,
    key: KeyArray | None = None,
) -> ArrayTree:
    """Prepare the state for the filter at the next time point.

    Converts the model inputs (and any stochasticity) into a unified state
    object which can be combined with a state (of the same form) from the
    previous time point with FilterCombine.

    state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))

    Args:
        model_inputs: The model inputs at the next time point.
        key: The key for the random number generator.
            Optional, as only used for stochastic inference methods

    Returns:
        The state prepared for FilterCombine,
            a NamedTuple with inference-specific fields.
    """
    ...

FilterCombine

Bases: Protocol

Protocol for combining the previous state with the state from FilterPrepare.

__call__(state_1, state_2)

Combine state from previous time point with state from FilterPrepare.

state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))

Parameters:

Name Type Description Default
state_1 ArrayTreeLike

The state from the previous time point.

 required
state_2 ArrayTreeLike

The state from FilterPrepare for the current time point.

 required

Returns:

Type Description
ArrayTree

The combined filter state, a NamedTuple with inference-specific fields.
Source code in cuthbert/inference.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def __call__(
    self,
    state_1: ArrayTreeLike,
    state_2: ArrayTreeLike,
) -> ArrayTree:
    """Combine state from previous time point with state from FilterPrepare.

    ```python
    state = FilterCombine(prev_state, FilterPrepare(model_inputs, key))
    ```

    Args:
        state_1: The state from the previous time point.
        state_2: The state from FilterPrepare for the current time point.

    Returns:
        The combined filter state, a NamedTuple with inference-specific fields.
    """
    ...