Suêtes: An end-to-end differentiable non-hydrostatic limited-area dynamical core

Abstract

We present Suêtes, a fully differentiable, non-hydrostatic, limited-area atmospheric dynamical core implemented in JAX. Suêtes includes both a semi-implicit semi-Lagrangian scheme and a Split-Explicit Runge-Kutta scheme with acoustic substepping, enabling applications from coarse regional simulations to convection-permitting dynamics. Its terrain-following geometry, spatial discretization, time integration, initial conditions, lateral boundary treatment, and included physical parameterizations are all compatible with reverse-mode automatic differentiation. A single reverse-mode pass can therefore compute sensitivities of forecast diagnostics to selected initial and boundary fields, physical parameters, topography, and computational geometry. We provide a series of numerical experiments---including a tracer-source inverse problem, a sensitivity analysis of a three-dimensional squall-line, adversarial perturbation construction an ERA5-driven downslope hurricane force wind event, topography optimization, and terrain-following coordinate optimization---demonstrating the versatility of Suêtes as a framework for gradient-based diagnosis and model development at regional and convection-permitting scale.