Learning vertical coordinates via automatic differentiation of a dynamical core

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.

💡 Research Summary

The paper introduces a novel framework that treats the vertical coordinate transformation in atmospheric models as a learnable, differentiable component and leverages automatic differentiation (AD) to compute exact geometric metric terms. Traditional terrain‑following coordinates such as sigma, hybrid, and SLEVE rely on analytically prescribed decay functions with fixed scale parameters, which often induce discretization errors—especially over steep topography—because the metric derivatives are approximated by finite differences.

To overcome these limitations, the authors develop a two‑dimensional non‑hydrostatic Euler solver on an Arakawa C‑grid using second‑order centered spatial differences and a third‑order strong‑stability‑preserving Runge‑Kutta time integrator. The solver is implemented in JAX, enabling end‑to‑end differentiability. The vertical coordinate is defined by

 z(x, ζ) = ζ + h(x)·B(ζ)

where h(x) is the terrain height and B(ζ) is a monotonic decay function. The new NEUVE (Neural Vertical Enhancement) coordinate parameterizes B(ζ) through an integral of a neural network‑generated density ρ₍ϕ₎(η) = log(1 + exp(f₍ϕ₎(η))) + ε, with f₍ϕ₎ being a three‑layer MLP (width = 64) using hyperbolic‑tangent activations. By integrating ρ₍ϕ₎ from the surface to a given ζ/H, strict monotonicity and the boundary conditions B(0)=1, B(H)=0 are guaranteed regardless of the learned weights.

The key methodological advance is the use of AD to obtain the Jacobian ∂z/∂x and ∂z/∂ζ, and consequently all metric coefficients, to machine precision. This eliminates the “geometric truncation error” that traditionally contaminates the pressure‑gradient term cₚ θ ∇π, which is the primary source of spurious horizontal and vertical motions over steep slopes.

Training proceeds by defining a physics loss based on the root‑mean‑square error (RMSE) between the model’s scalar density field µ (advected by a prescribed velocity field) and its analytical solution. A regularization term L_reg penalizes extreme coordinate distortions, and the total objective J(ϕ)=L_physics + λ L_reg is minimized using the Adam optimizer. At each iteration, a random topography is sampled (domain randomization), the corresponding NEUVE coordinate is constructed, metric terms are pre‑computed via AD, the forward model is integrated to a fixed horizon T, and gradients with respect to ϕ are back‑propagated through the entire time‑stepping loop.

Benchmark tests include linear advection, non‑linear density current, and a suite of terrain configurations (single bell‑shaped mountain, multi‑mountain, and randomly generated profiles). Results show that the learned NEUVE coordinate reduces the mean RMSE by a factor of 1.4–2 compared with Gal‑Chen, Hybrid (s = 1 km), and SLEVE (s₁ = 2 km, s₂ = 0.5 km) coordinates on the same grid resolution (Δx = 1 km, 40 vertical levels). More importantly, spurious vertical‑velocity striations that appear in the hybrid and SLEVE cases over slopes steeper than 30° are virtually eliminated in the NEUVE runs. Analysis of the learned decay profiles reveals a data‑driven adaptation: near the surface the coordinate follows the terrain closely, while higher up it transitions smoothly to near‑flat layers, effectively balancing the need for terrain conformity and numerical stability.

A sensitivity study confirms that a three‑layer, 64‑unit MLP provides sufficient expressive power, while the regularization weight λ ≈ 10⁻³ and ε = 0.05 prevent over‑distortion. The learned coordinate generalizes to unseen topographies, indicating that the optimization captures a robust representation of the “terrain‑vertical interaction” rather than over‑fitting to specific shapes.

The authors discuss several implications: (1) exact metric computation via AD removes a longstanding source of error in terrain‑following models; (2) treating the coordinate as a trainable component opens the door to joint optimization of discretization, physics, and sub‑grid parameterizations; (3) the framework is modular and can be extended to three‑dimensional global models, to include moisture, radiation, and chemistry, or to couple with data‑driven forecast components. Future work will explore large‑scale training on realistic orography, integration with operational NWP pipelines, and the impact on extreme‑event prediction.

In summary, this study demonstrates that automatic differentiation can be harnessed not only for calibrating physical parameterizations but also for designing the numerical grid itself, yielding more accurate, less grid‑dependent simulations of atmospheric flow over complex terrain.