Learning Vertical Coordinates via Automatic Differentiation of a Dynamical Core

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 that defines 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. By backpropagating simulation errors through the time integration to the coordinate parameters, our formulation finds a grid structure optimized for both the underlying physics and the numerical discretization. Using several standard tests, we demonstrate that these learned coordinates reduce the mean squared error by a factor of 1.4–2 in nonlinear statistical benchmarks, and eliminate spurious vertical velocity striations over steep topography.