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

Kotsuki, Shunji; Shiraishi, Kenta; Okazaki, Atsushi

doi:https://doi.org/10.48550/arXiv.2407.17781

Preprints

https://doi.org/10.48550/arXiv.2407.17781

Preprints

15 Oct 2024

| 15 Oct 2024

Status: this preprint is open for discussion and under review for Geoscientific Model Development (GMD).

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

Shunji Kotsuki, Kenta Shiraishi, and Atsushi Okazaki

Abstract. 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.

Received: 04 Oct 2024 – Discussion started: 15 Oct 2024

Shunji Kotsuki, Kenta Shiraishi, and Atsushi Okazaki

Status: open (extended)

Post a comment Subscribe to comment alert

CEC1:
'Comment on egusphere-2024-3102', Juan Antonio Añel, 24 Nov 2024 reply

Dear authors,
Unfortunately, after checking your manuscript, it has come to our attention that it does not comply with our "Code and Data Policy".

https://www.geoscientific-model-development.net/policies/code_and_data_policy.html

You have archived your code on GitHub (both ClimaX and the LETKF codes). However, GitHub is not a suitable repository for scientific publication. GitHub itself instructs authors to use other long-term archival and publishing alternatives, such as Zenodo. Therefore, the current situation with your manuscript is irregular.
Please, publish your code in one of the appropriate repositories and reply to this comment with the relevant information (link and a permanent identifier for it (e.g. DOI)) as soon as possible, as we can not accept manuscripts in Discussions that do not comply with our policy. Also, you must include the modified 'Code and Data Availability' section in a potentially reviewed manuscript, with the links and DOIs of both the ClimaX and LETKF repositories. Remember to include the licenses of each code in the new repositories, MIT and BSD-Clause 3 respectively.
Note that if you do not fix this problem, we will have to stop the Discussions process and reject your manuscript for publication in our journal.
Juan A. Añel

Geosci. Model Dev. Executive Editor

Reply

