Learning vertical coordinates via automatic differentiation of a dynamical core

TL;DR

This paper introduces a differentiable framework for learning terrain-following vertical coordinates in atmospheric models, reducing errors and spurious motions over steep topography by optimizing coordinate parameters through automatic differentiation.

Contribution

It develops a novel end-to-end differentiable solver with a neural network-based terrain-following coordinate, enabling automatic learning of optimal vertical grid structures.

Findings

Reduces mean squared error by 1.4 to 2 times in benchmarks

Eliminates spurious vertical velocity striations

Demonstrates improved accuracy over standard coordinates

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.