Notes for the Ensemble Kalman Filter (EnsKF).
Model ¶ x ( t ) = M x ( t − 1 ) + η ( t ) , η ∼ N ( 0 , Q ) y ( t ) = H x ( t ) + ϵ ( t ) , ϵ ∼ N ( 0 , R ) \begin{aligned}
\mathbf{x}(t) &= \mathbf{Mx}(t-1) + \boldsymbol{\eta}(t), &\boldsymbol{\eta} \sim \mathcal{N}(\mathbf{0}, \mathbf{Q}) \\
\mathbf{y}(t) &= \mathbf{Hx}(t) + \boldsymbol{\epsilon}(t), &\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{R}) \\
\end{aligned} x ( t ) y ( t ) = Mx ( t − 1 ) + η ( t ) , = Hx ( t ) + ϵ ( t ) , η ∼ N ( 0 , Q ) ϵ ∼ N ( 0 , R ) Forecast Step ¶ x f ( t ) = M x a ( t − 1 ) P f ( t ) = M P a ( t − 1 ) M ⊤ + Q \begin{aligned}
\mathbf{x}^f(t) &= \mathbf{Mx}^a(t-1) \\
\mathbf{P}^f(t) &= \mathbf{MP}^a(t-1) \mathbf{M}^\top + \mathbf{Q}
\end{aligned} x f ( t ) P f ( t ) = Mx a ( t − 1 ) = MP a ( t − 1 ) M ⊤ + Q Analysis Step (Forward) ¶ x a ( t ) = x f ( t ) + K a [ y ( t ) − H x f ( t ) ] P a ( t ) = [ I − K a H ] P f ( t ) \begin{aligned}
\mathbf{x}^a(t) &= \mathbf{x}^f(t) + \mathbf{K}^a[ \mathbf{y}(t) - \mathbf{H}\mathbf{x}^f(t)] \\
\mathbf{P}^a(t) &= [ \mathbf{I} - \mathbf{K}^a \mathbf{H}]\mathbf{P}^f(t)
\end{aligned} x a ( t ) P a ( t ) = x f ( t ) + K a [ y ( t ) − H x f ( t )] = [ I − K a H ] P f ( t ) where:
K a = P f ( t ) H ⊤ [ H P f ( t ) H ⊤ + R ] − 1 \mathbf{K}^a = \mathbf{P}^f(t) \mathbf{H}^\top [ \mathbf{HP}^f(t)\mathbf{H}^\top + \mathbf{R}]^{-1} K a = P f ( t ) H ⊤ [ HP f ( t ) H ⊤ + R ] − 1 Analysis Step (Backwards) ¶ x s ( t ) = x a ( t ) + K s [ x s ( t + 1 ) − x f ( t + 1 ) ] P s ( t ) = P a ( t ) + K s [ P s ( t + 1 ) − P f ( t + 1 ) ] ( K s ) ⊤ \begin{aligned}
\mathbf{x}^s(t) &= \mathbf{x}^a(t) + \mathbf{K}^s[ \mathbf{x}^s(t+1) - \mathbf{x}^f(t+1)] \\
\mathbf{P}^s(t) &= \mathbf{P}^a(t) + \mathbf{K}^s[ \mathbf{P}^s(t+1) - \mathbf{P}^f(t+1)](\mathbf{K}^s)^\top
\end{aligned} x s ( t ) P s ( t ) = x a ( t ) + K s [ x s ( t + 1 ) − x f ( t + 1 )] = P a ( t ) + K s [ P s ( t + 1 ) − P f ( t + 1 )] ( K s ) ⊤ where:
K s = P a ( t ) M ⊤ [ P f ( t + 1 ) ] − 1 \mathbf{K}^s = \mathbf{P}^a(t) \mathbf{M}^\top [ \mathbf{P}^f (t+1)]^{-1} K s = P a ( t ) M ⊤ [ P f ( t + 1 ) ] − 1 Model ¶ In man applications, there are cases where we cannot write the transition model and/or the emission model in a simple matrix operation. Instead, we use non-linear functions.
x ( t ) = m ( x , t − 1 ) + η ( t ) y ( t ) = h ( x , t ) + ϵ ( t ) \begin{aligned}
\mathbf{x}(t) &= \boldsymbol{m}(\mathbf{x}, t-1) + \boldsymbol{\eta}(t) \\
\mathbf{y}(t) &= \boldsymbol{h}(\mathbf{x}, t) + \boldsymbol{\epsilon}(t)
\end{aligned} x ( t ) y ( t ) = m ( x , t − 1 ) + η ( t ) = h ( x , t ) + ϵ ( t ) This gives us added flexibility and this also allows us to model more complexities found within the data. However, we cannot use the linear Kalman filter equations for inference.
Solutions : In general, you’ll find 3 classes of algorithms to deal with this: Extended Kalman Filters (EKF), Ensemble Kalman Filters (EnKF) and Particle Filters (PF).
Ensemble Kalman Filter ¶ In the EnKF, we use an ensemble to track the state of the system. So x ( t ) \mathbf{x}(t) x ( t ) N e N_e N e
x ( t ) = ( x 1 ( t ) , x 2 ( t ) , … , x N e ( t ) ) \mathbf{x}(t) = (\mathbf{x}_1(t), \mathbf{x}_2(t), \ldots, \mathbf{x}_{N_e}(t)) x ( t ) = ( x 1 ( t ) , x 2 ( t ) , … , x N e ( t )) x ( i ) f ( t ) = m ( x ( i ) a ( t − 1 ) ) + η ( i ) ( t ) , ∀ i = 1 , 2 , … , N e \mathbf{x}_{(i)}^f(t) = \boldsymbol{m}(\mathbf{x}_{(i)}^a(t-1)) + \boldsymbol{\eta}_{(i)}(t), \hspace{5mm} \forall i=1,2, \ldots, N_e x ( i ) f ( t ) = m ( x ( i ) a ( t − 1 )) + η ( i ) ( t ) , ∀ i = 1 , 2 , … , N e K = P f ( t ) H ⊤ ( H P f ( t ) H ⊤ + R ) − 1 \mathbf{K} = \mathbf{P}^f(t) \mathbf{H}^\top (\mathbf{H} \mathbf{P}^f(t) \mathbf{H}^\top + \mathbf{R})^{-1} K = P f ( t ) H ⊤ ( H P f ( t ) H ⊤ + R ) − 1 Forecast Step ¶ x ( i ) f ( t ) = m ( x ( i ) a ( t − 1 ) ) + η ( i ) ( t ) , ∀ i = 1 , 2 , … , N e x f ( t ) = 1 N e ∑ i = 1 N e x ( i ) ( t ) P f ( t ) = 1 N e − 1 ∑ i = 1 N e ( x ( i ) ( t ) − x ˉ ( t ) ) ( x ( i ) ( t ) − x ˉ ( t ) ) ⊤ K = P f ( t ) H ⊤ ( H P f ( t ) H ⊤ + R ) \begin{aligned}
\mathbf{x}_{(i)}^f(t) &= \boldsymbol{m}(\mathbf{x}_{(i)}^a(t-1)) + \boldsymbol{\eta}_{(i)}(t), &\forall i=1,2, \ldots, N_e \\
\mathbf{x}^f(t) &= \frac{1}{N_e}\sum_{i=1}^{N_e} \mathbf{x}_{(i)}(t) \\
\mathbf{P}^f(t) &= \frac{1}{N_e - 1} \sum_{i=1}^{N_e} \left( \mathbf{x}_{(i)}(t) - \bar{\mathbf{x}}(t) \right)\left( \mathbf{x}_{(i)}(t) - \bar{\mathbf{x}}(t) \right)^\top \\
\mathbf{K} &= \mathbf{P}^f(t)\mathbf{H}^\top \left( \mathbf{HP}^f(t)\mathbf{H}^\top + \mathbf{R} \right)
\end{aligned} x ( i ) f ( t ) x f ( t ) P f ( t ) K = m ( x ( i ) a ( t − 1 )) + η ( i ) ( t ) , = N e 1 i = 1 ∑ N e x ( i ) ( t ) = N e − 1 1 i = 1 ∑ N e ( x ( i ) ( t ) − x ˉ ( t ) ) ( x ( i ) ( t ) − x ˉ ( t ) ) ⊤ = P f ( t ) H ⊤ ( HP f ( t ) H ⊤ + R ) ∀ i = 1 , 2 , … , N e Deterministic Ensemble Kalman Filter ¶ x a ( t ) = x f ( t ) + K ( y ( t ) − h ( x f ( t ) ) ) P a ( t ) = ( I − K H ) P f ( t ) \begin{aligned}
\mathbf{x}^a(t) &= \mathbf{x}^f(t) + \mathbf{K}(\mathbf{y}(t) - \boldsymbol{h}(\mathbf{x}^f(t))) \\
\mathbf{P}^a(t) &= \left( \mathbf{I} - \mathbf{KH} \right)\mathbf{P}^f(t)
\end{aligned} x a ( t ) P a ( t ) = x f ( t ) + K ( y ( t ) − h ( x f ( t ))) = ( I − KH ) P f ( t ) Stochastic Ensemble Kalman Filter ¶ y ( i ) f ( t ) = h ( x ( i ) f ( t ) ) + ϵ ( i ) ( t ) x ( i ) a ( t ) = x ( i ) f ( t ) + K ( y ( t ) − y ( i ) f ( t ) ) x a ( t ) = 1 N e ∑ i = 1 N e x ( i ) a ( t ) P a ( t ) = 1 N e − 1 ∑ i = 1 N e ( x ( i ) a ( t ) − x a ( t ) ) ( x ( i ) a ( t ) − x a ( t ) ) ⊤ \begin{aligned}
\mathbf{y}^f_{(i)}(t) &= \boldsymbol{h}(\mathbf{x}_{(i)}^f(t)) + \boldsymbol{\epsilon}_{(i)}(t)\\
\mathbf{x}^a_{(i)}(t) &= \mathbf{x}_{(i)}^f(t) + \mathbf{K}\left( \mathbf{y}(t) - \mathbf{y}^f_{(i)}(t) \right) \\
\mathbf{x}^a(t) &= \frac{1}{N_e}\sum_{i=1}^{N_e} \mathbf{x}_{(i)}^a(t) \\
\mathbf{P}^a(t) &= \frac{1}{N_e - 1} \sum_{i=1}^{N_e} \left( \mathbf{x}_{(i)}^a(t) - \mathbf{x}^a(t) \right)\left( \mathbf{x}_{(i)}^a(t) - \mathbf{x}^a(t) \right)^\top \\
\end{aligned} y ( i ) f ( t ) x ( i ) a ( t ) x a ( t ) P a ( t ) = h ( x ( i ) f ( t )) + ϵ ( i ) ( t ) = x ( i ) f ( t ) + K ( y ( t ) − y ( i ) f ( t ) ) = N e 1 i = 1 ∑ N e x ( i ) a ( t ) = N e − 1 1 i = 1 ∑ N e ( x ( i ) a ( t ) − x a ( t ) ) ( x ( i ) a ( t ) − x a ( t ) ) ⊤ Computational Issues ¶ Sherman-Morrison Woodbury
( A B ⊤ ( C + B A B ⊤ ) ) − 1 = ( A − 1 + B ⊤ C B ) − 1 B ⊤ C − 1 (AB^\top (C + BAB^\top))^{-1} = (A^{-1} + B^\top CB)^{-1}B^\top C^{-1} ( A B ⊤ ( C + B A B ⊤ ) ) − 1 = ( A − 1 + B ⊤ CB ) − 1 B ⊤ C − 1 Kalman Gain (Observation State)
K = P H ⊤ ( H P H ⊤ + R ) − 1 \mathbf{K} = \mathbf{PH}^\top (\mathbf{HPH}^\top + \mathbf{R})^{-1} K = PH ⊤ ( HPH ⊤ + R ) − 1 Reduced Space
K = S ( I + ( H S ) ⊤ R − 1 H S ) − 1 ( H S ) ⊤ R − 1 \mathbf{K} = \mathbf{S}(\mathbf{I} + (\mathbf{HS})^\top \mathbf{R}^{-1}\mathbf{HS})^{-1}(\mathbf{HS})^\top \mathbf{R}^{-1} K = S ( I + ( HS ) ⊤ R − 1 HS ) − 1 ( HS ) ⊤ R − 1 where S = x N − 1 \mathbf{S} = \frac{\mathbf{x}}{\sqrt{N-1}} S = N − 1 x
Inverting Symmetric Definite Matrices ¶ Standard
C_inv = linalg.pinv(C, hermitian=True)SVD
P = U Λ V ⊤ P − 1 = V Λ − 1 U ⊤ \begin{aligned}
\mathbf{P} &= \mathbf{U\Lambda V}^\top \\
\mathbf{P}^{-1} &= \mathbf{V} {\Lambda}^{-1} \mathbf{U}^\top
\end{aligned} P P − 1 = UΛV ⊤ = V Λ − 1 U ⊤ U, D, VH = linalg.svd(C, full_matrices=False, hermitian=True)
C_inv = VH.T @ diag(1/D) @ U.TNote : This is useful for data assimilation with square roots
Cholesky
P = L L ⊤ L U = I L ⊤ P − 1 = U \begin{aligned}
\mathbf{P} &= \mathbf{LL}^\top \\
\mathbf{LU} &= \mathbf{I} \\
\mathbf{L}^\top \mathbf{P}^{-1} &= \mathbf{U}
\end{aligned} P LU L ⊤ P − 1 = LL ⊤ = I = U L = linalg.cholesky(C)
U = solve_triangular(L, eye(L.shape[1]), lower=True)
C_inv = solve_triangular(L.transpose(1, 2), U, lower=False)Note : This is useful for optimal interpolation.
Analysis Step (SVD) ¶ Forecast Ensemble Anomalies
S f = 1 N − 1 ( x f − x ˉ f ) \mathbf{S}^f = \frac{1}{N-1}\left( \mathbf{x}^f - \bar{\mathbf{x}}^f \right) S f = N − 1 1 ( x f − x ˉ f ) Forecast Ensemble Covariance
P f = S f S f ⊤ \mathbf{P}^f = \mathbf{S}^f\mathbf{S}^{f^\top} P f = S f S f ⊤ Kalman Gain (Ensemble Space)
K = S f ( I + ( H S ) ⊤ R − 1 H S ) − 1 ( H S f ) ⊤ R − 1 \mathbf{K} = \mathbf{S}^f\left( \mathbf{I} + (\mathbf{HS})^\top \mathbf{R}^{-1}\mathbf{HS} \right)^{-1}( \mathbf{HS}^f)^\top \mathbf{R}^{-1} K = S f ( I + ( HS ) ⊤ R − 1 HS ) − 1 ( HS f ) ⊤ R − 1 S @ inv(eye(Ne) + (H @ S).T @ inv(R) @ H @ S) @ (H @ S).T @ inv(R)[Ns x D] = [Ns x D] @ []
SVD of Signal/Noise
( H S ) ⊤ R − 1 H S = U Λ U ⊤ (\mathbf{HS})^\top\mathbf{R}^{-1}\mathbf{HS} = \mathbf{U\Lambda U}^\top ( HS ) ⊤ R − 1 HS = UΛU ⊤ Kalman Gain with SVD
K = S f U ( I + Λ ) − 1 / 2 U ⊤ \mathbf{K} = \mathbf{S}^f \mathbf{U}(\mathbf{I}+\mathbf{\Lambda})^{-1/2}\mathbf{U}^\top K = S f U ( I + Λ ) − 1/2 U ⊤ Analysis Ensemble of Anomalies
S a = S f U ( I + Λ ) − 1 / 2 U ⊤ \mathbf{S}^a = \mathbf{S}^f \mathbf{U}(\mathbf{I} + \mathbf{\Lambda})^{-1/2} \mathbf{U}^\top S a = S f U ( I + Λ ) − 1/2 U ⊤ Analysis Ensemble of Covariances
P a = P f − K H P f = S a S a ⊤ \mathbf{P}^a = \mathbf{P}^f - \mathbf{KHP}^f = \mathbf{S}^a\mathbf{S}^{a\;\top} P a = P f − KHP f = S a S a ⊤ Analysis Mean
x ˉ a = x ˉ f + K ( y obs − H x ˉ f ) \bar{\mathbf{x}}^a = \bar{\mathbf{x}}^f + \mathbf{K}(\mathbf{y}_\text{obs} - \mathbf{H}\bar{\mathbf{x}}^f) x ˉ a = x ˉ f + K ( y obs − H x ˉ f ) Analysis Ensemble
x a = x ˉ a + N − 1 S a \mathbf{x}^a = \bar{\mathbf{x}}^a + \sqrt{N - 1}\;\mathbf{S}^a x a = x ˉ a + N − 1 S a Perturbed Observations ¶ Sequential Data Assimilation with Non-Linear Quasi-Geostrophic Model using Monte Carlo methods to forecast error statistics - Evensen (1994) Data Assimilation using Ensemble Kalman Filter Technique - Houtekamer & Mitchell (1998) Inputs
We have some ensemble members.
X e ∈ R N e × D x \mathbf{X}_{e} \in \mathbb{R}^{N_e \times D_x} X e ∈ R N e × D x y ∈ R D y \mathbf{y} \in \mathbb{R}^{D_y} y ∈ R D y Sample the Observations given the state
Here, we propagate these samples through our emission function, H \boldsymbol{H} H
Y e = H ( X e ; θ ) \begin{aligned}
\mathbf{Y}_e &= \boldsymbol{H}(\mathbf{X}_e; \boldsymbol{\theta})
\end{aligned} Y e = H ( X e ; θ ) where Y e ∈ R N e × D y \mathbf{Y}_e \in \mathbb{R}^{N_e \times D_y} Y e ∈ R N e × D y
# emission function
Y_ens = H(X_ens, t, rng)[Ns x Dy]
Calculate 1st Moment
Now we need to calculate the 1st moment (the mean) to approximate the Gaussian distribution. We need to do this for both the predicted state, X e \mathbf{X}_e X e Y e \mathbf{Y}_e Y e
x ˉ = 1 N e ∑ i N e X e ( i ) y ˉ = 1 N e ∑ i N e Y e ( i ) \begin{aligned}
\bar{\mathbf{x}} &= \frac{1}{N_e}\sum_{i}^{N_e} \mathbf{X}_{e\;(i)} \\
\bar{\mathbf{y}} &= \frac{1}{N_e}\sum_{i}^{N_e} \mathbf{Y}_{e\;(i)} \\
\end{aligned} x ˉ y ˉ = N e 1 i ∑ N e X e ( i ) = N e 1 i ∑ N e Y e ( i ) where x ˉ ∈ R D \bar{\mathbf{x}} \in \mathbb{R}^{D} x ˉ ∈ R D
Now we need to calculate the 1st and 2nd moments (mean, covariance) to approximate the Gaussian distribution. We need to do this for both the predicted state, X e \mathbf{X}_e X e Y e \mathbf{Y}_e Y e
x ˉ = X e − \begin{aligned}
\bar{\mathbf{x}} = \mathbf{X}_e -
\end{aligned} x ˉ = X e − # mean of state ensembles (Dx)
x_mu = mean(X_ens)
# mean of obs ensembles (Dy)
y_mu = mean(Y_ens)Calculate Residuals
We calculate the errors (residuals) for our ensemble of predictions, Y e \mathbf{Y}_e Y e y obs \mathbf{y}_\text{obs} y obs
Resources ¶ Auto-differentiable Ensemble Kalman Filters - Chen et al (2021) - arxiv | PyTorch