JAX-SCM v1.0: a modern atmospheric single-column model for boundary layer research

Abstract: We present JAX-SCM v1.0, an open-source atmospheric single-column model for boundary layer research, implemented in Python using the JAX computing library. The model solves for horizontal wind, potential temperature, and specific humidity, combined with prognostic turbulent kinetic energy and turbulent statistics parameterized by the Mellor-Yamada-Nakanishi-Niino level-2.5 (MYNN-2.5) turbulence closure. We verify the implementation against three well-established benchmark cases covering neutral (turbulent Ekman layer), stable (GABLS1), and convective (Wangara Day 33) conditions. Close agreement with reference solutions is demonstrated across all regimes. By building on JAX, the model benefits from just-in-time compilation and native GPU support. While JAX-SCM is not yet fully differentiable, basing it on JAX also lays the foundation for future integration with machine learning components. The model is designed for simplicity and modularity, lowering the barrier to entry for users and developers alike.

• The paper introduces JAX-SCM v1.0, an open-source Python single-column model that implements the NN09 formulation of MYNN-2.5 turbulence and evolves wind, temperature, moisture, and TKE.
• The model uses staggered-grid discretization, explicit Adams–Bashforth and semi-implicit Crank–Nicolson solvers, and JAX features that support transparent CPU/GPU execution and scalable ensembles.
• Verification across neutral Ekman, GABLS1 stable, and Wangara convective cases shows close agreement with LES, intercomparison, and published reference results, while condensation and full automatic differentiation remain future work.

JAX-SCM v1.0 is an open-source atmospheric single-column model (SCM) implemented in Python with the JAX computing library, targeting boundary layer research (2605.24544). The model integrates the one-dimensional evolution of horizontal wind, potential temperature, specific humidity, and prognostic turbulent kinetic energy (TKE), with turbulence parameterized by the level-2.5 Mellor–Yamada–Nakanishi–Niino (MYNN-2.5) closure. The paper's central contribution is twofold: a verified numerical implementation of the NN09 formulation of MYNN-2.5, and a software design that positions the model for just-in-time (JIT) compilation, transparent GPU execution, and eventual automatic differentiation (AD).

Motivation and positioning

Single-column models serve as inexpensive, controlled test beds for developing and tuning parameterizations, avoiding both the cost of 3D simulations and the difficulty of attributing errors to individual processes in full models. The author argues that existing SCMs are predominantly implemented in Fortran, which impedes adoption of modern hardware (GPUs) and integration with machine learning tooling. The choice of JAX is motivated by prior results: a JAX-based ocean model reported on-par CPU performance with its Fortran counterpart, roughly a 2x GPU speedup over Fortran, and successful use of AD in JAX-based atmospheric and fluid solvers. The author is careful not to claim current differentiability — this is stated explicitly as future work.

Model formulation

The governing equations are forced 1D vertical diffusion equations for $U$, $V$, $\Theta$, $Q_v$, and $q^2$ (twice the TKE), including Coriolis forcing, geostrophic wind, and thermal-wind advection of temperature under the hydrostatic assumption. Turbulent fluxes follow gradient diffusion with eddy diffusivities $K_\Phi = L_M q S_\Phi$, where $L_M$ is the harmonic mean of surface, turbulent, and buoyancy length scales and $S_\Phi$ are MYNN-2.5 stability functions. The NN09 equations are implemented unchanged, with one deliberate simplification: the partial-condensation scheme is omitted, so no condensation or cloud formation is represented even though moisture is a prognostic variable.

Spatial discretization uses a staggered grid with second-order central differences on half levels, plus 1-2-1 filtering of $L_M$, $K_m$, and $V$0 each step to suppress small-scale oscillations. Two time integrators are provided: an explicit second-order Adams–Bashforth scheme, whose stability constraint $V$1 becomes restrictive in convective regimes, and a semi-implicit Crank–Nicolson solver that treats flux divergences implicitly and is unconditionally stable. A notable numerical detail is the treatment of TKE dissipation: following Janjić's finding that explicit treatment of $V$2 causes instability, dissipation is quasi-linearized as $V$3 and added semi-implicitly, keeping the implicit system linear and solvable with a tridiagonal solver.

