Learning vertical coordinates via automatic differentiation of a dynamical core

TL;DR

The paper tackles spurious discretization errors introduced by terrain-following vertical coordinates in high-resolution atmospheric simulations with non-hydrostatic dynamics.It introduces NEUVE, a learnable vertical coordinate parameterization embedded in an end-to-end differentiable 2D Euler solver, leveraging automatic differentiation to compute exact geometric terms and a solver-in-the-loop optimization to minimize discretization error.Across linear advection, rising bubble, and density current tests, NEUVE reduces mean-squared errors and suppresses spurious vertical velocity artifacts, while revealing how the learned grid concentrates discretization accuracy where it matters (e.g., near shear layers).The approach enables regime-specific grid optimization, offers potential computational savings, and suggests a path toward progressively replacing heuristic solver components with differentiable, machine-learned counterparts in dynamical cores.

Abstract

Terrain-following coordinates in atmospheric models often imprint their grid structure onto the solution, particularly over steep topography, where distorted coordinate layers can generate spurious horizontal and vertical motion. Standard formulations, such as hybrid or SLEVE coordinates, mitigate these errors by using analytic decay functions controlled by heuristic scale parameters that are typically tuned by hand and fixed a priori. In this work, we propose a framework to define a parametric vertical coordinate system as a learnable component within a differentiable dynamical core. We develop an end-to-end differentiable numerical solver for the two-dimensional non-hydrostatic Euler equations on an Arakawa C-grid, and introduce a NEUral Vertical Enhancement (NEUVE) terrain-following coordinate based on an integral transformed neural network that guarantees monotonicity. A key feature of our approach is the use of automatic differentiation to compute exact geometric metric terms, thereby eliminating truncation errors associated with finite-difference coordinate derivatives. By coupling simulation errors through the time integration to the parameterization, our formulation finds a grid structure optimized for both the underlying physics and numerics. Using several standard tests, we demonstrate that these learned coordinates reduce the mean squared error by a factor of 1.4 to 2 in non-linear statistical benchmarks, and eliminate spurious vertical velocity striations over steep topography.

Figures (10)

• Figure 1: (a) Vertical decay functions $B(\zeta)$ for Gal-Chen and Somerville (dotted), Hybrid (dashed), SLEVE (dash-dot), and the learned NEUVE (solid red). (b-e) Visual comparison of the resulting computational grids over a sample mountain profile for (b) Gal-Chen and Somerville, (c) Hybrid, (d) SLEVE, and (e) NEUVE coordinates.
• Figure 2: Statistical evaluation of coordinate performance on the passive scalar advection test case. (a) Probability density of $L_2$ errors across the stochastic topography test ensemble for NEUVE (solid), Hybrid (dashed), and SLEVE (dotted) coordinates. (b) Sensitivity of $L_2$ error to maximum mountain height ($H_{max}$). (c) Sensitivity of $L_2$ error to topographic roughness length ($z_0$). Shaded regions in (b) and (c) indicate the error variance across the batch. (d--e) The improvement factor, defined as the ratio of the SLEVE $L_2$ error to NEUVE $L_2$ error, plotted against (d) $H_{max}$ and (e) $z_0$. The classic Gal-Chen and Somerville coordinate is excluded from these comparisons as its errors are orders of magnitude larger, which would distort the graphical scale.
• Figure 3: Visual comparison of advection accuracy across coordinate formulations for a randomized topography sample within the extra advection test case. The left column displays the final state ($t=T$) of the passive scalar contour after advecting over a rough surface profile, while the right column shows the corresponding error field defined as the difference between the model and the exact solution. The rows, from top to bottom, correspond to the Gal-Chen and Somerville, Hybrid, SLEVE, and NEUVE coordinate solutions.
• Figure 4: Statistical evaluation of coordinate performance on the rising thermal bubble test case. (a) Probability density of $L_2$-errors across the stochastic topography test ensemble ($N=500$) for NEUVE (solid), Hybrid (dashed), and SLEVE (dotted) coordinates. (b) Sensitivity of $L_2$-error to maximum mountain height ($H_{max}$). (c) Sensitivity of $L_2$-error to topographic roughness length ($z_0$). The improvement factor, defined as the ratio of the SLEVE $L_2$-error to NEUVE $L_2$-error, plotted against (d) $H_{max}$ and (e) $z_0$.
• Figure 5: Rising thermal bubble test over a jagged mountain range. Same as Figure \ref{['fig:bubble_large']} but evaluated over high-frequency multi-peak topography ($h_c=100$ m, $a_c=400$ m, $\lambda_c=150$ m). The SLEVE and Hybrid parameters are tuned with $s_1 = 400$, $s_2 = 130$, and $s=400$.