Data assimilation algorithms are necessary to track or estimate the hidden state of chaotic systems through partial and noisy observations. Variational data assimilation aims to find a the best trajectory of the dynamical system which minimizes certain cost function.

This repositiory contains a new code for solving the weak-constraint 4dvar or simply weak-4dvar data assimilation for QG model. All the code in this repo is implemented in pytorch with handing experiment configurations using hydra.

Given the sequence of observations Observations sequence $Y^i=\left(y^i_0,y^i_1,...y^i_n\right)$ on $\left(\Omega_i\right) \in \Omega$, find the optimal trajectory $X^i=\left(x^i_0,x^i_1,...x^i_n\right)$ that minimizes the following cost function. The weak-4dvar cost function is:

$$\mathcal{J}(x^i_0,x^i_1,...x^i_n)=\sum_{k=1} ||x_k - \mathcal{M}(x_{k-1})||^2+ ||y_i-\mathcal{H}(x_i)||^2$$

The dynamical systems $\mathcal{M}$ takes the system state $x_k$ to $x_{k+1}$. Where as the above weak formulation of the 4dvar problem accounts for additional model errors in the dynamical system.

The first term minimizes the depatures from a pure model trajectory since the aim to find a trajectory close to the model trajectory while accounting for the model error- a part we refer to as the dynamical cost. The second term makes the trajectory fit to the obsrvations while accounting for the observation error and is referred as observation cost.

For testing algorithms, we extensive experiments on toy models such as L63/ L96 which are ODEs, have been shown in literature.
QG model is interesting PDE model of intermediate complexity for studying performace of state estimation problem in machine learning, deep learning problem in data assimilation. The vorticity is the dynamical variable and the observations are in the stream function space with some noise added to them.

Domain $=\left[0,2 \pi\right)^2$, $\omega $: vorticity, $\psi$: stream-function.

$$ \frac{\partial \omega}{\partial t}+ J(\psi, \omega)= \nu \nabla^2 \omega - \mu \omega - \beta \partial_x \psi ,, \quad \omega = \nabla^2 \psi ,, \quad $$

The current implementation has Quasi-geostrophic model on a grid of $128$ X $128$ dimension. The ground truth is a $1024$ X $1024$ simulated voriticity and stream function coarse grained on the low resolution grid. The PDE has PBC and we use psueod-spectral methods to solve them( all in pytorch.)

These observations are available on masks which have been obtained from Nadir Satelite altimetry tracks. The nature of these observations are realistic- they are really quite sparse! Below is a snapshot of vorticity field, the stream function and the kind of observations we have for any time. The observations and the masks are different at different times.

The paramters for any numerical experiments is loaded as a 'config.yaml' file and hydra is used to initialise experiments. The code is both CPU and GPU compatible. Here is a result from one of the initial conditions, (discussed further below ):

We solve the weak-4dvar problem for 10-day assimilation window. The
The default implementation performs optimization of the loss function of the weak-4dvar using SGD based algorithm, although this can be easily changed to a different optimizer that is written or available in pytorch. We have explored different initial conditions at the moment:

1. True i.c. - just to check for consistency and stability.
2. Blurred i.c. - to test stability and convergence to the true solution.
3. Gaussian random field- a prior with temporal correlation on the initial state.
4. Coherent-shifted field- to mimic psition based errors.

The last kind of error is the focus of my experiments. How I am simulating these kind of errors is part of another repository here-. But this kind of IC is specific to our study and not necessary for the weak-4dvar algorithm in general.

The folder structure is a as follows:

To experiment with new models, two things need to be worked upon-

1. The dataset and the dataloader within it for the observations of the system.
2. Create a pytorch implementation of 'your_dynamical_system.py' using pytorch's nn module.
3. The neural ode package which will be able to handle derivative computation via the adjoint implementation.

At this moment, the code is purely designed for my own experiments. If you find it useful and want to collaborate on an interesting idea in the space of data assimilation and dynamical systems, reach out to me via email or linkedin :).