Variational Autoencoder-Enhanced Variational Methods for Data Assimilation

Data assimilation (DA) is a statistical approach used to estimate the states of physical systems by integrating prior model predictions (background states x_b​) with observational data (y). This integration produces an accurate estimate, called analysis states (​x_a), by sampling or maximizing the posterior likelihood p(xx_b, y). In weather forecasting, background states are generated by imperfect models, and the likelihood p(xx_b) is often unknown. Observations, sourced from diverse instruments, are mapped to model space using observation operators (H). Effective DA algorithms must accurately estimate p(xx_b) while accommodating various observation operators, including those involving sparse, noisy or irregular data.

Traditional DA methods, such as variational assimilation, assume that the background error (​x - x_b) follows a Gaussian distribution independent of x_b​. This allows explicit computation of p(xx_b, y) and optimization via techniques like gradient descent. While robust to various observation operators, these methods depend heavily on expert knowledge to construct error correlations and are limited by their Gaussian assumptions.

Generative neural networks, particularly diffusion models, have emerged as alternatives for modeling p(xx_b). Notable examples include SDA and DiffDA, which use diffusion models to learn background distributions. SDA incorporates observations via diffusion posterior sampling, while DiffDA employs the repaint technique. These approaches improve on traditional methods by capturing more complex distributions but often struggle with sparse, irregular observations. For instance, DiffDA assumes grid-aligned data, while SDA relies on assumptions that can reduce accuracy in real-world scenarios.

In this research, we aim to develop a neural network-based data assimilation algorithm that not only captures the non-Gaussian characteristics of the conditional background distribution for enhanced accuracy but also effectively assimilates data under real-world observations (sparse, noisy and outside of the grid). We introduce VAE-Var, a novel data assimilation algorithm in which a variational autoencoder is first employed to learn the conditional background distribution and then the decoder component is utilized to construct a variational cost function, which, when optimized, yields the analysis states.

Key advantages of VAE-Var include:

• This algorithm inherits the framework of traditional variational assimilation by explicitly modeling the posterior probability function p(xx_b, y) and maximizing it to derive the analysis states. As a result, compared to other neural network data assimilation methods such as SDA and DiffDA, VAE-Var can better handle different types of observation operators, particularly irregular observations that do not fall on the grid points of the physical field.
• Unlike traditional variational assimilation algorithms, VAE-Var alleviates the dependence on expert knowledge for constructing the conditional background distribution, enabling the model to effectively capture non-Gaussian structures. This makes VAE-Var perform better in sparse observational settings.

Experiments with the FengWu weather forecasting system at 0.25° resolution show that VAE-Var achieves higher accuracy than DiffDA and traditional algorithms (interpolation and 3DVar) in sparse observational settings. When integrated with FengWu, VAE-Var reliably assimilates real-world GDAS prepbufr observations over a one-year period.