Skip to content

Kalman Filter and Smoother

cuthbert.gaussian.kalman

Implements the square-root, parallel-in-time Kalman filter for linear Gaussian SSMs.

See Yaghoobi et. al. (2025).

build_filter(get_init_params, get_dynamics_params, get_observation_params)

Builds an exact Kalman filter object for linear Gaussian SSMs.

Parameters:

Name Type Description Default
get_init_params GetInitParams

Function to get m0, chol_P0 to initialize filter state, given model inputs sufficient to define p(x_0) = N(m0, chol_P0 @ chol_P0^T).

 required
get_dynamics_params GetDynamicsParams

Function to get dynamics parameters, F, c, chol_Q given model inputs sufficient to define p(x_t | x_{t-1}) = N(F @ x_{t-1} + c, chol_Q @ chol_Q^T).

 required
get_observation_params GetObservationParams

Function to get observation parameters, H, d, chol_R, y given model inputs sufficient to define p(y_t | x_t) = N(H @ x_t + d, chol_R @ chol_R^T).

 required

Returns:

Type Description
Filter

Filter object for exact Kalman filter. Suitable for associative scan.
Source code in cuthbert/gaussian/kalman.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def build_filter(
    get_init_params: GetInitParams,
    get_dynamics_params: GetDynamicsParams,
    get_observation_params: GetObservationParams,
) -> Filter:
    """Builds an exact Kalman filter object for linear Gaussian SSMs.

    Args:
        get_init_params: Function to get m0, chol_P0 to initialize filter state,
            given model inputs sufficient to define p(x_0) = N(m0, chol_P0 @ chol_P0^T).
        get_dynamics_params: Function to get dynamics parameters, F, c, chol_Q
            given model inputs sufficient to define
            p(x_t | x_{t-1}) = N(F @ x_{t-1} + c, chol_Q @ chol_Q^T).
        get_observation_params: Function to get observation parameters, H, d, chol_R, y
            given model inputs sufficient to define
            p(y_t | x_t) = N(H @ x_t + d, chol_R @ chol_R^T).

    Returns:
        Filter object for exact Kalman filter. Suitable for associative scan.
    """
    return Filter(
        init_prepare=partial(
            init_prepare,
            get_init_params=get_init_params,
            get_observation_params=get_observation_params,
        ),
        filter_prepare=partial(
            filter_prepare,
            get_dynamics_params=get_dynamics_params,
            get_observation_params=get_observation_params,
        ),
        filter_combine=filter_combine,
        associative=True,
    )

build_smoother(get_dynamics_params, store_gain=False, store_chol_cov_given_next=False)

Builds an exact Kalman smoother object for linear Gaussian SSMs.

Parameters:

Name Type Description Default
get_dynamics_params GetDynamicsParams

Function to get dynamics parameters, F, c, chol_Q given model inputs sufficient to define p(x_t | x_{t-1}) = N(F @ x_{t-1} + c, chol_Q @ chol_Q^T).

 required
store_gain bool

Whether to store the gain matrix in the smoother state.

 False
store_chol_cov_given_next bool

Whether to store the chol_cov_given_next matrix in the smoother state.

 False

Returns:

Type Description
Smoother

Smoother object for exact Kalman smoother. Suitable for associative scan.
Source code in cuthbert/gaussian/kalman.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def build_smoother(
    get_dynamics_params: GetDynamicsParams,
    store_gain: bool = False,
    store_chol_cov_given_next: bool = False,
) -> Smoother:
    """Builds an exact Kalman smoother object for linear Gaussian SSMs.

    Args:
        get_dynamics_params: Function to get dynamics parameters, F, c, chol_Q
            given model inputs sufficient to define
            p(x_t | x_{t-1}) = N(F @ x_{t-1} + c, chol_Q @ chol_Q^T).
        store_gain: Whether to store the gain matrix in the smoother state.
        store_chol_cov_given_next: Whether to store the chol_cov_given_next matrix
            in the smoother state.

    Returns:
        Smoother object for exact Kalman smoother. Suitable for associative scan.
    """
    return Smoother(
        convert_filter_to_smoother_state=partial(
            convert_filter_to_smoother_state,
            store_gain=store_gain,
            store_chol_cov_given_next=store_chol_cov_given_next,
        ),
        smoother_prepare=partial(
            smoother_prepare,
            get_dynamics_params=get_dynamics_params,
            store_gain=store_gain,
            store_chol_cov_given_next=store_chol_cov_given_next,
        ),
        smoother_combine=smoother_combine,
        associative=True,
    )

cuthbert.gaussian.types

Provides shared types for Gaussian representations in state-space models.

GetInitParams

Bases: Protocol

Protocol for defining the initial distribution of a linear Gaussian SSM.

__call__(model_inputs)

Get initial parameters (m0, chol_P0) from model inputs.

Source code in cuthbert/gaussian/types.py 
13
14
15
def __call__(self, model_inputs: ArrayTreeLike) -> tuple[Array, Array]:
    """Get initial parameters (m0, chol_P0) from model inputs."""
    ...

GetDynamicsParams

Bases: Protocol

Protocol for defining the dynamics model of a linear Gaussian SSM.

__call__(model_inputs)

Get dynamics parameters (F, c, chol_Q) from model inputs.

Source code in cuthbert/gaussian/types.py 
21
22
23
def __call__(self, model_inputs: ArrayTreeLike) -> tuple[Array, Array, Array]:
    """Get dynamics parameters (F, c, chol_Q) from model inputs."""
    ...

GetObservationParams

Bases: Protocol

Protocol for defining the observation model of a linear Gaussian SSM.

__call__(model_inputs)

Get observation parameters (H, d, chol_R, y) from model inputs.

Source code in cuthbert/gaussian/types.py 
29
30
31
32
33
def __call__(
    self, model_inputs: ArrayTreeLike
) -> tuple[Array, Array, Array, Array]:
    """Get observation parameters (H, d, chol_R, y) from model inputs."""
    ...