Best Practices for Fortran-Python Bridges to Integrate Neural Networks in Earth System Models

In recent years, machine learning (ML) based parameterizations have become increasingly common in Earth System Models (ESM). Sub-grid scale physical processes that would be computationally too expensive, e.g., atmospheric chemistry and cloud microphysics, can be emulated by ML algorithms such as neural networks.

Neural networks are trained first on simulations of the sub-grid scale process that is to be emulated. They are then used in so-called inference mode to make predictions during the ESM run, replacing the original parameterization. Training usually requires GPUs, while inference may be done on CPU architectures.

At first, neural networks are evaluated offline, i.e., independently of the ESM on appropriate datasets. However, their performance can ultimately only be evaluated in an online setting, where the ML algorithm is coupled to the ESM, including nonlinear interactions.

We want to shorten the time spent in neural network development and offline testing and move quickly to online evaluation of ML components in our ESM of choice, ICON (Icosahedral Nonhydrostatic Weather and Climate Model). Since ICON is written in Fortran, and modern ML algorithms are developed in the Python ecosystem, this requires efficient bridges between the two programming languages. The Fortran-Python bridge must be flexible to allow for iterative development of the neural network. Changes to the ESM codebase should be as few as possible, and the runtime overhead should not limit development.

In our contribution we explore three strategies to call the neural network inference from within Fortran using (i) embedded Python code compiled in a dynamic library, (ii) pipes, and (iii) MPI using the ICON coupler YAC. We provide quantitative benchmarks for the proposed Fortran-Python bridges and assess their overall suitability in a qualitative way to derive best practices. The Fortran-Python bridge enables scientists and developers to evaluate ML components in an online setting, and can be extended to other parameterizations and ESMs.