The Ai2 Climate Emulator
AI-driven climate modeling reaches for the future
Chris Bretherton / December 10, 2024
Could AI help the world deal with climate change?
Fossil fuel combustion powers the world economy. The resulting CO2 is driving rapid global warming and new extreme heat waves, floods, droughts, melting ice sheets, acidifying oceans and rising sea levels. Over the last 60 years, climate models have become our best tool for predicting future climate change and its many consequences. They translate physical equations into complex computer models to simulate the earth system. The atmospheric component of a physically-based climate model works by simulating day-to-day weather and its interactions with land, ocean, and ice. Recently, fast AI-driven weather forecast models have become more accurate than the best physics-based computer models. Can similar technology be used to develop fast, easy-to-use AI climate models for reliably simulating decades to centuries across a range of climates? Such models could enable a much broader use of climate simulations to inform planning for climate change. The Ai2 Climate Emulator (ACE) is leading the way toward this challenging research and engineering goal.
Why is climate modeling a grand challenge for AI?
Several aspects of climate modeling make it challenging for AI. AI-based weather models work by updating the ‘state’ - the temperature, moisture and winds throughout the global atmosphere in steps of 6-24 hours. Almost all such models learn those updates from regular global observationally based state estimates created by a model-data fusion process called reanalysis. These estimates are most accurate for the last 45 years, due to global coverage of weather satellite measurements. The weather model is then ‘rolled out’ from a known present state to a desired forecast time. Skillful AI weather models such as GraphCast, Pangu and FourCastNet have used diverse ML architectures such as graph neural nets, transformers and neural operators that can efficiently learn the effects of winds, storm systems, mountains on a spherical earth using a grid of points 25-150 km apart.
However, adapting AI weather models to climate prediction faces three fundamental challenges:
- Most such models cannot be accurately rolled out more than a few weeks before they predict highly unrealistic atmospheric states, so they are not useful for climate change prediction.
- Future climates will lie outside the sample of historically observed climates on which AI weather models have been trained.
- The AI weather model must be ‘coupled’ to AI or physics-based models of ocean, land, and ice.
Last year: idealized stationary climate
During 2023, Ai2’s Climate Modeling group developed a first version of ACE to start addressing these challenges. ACE was based on NVIDIA’s open-source Spherical Fourier Neural Operator (SFNO) architecture. The ocean has a much larger heat capacity than the atmosphere or land, so it evolves more slowly. Thus ACE, unlike AI weather emulators, treated the ocean temperature as known, which helps stabilize long rollouts of AI weather prediction. The outputs are chosen to enable testing of global air mass, moisture, and energy conservation properties, important for physical interpretability. To generalize into future climates, ACE can be trained on outputs from a reference physics-based climate model that can be run across a range of climates. ACE is thus an emulator of that reference model, whose flaws ACE inherits. Because the ML optimizes six-hourly updates of atmospheric state, and there are almost 1500 6-hourly steps in a year, there is no guarantee that this process will naturally produce an accurate climate, but it did. Our first version of ACE, trained with a realistic repeating annual cycle of present-day ocean temperature, accurately emulated decades-long simulations of its reference model on a 100 km grid (Watt-Meyer et al. 2023) at 800 simulated years per day on a single A100 GPU – 100x faster and 1000x more energy efficient than the reference physics-based model. We also showed that ACE can easily be trained on other reference models without hyperparameter changes and can realistically emulate the frequency of extreme rainfall events, as shown in Fig. 1 (Duncan et al. 2024).
This year: historical and future climate change
The new ACE version 2 (ACE2) ups the ante - it is the world’s first climate emulator accurate for climate variability and climate change. ACE2 incorporates numerous technical improvements, including exact ML-based conservation of air mass and moisture and a higher-capacity model. It has been trained on observed climate (1940-2020), as detailed by Watt-Meyer et al. (2024). This includes natural year-to-year variability in ocean temperatures associated with such phenomena as El Niño, as well as the bulk of the global warming experienced to date.
When trained on historical reanalysis, with a 15-year holdout period (1996-2010) for validation and testing, ACE2 replicates observed trends in atmospheric and land temperature (Fig. 2) and the response of rainfall and other atmospheric features to El Niño and La Niña events. ACE2 is already being stress-tested by external academic groups and by one of our institutional collaborators for applications we did not foresee.
ACE2 can also be coupled to a simple ‘slab’ thermodynamic ocean model (SOM) which heats or cools in response to heat exchange with the overlying atmosphere, a configuration we refer to as ACE2-SOM (Clark et al., 2024). The same SOM is implemented in the reference Fortran climate model and in PyTorch for coupling with ACE2. In this configuration, the atmospheric concentration of CO2, the greenhouse gas most important for climate change, can be varied and the response of the reference model and the emulator to changes in CO2 can be compared - an important demonstration of the potential utility of ACE2 for climate change simulations.
ACE2-SOM was trained to simulate 6-hourly atmospheric state evolution in 10-year NOAA/GFDL SHiELD+SOM reference model simulations with present-day, 2x and 4x CO2 concentrations, in which the climate has been allowed to equilibrate to the CO2 concentration. For comparison, in 1800, CO2 concentration was ⅔ as large as today, and by 2100, CO2 may double again from present-day concentrations unless the world’s power supply is aggressively decarbonized. ACE2-SOM accurately reproduces the geographic patterns of temperature and precipitation change from 1xCO2 to an out-of-sample 3xCO2 climate simulated by the reference model (Fig. 3). It also reproduces the rate of surface warming in an out-of-sample ramp simulation in which CO2 is increased by 2% per year over 70 years from 1x to 4x, suggesting ACE2-SOM is useful for simulating a warming climate, not just an equilibrium climate.
Try it out yourself
Pre-trained ACE models and the data necessary for inference are available in our ACE Hugging Face collection. Although ACE models run fastest on the latest GPUs, they can also be run on your laptop and on a Macbook they are still faster than comparable physics-based models at the same horizontal resolution and requiring hundreds of CPU cores. The code used for training and evaluation is available with documentation at https://github.com/ai2cm/ace.
The future of ACE
With ACE2, we’ve come a long way toward addressing the first two challenges that we laid out for applying AI to climate modeling (stability and accuracy across a range of climates), and part way toward addressing the third challenge (coupling to the ocean, ice and land) using a simple but informative slab ocean model. Our goal for next year is coupling ACE2 to a realistic AI ocean emulator developed by external partners that can simulate near-surface currents and ocean temperatures, and the slow uptake of heat into the deep ocean. This coupled emulator will be attractive for many core tasks of climate modeling, from quickly generating many possible realizations of future weather and climate to look for disruptive extreme events, to academic use as a climate model designed to be conveniently run on a laptop or cloud computing resources.