Scientists from NVIDIA, in collaboration with Lawrence Berkeley National Laboratory (Berkeley Lab), released a machine learning tool called Huge Ensembles (HENS) for extreme-weather prediction that brings supercomputer-class forecasting but at significantly less computational power and cost. Available as open source code or ready-to-run model, it forecasts low-likelihood, high-impact events—from prolonged heat waves to 100-year hurricanes. The technology could help climate scientists, city officials, and emergency managers quickly test scenarios and update response plans with minimal computing resources.
The two-part study published in the journal Geoscientific Model Development, introduces a method called HENS to produce 27,000 years of data and is one of the largest and most reliable ensembles of weather and climate simulations available.
Using NVIDIA PhysicsNeMo, an open source Python framework for building, training, and fine-tuning physics AI models at scale, and Makani open source frameworks, the researchers trained global weather models to refine the HENS methodology.
“Twenty-seven thousand years of simulations is a goldmine for studying the statistics and drivers of extreme weather events,” said Ankur Mahesh, co-author on the study and a graduate student researcher in Berkeley Lab’s Earth and Environmental Sciences Area. “This large sample size is truly at a scale that has not been seen before.”
According to the study, HENS can predict weather faster than other methods, taking minutes instead of hours. It also extends the forecast window, predicting extreme weather events from six hours to 14 days into the future at a resolution of 15 miles (25 kilometers). It can help researchers study weather patterns at high resolution over many decades to identify new clues leading up to an extreme event.
“With HENS, we now have the luxury of going after low-likelihood, high-impact extreme events predicted over years and decades instead of single near-term events,” said senior co-author Bill Collins, a faculty senior scientist in Berkeley Lab’s Earth and Environmental Sciences Area and a professor at UC Berkeley.
This new approach also requires far less energy and people hours than other methods, and saves energy by retraining models on new data—a technique to ensure accuracy—more quickly than other methods, Collins added.
Training HENS: PhysicsNeMo and 40 years of climate data
HENS employs an AI model trained using PhysicsNeMo on 40 years of ERA5 data, one of the best historical atmospheric state sources. Once trained, the model offers a far cheaper computational approach for forecasts, said Shashank Subramanian, a machine learning engineer in the Department of Energy’s National Energy Research Scientific Computing Center (NERSC) at Berkeley Lab and study co-author who helped Mahesh develop and test the training and evaluation workflows.
“HENS is a game changer. Until today, generating 1,000- or 10,000-member ensembles of simulations was simply impractical because of prohibitive compute and data storage costs,” said co-author Michael Pritchard, director of climate simulation research at NVIDIA and a professor at UC Irvine. “Thanks to this team’s careful work calibrating novel AI simulation technology, it is now fit for purpose to generate massive ensembles including realistic heat wave counterfactuals at orders-of-magnitude faster completion than traditional numerical simulations.”
How can you improve weather prediction accuracy using HENS?
To capture the range of possible future weather outcomes, national weather services run multiple different simulations, or “ensemble members,” each with small changes to the initial conditions. These numerical models are based on laws of physics such as the conservation of mass, the conservation of momentum, and the conservation of energy. There is a lot of trust in these physics-based simulations, but they are also very computationally expensive because they require a supercomputer.
Due to this expense, traditional weather models can only have 50 ensemble members. To find extreme weather, the initial conditions of a model need to be perturbed thousands of times and require hundreds of supercomputing hours.
The researchers used HENS to create 7,424 ensemble members based on initial weather conditions from each day of summer 2023, the hottest on record at the time—nearly 150x more members than what’s possible with conventional models—each ensemble member represents an alternate weather trajectory, or a different way the weather could have unfolded.
“This allowed us to get a better estimate of the tail of the distribution and to understand extreme events that could have occurred that summer,” Mahesh said.
The predictions made by HENS have uncertainties that are over 10 times smaller than those from traditional models. It is able to catch 96% of rare but severe extreme weather events that other models usually miss. Together, these strengths have allowed the team to create an enormous dataset, about 27,000 years’ worth of climate data (20 petabytes).
During rigorous validation experiments at NERSC, Mahesh and team weighed the ensemble predictions on a wide range of diagnostic metrics, showing that HENS is very close to the gold standard.
What’s next?
In future work, Mahesh said that the team plans to study the 27,000-year simulations with the hope of uncovering new insight into the drivers behind the low-likelihood high-impact events, such as catastrophic heat waves, hurricanes, and atmospheric rivers, that have devastated communities in recent years. They also aim to further reduce the computational requirements for running HENS.
Related Resources
- HENS open source code and fully-trained model on GitHub
- HENS research papers:
NERSC is a DOE Office of Science user facility at Berkeley Lab. This work was supported by the DOE Office of Science.
