Tensor-Var: Efficient Four-Dimensional Variational Data Assimilation

We sincerely thank reviewer wFvT for these insightful questions. Below, we address each point in detail.

• Choice of kernel functions

  Thank you for raising the question about choices of kernel functions in Tensor-Var. We would like to clarify that we have included an ablation study in Section 4.4 Table 3 on the Lorenz96 systems with $n_s=40$ and $n_s=80$. To further address the reviewer's comments, we also tested other kernel features, including Matern-1.5 and Matern-2.5. The results are provided in the following Table:

  Table: Comparison of different features with NRMSE (%) as the metric for state estimation accuracy.

  Features/Experiments	Lorenz-96 $n_s=40$	Lorenz-96 $n_s=80$
  $d_s=20$	16.7±2.1	17.3±2.6
  $d_s=40$	14.1±1.3	16.8±2.2
  $d_s=60$	8.3±0.9	11.4±0.9
  $d_s=120$	9.7±0.7	9.0±0.9
  Gaussian Kernel	8.4±0.5	14.3±1.4
  Matern-1.5	9.6±0.8	13.2±1.1
  Matern-2.5	7.2±0.2	9.7±0.7

In a low dimensional system ($n_s=40$), Matern-2.5 achieves the lowest mean and standard deviation. In the higher-dimensional setting ($n_s=80$), the deep feature with $d_s=120$ outperforms the other approaches, albeit with a slightly higher standard deviation compared to Matern-2.5. In general, the results suggest that fixed kernel features may be more reliable at lower dimensions, but deep features become advantageous as dimensionality increases.

For high-dimensional weather forecasting problems, fixed kernel features tend to be poorly scaled and are inadequate for processing structured data. Thus, we propose deep features to address scalability challenges and extend Tensor-Var for broad applications. 

• Computational Complexity and Scalability

  We thank the reviewer for their insightful question regarding the scalability of Tensor-Var. The computational complexity of Tensor-Var primarily depends on the feature dimension rather than the system dimension, as DA optimisation is performed within the feature space. As shown in Table 1 (Section 4.1), the evaluation time exhibits only a slight increase with the increasing system dimension.

  Moreover, the experiments presented in Section 4.2 and Section 4.3 differ in spatial resolution, with Section 4.3 using a resolution of $240*121$, approximately 14 times higher than that of Section 4.2. Nonetheless, the computation time does not increase substantially.

  One potential bottleneck arises from solving the least square problems to compute the two CME operators in high-dimensional feature spaces during training. Incorporating more efficient solvers, such as the Krylov subspace method [Liesen et al. (2013)], could address this challenge in future work.

We hope these responses adequately address the reviewer’s concerns, and we look forward to any further comments you may have.

References

• Liesen, Jörg, and Zdenek Strakos. Krylov subspace methods: principles and analysis. Numerical Mathematics and Scie, 2013.

We greatly appreciate reviewer's constructive feedback and recognition of the strengths of our work. Below, we address each of the questions in detail.

• Compare with sota ML-based DA methods.

  We thank reviewer for raising the questions regarding the comparison with recent ML-based methods.

  • Most existing ML-based methods, such as GraphCast [Lam et al. (2023)], primarily focus on weather forecasting rather than data assimilation. Therefore, direct comparisons between our DA framework and these forecasting methods are not straightforward.

  • Nevertheless, our proposed DA framework can be seamlessly integrated with advanced AI-based weather forecasting models. By incorporating real-time observations, our approach consistently corrects the AI-based weather forecasting models, providing a more reliable starting point for subsequent forecasts. Furthermore, GNN-based approaches can be directly integrated into our proposed DA framework as kernel features, enabling the effective assimilation of spatially irregular data. This integration would further improve the model's ability to handle complex spatial information within the DA context.

  • We recognise the importance of comparing our method with state-of-the-art ML-based DA approaches, as also highlighted by Reviewer dYtv. To address this concern, our revised manuscript has added the new results from a state-of-the-art machine learning-based data assimilation method, FengWu 4D-Var [Xiao et al. (2023)]. The results and comparison details are provided in our response to Reviewer dYtv under the question Compare with sota ML-based DA methods & Baseline selections and are also presented in the accessible anonymous Google Slides. We believe this addition improves the credibility of our study and provides a more comprehensive evaluation.

We hope this response adequately addresses the reviewer’s concerns. We sincerely appreciate the reviewer’s invested time and effort in carefully reviewing our work and providing valuable, constructive feedback.

References

