PhyDA: Physics-Guided Diffusion Models for Data Assimilation in Atmospheric Systems

TL;DR

This work addresses the challenge of physically plausible data assimilation in atmospheric systems by integrating domain physics into diffusion-based generative modeling. It introduces PhyDA, a physics-guided diffusion framework that combines a Physically Regularized Diffusion Objective (PRDO) with a Virtual Reconstruction Encoder (VRE) to enforce PDE-consistent reconstructions from sparse observations. The approach is grounded in an energy-based interpretation of the diffusion score and leverages an Adaptive Physics-Proxy Operator for efficient PDE handling, achieving superior SpecDiv, MSE, MAE, and RMSE on ERA5 data compared to strong baselines. The results demonstrate that coupling physics with generative diffusion improves both numerical accuracy and physical realism, suggesting a practical path toward more reliable real-world data assimilation systems; future work will explore multimodal observations and end-to-end assimilation-prediction pipelines.

Abstract

Data Assimilation (DA) plays a critical role in atmospheric science by reconstructing spatially continous estimates of the system state, which serves as initial conditions for scientific analysis. While recent advances in diffusion models have shown great potential for DA tasks, most existing approaches remain purely data-driven and often overlook the physical laws that govern complex atmospheric dynamics. As a result, they may yield physically inconsistent reconstructions that impair downstream applications. To overcome this limitation, we propose PhyDA, a physics-guided diffusion framework designed to ensure physical coherence in atmospheric data assimilation. PhyDA introduces two key components: (1) a Physically Regularized Diffusion Objective that integrates physical constraints into the training process by penalizing deviations from known physical laws expressed as partial differential equations, and (2) a Virtual Reconstruction Encoder that bridges observational sparsity for structured latent representations, further enhancing the model's ability to infer complete and physically coherent states. Experiments on the ERA5 reanalysis dataset demonstrate that PhyDA achieves superior accuracy and better physical plausibility compared to state-of-the-art baselines. Our results emphasize the importance of combining generative modeling with domain-specific physical knowledge and show that PhyDA offers a promising direction for improving real-world data assimilation systems.

Figures (5)

• Figure 1: Figure (a) provides a systematic framework that calibrates background states with sparse observations to produce analysis. Figure (b) shows the challenge caused by observational sparsity. The absence of neighbor observations hinders the calculation of spatial gradients which is common in physical PDE calculation.
• Figure 2: Overview of the PhyDA framework. Sparse observations and gridded atmospheric variables (background) are first normalized. (a) The background is further encoded by a variational autoencoder into a unified latent space where a score-based diffusion model simulates stochastic trajectories. (c) During training, a Physically Regularized Diffusion Objective enforces consistency with physical laws by minimizing PDE residuals while an Adaptive Physics-Proxy Operator is adopted to simplify the calculation of partial derivative in complex PDEs. (b) A Virtual Reconstruction Encoder extracts hierarchical features from sparse inputs for conditioning the diffusion process. The model reconstructs physically consistent analysis states that align with both observations and governing physical PDEs.
• Figure 3: Ablation study.
• Figure 4: Realistic case where we use PhyDA to assimilate t@250hPa.
• Figure 5: Model performance with multiple observational sparsity ($\{10\%, 5\%, 2\%, 1\%, 0.5\%\}$).