Citation: https://doi.org/10.5194/egusphere-2024-3102-CEC1
- CC1:
  'Reply on CEC1', Shunji Kotsuki, 28 Nov 2024 reply
  
  Thank you for your message.
  
  We would like to confirm our understanding of your concerns. As mentioned in the manuscript, we have already archived the code and data on Zenodo, and the DOI is provided in the "Code and Data Availability" section of the preprint paper (https://arxiv.org/pdf/2407.17781). The archived files on zenodo include all necessary codes and docker containers (https://zenodo.org/records/13884167). Also, the license of each code is also described in the zonodo (Creative Commons Attribution 4.0 International).
  
  The statement "The default ClimaX was obtained from https://microsoft.github.io/ClimaX (last access: October 1, 2024). The LETKF system was obtained from https://github.com/skotsuki/speedy-lpf (last access: July 12, 2024)." describes the original source codes of the ClimaX and LETKF. We developed the data assimilation system (ClimaX-LETKF) by modifying the original codes. Again, the developed system is archived on the zonodo.
  Could you please clarify if there is any specific aspect of the current manuscript that does not comply with the journal's policy?
  
  Thank you for your attention, and we look forward to your response.
  
  Best regards,
  Shunji Kotsuki
  
  Reply
  
  Citation: https://doi.org/10.5194/egusphere-2024-3102-CC1
  - CEC2: 'Reply on CC1', Juan Antonio Añel, 28 Nov 2024 reply
    
    Dear authors,
    I refer you to my previous comment. You must publish in an adequate repository the original ClimaX and LETKF codes, not only the adapted versions that you use. It is my understanding that they are not included in the current repository in Zenodo. In any case, please, remove the mention to the GitHub webpages to avoid confusion.
    I hope this clarifies the situation.
    Juan A. Añel
    Geosci. Model Dev. Executive Editor
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CEC2
    
    CC2: 'Reply on CEC2', Shunji Kotsuki, 29 Nov 2024 reply
    
    Thanks again for your message.
    We understand your concern regarding the archival of the original ClimaX and LETKF codes. However, these codes are not our intellectual property but are developed by third parties such as Microsoft. Therefore, it would be inappropriate for us to re-archive them in Zenodo under our name, in my understanding.
    The modified system we developed (ClimaX-LETKF) is fully archived in Zenodo with a DOI and includes all necessary information for reproduction, as per the journal's policy.
    Based on the considerations above, we plan to update the "code and data availability section" as follows, with the removal of GitHub webpages:
    "The exact version of the model used to produce the results used in this paper is archived on Zenodo (https://zenodo.org/records/13884167), as are input data and scripts to run the model and produce the plots for all the simulations presented in this paper."
    Please let us know if this update aligns with the journal's requirements. If you agree with the proposed "code and data availability section", I will update the section of the preprint paper on arXiv.
    Thank you for your guidance.
    Sincerely,
    Shunji Kotsuki
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CC2
    
    CEC3: 'Reply on CC2', Juan Antonio Añel, 29 Nov 2024 reply
    
    Dear authors,
    Your understanding regarding the ClimaX and LETKF codes is wrong. As I stated in my first comment they are released under the MIT and BSD Clause 3 licenses, which explicitly allow you to copy, use and redistribute the codes. Actually, the alarming issue is that you have used the codes without paying attention to their licensing. If you were not allowed to copy and store the codes in other servers, probably you could not use them for your research too.
    Therefore, again, I have to insist on my requirement to create new repositories for them and reply to this comment with the information about them.
    Regards,
    Juan A. Añel
    Geosci. Model Dev. Executive Editor
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CEC3
    
    CC3: 'Reply on CEC3', Shunji Kotsuki, 02 Dec 2024 reply
    
    Dear Dr. Añel,
    Thank you very much for your detailed feedback and kind guidance. We have addressed the issues you raised as follows:
    We have uploaded the relevant codes to an appropriate repository, ensuring long-term archival.
    ClimaX version 0.3.1: https://zenodo.org/records/14258100
    
    LETKF: https://zenodo.org/records/14258014
    
    We have updated the "Code and Data Policy" section in the preprint manuscript accordingly and have submitted the revised version to arXiv. Once the updated preprint is available, we will notify you immediately. Here is the updated "Code and Data Policy" submitted to arXiv.
    
    The data assimilation system, experimental data, and visualization scripts used in this manuscript are archived on Zonodo (https://zenodo.org/records/13884167; doi: 10.5281/zenodo.13884167). The original ClimaX version 0.3.1 and LETKF codes are also archived on Zenodo; ClimaX version 0.3.1 (https://zenodo.org/records/14258100, doi: 10.5281/zenodo.14258099) and LETKF (https://zenodo.org/records/14258014, doi: 10.5281/zenodo.14258014).
    
    Please let us know if any further revisions or clarifications are needed.
    
    Thank you again for your patience and guidance.
    Sincerely,
    
    Shunji Kotsuki
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CC3
    
    CEC4: 'Reply on CC3', Juan Antonio Añel, 02 Dec 2024 reply
    
    Dear authors,
    Many thanks for addressing the mentioned issues. Now we can consider the current version of your manuscript in compliance with the journal's policy. Please, pay attention to potential modifications that you make to the code for revised versions that could follow-up during the review process.
    Regards,
    Juan A. Añel
    Geosci. Model Dev. Executive Editor
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CEC4
    
    CC4: 'Reply on CEC4', Shunji Kotsuki, 03 Dec 2024 reply
    
    Thank you once again for your kind and thorough feedback. We truly appreciate your careful attention to our work.
    
    Additionally, we would like to inform you that we have issued the new preprint with the modified code and data policy (cf., https://arxiv.org/abs/2407.17781).
    
    Reply
    
    Citation: https://doi.org/10.5194/egusphere-2024-3102-CC4
RC1: 'Comment on egusphere-2024-3102', Anonymous Referee #1, 24 Dec 2024 reply

The authors apply an ensemble Kalman filter (EnKF) to a machine learning (ML)–based model for weather prediction. This is an interesting and important study, which I believe deserves to be published after some revisions. In particular, the literature review is insufficient, and the paper is lacking some details.
Issues:

1. "On the other hand, recent studies have started to solve the inverse problem inherent in data assimilation by deep neural networks". Some other papers can also be cited along these lines: McCabe & Brown (2021), Luk et al. (2024), and Bocquet et al. (2024).

2. "However, no study has succeeded in employing ensemble Kalman filtering with AI models." This statement should be made more specific: this is perhaps the first application of an EnKF to a global ML model of the atmosphere. However, there is older work applying EnKFs to ML models which should also be mentioned: Hamilton et al. (2016), Penny et al. (2022), and Chattopadhyay et al. (2022, 2023).

3. The authors say that the ClimaX model was retrained to predict additional variables. More detail would be helpful here. Was the entire training process repeated with the additional outputs, or did only some new part of the network have to be trained? Given the fact that it has been mostly tech companies training these ML models partly due to the computational limitations in academia, what were the computational resources involved in the training?

4. "at larger localization scales (L_h = 700 and 800 km), analysis RMSEs tended to be higher than the first guess RMSEs". I am somewhat confused by this claim. For L_h = 800 km this makes sense, since the filter appears to diverge. But for L_h = 700 km, if the filter is stable, then is it not contradictory for the analysis RMSEs to be consistently higher than the first guess ones (since the error would just keep growing)? Of course, it can be the case that the RMSEs of some variables worsen with analysis, but I believe that there should be an overall reduction for the filter to be stable (unless the system is not chaotic).

5. "Our optimal localization scale for the 20-member ClimaX-LETKF was 600 km, which is significantly shorter than the 900-km scale of the 20-member LETKF experiment coupled with a dynamical NWP model". In Miyoshi and Kondo (2013), they find an optimal radius of 700 km in a similar setup, which seems not too different than what the authors get here with ClimaX. Unless there are systematic comparisons done it seems difficult to conclude that the optimal radius is significantly lower than for a dynamical model, and draw the resulting conclusions.

6. "Another important property is that the ClimaX is less chaotic than dynamical NWP models, as indicated by the estimated inflation factor β diagnosed by observation-space statistics": I do not think this is a justified conclusion. It could be that ClimaX has higher model error, which is compensated for using the inflation.

7. "It should be noted that we were unable to conduct observation system simulation experiments (k.a. OSSEs), which requires a natural run by ClimaX." I think that this paragraph is quite important and that it could go earlier in the paper (along with Figure 7). As it stands it is somewhat hidden in the discussion section.

8. "In other words, this suggests that ClimaX is unable to return to a meteorologically plausible attractor (or trajectory) while data assimilation enables the ClimaX to synchronize with the real atmosphere." This fact is actually proved rigorously in a simpler setting in Theorem 1 in Adrian et al. (2024): "our theory rigorously shows that if we have long-term filter accuracy with the true dynamics model F and a surrogate model F_s that provides accurate short-term forecasts, we can achieve long-term filter accuracy with the surrogate dynamics." This should be mentioned.

9. In the discussion section, the authors discuss two major improvements in ML forecast models that would improve ensemble data assimilation: accurate covariances, and accurate error growth rates. In this connection, it would be helpful to mention that a major area of research is in models that are trained to produce statistically accurate ensembles by using generative models (Price et al., 2024) or by training on probabilistic cost functions (Kochkov et al., 2024). Also, there has previously been research on enforcing the error growth rate during training as measured by the Lyapunov exponent (Platt et al., 2023).

10. The quality of Figure 3 is quite low.
Minor issues:

1. "Since Google DeepMind issued the first artificial intelligence (AI) weather prediction model, GraphCast". I believe that FourCastNet was earlier in 2022, which seems to be the case based on the arXiv dates. In any case, there are many older attempts to forecast weather based on statistical/machine learning techniques. It would be more accurate to say that FourCastNet was the first ML weather prediction model competitive in skill with dynamical forecast models.

2. "A number of AI weather prediction models have been proposed even by private information and technology (IT) companies". In fact, all the leading models I am aware of have come from private companies, with the exception of the ECMWF's AI Forecasting System (AIFS).

3. "Nguyen" is misspelled as "Nguen" several times in the paper. Also, the authors cite "Nguen et al. (2023, 2024)" but there is no 2024 paper in the bibliography, just two 2023 papers. The second one has a 2024 version on arXiv, so perhaps this is what the authors meant to cite.

4. "grids" should be replaced with "grid points" throughout the manuscript.

5. "Since AI models require significantly lower computational costs compared to dynamical NWP models, ensemble-based methods, such as ensemble Kalman filters (EnKFs) and particle filters, also offer benefits for AI models." I think, given this justification, that it would be more accurate to say that the AI models offer benefits for ensemble-based methods.

6. "We employed a series of data assimilation experiments over a year of 2017, which is not used for training and validation of the ClimaX. The ensemble size is 20, and their initial conditions were taken from WeatherBench data in 2006." Are you saying that the initial ensemble (for January 1, 2017) was taken from 2006 data? Did this come from consecutive days in 2006?

7. Figure 3 caption, typo "WeatharBench".

8. "the experiment with L_h = 400 km kept reducing the RMSEs over a year, indicating that a too small localization scale is suboptimal". It would be good to clarify that the RMSEs kept reducing but are still higher than those of the other experiments, with the exception of L_h = 800 km.

9. "Negative and positive values indicate improvements and degradations due to data assimilation": needs a "respectively" at the end.

10. "Compared to our study, Kotsuki et al. (2017) estimated much smaller inflation factor for a global ensemble data assimilation system using a dynamical model": Please state the average inflation factors in Kotsuki et al. (2017) compared to yours.

11. "The estimated inflation factor of Kotsuki et al. (2017) was substantially smaller than our study." This sentence is essentially repeating the previous one.

12. "In addition, the ensemble-based error covariance was reasonable in sparsely observed regions, even according to AI weather prediction models." I'm not sure what is meant by "even according to AI weather prediction models".

13. In Figure 4 the acronym "WB" (I suppose meaning WeatherBench) is not defined.
References:

– Adrian, M., Sanz-Alonso, D., & Willett, R. (2024). Data Assimilation with Machine Learning Surrogate Models: A Case Study with FourCastNet (No. arXiv:2405.13180; Version 1). arXiv. https://doi.org/10.48550/arXiv.2405.13180

– Bocquet, M., Farchi, A., Finn, T. S., Durand, C., Cheng, S., Chen, Y., Pasmans, I., & Carrassi, A. (2024). Accurate deep learning-based filtering for chaotic dynamics by identifying instabilities without an ensemble. Chaos: An Interdisciplinary Journal of Nonlinear Science, 34(9), 091104. https://doi.org/10.1063/5.0230837

– Chattopadhyay, A., Mustafa, M., Hassanzadeh, P., Bach, E., & Kashinath, K. (2022). Towards physics-inspired data-driven weather forecasting: Integrating data assimilation with a deep spatial-transformer-based U-NET in a case study with ERA5. Geoscientific Model Development, 15(5), 2221–2237. https://doi.org/10.5194/gmd-15-2221-2022

– Chattopadhyay, A., Nabizadeh, E., Bach, E., & Hassanzadeh, P. (2023). Deep learning-enhanced ensemble-based data assimilation for high-dimensional nonlinear dynamical systems. Journal of Computational Physics, 477, 111918. https://doi.org/10.1016/j.jcp.2023.111918

– Hamilton, F., Berry, T., & Sauer, T. (2016). Ensemble Kalman Filtering without a Model. Physical Review X, 6(1), 011021. https://doi.org/10.1103/PhysRevX.6.011021

– Kochkov, D., Yuval, J., Langmore, I., Norgaard, P., Smith, J., Mooers, G., Klöwer, M., Lottes, J., Rasp, S., Düben, P., Hatfield, S., Battaglia, P., Sanchez-Gonzalez, A., Willson, M., Brenner, M. P., & Hoyer, S. (2024). Neural general circulation models for weather and climate. Nature, 632(8027), 1060–1066. https://doi.org/10.1038/s41586-024-07744-y

– Luk, E., Bach, E., Baptista, R., & Stuart, A. (2024). Learning Optimal Filters Using Variational Inference (No. arXiv:2406.18066). arXiv. http://arxiv.org/abs/2406.18066

– McCabe, M., & Brown, J. (2021). Learning to Assimilate in Chaotic Dynamical Systems. Advances in Neural Information Processing Systems, 34, 12237–12250. https://proceedings.neurips.cc/paper/2021/hash/65cc2c8205a05d7379fa3a6386f710e1-Abstract.html

– Miyoshi, T., & Kondo, K. (2013). A Multi-Scale Localization Approach to an Ensemble Kalman filter. Sola, 9, 170–173. https://doi.org/10.2151/sola.2013-038

– Miyoshi, T., Kondo, K., & Imamura, T. (2014). The 10,240-member ensemble Kalman filtering with an intermediate AGCM. Geophysical Research Letters, 41(14), 5264–5271. https://doi.org/10.1002/2014GL060863

– Penny, S. G., Smith, T. A., Chen, T.-C., Platt, J. A., Lin, H.-Y., Goodliff, M., & Abarbanel, H. D. I. (2022). Integrating Recurrent Neural Networks With Data Assimilation for Scalable Data-Driven State Estimation. Journal of Advances in Modeling Earth Systems, 14(3), e2021MS002843. https://doi.org/10.1029/2021MS002843

– Platt, J. A., Penny, S. G., Smith, T. A., Chen, T.-C., & Abarbanel, H. D. I. (2023). Constraining chaos: Enforcing dynamical invariants in the training of reservoir computers. Chaos: An Interdisciplinary Journal of Nonlinear Science, 33(10), 103107. https://doi.org/10.1063/5.0156999

– Price, I., Sanchez-Gonzalez, A., Alet, F., Andersson, T. R., El-Kadi, A., Masters, D., Ewalds, T., Stott, J., Mohamed, S., Battaglia, P., Lam, R., & Willson, M. (2024). Probabilistic weather forecasting with machine learning. Nature, 1–7. https://doi.org/10.1038/s41586-024-08252-9

Reply

Citation: https://doi.org/10.5194/egusphere-2024-3102-RC1

Shunji Kotsuki, Kenta Shiraishi, and Atsushi Okazaki

Data sets

Experimental data, source codes and scripts used in Kotsuki et al. (2024) submitted to GMD Shunji Kotsuki https://zenodo.org/records/13884167

Shunji Kotsuki, Kenta Shiraishi, and Atsushi Okazaki

Viewed

Since the preprint corresponding to this journal article was posted outside of Copernicus Publications, the preprint-related metrics are limited to HTML views.

Total article views: 193 (including HTML, PDF, and XML)

HTML	PDF	XML	Total	BibTeX	EndNote
193	0	0	193	0	0

HTML: 193
PDF: 0
XML: 0
Total: 193
BibTeX: 0
EndNote: 0

Views and downloads (calculated since 15 Oct 2024)

Month	HTML	XML
Oct 2024	65	65
Nov 2024	63	63
Dec 2024	28	28
Jan 2025	37	37

Cumulative views and downloads (calculated since 15 Oct 2024)

Month	HTML	XML
Oct 2024	65	65
Nov 2024	63	63
Dec 2024	28	28
Jan 2025	37	37

Viewed (geographical distribution)

Since the preprint corresponding to this journal article was posted outside of Copernicus Publications, the preprint-related metrics are limited to HTML views.

Total article views: 182 (including HTML, PDF, and XML) Thereof 182 with geography defined and 0 with unknown origin.

Country	#	Views	%

Latest update: 20 Jan 2025

Short summary

Artificial intelligence (AI) is playing a bigger role in weather forecasting, often competing with physical models. However, combining AI models with data assimilation, a process that improves weather forecasts by incorporating observation data, is still relatively unexplored. This study explored coupling ensemble data assimilation with an AI weather prediction model ClimaX, succeeded in employing weather forecasts stably by applying techniques conventionally used for physical models.


Total:	0
HTML:	0
PDF:	0
XML:	0