TorchDA: A Python package for performing data assimilation with deep learning forward and transformation functions

Data assimilation techniques are often confronted with challenges handling complex high dimensional physical systems, because high precision simulation in complex high dimensional physical systems is computationally expensive and the exact observation functions that can be applied in these systems are difficult to obtain. It prompts growing interest in integrating deep learning models within data assimilation workflows, but current software packages for data assimilation cannot handle deep learning models inside. This study presents a novel Python package seamlessly combining data assimilation with deep neural networks to serve as models for state transition and observation functions. The package, named TorchDA, implements Kalman Filter, Ensemble Kalman Filter (EnKF), 3D Variational (3DVar), and 4D Variational (4DVar) algorithms, allowing flexible algorithm selection based on application requirements. Comprehensive experiments conducted on the Lorenz 63 and a two-dimensional shallow water system demonstrate significantly enhanced performance over standalone model predictions without assimilation. The shallow water analysis validates data assimilation capabilities mapping between different physical quantity spaces in either full space or reduced order space. Overall, this innovative software package enables flexible integration of deep learning representations within data assimilation, conferring a versatile tool to tackle complex high dimensional dynamical systems across scientific domains.数据同化（Data assimilation）技术在处理复杂高维物理系统时往往面临诸多挑战：一方面，对复杂高维物理系统开展高精度模拟的计算成本极高；另一方面，难以获取可适用于此类系统的精确观测函数（exact observation functions）。这使得学界对在数据同化工作流中集成深度学习模型的兴趣与日俱增，但当前的数据同化软件包均无法在其内部处理深度学习模型。本研究提出了一款全新的Python软件包，可将数据同化与深度神经网络无缝结合，以作为状态转移与观测算子（state transition and observation functions）的建模工具。该软件包命名为TorchDA，实现了卡尔曼滤波（Kalman Filter）、集合卡尔曼滤波（Ensemble Kalman Filter, EnKF）、三维变分（3D Variational, 3DVar）以及四维变分（4D Variational, 4DVar）算法，可根据应用需求灵活选择算法。基于洛伦兹63系统（Lorenz 63）与二维浅水波系统开展的全面实验表明，相较于未进行数据同化的独立模型预测结果，本软件包的性能得到了显著提升。针对浅水波系统的分析验证了该软件包的数据同化能力，可在全空间或降阶空间（reduced order space）中实现不同物理量空间之间的映射。总体而言，这款创新性软件包实现了深度学习表征在数据同化框架下的灵活集成，为跨科学领域处理复杂高维动力学系统提供了一款通用工具。