PnP-DA: Towards Principled Plug-and-Play Integration of Variational Data Assimilation and Generative Models

TL;DR

The paper introduces PnP-DA, a novel plug-and-play data assimilation method that combines gradient-based analysis updates with pretrained generative models, improving forecast accuracy in chaotic systems.

Contribution

It presents a new algorithm that relaxes Gaussian assumptions in data assimilation by integrating generative models without backpropagation, enhancing forecast accuracy.

Findings

Consistently reduces forecast errors in chaotic testbeds.

Outperforms classical variational methods across various conditions.

Effectively handles sparse and noisy observations.

Abstract

Earth system modeling presents a fundamental challenge in scientific computing: capturing complex, multiscale nonlinear dynamics in computationally efficient models while minimizing forecast errors caused by necessary simplifications. Even the most powerful AI- or physics-based forecast system suffer from gradual error accumulation. Data assimilation (DA) aims to mitigate these errors by optimally blending (noisy) observations with prior model forecasts, but conventional variational methods often assume Gaussian error statistics that fail to capture the true, non-Gaussian behavior of chaotic dynamical systems. We propose PnP-DA, a Plug-and-Play algorithm that alternates (1) a lightweight, gradient-based analysis update (using a Mahalanobis-distance misfit on new observations) with (2) a single forward pass through a pretrained generative prior conditioned on the background forecast via a conditional Wasserstein coupling. This strategy relaxes restrictive statistical assumptions and leverages rich historical data without requiring an explicit regularization functional, and it also avoids the need to backpropagate gradients through the complex neural network that encodes the prior during assimilation cycles. Experiments on standard chaotic testbeds demonstrate that this strategy consistently reduces forecast errors across a range of observation sparsities and noise levels, outperforming classical variational methods.