[Paper Review] DiffDA: a Diffusion Model for Weather-scale Data Assimilation

Langwen Huang, Lukas Gianinazzi+3 more|arXiv (Cornell University)|Jan 11, 2024

Meteorological Phenomena and Simulations10 citations

TL;DR

DiffDA uses a denoising diffusion model, built on GraphCast, to assimilate high-resolution atmospheric data with predicted states and sparse observations, enabling real-time-like reanalysis and forecast-ready initial conditions with up to 24 hours lead-time loss.

ABSTRACT

The generation of initial conditions via accurate data assimilation is crucial for weather forecasting and climate modeling. We propose DiffDA as a denoising diffusion model capable of assimilating atmospheric variables using predicted states and sparse observations. Acknowledging the similarity between a weather forecast model and a denoising diffusion model dedicated to weather applications, we adapt the pretrained GraphCast neural network as the backbone of the diffusion model. Through experiments based on simulated observations from the ERA5 reanalysis dataset, our method can produce assimilated global atmospheric data consistent with observations at 0.25 deg (~30km) resolution globally. This marks the highest resolution achieved by ML data assimilation models. The experiments also show that the initial conditions assimilated from sparse observations (less than 0.96% of gridded data) and 48-hour forecast can be used for forecast models with a loss of lead time of at most 24 hours compared to initial conditions from state-of-the-art data assimilation in ERA5. This enables the application of the method to real-world applications, such as creating reanalysis datasets with autoregressive data assimilation.

Motivation & Objective

• Demonstrate a machine learning based data assimilation method capable of handling high-resolution atmospheric data.
• Integrate a pretrained weather forecast model as the backbone of a diffusion-based assimilator.
• Show conditioning on predicted states during training and inference, plus conditioning on sparse observations at inference.
• Enable autoregressive data assimilation to generate reanalysis data compatible with ERA5 in principle.

Proposed method

• Adapt GraphCast as a denoising diffusion model for data assimilation.
• Condition the diffusion model on the forecasted state  x̂ during training and inference.
• Condition the diffusion process on sparse observations at inference time using a soft masking and interpolation strategy.
• Use a two-phase conditioning approach to allow post-processing when observations are unavailable.
• Train with a diffusion objective to learn p(x^0 | x̂) and sample via reverse diffusion steps.
• Provide flexibility to plug in alternative forecast backbones beyond GraphCast.

Figure 1: Diagram of numerical weather forecast pipeline. It consists of data assimilation, forecast and post-processing. Data assimilation produces gridded values from sparse observations and predicted gridded values from previous time steps. Forecast takes in gridded values and produces prediction

Experimental results

Research questions

• RQ1 Can a diffusion-based model assimilate high-resolution atmospheric fields conditioned on both predicted states and sparse observations?
• RQ2 Does conditioning on the predicted state during training and inference improve assimilation closeness to ground truth?
• RQ3 Can the approach produce reanalysis-like data and remain forecast-ready with acceptable lead-time loss?
• RQ4 How does the method perform with autoregressive data assimilation and varying amounts of observations?

Key findings

• Assimilated data converge toward ground truth as more observations are used.
• When used as input to a forecast model, the 48h forecast error with assimilated data shows lead-time loss of at most 24 hours compared to using ground-truth initial conditions.
• The method enables generating reanalysis-like data with autoregressive assimilation cycles.
• Conditioning on predicted state during training and inference provides post-processing capabilities even without observations at inference time.
• A soft masking strategy for observations improves conditioning effectiveness over a hard mask.
• The approach scales to 0.25-degree resolution with 13 vertical levels using GraphCast as backbone.

Figure 2: Architecture of the diffusion based data assimilation method. We take advantage of the input and output shape of the pretrained GraphCast model which takes the state of the atmosphere at two time steps as input. In each iteration of the denoising diffusion process in our method, the adapte