• Lam, Remi, et al. "Learning skillful medium-range global weather forecasting." Science 382.6677 (2023): 1416-1421.
• Xiao, Yi, et al. "Fengwu-4dvar: Coupling the data-driven weather forecasting model with 4d variational assimilation." arXiv preprint arXiv:2312.12455 (2023).

We greatly appreciate valuable and constructive suggestions from Reviewer dYtv. Your efforts have significantly contributed to improving the quality of our paper. Below, we provide detailed responses to each of your comments.

• Higher-resolution experiment

  We would like to clarify that the experiments in Section 4.3 are conducted at a resolution of 240×121, which is roughly 14 times higher than the 64×32 resolution used in Section 4.2.

• Compare with sota ML-based DA methods & Baseline selections

  We understand the reviewer’s concern about comparing with state-of-the-art ML-based DA methods, particularly FengWu-4DVar. To address this, we conducted additional evaluations of FengWu-4DVar in two weather data assimilation experiments. We used the official implementation from the GitHub repository. Training and evaluation settings follow our implementation as described in Sections 4.2 and 4.3.

  • We have summarized the updated results for Figure 2 (page 7) for Section 4.2 into the following table for clarity.

  • For the experiment with higher resolution, we have also added the additional results of FengWu-4DVar into Figure 3 (page 7). Both updated figures (2 and 3) are available in the anonymous Google Slides, providing a clearer demonstration of the results.

    Table: Assimilation NRMSE (%) across variables and evaluation time (s).

    Method	z500	t850	q700	u850	v850	Evaluation Time
    Latent 3D-Var	3.10±0.28	3.83±0.28	11.4±0.65	6.40±0.38	6.36±0.36	1.84±0.06
    Latent 4D-Var	3.02±1.12	3.89±1.29	13.79±3.17	6.57±1.07	6.31±0.83	4.90±0.26
    FengWu 4D-Var	2.80±0.23	2.88±0.16	5.00±0.43	5.51±0.19	5.40±0.24	1.81±0.11
    Tensor-Var	1.79±0.11	1.94±0.09	4.77±0.19	3.19±0.14	3.47±0.15	0.83±0.02

    The results demonstrate that Tensor-Var outperforms FengWu-4DVar in terms of both accuracy and computational efficiency. Similar results are observed in higher resolution experiments. We provide result figures along with detailed discussions in the shared Google Slides for the reviewer's potential interest.

• Prediction Model Details

  In Tensor-Var, predictions are made using the fitted linear operator $\mathcal{C}_{S^+|S}$ without the need for additional learnable models. To improve clarity, we have included a demonstrative figure of the fitting procedure in the revised manuscript's Appendix C and made it available in the Google Slides.

• Stability Test

  We thank reviewer for highlighting the importance of long-term stability. In response, we added a long-term rollout stability test with the following key settings:

  • Spatial Resolution: 240×121.
  • Meteorological Variables: z500, t850, q700, u850, v850.
  • Time Step: 6 hours per step.
  • Assimilation Window Length: 5 steps.
  • Test Horizon: One year (January 1, 2018 - December 31, 2018).
  • Leading Time: 7 days.
  • others: Training and evaluation settings follow Section 4.3.

  The result figures can be found in the Google Slides. These results demonstrate that Tensor-Var consistently controls the estimation error over one-year horizon. To benchmark performance, we also include the results of FengWu 4D-Var. Our experiments show that Tensor-Var achieves better error control with DA compared to FengWu 4D-Var. However, Tensor-Var rollouts unstablely after 800 steps, where FengWu demonstrates more stable long-term rollouts.

  Tensor-Var’s long-term instability likely stems from linear dynamics, which only capture short-term evolution accurately. In contrast, transformer model in FengWu better handles long-term evolution. We value the reviewer’s insight and will work on improvng Tensor-Var’s stability while maintaining DA efficiency and accuracy.

• Different metrics

  We thank reviewer for suggesting using different metrics. To address this, we evaluated FengWu 4D-Var (considered cutting-edge) and Tensor-Var using the anomaly correlation coefficient (ACC) at a resolution of $240*121$. Other settings follow Section 4.3 Figure 3. The results demonstrate that Tensor-Var outperforms FengWu 4D-Var in ACC as well. The corresponding figures can be found in the shared Google Slides.

• Clarification of Figure 10

  We thank reviewer for carefully reviewing our manuscript. The reported metric is normalized RMSE (NRMSE), and we have clarified this in the revised manuscript in Appendix D.3.

