Online learning of subgrid-scale models for quasi-geostrophic turbulence in planetary interiors

TL;DR

This work addresses the challenge of representing subgrid-scale fluxes in quasi-geostrophic turbulence within axisymmetric bounded domains relevant to planetary interiors. It introduces an online learning framework using a differentiable pseudo-spectral solver, with a Galerkin-based coarse-graining scheme that preserves boundary conditions, and a two-path spectral-space neural network that provides SGS corrections during simulation. The online SGS model achieves long-time stability and accurately reproduces DNS statistics, spectra, zonal jets, and fast Rossby-wave dynamics across three geometries, outperforming hyperdiffusive and no-SGS baselines. The results demonstrate the viability of ML-driven SGS closures for planetary interior dynamics and motivate extensions toward geodynamo models and more turbulent regimes.

Abstract

The use of machine learning to represent subgrid-scale (SGS) dynamics is now well established in weather forecasting and climate modelling. Recent advances have demonstrated that SGS models trained via ``online'' end-to-end learning -- where the dynamical solver operating on the filtered equations participates in the training -- can outperform traditional physics-based approaches. Most studies, however, have focused on idealised periodic domains, neglecting the mechanical boundaries present e.g. in planetary interiors. To address this issue, we consider two-dimensional quasi-geostrophic turbulent flow in an axisymmetric bounded domain that we model using a pseudo-spectral differentiable solver, thereby enabling online learning. We examine three configurations, varying the geometry (between an exponential container and a spherical shell) and the rotation rate. Flow is driven by a prescribed analytical forcing, allowing for precise control over the energy injection scale and an exact estimate of the power input. We evaluate the accuracy of the online-trained SGS model against the reference direct numerical simulation using integral quantities and spectral diagnostics. In all configurations, we show that an SGS model trained on data spanning only one turnover time remains stable and accurate over integrations at least a hundred times longer than the training period. Moreover, we demonstrate the model's remarkable ability to reproduce slow processes occurring on time scales far exceeding the training duration, such as the inward drift of jets in the spherical shell. These results suggest a promising path towards developing SGS models for planetary and stellar interior dynamics, including dynamo processes.

Figures (11)

• Figure 1: (a) Exponential and (b) spherical container geometries. The shape is prescribed through the half-height $h(s)$, where the cylindrical radius $s$ varies between $s_i$ and $s_o$. The container rotates at constant angular velocity $\Omega$ about the $z$-axis.
• Figure 2: Cartesian forcing pattern $\mathcal{F}$ as described by lemasquerier2023zonal producing inlet-outlet vorticity sources. Vortex spacing $\Delta_{\mathcal{F}} = 0.08$, radius $\ell_{\mathcal{F}} = 0.04$ and amplitude $a_{\mathcal{F}} = 2 \times 10^{10}$ are the same for all configurations studied. The magnitude of the pumps on this figure is normalised between 1 and -1 for cyclonic and anticyclonic vortices, respectively.
• Figure 3: Illustration of the coarse-graining process for an analytical function $z(x) = (x^{2} - 1) \sin (15 \pi x)$ with $N_{s} = 300$ and $\overline{N_{s}} = 60$ using Chebyshev coefficients and using the boundary-preserving Dirichlet-Neumann Galerkin basis. The zoom inset clearly demonstrates that the Chebyshev truncation does not preserve the boundary conditions $z(\pm 1) = 0$.
• Figure 4: Illustration of the spectral-space architecture. The architecture is split into two paths for the axisymmetric velocity and the vorticity equations. The model applies three (residual) blocks of increasing number of features $F \in \{32, 64, 128\}$ containing each a Conv layer, a Linear layer and a non-linear mod ReLU, or $|\mathrm{ReLU}| = \mathrm{ReLU}(z + b) z / |z|$ activation arjovsky2016unitary followed by a Linear layer. The last Linear layer projects the high-dimensional features $F$ into a single feature. The total number of trainable parameters for this architecture is approximately 666k.
• Figure 5: Vorticity $\omega$ (left) and azimuthal velocity $u_{\varphi}$ (right) from DNS for the three configurations. Top: configuration (i), exponential container with $E = 2 \times 10^{-7}$ and $\beta = -1 / s_{o}$. Middle: configuration (ii), spherical shell with $E = 3 \times 10^{-7}$. Bottom: configuration (iii), spherical shell with $E = 10^{-6}$. In each panel, the left and right halves represent direct-resolution and coarse-grained fields, respectively. See text for details.