# Ensemble data assimilation to diagnose AI-based weather prediction model: A case with ClimaX version 0.3.1

arXiv: https://arxiv.org/abs/2407.17781 · HTML: https://mlsanity.com/paper/2407.17781/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 46 references that are in the map.
- First posted on arXiv 25 Jul 2024.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [WeatherBench: A benchmark dataset for data-driven weather forecasting](https://mlsanity.com/paper/2002.00469/) (2020)
- [Neural General Circulation Models for Weather and Climate](https://mlsanity.com/paper/2311.07222/) (2023)
- [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://mlsanity.com/paper/2308.15560/) (2023)
- [Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere](https://mlsanity.com/paper/2306.03838/) (2023)
- [Scaling transformer neural networks for skillful and reliable medium-range weather forecasting](https://mlsanity.com/paper/2312.03876/) (2023)
- [Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models](https://mlsanity.com/paper/2102.05107/) (2021)
- [Combining data assimilation and machine learning to infer unresolved scale parametrisation](https://mlsanity.com/paper/2009.04318/) (2020)
- [FengWu-4DVar: Coupling the Data-driven Weather Forecasting Model with 4D Variational Assimilation](https://mlsanity.com/paper/2312.12455/) (2023)
- [Integrating Recurrent Neural Networks with Data Assimilation for Scalable Data-Driven State Estimation](https://mlsanity.com/paper/2109.12269/) (2021)
- [Towards an end-to-end artificial intelligence driven global weather forecasting system](https://mlsanity.com/paper/2312.12462/) (2023)
- [Deep learning-enhanced ensemble-based data assimilation for high-dimensional nonlinear dynamical systems](https://mlsanity.com/paper/2206.04811/) (2022)
- [Learning to Assimilate in Chaotic Dynamical Systems](https://mlsanity.com/paper/2111.01058/) (2021)
- [Data Assimilation with Machine Learning Surrogate Models: A Case Study with FourCastNet](https://mlsanity.com/paper/2405.13180/) (2024)
- [Learning Optimal Filters Using Variational Inference](https://mlsanity.com/paper/2406.18066/) (2024)

## Built on by (most cited first, 1 of 1)

- [Bridging Artificial Intelligence and Data Assimilation: The Data-driven Ensemble Forecasting System ClimaX-LETKF](https://mlsanity.com/paper/2512.14444/) (2025)

## Abstract (arXiv)

Artificial intelligence (AI)-based weather prediction research is growing rapidly and has shown to be competitive with the advanced dynamic numerical weather prediction models. However, research combining AI-based weather prediction models with data assimilation remains limited partially because long-term sequential data assimilation cycles are required to evaluate data assimilation systems. This study proposes using ensemble data assimilation for diagnosing AI-based weather prediction models, and marked the first successful implementation of ensemble Kalman filter with AI-based weather prediction models. Our experiments with an AI-based model ClimaX demonstrated that the ensemble data assimilation cycled stably for the AI-based weather prediction model using covariance inflation and localization techniques within the ensemble Kalman filter. While ClimaX showed some limitations in capturing flow-dependent error covariance compared to dynamical models, the AI-based ensemble forecasts provided reasonable and beneficial error covariance in sparsely observed regions. In addition, ensemble data assimilation revealed that error growth based on ensemble ClimaX predictions was weaker than that of dynamical NWP models, leading to higher inflation factors. A series of experiments demonstrated that ensemble data assimilation can be used to diagnose properties of AI weather prediction models such as physical consistency and accurate error growth representation.