# Learning vertical coordinates via automatic differentiation of a dynamical core

arXiv: https://arxiv.org/abs/2512.17877 · HTML: https://mlsanity.com/paper/2512.17877/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 45 references that are in the map.
- First posted on arXiv 19 Dec 2025.

## Builds on (most cited first, 12 of 12)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Neural General Circulation Models for Weather and Climate](https://mlsanity.com/paper/2311.07222/) (2023)
- [WeatherBench 2: A benchmark for the next generation of data-driven global weather models](https://mlsanity.com/paper/2308.15560/) (2023)
- [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://mlsanity.com/paper/2007.00016/) (2020)
- [Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling](https://mlsanity.com/paper/2309.15214/) (2023)
- [Kilometer-Scale Convection Allowing Model Emulation using Generative Diffusion Modeling](https://mlsanity.com/paper/2408.10958/) (2024)
- [Can AI weather models predict out-of-distribution gray swan tropical cyclones?](https://mlsanity.com/paper/2410.14932/) (2024)
- [Differentiable Programming for Earth System Modeling](https://mlsanity.com/paper/2208.13825/) (2022)
- [Differentiable Programming for Differential Equations: A Review](https://mlsanity.com/paper/2406.09699/) (2024)
- [Validating Deep Learning Weather Forecast Models on Recent High-Impact Extreme Events](https://mlsanity.com/paper/2404.17652/) (2024)
- [Numerical models outperform AI weather forecasts of record-breaking extremes](https://mlsanity.com/paper/2508.15724/) (2025)
- [Constructing Extreme Heatwave Storylines with Differentiable Climate Models](https://mlsanity.com/paper/2506.10660/) (2025)

## Built on by (most cited first, 1 of 1)

- [A solver-in-the-loop framework for end-to-end differentiable coastal hydrodynamics](https://mlsanity.com/paper/2604.07129/) (2026)

## Abstract (arXiv)

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.