Neural variational Data Assimilation with Uncertainty Quantification using SPDE priors
This work addresses uncertainty quantification in data assimilation for geophysical applications, offering an incremental improvement over existing neural methods by incorporating SPDE-based priors.
The paper tackles the problem of spatio-temporal interpolation in geophysical data by extending a neural variational data assimilation framework to include uncertainty quantification using Stochastic Partial Differential Equations (SPDE) priors, achieving performance matching Optimal Interpolation in Gaussian cases and outperforming it for nonlinear dynamics with fast parameter estimation.
The spatio-temporal interpolation of large geophysical datasets has historically been addressed by Optimal Interpolation (OI) and more sophisticated equation-based or data-driven Data Assimilation (DA) techniques. Recent advances in the deep learning community enables to address the interpolation problem through a neural architecture incorporating a variational data assimilation framework. The reconstruction task is seen as a joint learning problem of the prior involved in the variational inner cost, seen as a projection operator of the state, and the gradient-based minimization of the latter. Both prior models and solvers are stated as neural networks with automatic differentiation which can be trained by minimizing a loss function, typically the mean squared error between some ground truth and the reconstruction. Such a strategy turns out to be very efficient to improve the mean state estimation, but still needs complementary developments to quantify its related uncertainty. In this work, we use the theory of Stochastic Partial Differential Equations (SPDE) and Gaussian Processes (GP) to estimate both space-and time-varying covariance of the state. Our neural variational scheme is modified to embed an augmented state formulation with both state and SPDE parametrization to estimate. We demonstrate the potential of the proposed framework on a spatio-temporal GP driven by diffusion-based anisotropies and on realistic Sea Surface Height (SSH) datasets. We show how our solution reaches the OI baseline in the Gaussian case. For nonlinear dynamics, as almost always stated in DA, our solution outperforms OI, while allowing for fast and interpretable online parameter estimation.