Skip to content

TUTORIAL: Real-time data assimilation

import numpy as np
from scipy import linalg
from romda.models.physical import VdP
from romda.estimators import EnSRKF
rng = np.random.default_rng(0)


def EnKF(Af, d, Cdd, M):
    """Ensemble Kalman Filter as derived in Evensen (2009) eq. 9.27.
        Inputs:
            Af: forecast ensemble at time t
            d: observation at time t
            Cdd: observation error covariance matrix
            M: matrix mapping from state to observation space
        Returns:
            Aa: analysis ensemble (or Af is Aa is not real)
    """
    m = np.size(Af, 1)

    psi_f_m = np.mean(Af, 1, keepdims=True)
    Psi_f = Af - psi_f_m

    # Create an ensemble of observations
    D = rng.multivariate_normal(d, Cdd, m).transpose()

    # Mapped forecast matrix M(Af) and mapped deviations M(Af')
    Y = np.dot(M, Af)
    S = np.dot(M, Psi_f)

    # Matrix to invert
    C = (m - 1) * Cdd + np.dot(S, S.T)
    Cinv = linalg.inv(C)

    X = np.dot(S.T, np.dot(Cinv, (D - Y)))

    Aa = Af + np.dot(Af, X)

    return Aa

1. State and parameter estimation

Motivation: uncertain model parameters

from romda.observations import Observations

alpha_true = dict(beta = 70.,
                  kappa = 4.0,
                  zeta = 60.,
                  Nq=2,
                  dt=2e-4,
                  psi0=rng.random(2))

truth = Observations(VdP, 
                     t_start=1., 
                     t_stop=1.5, 
                     t_max=2.5, 
                     Nt_obs=30, 
                     std_obs = 0.1,
                     noise_type = 'gaussian, additive',
                     **alpha_true
                    )


# truth.y_true.shape
truth.plot_truth(truth, Nq=2, window=.2, f_max=2000)
ensemble = EnSRKF(parent_model=VdP, 
                       Nq=2,                
                       m=20,               # Number of ensemble members
                       std_phi=.8,        # Initial uncertainty in the state
                       psi0=np.array([1, 1000]),
                       # Initial guess on the parameters
                       beta = 67.,
                       kappa = 4.0,
                       zeta = 65.,
                    )
filter_ens = ensemble.copy()

# Define measurement error covariance matrix

var_d  = [.1, 100]
Cdd = np.eye(filter_ens.model.Nq) * var_d


# ----------------------------------------------------------------------------

for d, t_d in zip(truth.y_obs, truth.t_obs):
    # Parallel forecast to next observation

    filter_ens.forecast_step(t_end=t_d)
    filter_ens.assimilated_data = (d, t_d)

    # Create augmented state matrix
    Af = np.vstack([filter_ens.current_state,
                   filter_ens.model.get_observables()]) # augmented state matrix [phi; alpha; q] x m

    # Perform assimilation 
    Aa = EnKF(Af, d, Cdd, filter_ens.model.M)  # Analysis step

    # Update the initial condition for the next forecast
    filter_ens.update_history(Aa[:-filter_ens.model.Nq, :],
                              t=t_d,
                              modify_saved_states=True)

# ----------------------------------------------------------------------------

#Forecast the ensemble further without assimilation
t_extra = filter_ens.model.hist_t[-1] + filter_ens.model.t_CR

filter_ens.forecast_step(t_extra, averaged=True, close=True)
filter_ens.visualize_history(plot_members=True, truth=truth)

Therefore, we need to account for possible errors in the parameters. As discussed in the previois tutorial, we model aleatoric errors as Gaussian processes.


1.1 Augmented State-Space Formulation

To estimate both the state \(\boldsymbol{\phi}\) and parameters \(\boldsymbol{\alpha}\), we define an augmented state vector:

\[ \boldsymbol{\psi} = \begin{bmatrix} \boldsymbol{\phi} \\ \boldsymbol{\alpha} \\ \mathbf{q} \end{bmatrix} \]

We can treat the parameters as state variables which are constant in time, such that the augmented state-space formulation reads