• Suggested References

  We thank for the suggested references. After careful review, we have properly cited them in the revised manuscript and discussed their relevance.

In the end, we would like to thank reviewer dYtv again for the valuable feedback and invested efforts. We believe that the revisions address your concerns and improve the overall confidence to our work.

We sincerely thank reviewer Qu2S for your thoughtful review and constructive feedback on our manuscript. Your comments have been invaluable in improving the clarity and presentation of our work. Below, we address each of your comments.

• Reorganize the method section

  Thanks reviewer’s suggestion, we have improved clarity in the methods section as follows:

  • Added a dedicated Related Work section to clearly distinguish the literature review from our methodological contributions.
  • Included a diagram illustrating the training procedure for CME opeartor $\mathcal{C}_{S^+|S}$ (can be accessed in the anonymous Google Slides)
  • Moved the algorithm (previously Algorithm 1) to the main text (end of Section 3.1), revising it to highlight key steps.

• Clarification of Notations

  Thank you for carefully reviewing our manuscript. We would like to clarify the following:

  • We do not state the CME operator as a mapping from the observation space $\mathcal{O}$ to latent space $\mathcal{S}$. Instead, we defined it as $\mathcal{C}_{S|O}\colon\mathbb{H}_O\rightarrow\mathbb{H}_S$, where $\mathbb{H}$ with subscript denotes the corresponding feature (latent) spaces.

    The forward CME operator is consistently expressed as $\mathbb{H}_{S^+}\rightarrow\mathbb{H}_{S}$ later in the manuscript.

  • Algorithm 2 (Appendix C) focuses on Tensor-Var with predetermined kernel features, without learnable parameters $\theta$ as in deep features.

  • We realize the notations in our manuscript may be unclear. To improve clarity, we have added a dedicated table in the Appendix of the updated manuscript. You can also find it in the shared Google Slides above.

• Predetermined kernel features compare to deep features

  We would like to clarify that we conducted an ablation study to compare different feature dimensions and the Gaussian kernel (Section 4.4, Table 3). Details are in Appendix D.5. Additionally, we tested other kernel features, including Matern-1.5 and Matern-2.5.

  Table: Comparison of features using NRMSE (%) for state estimation accuracy.

  Features/Experiments	Lorenz-96 $n_s=40$	Lorenz-96 $n_s=80$
  $d_s=20$	16.7±2.1	17.3±2.6
  $d_s=40$	14.1±1.3	16.8±2.2
  $d_s=60$	8.3±0.9	11.4±0.9
  $d_s=120$	9.7±0.7	9.0±0.9
  Gaussian Kernel	8.4±0.5	14.3±1.4
  Matern-1.5	9.6±0.8	13.2±1.1
  Matern-2.5	7.2±0.2	9.7±0.7

  In $n_s=40$, Matern-2.5 shows the lowest mean and standard deviation (std). In $n_s=80$, the deep feature with $d_s=120$ outperforms other methods, despite a slightly higher std than Matern-2.5. These results indicate fixed kernels are more reliable at low dimensions, while deep features excel as dimensionality increases. For high-dimensional weather forecasting, fixed kernel features often scale poorly and struggle with structured data. Thus, we propose deep features to extend Tensor-Var for broader applications.

• Comparing Training Time Across Methods

  We thank reviewer for the interest in training time. To address this, we measured the training time (per epoch) under identical configurations (NVIDIA RTX 4090) for the experiments in Section 4.2. These results are summarised in the following table and are included in the revised manuscript (Appendix D).

  Methods	Latent 3D-Var	Latent 4D-Var	FengWu 4D-Var	Tensor-Var
  Training time (min)	0.75±0.06	0.88±0.41	11.05±1.65	5.32±1.07 (forward), 4.17±0.31 (inverse)

  Latent 3D and 4D-Var are computationally efficient as they only consist of convolutional and fully connected layers, requiring only one forward pass in training. FengWu 4D-Var, with its transformer-based encoder-decoder, is more time-consuming due to self-attention operations. Tensor-Var’s forward and inverse models are slower than latent DA due to solving least square problems in feature spaces, while the inverse model is faster as no decoding process.

• Suggested References

  We appreciate suggested references. After careful review, we agree that they are relevant to our study. We have properly cited them in the updated manuscript.

We thank reviewer Qu2S again for invested efforts and constructive feedback, which helped improve the manuscript. We have carefully addressed all concerns and made substantial revisions to improve clarity. We look forward to any further feedback.