Coupling of NEMO to a neural network emulator of PISCES

The Pelagic Interaction Scheme for Carbon and Ecosystem Studies (PISCES) is a marine biogeochemical model that is used in several IPCC-Class Earth System models. PISCES simulates the distribution of nutrients (four macronutrients and one micronutrient) that regulate the growth of two phytoplankton classes (nanophytoplankton and diatoms). It also simulates the ocean carbon cycle with a complete representation of the marine carbonate systems. PISCES includes 24 state variables, and increases the runtime of NEMO, the physical ocean model with which it is coupled, by a factor of 3.4, indicating a high computational cost.

PISCES-AI has been developed as a U-Net based machine learning PISCES emulator, which takes a small number of input variables (TOS, ZOS, SOS, PAR, atmospheric CO₂), and predicts two output variables: surface chlorophyll and the difference in partial pressure of CO₂ between the atmosphere and the ocean. These are the only outputs which have a direct influence on climate simulations by Earth system models. Previous work has shown the predictive power of PISCES-AI across multiple timescales, and in an out-of-domain setting.

In this work, we couple the AI emulator of PISCES to NEMO, using Eophis and Morays for Python-Fortran interaction. We evaluate its performance, as well as its computational efficiency, to give a holistic picture of the challenges and opportunites for AI emulation of ocean biogeochemistry. With a particular interest in the computational speed, we find that inference for a single time-step to be around 10ms, with a much larger preliminary bottleneck due to CPU-GPU transfer (200ms per timestep). Even with this bottleneck, with our implementation we obtain a speed-up of factor 3 compared to PISCES, and we explore ways in which the data transfer bottleneck could be reduced.