\[ \left\{ \begin{aligned} \mathrm{d} \begin{bmatrix} \boldsymbol{\phi} \\ \boldsymbol{\alpha} \\ \mathbf{q} \end{bmatrix} &= \begin{bmatrix} \mathcal{F}(\boldsymbol{\phi} + \boldsymbol{\epsilon}_\phi, \boldsymbol{\alpha} + \boldsymbol{\epsilon}_\alpha) \\ \mathbf{0}_{N_\alpha} \\ \mathbf{0}_{N_q} \end{bmatrix} \mathrm{d}t \\ \mathbf{q} &= \mathcal{M}(\boldsymbol{\phi}) + \boldsymbol{\epsilon}_q \end{aligned} \right. \quad \leftrightarrow \quad \left\{ \begin{aligned} \mathrm{d}\boldsymbol{\psi} &= \mathbf{F}(\boldsymbol{\psi} + \boldsymbol{\epsilon}_\psi)\mathrm{d}t \\ \mathbf{q} &= \mathbf{M} \boldsymbol{\psi} + \boldsymbol{\epsilon}_q \end{aligned} \right. \]

Here: - \(\mathbf{F}\): augmented nonlinear operator
- \(\boldsymbol{\epsilon}_\psi\): augmented uncertainty
- \(\mathbf{M} = \begin{bmatrix} \mathbf{0} & \mathbb{I}_{N_q} \end{bmatrix}\): linear measurement operator

This augmented state-space approach effectively linearises the nonlinear observation operator around the point of observation, which simplifies the derivation of data assimilation methods.

1.2 Augmented ensemble statistics

\[ \begin{aligned} \mathbb{E}(\boldsymbol{\psi})\approx\bar{\boldsymbol{\psi}}=\dfrac{1}{m}\sum^m_{j=1}{\boldsymbol{\psi}_j} \quad \text{and} \quad \mathbf{C}_{\psi\psi} = \begin{bmatrix} \mathbf{C}_{\phi\phi} & \mathbf{C}_{\phi\alpha}& \mathbf{C}_{\phi q} \\ \mathbf{C}_{\alpha \phi} & \mathbf{C}_{\alpha\alpha}& \mathbf{C}_{\alpha q} \\ \mathbf{C}_{q \phi} & \mathbf{C}_{q \alpha}& \mathbf{C}_{q q} \\ \end{bmatrix} \approx\dfrac{1}{m-1}\sum^m_{j=1}(\boldsymbol{\psi}_i-\bar{\boldsymbol{\psi}})\otimes(\boldsymbol{\psi}_i-\bar{\boldsymbol{\psi}}). \end{aligned} \]

Each ensemble member \(j\) is forecast independently in time with \(\mathbf{F}(\boldsymbol{\psi}_j)\) to obtain an ensemble of forecast states \(\boldsymbol{\psi}_j^\text{f}\). When a sensor provides noisy data \(\mathbf{d}\), real-time data assimilation statistically combines the noisy data and the forecast ensemble to improve our knowledge in the system's parameters and states (i.e., to compute an analysis ensemble \(\boldsymbol{\psi}_j^\mathrm{a}\)). Mathematically, we aim to minimize the cost function

\[ \begin{aligned} \mathcal{J}(\boldsymbol{\psi}_j) = &\left\|\boldsymbol{\psi}_j-\boldsymbol{\psi}_j^\mathrm{f}\right\|^2_{\mathbf{C}^{\mathrm{f}^{-1}}_{\psi\psi}} + \left\|{\boldsymbol{y}}_j-\boldsymbol{d}_j\right\|^2_{\mathbf{C}^{-1}_{dd}}, \quad \mathrm{for} \quad j=0,\dots,m-1, \end{aligned} \]

where \(\left\|\cdot\right\|^2_{\mathbf{C}^{-1}}\) is the L2-norm weighted by the semi-positive definite matrix \(\mathbf{C}^{-1}\). The ensemble Kalman filter (EnKF) minimize the cost function to obtain an analysis ensemble \(\boldsymbol{\psi}_j^\mathrm{a}\) from the forecast ensemble \(\boldsymbol{\psi}_j^\mathrm{f}\) and the observations \(\mathbf{d}\) as

\[ \begin{aligned} \boldsymbol{\psi}_j^\mathrm{a} &= \boldsymbol{\psi}_j^\mathrm{f}+\mathbf{K}\left[\mathbf{d}_j - \mathbf{M}\boldsymbol{\psi}_j^\mathrm{f}\right], \quad j=0,\dots,m-1, \end{aligned} \]

where \(\mathbf{K}=\mathbf{C}_{\psi\psi}^\mathrm{f}\mathbf{M}^\mathrm{T}\left(\mathbf{C}_{dd}+\mathbf{M}\mathbf{C}_{\psi\psi}^\mathrm{f}\mathbf{M}^\mathrm{T}\right)^{-1}\) is the Kalman gain matrix.


2. Test case: twin experiment on the Van der Pol oscillator

The time evolution of the Van der Pol oscillator is governed by the second-order differential equation \begin{aligned} \ddot{\eta} + \omega^2{\eta} = \dot{q} - \zeta\dot{\eta}, \end{aligned} where \(\omega\) is the angular oscillating frequency, \(\zeta\) is the damping coefficient, and \(\dot{q}\) is a forcing term. One application of this system is thermoacoustic systems, in which \(\eta\) represents the acoustic velocity, and \(\dot{q}\) is the heat release rate, which can be modelled as
\dot{q} = \beta\dot{\eta}\left(1 - \dfrac{\kappa\eta^2}{\beta + \kappa\eta^2}\right), where \(\kappa\) is the nonlinearity coefficient, and \(\beta\) is the forcing strength. Using this heat release law, the Van der Pol oscillator can be written as the system of ordinary differential equations \left{ \begin{array}{rcl} \dfrac{\mathrm{d}\eta}{\mathrm{d} t} &=& \mu\ \dfrac{\mathrm{d}\mu}{\mathrm{d} t} &=& -\omega^2\eta + \mu \left(\beta - \zeta -\dfrac{\beta\kappa\eta^2}{\beta + {\kappa}\eta^2}\right). \end{array} \right. In state-space notation, the state vector is \(\boldsymbol{\phi} = [\mu; \eta]\), the model parameters are \(\boldsymbol{\alpha} = [\zeta; \beta; \kappa]\); and the model estimate (i.e., the measurable quantity) is the acoustic velocity \(\eta\). Therefore, the augmented vector to estimate is

\[ \boldsymbol{\psi} = \begin{bmatrix} \boldsymbol{\phi}\\ \boldsymbol{\alpha}\\ \mathbf{q} \end{bmatrix} = \begin{bmatrix} {\mu}\\ {\eta}\\ \zeta\\ \beta\\ \kappa\\ \eta \end{bmatrix} \]

Create forecast ensemble with uncertain parameters

fixed_params = dict(beta = 70.,
                  kappa = 4.0,
                  zeta = 60.)


alpha0 = dict(beta = (55., 70.),
              # kappa=(3.,4.5),
            #   zeta=(60., 62.),
              )

for key in alpha0.keys():
    fixed_params.pop(key)


forecast_params = dict(Nq = 2,           
                       m=10,               # Number of ensemble members
                       std_phi=0.1,        # Initial uncertainty in the state
                       std_alpha=alpha0,       # Initial uncertainty in the parameters
                       **fixed_params
                       )


ensemble = EnSRKF(parent_model=VdP,  **forecast_params) # type: ignore

Apply data assimilation

rng = np.random.default_rng(0)


filter_ens = ensemble.copy()
pm = filter_ens.model

# Define measurement error covariance matrix
var_d  = [.1, 100]
Cdd = np.eye(filter_ens.model.Nq) * var_d


# ----------------------------------------------------------------------------

for d, t_d in zip(truth.y_obs, truth.t_obs):
    # Parallel forecast to next observation

    filter_ens.forecast_step(t_end=t_d)

    # Create augmented state matrix
    Af = np.vstack([pm.current_state, 
                    pm.get_observables()]) # augmented state matrix [phi; alpha; q] x m

    # Perform assimilation 
    Aa = EnKF(Af, d, Cdd, pm.M)  # Analysis step

    Aa_m = np.mean(Aa, axis=-1, keepdims=True)
    Aa = Aa_m + 1.01 * (Aa - Aa_m)

    # Update the initial condition for the next forecast
    filter_ens.update_history(Aa[:-pm.Nq, :],
                              t=t_d,
                              modify_saved_states=True)

    filter_ens.assimilated_data = (d, t_d)
# ----------------------------------------------------------------------------


#Forecast the ensemble further without assimilation
t_extra = filter_ens.model.hist_t[-1] + filter_ens.model.t_CR

filter_ens.forecast_step(t_extra, averaged=True, close=True)
filter_ens.visualize_state(time_indices=[0, -1])
filter_ens.visualize_history(plot_members=True, truth=truth, reference_a=alpha_true)



Exercises

How would you expect the solution to change if we decrease/increase the number of uncertain parameters?

How would you expect the solution to change if we decrease/increase the ensemble size (m) or the observation noise (std_d)?


What if the dynamics are chaotic? In tutorial 15 we test the framework in a prototypical chaotic system, the Lorenz 63.