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

This study explores the use of ensemble DA,  in particular, LETKF to diagnose and improve AI-based weather prediction models. The authors use Microsoft’s ClimaX v0.3.1 as a test case. Overall, I found the research idea interesting and potentially impactful for the field. However, I believe the presentation of the methodology could be improved to highlight the contribution of the author's work and make the paper easier to follow.

1. Introduction: Many recent studies have combined data assimilation with AI-based weather prediction models, although most rely on variational data assimilation as mentioned by the authors. The authors should further summarize the existing literature and clearly highlight their contribution in relation to these existing approaches. In general, many efforts have also been given to combine DA (both variational and Kalman-type) with machine learning predictive (or surrogate) models (not necessary for weather prediction). It would be beneficial to explore how the methods developed in this paper could be applied to other fields where AI predictive models and real-time observations are available.
2. Introduction: It seems the main objective of using LETKF in this paper is to diagnose AI-based models but can it be used to make real-time update of the AI weather model based on real-time observations? to demonstrate that, one would probably need numerical experiments with higher resolution data?
3. Methodology: in Section 2.2, the authors are encouraged to define the state and observation variables, as well as the observation operator, in the context of the predictive climate model X used in this paper. It seems that these important information is only available in table 1?
4. Methodology: How sensitive are the results to localization radius, inflation methods, or ensemble size? Could variational DA (with the help of autodifferentiation codes as mentioned by the authors) methods with AI weather models outperform LETKF?
5. Results: How does the relatively coarse resolution (5.625°) affect the diagnosis of localized weather phenomena? Would ClimaX + LETKF behave differently if applied at higher spatial resolution?
6. Could we use ensemble-based DA (such as the one developed in this paper) to generate synthetic training data for AI models in under-observed regions?