A Framework for Hybrid Physics-AI Coupled Ocean Models

TL;DR

This work focuses on the ocean and sea-ice components of a state-of-the-art climate model by implementing a spectrum of data-driven parameterizations, ranging from complex deep learning models to more interpretable equation-based models.

Abstract

Climate simulations, at all grid resolutions, rely on approximations that encapsulate the forcing due to unresolved processes on resolved variables, known as parameterizations. Parameterizations often lead to inaccuracies in climate models, with significant biases in the physics of key climate phenomena. Advances in artificial intelligence (AI) are now directly enabling the learning of unresolved processes from data to improve the physics of climate simulations. Here, we introduce a flexible framework for developing and implementing physics- and scale-aware machine learning parameterizations within climate models. We focus on the ocean and sea-ice components of a state-of-the-art climate model by implementing a spectrum of data-driven parameterizations, ranging from complex deep learning models to more interpretable equation-based models. Our results showcase the viability of AI-driven parameterizations in operational models, advancing the capabilities of a new generation of hybrid simulations, and include prototypes of fully coupled atmosphere-ocean-sea-ice hybrid simulations. The tools developed are open source, accessible, and available to all.

Figures (10)

• Figure 1: Schematic of learning subgrid processes from high-resolution simulation. Here, $x_n$ represents the physical variables (e.g., velocities at a grid point) simulated by a high-resolution model. A coarse-graining procedure, which relies on spatial filtering, is applied to the high-resolution data to explicitly compute the subgrid forcing and input variables on a coarse-resolution mesh. The overbar represents a spatial average, and the prime represents the fluctuation from the mean. A neural network is then trained to predict the subgrid processes (e.g., $\overline{x'_1 x'_2}$) from coarse-resolution variables (e.g., $\overline{x_1}$).
• Figure 2: Schematic of Structural Error vs. Complexity. Panel a--d are organized on a structural error vs complexity curve, which generally transitions from high error and low complexity (e.g., small number of free parameters) to low error and high complexity. Here, each panel uses the case of parameterizing the momentum stress tensor ($\mathbf{T}$) in different ways as an example. Panel a) illustrates a more traditional approach to flux parameterization, where the stress tensor is simply a function of the velocity gradient, for anti-viscosity chang2023remote. Panel b) shows equation discovery methods, where the basis functions $\bm{\varphi}$ may be chosen or discovered functions of the velocity field, such as vorticity $\bm{\varphi}_1 = \bm{\zeta} = \overline{\bm{\nabla}}\times\overline{\mathbf{u}}$, shearing deformation $\bm{\varphi}_2=\mathbf{D} = \partial_y\overline{u} + \partial_x\overline{v}$, and the stretching deformation $\bm{\varphi}_3 = \widetilde{\bm{D}} = \partial_x\overline{u} - \partial_y\overline{v}$zanna2020dataross2022benchmarking. Panel c) shows physics-aware ANNs, where as an example the inputs are also hand chosen (e.g., $\bm{\zeta}$, $\mathbf{D}$, $\widetilde{\bm{D}}$), and the outputs from the network are dimensionalized using some physics-informed function $g$ of an appropriate norm of the input fields, $||\bm{\nabla}\overline{\mathbf{u}}||$ and the grid-spacing $\Delta$. Panel d) shows CNNs where the input fields are chosen to be some variables, e.g., filtered velocity fields, with minimal inductive bias.
• Figure 3: Schematic of an ANN module, written in Fortran, in MOM6 being used for multiple parameterizations within the GCM. The ANN module is created to perform the inference for multi-layer ANN models trained in PyTorch or other ML libraries. The weights, biases, and others parameters are stored in a NetCDF file, which communicate the information to the Fortran module. For ocean vertical mixing, the module computes the vertical profile of diffusivity, while for the lateral momentum mixing parametrization it computes the full eddy momentum flux.
• Figure 4: Summer Mixed Layer Depth (MLD) in meters: Changes induced by ML parameterizations for ePBL and lateral momentum fluxes. (A, E) Summer MLD, found using the potential anomaly criterion ReichlMLD, in the control run for ePBL with JRA forcing and for mesoscale parameterization with CORE IAF, respectively; (B) Bias of control with respect to ARGO; (C) Difference between ePBL_NN and control; (D) Difference between ePBL_EQ and control; (F) Difference between ANN and control; (G) Difference between ZB20 and control.
• Figure 5: Winter Mixed Layer Depth: Changes induced by ML parameterizations for ePBL and lateral momentum fluxes. Panels are the same as Fig. \ref{['fig:mld-summer']} but for winter MLD. Note that control biases exhibit the same patterns under different forcing conditions (JRA vs. CORE IAF, but with varying magnitudes.