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

**Authors:** Shunji Kotsuki, Kenta Shiraishi, Atsushi Okazaki  
**arXiv:** [2407.17781](https://arxiv.org/abs/2407.17781)  
**Date:** 2024-07-25  
**Summary page:** https://www.summarizepaper.com/en/paper/2407.17781v4/

## Key points

- AI-based weather prediction research is competitive with traditional dynamic numerical weather prediction models
- Integration of AI-based weather prediction models with data assimilation techniques has been limited due to long-term sequential data assimilation cycles
- Novel approach introduced using ensemble data assimilation to diagnose AI-based weather prediction models, specifically ClimaX
- ClimaX model is a ViT-based AI weather prediction model designed for global atmospheric forecasting
- Utilizes variable tokenization and aggregation for flexibility and generality in architecture
- Low-resolution version of ClimaX (version 0.3.1) employed in the study with specific grid configurations and vertical model levels
- Model trained on five variables initially, updated for data assimilation to predict additional variables for accurate forecasts
- Training curves showed improved anomaly correlation coefficients and reduced root mean square errors after updating the model
- Local Ensemble Transform Kalman Filter (LETKF) adapted for use with ClimaX model as a widely used data assimilation method in operational NWP centers like ECMWF, DWD, and JMA
- Ensemble data assimilation effectively utilized to evaluate AI-based weather prediction models like ClimaX by assessing physical consistency, error growth representation, and forecast accuracy

## Summary

, , , , 


Artificial intelligence (AI)-based weather prediction research is rapidly advancing and has proven to be competitive with traditional dynamic numerical weather prediction models. However, the integration of AI-based weather prediction models with data assimilation techniques has been limited due to the need for long-term sequential data assimilation cycles for evaluation. This study introduces a novel approach using ensemble data assimilation to diagnose AI-based weather prediction models, marking the first successful implementation of ensemble Kalman filter with an AI-based model called ClimaX. The ClimaX model is a ViT-based AI weather prediction model designed for global atmospheric forecasting. It utilizes variable tokenization and aggregation to enhance flexibility and generality in its architecture. In this study, the low-resolution version of ClimaX (version 0.3.1) was employed, featuring 64 zonal grids and 32 meridional grids at a spatial resolution of 5.625° × 5.625°, with seven vertical model levels ranging from 900 hPa to 50 hPa. By default, ClimaX is trained on five variables: geopotential at 500 hPa, temperature at 850 hPa, temperature at 2 m, zonal wind at 10 m, and meridional wind at 10 m. The model was updated for data assimilation purposes to predict additional variables required for accurate forecasts. Surface pressure was diagnosed based on geopotential and surface elevation inputs. Training curves comparing the default and updated ClimaX models against WeatherBench data demonstrated improved anomaly correlation coefficients and reduced root mean square errors after training. The study also introduced the Local Ensemble Transform Kalman Filter (LETKF) as a widely used data assimilation method in operational NWP centers like ECMWF, DWD, and JMA. The LETKF was adapted for use with the ClimaX model by replacing the SPEEDY weather prediction model in the existing system with ClimaX. Overall, this research showcases how ensemble data assimilation can be effectively utilized to evaluate AI-based weather prediction models like ClimaX by assessing physical consistency, error growth representation, and forecast accuracy. The findings suggest that ensemble data assimilation enhances the stability of AI-based weather predictions while providing valuable insights into their performance compared to traditional dynamical NWP models.