Emulating the E3SM Atmosphere with AI

Key Points

• The Ai2 Climate Emulator (ACE), based on NVIDIA’s SFNO architecture, is trained to match the weather and climate of the EAMv2 global atmosphere model using 40 yrs of 6-hourly EAMv2 output with 110 km grid spacing and prescribed sea surface temperature (SST)
• ACE can reproduce E3SM’s time-mean climatology, its extreme events and its tropical waves
• ACE is 100x faster than E3SMv2. This is already being exploited to create large ensembles of simulations for exploring the response of global radiation balance to the spatial pattern of SST
• This emulator for the E3SM atmosphere is an important step towards emulating the full coupled Earth system, which could transform how we do climate science

Background

Lawrence Livermore National Laboratory (LLNL) and the climate modeling group of the Allen Institute for Artificial Intelligence (Ai2), a nonprofit research institute in Seattle, have been partnering on AI tools that enable impactful new climate-change applications of current and future versions of E3SM. There are many ways AI is being used to this end; our approach is inspired by the recent success of ML-based global weather forecasts trained on reanalyses, for example, GraphCast and Pangu. Those models are typically trained to step the full atmospheric state in a global set of grid columns forward in steps of 6-24 hours. They use diverse ML architectures such as graph neural nets, transformers, or U-Nets that combine column-local and nonlocal information processing. This methodology is called full-model emulation. These models can be 100 times faster than physics-based global weather forecast models of the same grid resolution, because they use fewer vertical levels and less physical variables and because they are not subject to the same stability-based limits on time step. They are less susceptible than physics-based models to systematic forecast bias.

It is natural to ask whether such full-model emulators could also be used as the atmospheric component of a climate model such as E3SM. Challenges emerge. First, ML weather emulators are trained with little or no information about ocean evolution, and typically drift from the target climate or even become unstable after periods of weeks to months. Second, future climates are out of sample of observational reanalyses, and ML is not generally as trustworthy for extrapolation as models based on physical principles.

Vision of ACE

Nevertheless, the speed and efficiency of full-model emulation are enticing, motivating the Ai2 Climate Emulator (ACE). Unlike other emulators that were trained on observational reanalyses, ACE aims to accurately emulate the weather and climate of a physically-based climate model so it can be used for climate projection. ACE can be trained using a modest set of climate model simulations that span any desired range of climates; the hope is that ACE will also accurately emulate intermediate or changing climates within this range. Realizing this hope is still a work in progress, involving enforcing global constraints such as mass, humidity and energy conservation, among other things. To ensure trustworthiness, the trained emulator must be challenged by a portfolio of difficult out-of-sample test cases, e.g. taken from the CMIP Deck.

Since ACE is a ‘model of a model’, it can be faster, but not better – it inherits all the biases of the reference model on which it is trained. Furthermore, it must be retrained or fine-tuned to reflect changes in the reference model. However, ACE training is quick — typically 1-3 days on a few GPUs, and ACE’s inference speed enables attractive new uses. One application we are exploring is an emulator trained on decade+ km-scale SCREAM simulations across multiple climates, coarsened to an affordable emulator grid resolution and downscaled back to km-scale resolution where desired. A second application is emulating “SST Green’s functions” (generated from an array of EAM simulations with localized patches of warm and cold sea surface temperature (SST) anomalies centered on locations spanning the entire global ocean) for understanding how the spatial pattern of SSTs affects cloud feedbacks on global warming. Another important, challenging development direction is to develop a coupled ML emulator of atmospheric and ocean dynamics (including sea-ice and land surface processes) suitable for certain climate change applications, such as large ensembles of CMIP-type simulations.

The ACE vision is complementary to the goal of physical earth system modeling. A more robust and less biased reference physics-based model (possibly using ML-based parameter optimization) can enable a more accurate and useful full model emulator. The emulator can be quasi-automatically trained at a variety of useful grid resolutions rather than requiring the tedious parameter retuning typically required for a physical model. Physical model development and tuning can then focus on the finest target resolution that is sufficiently computationally affordable to allow a suite of simulations covering enough climate scenarios to train the emulator. This enables DOE scientists to focus more on better encoding of earth system physics and less on complex models of subgrid variability, driving faster advances in the reliability of E3SM.

