Skip to content

Ensemble Rauch-Tung-Striebel smoother

The EnRTS is an ensemble RTS-like smoothing algorithm, similar in form to the RTS smoothing algorithm. It forms a backwards recursion using a gain \(J_t\):

\[ x_{t \mid T}^{(i)} = x_{t \mid t}^{(i)} + J_t \left(x_{t+1 \mid T}^{(i)} - x_{t+1 \mid t}^{(i)}\right), \]

where \(J_t\) is determined by ensembles to be an empirical version of the Kalman gain.

The computation of \(J_t\) involves the ensemble of \(x_{t+1 \mid t}\); to avoid duplicate computation, the EnRTS therefore requires store_predicted_ensemble=True in the forward filtering run.

See Raanes (2016) for more info on the EnRTS algorithm and its equivalence to the ensemble Kalman smoother (EnKS), a simliar ensemble smoothing algorithm which consists of a forward pass of increasing dimension.

cuthbert.ensemble_kalman.ensemble_rts_smoother

Implements the high-level Ensemble Rauch-Tung-Striebel smoother (EnRTS).

EnRTSState

Bases: NamedTuple

Ensemble Rauch-Tung-Striebel smoother state.

ensemble instance-attribute

predicted_ensemble instance-attribute

model_inputs instance-attribute

n_particles property

Number of particles.

mean property

Ensemble mean.

chol_cov property

Generalised Cholesky factor of the ensemble sample covariance.

build_smoother()

Build an Ensemble Rauch-Tung-Striebel smoother object.

Filtered states must come from an EnKF built with store_predicted_ensemble=True.

Returns:

Type Description
Smoother

Smoother object for the EnRTS.
Source code in cuthbert/ensemble_kalman/ensemble_rts_smoother.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def build_smoother() -> Smoother:
    """Build an Ensemble Rauch-Tung-Striebel smoother object.

    Filtered states must come from an EnKF built with
    ``store_predicted_ensemble=True``.

    Returns:
        Smoother object for the EnRTS.
    """
    return Smoother(
        convert_filter_to_smoother_state=convert_filter_to_smoother_state,
        smoother_prepare=smoother_prepare,
        smoother_combine=smoother_combine,
        associative=False,
    )

smoother_prepare(filter_state, model_inputs, key=None)

Prepare a state for an EnRTS step.

Parameters:

Name Type Description Default
filter_state EnKFState

EnKF state at time t.

 required
model_inputs ArrayTreeLike

Model inputs for the transition from t to t + 1.

 required
key KeyArray | None

JAX random key; unused.

 None

Returns:

Type Description
EnRTSState

Prepared EnRTS state.

Raises:

Type Description
ValueError

If the EnKF did not store predicted ensembles.
Source code in cuthbert/ensemble_kalman/ensemble_rts_smoother.py 
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
def smoother_prepare(
    filter_state: EnKFState,
    model_inputs: ArrayTreeLike,
    key: KeyArray | None = None,
) -> EnRTSState:
    """Prepare a state for an EnRTS step.

    Args:
        filter_state: EnKF state at time t.
        model_inputs: Model inputs for the transition from t to t + 1.
        key: JAX random key; unused.

    Returns:
        Prepared EnRTS state.

    Raises:
        ValueError: If the EnKF did not store predicted ensembles.
    """
    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
    predicted_ensemble = filter_state.predicted_ensemble
    if predicted_ensemble is None:
        raise ValueError(
            "EnRTS requires an EnKF built with store_predicted_ensemble=True."
        )
    return EnRTSState(
        ensemble=filter_state.ensemble,
        predicted_ensemble=predicted_ensemble,
        model_inputs=model_inputs,
    )

smoother_combine(state_1, state_2)

Combine a prepared state with the next EnRTS state.

Parameters:

Name Type Description Default
state_1 EnRTSState

Prepared state at time t.

 required
state_2 EnRTSState

Smoothed state at time t + 1.

 required

Returns:

Type Description
EnRTSState

Smoothed state at time t.
Source code in cuthbert/ensemble_kalman/ensemble_rts_smoother.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def smoother_combine(
    state_1: EnRTSState,
    state_2: EnRTSState,
) -> EnRTSState:
    """Combine a prepared state with the next EnRTS state.

    Args:
        state_1: Prepared state at time t.
        state_2: Smoothed state at time t + 1.

    Returns:
        Smoothed state at time t.
    """
    ensemble, _ = enkf_lib.smoother_update(
        state_1.ensemble,
        state_2.predicted_ensemble,
        state_2.ensemble,
    )
    return EnRTSState(
        ensemble=ensemble,
        predicted_ensemble=state_1.predicted_ensemble,
        model_inputs=state_1.model_inputs,
    )

convert_filter_to_smoother_state(filter_state, model_inputs=None, key=None)

Convert the final EnKF state to an EnRTS state.

Requires filter_state to contain predicted ensembles (via store_predicted_ensemble=True in the filter).

Parameters:

Name Type Description Default
filter_state EnKFState

Final EnKF state.

 required
model_inputs ArrayTreeLike | None

Model inputs used to define the output tree structure.

 None
key KeyArray | None

JAX random key - not used.

 None

Returns:

Type Description
EnRTSState

Final EnRTS state with dummy model inputs.

Raises:

Type Description
ValueError

If the EnKF did not store predicted ensembles.
Source code in cuthbert/ensemble_kalman/ensemble_rts_smoother.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def convert_filter_to_smoother_state(
    filter_state: EnKFState,
    model_inputs: ArrayTreeLike | None = None,
    key: KeyArray | None = None,
) -> EnRTSState:
    """Convert the final EnKF state to an EnRTS state.

    Requires `filter_state` to contain predicted ensembles (via `store_predicted_ensemble=True` in the filter).

    Args:
        filter_state: Final EnKF state.
        model_inputs: Model inputs used to define the output tree structure.
        key: JAX random key - not used.

    Returns:
        Final EnRTS state with dummy model inputs.

    Raises:
        ValueError: If the EnKF did not store predicted ensembles.
    """
    if model_inputs is None:
        model_inputs = filter_state.model_inputs

    predicted_ensemble = filter_state.predicted_ensemble
    if predicted_ensemble is None:
        raise ValueError(
            "EnRTS requires an EnKF built with store_predicted_ensemble=True."
        )

    return EnRTSState(
        ensemble=filter_state.ensemble,
        predicted_ensemble=predicted_ensemble,
        model_inputs=dummy_tree_like(model_inputs),
    )