Improving AI weather prediction models using global mass and energy conservation schemes

Yingkai Sha, John S. Schreck, William Chapman, David John Gagne II

NSF National Center for Atmospheric Research, Boulder, Colorado, USA

Abstract

Artificial Intelligence (AI) weather prediction (AIWP) models are powerful tools for medium-range forecasts but often lack physical consistency, leading to outputs that violate conservation laws. This study introduces a set of novel physics-based schemes designed to enforce the conservation of global dry air mass, moisture budget, and total atmospheric energy in AIWP models. The schemes are highly modular, allowing for seamless integration into a wide range of AI model architectures. Forecast experiments are conducted to demonstrate the benefit of conservation schemes using FuXi, an example AIWP model, modified and adapted for 1.0-degree grid spacing. Verification results show that the conservation schemes can guide the model in producing forecasts that obey conservation laws. The forecast skills of upper-air and surface variables are also improved, with longer forecast lead times receiving larger benefits. Notably, large performance gains are found in the total precipitation forecasts, owing to the reduction of drizzle bias. The proposed conservation schemes establish a foundation for implementing other physics-based schemes in the future. They also provide a new way to integrate atmospheric domain knowledge into the design and refinement of AIWP models.

Resources

• NSF NCAR Research Data Archive, ERA5 Reanalysis (0.25 Degree Latitude-Longitude Grid)

  • Note: this repository visits RDA internally, e.g., /glade/campaign/collections/rda/data/d633000/e5.oper.an.pl/197901/

• Google Research, Analysis-Ready, Cloud Optimized (ARCO) ERA5 [link]

• IFS-HRES and ERA5 climatology from the Weatherbench2 project [link]

• NSF NCAR MILES Community Research Earth Digital Intelligence Twin (CREDIT) [link]

Navigation

• Derivations of conservation schemes: mass conservation, energy conservation, Pytorch integration
• Results: Dry air mass conservation, Moisture budget conservation, Total atmospheric energy conservation, SEEPS

Model weights

Model name	Weights	Notes
FuXi-base	link	The 1.0 degree FuXi baseline without conservation schemes
FuXi-physics	link	The 1.0 degree FuXi run with conservation schemes

Citation

Sha, Y., J. Schreck, W. Chapman, D. J. Gagne II, 2025a: Improving AI weather prediction models using global mass and energy conservation schemes. In review: Journal of Advances in Modeling Earth Systems. Pre-print: https://arxiv.org/abs/2501.05648

@article{sha2025improving,
  title={Improving AI weather prediction models using global mass and energy conservation schemes},
  author={Sha, Yingkai and Schreck, John S and Chapman, William and Gagne II, David John},
  journal={arXiv preprint arXiv:2501.05648},
  year={2025}
}