Surface coupling uses Monin–Obukhov similarity theory with the original Businger–Dyer relations in Paulson's integrated form. Users prescribe either a sensible heat flux or a surface temperature; the latter option is highlighted as essential for stably stratified conditions because of the well-known duality of the heat flux in the stable boundary layer. Since friction velocity and Obukhov length depend on each other, they are obtained iteratively. Near-surface wind gradients required for TKE shear production are taken from MOST rather than extrapolated from the first interior grid level.

Verification against three benchmark cases

Verification spans neutral, stable, and convective regimes in order of increasing complexity, comparing against reference ensembles or the closure's original publication.

Neutral Ekman layer (A94): driven by constant geostrophic wind ($V$4 m s$V$5), run for 10 inertial periods at $V$6 m over $V$7 m. Column-integrated normalized TKE, wind non-stationarity metrics, and time-averaged momentum flux profiles all agree closely with the LES ensemble of Andren et al. Notably, the MYNN-2.5-parameterized momentum fluxes match the partially resolved LES fluxes well — encouraging, since SCM curves are smoother by construction.

GABLS1 stable boundary layer: constant geostrophic forcing ($V$8 m s$V$9) with surface cooling of 0.25 K h$\Theta$0 over 9 hours in a 400 m domain. JAX-SCM reproduces the expected low-level jet near $\Theta$1 m, realistic stratification, and a surface heat flux of approximately $\Theta$2 K m s$\Theta$3, matching the intercomparison ensemble across mean and turbulent quantities, including eddy diffusivities and boundary layer height.

Wangara day 33 (convective): initialized from observed 09:00 soundings and forced with time-dependent sinusoidal surface heat and moisture fluxes; this case exercises all model components including humidity. Because neither the campaign data nor NN09 specify an initial TKE profile, the author estimates a surface value from an observed friction velocity ($\Theta$4 m² s⁻²) with a cubic power-law decay — an assumption acknowledged as a likely source of early-simulation discrepancies. Profiles of temperature, moisture, their fluxes, TKE, and master length scale match NN09 closely; mixed-layer parameters (inversion height $\Theta$5, convective velocity scale $\Theta$6) are captured well. One deviation is flagged: the entrainment-flux ratio $\Theta$7 deviates noticeably at 10:00 LST, attributed to spin-up and vertical discretization affecting resolution of the thin entrainment layer; the author judges this unproblematic given accurate $\Theta$8 and $\Theta$9. This case demonstrates real-world applicability beyond synthetic benchmarks.

Technical design

The Python interface emphasizes composability: forcings are arbitrary functions of simulation time, initial conditions can be ingested via pandas from diverse formats, and all settings are collected in a tree-structured Simulation object that prevents configuration mismatches. Output integrates with xarray, supporting netCDF and Zarr backends. JIT compilation to XLA is transparent, so identical code runs on CPU and GPU; the paper positions GPU acceleration as useful primarily for parameter sweeps and ensemble simulations rather than single runs.

On differentiability, the paper is candid about a concrete obstacle: the turbulent velocity scale $Q_v$0 has a derivative that diverges as $Q_v$1, breaking gradients. Physical parameterizations were not designed with AD in mind, so achieving full differentiability requires targeted modifications. Until then, the AD-enabled use cases — gradient-based closure tuning, data assimilation, and embedding learned components in the solver — remain prospective rather than demonstrated.

Limitations and open questions

Several limitations are conceded directly. The absence of condensation excludes cloud-parameterization development, historically a central SCM use case, pending a future partial-condensation scheme. Full differentiability is not yet achieved, and the wall-clock advantage of Crank–Nicolson over Adams–Bashforth is described as situation-dependent rather than established. The Wangara verification relies on an assumed initial TKE profile and digitized reference data, leaving residual attribution of small differences between JAX-SCM and NN09 to numerical implementation details unresolved. Whether the similarity functions should extend beyond Businger–Dyer relations remains open, though the modular interface permits testing alternatives.

Conclusion

JAX-SCM v1.0 delivers a verified, operationally relevant turbulence closure (MYNN-2.5/NN09) in a modern, accessible Python package, validated across neutral, stable, and convective regimes with close agreement to canonical references. Its principal limitation — incomplete differentiability — is also its defining open question: demonstrating that gradient-based optimization or hybrid physics–ML coupling yields tangible benefits for boundary layer modeling would substantiate the design motivation. As released, it stands as an immediately usable tool for SCM-based boundary layer research.