Implementation of tangent linear and adjoint models for neural networks based on a compiler library tool

TL;DR

This paper introduces TorchNWP, a compilation library that efficiently couples neural networks with traditional numerical models, enabling improved data assimilation and weather prediction accuracy through tangent linear and adjoint models.

Contribution

It presents a unified interface for integrating deep learning models into numerical models, supporting tangent linear and adjoint models, and enabling deployment on heterogeneous platforms with minimal code modifications.

Findings

Efficient coupling of neural networks with numerical models demonstrated in weather prediction.

Implementation of tangent linear and adjoint models improves data assimilation accuracy.

Supports deployment on heterogeneous computing platforms.

Abstract

This paper presents TorchNWP, a compilation library tool for the efficient coupling of artificial intelligence components and traditional numerical models. It aims to address the issues of poor cross-language compatibility, insufficient coupling flexibility, and low data transfer efficiency between operational numerical models developed in Fortran and Python-based deep learning frameworks. Based on LibTorch, it optimizes and designs a unified application-layer calling interface, converts deep learning models under the PyTorch framework into a static binary format, and provides C/C++ interfaces. Then, using hybrid Fortran/C/C++ programming, it enables the deployment of deep learning models within numerical models. Integrating TorchNWP into a numerical model only requires compiling it into a callable link library and linking it during the compilation and linking phase to generate the executable. On this basis, tangent linear and adjoint model based on neural networks are implemented at the C/C++ level, which can shield the internal structure of neural network models and simplify the construction process of four-dimensional variational data assimilation systems. Meanwhile, it supports deployment on heterogeneous platforms, is compatible with mainstream neural network models, and enables mapping of different parallel granularities and efficient parallel execution. Using this tool requires minimal code modifications to the original numerical model, thus reducing coupling costs. It can be efficiently integrated into numerical weather prediction models such as CMA-GFS and MCV, and has been applied to the coupling of deep learning-based physical parameterization schemes (e.g., radiation, non-orographic gravity wave drag) and the development of their tangent linear and adjoint models, significantly improving the accuracy and efficiency of numerical weather prediction.