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

Sibo Cheng, Jinyang Min, Che Liu, Rossella Arcucci

Abstract: 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.Program summaryProgram Title: TorchDACPC Library link to program files: https://doi.org/10.17632/bm5d7xk6gw.1Developer's repository link: https://github.com/acse-jm122/torchdaLicensing provisions: GNU General Public License version 3Programming language: Python3External routines/libraries: Pytorch.Nature of problem: Deep learning has recently emerged as a potent tool for establishing data-driven predictive and observation functions within data assimilation workflows. Existing data assimilation tools like OpenDA and ADAO are not well-suited for handling predictive and observation models represented by deep neural networks. This gap necessitates the development of a comprehensive package that harmonizes deep learning and data assimilation.Solution method: This project introduces TorchDA, a novel computational tool based on the PyTorch framework, addressing the challenges posed by predictive and observation functions represented by deep neural networks. It enables users to train their custom neural networks and effortlessly incorporate them into data assimilation processes. This integration facilitates the incorporation of real-time observational data in both full and reduced physical spaces.