This research documents an early but significant step toward this vision. An early version of ACE was trained on 40 years of a reference EAMv2 simulation forced by a repeating annual cycle of SST and insolation, with constant greenhouse gas concentrations. The goal was to test whether ACE accurately emulated EAMv2’s mean climate, with a focus on precipitation patterns and extremes; this research shows that the emulator is quite successful in this context. Since then, the approach has been extended to AMIP-style training and testing as a step toward future coupled emulation.

Figure 3. Hovmöller diagrams of daily mean tropical-mean precipitation over two typical years, bandpassed to retain 20-100 day periods. Black horizontal lines mark month boundaries. ACE trained to emulate EAMv2 has a similar character as EAMv2.

ACE technical details

ACE is built on NVIDIA’s open-source Spherical Fourier Neural Operator (SFNO) architecture, also used by recent versions of FourCastNet. It combines the terrain-following levels of the reference model (72 for EAMv2) into 8 layers, enabling accurate comparison of column-integrated quantities such as precipitable water, as well as terms in their budget equations, between ACE and the reference model. The temperature, humidity and horizontal wind components, together with surface pressure and land surface temperature, are stepped forward in each model layer in 6-hour increments on a 1 degree lat/lon grid. The 6-hour time step is computationally efficient but captures the diurnal cycle. ACE also diagnoses surface radiative and turbulent fluxes, top-of-atmosphere radiative fluxes, and column horizontal moisture flux convergence at each update step. A root mean square error (RMSE)-based loss function combines all predicted quantities at each update step. Forty years of 6-hourly model outputs from EAMv2 are used for training. For seasonal-cycle emulation, ACE can be satisfactorily trained as little as a decade of reference model output. ACE’s excellent training and inference efficiency were noted in the third key point.

Key results from this research

Previously, Ai2 had trained ACE on a 1 degree resolution NOAA atmosphere model, FV3GFS, used at finer grid resolution for U.S. operational weather forecasting, and documented high skill in emulating the mean climate of FV3GFS. Would the same approach perform as well on EAMv2?

Indeed, this was the case for all emulated fields. An impressive result was that the mean climate of the emulator was much closer to the mean climate of EAMv2 than the climate bias of EAMv2 vs. observational climatology, as shown in Fig. 1 for precipitation. Another exciting new result, shown in Fig. 2, was that the probability density function (PDF) of daily-mean precipitation, aggregated across all grid columns and seasons, was nearly identical in the ACE emulator and the reference model, even in the extreme tails. Lastly, the emulator closely matched the statistical character of an EAMv2 simulation of MJO-like tropical convective variability, as illustrated in a Hovmöller diagram of 20 to 100 day period bandpass precipitation around the equatorial belt (Fig. 3).

These results have provided a strong foundation for pushing ACE further toward our symbiotic vision of an accurate ML-based coupled emulator trained on high-fidelity fine-resolution multi-climate atmospheric and ocean model output that the E3SM project will be increasingly generating in the next few years.

Code and Data Availability

• ACE model weights (2.5 GB) and the EAMv2 10-year validation set (165 GB) are available to download over HTTP from the E3SM project’s NERSC science gateway at https://portal.nersc.gov/archive/home/projects/e3sm/www/e3smv2-fme-dataset. Documentation, inference code, and an example configuration for running ACE are available in the following repository: https://github.com/ai2cm/ace

Reference

• Duncan, J. P. C., E. Wu, J.-C. Golaz, P. M. Caldwell, O. Watt-Meyer, S. K. Clark, J. McGibbon, G. Dresdner, K. Kashinath, B. Bonev, M. S. Pritchard, and C. S. Bretherton, 2024: Application of the AI2 Climate Emulator to E3SMv2’s global atmosphere model, with a focus on precipitation fidelity. J. Geophys. Res.: Machine Learning and Comp., https://doi.org/10.1029/2024JH000136

This article is a part of the E3SM “Floating Points” Newsletter, to read the full Newsletter check:

• E3SM Floating Points, Nov ’24: Celebrating a Decade of Earth System Modeling