NORi: A Novel, Physically-Principled Approach to Parameterization of Upper Ocean Turbulence using Neural Ordinary Differential Equations

Given our current computational resources, a state-of-the-art ocean model can achieve a grid resolution in the order of 10 km for realistic global simulations, meaning that small-scale convection and wind-driven mixing near the surface of the ocean with length scales of roughly 1 m cannot be explicitly resolved. However, these microturbulent processes play a fundamental role in setting the structure of the ocean stratification, govern air-sea fluxes exchange as well as tracer transport with the interior of the ocean. Therefore, we use parameterizations, models which approximate small-scale processes using large-scale variables, to represent their effects in climate simulations.

In this work, we propose NORi: a novel, physically-principled, and data-driven approach to parameterizing ocean vertical mixing using neural ordinary differential equations (NODEs). NORi uses neural ODEs (NO) to augment a simple eddy-diffusivity closure based on the local gradient Richardson number (Ri). The Ri-based diffusivity closure captures local convective- and shear-driven mixing, while the neural ODEs augment the base model with nonlocal entrainment fluxes due to convection using neural networks. NORi is designed for realistic seawater's nonlinear equation of state using TEOS-10 and explicitly represents temperature and salinity fluxes. When compared against high fidelity large-eddy simulations (LES) of different convective strengths, background stratifications, and shear conditions, NORi demonstrates excellent prediction and generalization capabilities. By design, NORi automatically satisfies tracer invariance and conservation. It also exhibits high numerical stability and accuracy owing to its online training paradigm, where neural networks are calibrated against time-integrated field variables of interest rather than on instantaneous, time-independent turbulent fluxes. When compared against other parameterizations, NORi produces deeper mixed layers which are in better agreement with the LES solution. NORi is implemented straightforwardly into Oceananigans.jl, the fastest ocean model to date without intermediate wrappers. This can be achieved owing to the cutting-edge paradigm of the Julia programming language as well as the simple, modern and flexible interface of Oceananigans.jl. Using large-scale simulations, we demonstrate that NORi is numerically stable for at least 100 years despite being trained with only a 2-day integration, is computationally efficient, and produces realistic fields which are comparable to existing parameterization.