ResCu in JCM

A hybrid convection scheme trained online inside a differentiable GCM

The convection and cloud networks of HyCAM, trained on super-parameterised CAM, were ported into JAX-GCM (JCM) as its moist physics and trained online, end to end through the differentiable dynamical core, against ERA5 and IMERG. JCM is the differentiable atmospheric model of the Climate Analytics Lab (Duncan Watson-Parris, UC San Diego); this page shows my work on top of it.

Setup

JCM host, T63, 1800 s step, every column, every step dynamical core (Dinosaur) spectral, differentiable host physics, fixed: RRTMGP radiation surface fluxes, drag vertical diffusion gravity-wave drag ResCu convection net (trained) NN cloud scheme in: column state and forcing out: heating, moistening, rain rollout, 12 to 48 hours from a balanced ERA5 state model output: daily precipitation u, v, T, q every 6 hours loss against observations IMERG precipitation ERA5 wind, temperature, humidity gradient through the core and the fixed physics, back to the ResCu weights (Adam) Only the networks are trained; the host physics is differentiated but not changed.

The host model needed three fixes before training was meaningful: a repaired orographic drag, balanced initial states from ERA5, and a Lanczos-tapered orography that removes the ringing of the spectral truncation.

Online training works

Day-1 precipitation averaged over 364 starts of 2014: IMERG, online-trained ResCu, original ResCu, the two differences against IMERG and the zonal means
Day-1 precipitation forecasts averaged over 364 daily starts of 2014: IMERG, the online-trained net and the original net; the online-trained net cuts the error against IMERG by a third and improves every one of the 364 starts.

The model stays stable

Precipitation of a year-long free run with the online-trained net, with the interactive surface and with prescribed fluxes, against IMERG, with the differences and zonal means
A year-long free run with the online-trained net, days 60 to 360, with the interactive surface and with prescribed ERA5 surface fluxes, against IMERG; the ITCZ, the monsoons and the storm tracks stay in place for the whole year.