The Quest to Make Climate Models Faster, Smarter and More Accurate

Infamous for its ferocious waves, violent winds, and copious shipwrecks, the Drake Passage cuts between the tip of South America and the outstretched arm of Antarctica. The passage is the bottleneck through which the largest ocean current must pass, and it exerts outsized influence on global climate dynamics. As such, it is an ideal test ground for the Differentiable programming in Julia for Earth system modeling (DJ4Earth) initiative. In a recent paper published in the American Geophysical Union’s Journal of Advances in Modeling Earth Systems, a team of researchers introduced a new framework for simulating components of the Earth and, to demonstrate the framework’s versatility, tested it on four different Earth component models — among them, an idealized ocean model of the Drake Passage.

When modeling a patch of ocean like the Southern Ocean around Drake Passage, a universal problem of computational science emerges: how much detail can a model capture without breaking the bank? 

“There is a fundamental tension within ocean modeling. As you increase the resolution of the ocean grid, the model is expected to become more accurate, but the computational cost also shoots up,” explains Joseph Kump, a fifth-year graduate student at the Oden Institute for Computational Engineering and Sciences who led the ocean modeling section of the study. “There are so many fine-grained processes that affect the ocean on a large scale, but you can’t accurately represent them within 100-kilometer (62-mile) resolution.” 

To mitigate this issue, large models like Oceananigans.jl (the software used to model the ocean around Drake Passage, whose name is a wordplay on “Ocean” and “shenanigans”) turn to parameterization. Parameterization essentially lumps all of the small processes together and approximates their collective effect. One example of a tiny but mighty phenomenon that gets swept up into a larger parameter is mesoscale eddies. 

Eddies are considered the ‘weather’ of the ocean and are the most energetic component of the oceanic environment. Their swirling patterns outside the main ocean current help move energy, nutrients, and warm and cool water temperatures around the globe. 

“The largest eddies, known as geostrophic or mesoscale eddies, span roughly 5 to 200 kilometers (3 to 124 miles) and add a lot of kinetic energy and turbulence to the ocean system, which influences the bigger flows that you can explicitly model,” explained Kump. “If you don't consider mesoscale eddies, then you're ignoring a major conduit of larger processes within the ocean.” This gap between mesoscale eddies’ small size and their large influence is what makes parameterization (rather than explicit modeling) the practical choice for global ocean general circulation models. 

However, parameterization invites another huge challenge. Since parameters represent the interplay of multiple physical phenomena, like instabilities surrounding powerful ocean currents, the values of parameters cannot be directly measured. As a result, parameterization introduces substantial uncertainty into the models.

Patrick Heimbach, professor of Earth and Planetary Sciences and lead of the Computational Research in Ice and Ocean Systems Group at the Oden Institute, explains, “Model biases appear when you run these simulations and see persistent differences between what the model outputs, like sea-surface temperature over a large region, and what observations tell us. We think a major reason for these discrepancies is that many of the parameters in these models haven’t been calibrated well.” Calibrating parameters using observational data or high-fidelity simulations often requires ad-hoc techniques that are both complex and unreliable. Heimbach, who also serves as Kump’s advisor, added, “If the parameters aren’t well-tuned, the model simply can’t reduce those biases by itself.”

This was a big motivator for the development of the DJ4Earth automatic differentiation (AD) framework. AD enables efficient computation of gradients: if the cost function is the gap between observational data and a model’s predictions, the gradient quantifies how that cost function changes when you adjust a parameter’s value, i.e., the cost function’s sensitivity to parameters. Typically, calculating the gradient requires changing one parameter at a time and rerunning the full simulation to measure its impact on the result, a tedious and expensive process for models with millions of lines of code and large sets of parameters. 

To address the need for quick and efficient gradient calculations, Billy Moses, assistant professor at the University of Illinois Urbana-Champaign and a former J. Tinsley Oden Faculty Fellow at the Oden Institute, led a team that built two complementary tools called Enzyme and Reactant. 

The tool that computes the gradients is Enzyme. Rather than working directly on human-written code, Enzyme converts a variety of human-written codes into a single representation called intermediate representation (IR) of the LLVM [HP1] compiler and computes gradients on this instead. The advantage is that many languages can be transformed into LLVM IR, making AD accessible on many codebases. 

Reactant, the second tool, focuses on hardware compatibility. Moses noted, “A lot of modern machines aren't tailored to scientific computing. There's been a big trend toward building chips for AI accelerators, which is great if you're a billion-dollar company with the resources to design chips for your exact workload, but is less helpful if you're not. So, the second big thing we built is a tool, inside the compiler, that takes existing workflows and automatically runs them on these AI chips, as well as distributed versions of other devices. That's Reactant."

With these tools, the DJ4Earth team can calibrate parameters effectively and quickly. But beyond parameter calibration, AD also allows them to easily integrate machine learning into Earth system models. “We also are very interested in applying machine learning —  not to replace the physics-based ocean model as a whole — but to augment it,” Kump added. 

He outlined two strategies. The first completely swaps parameterization for a small AI system that achieves similar results more efficiently. The second uses a mix of AI and conventional parameterization methods. 

“What we're more interested in doing is embedding even smaller neural networks within the parameterization itself to improve how we refine those free parameters.” Either way, AD is what makes the integration with machine learning possible: training a neural network means learning which adjustments improve its performance and final results. 

In addition to the Drake Passage ocean model, the team tested their framework on shallow water, ice sheet, and atmospheric models to showcase that it can successfully work on diverse applications. Looking ahead, the team hopes to combine these four parts into a single model where these systems can interact with each other.

This research has “successfully demonstrated that scientific codes can be differentiated, giving researchers predictive power across different applications, like anticipating a hurricane's path or identifying where more observational data is needed to improve weather forecasts,” said Moses. He added that for decades, scientific computing greatly benefited from Moore’s Law, the notion that every two years, computational power on a general-purpose chip doubles. But as that trend has stalled, progress increasingly depends on specialized hardware designed for AI, like TPUs, custom chips made by Google. 

Thanks to the new DJ4Earth framework, oceanographers can simulate the global ocean circulation or regional subdomains such as the Drake Passage on TPUs, a feat that had not been accomplished before. “What’s exciting is that by making these models differentiable and performance-portable, we can run them on those new architectures,” Moses said, adding, “It means the scientific community can keep advancing even as the traditional hardware plateaus.”

The work was made possible through an NSF Cyberinfrastructure for Sustained Software Innovation (CSSI) grant and involved tight collaboration between several institutions, among which MIT, Dartmouth College, UIUC, Argonne National Laboratory, the University of Oxford (UK), and the Potsdam Institute for Climate Impact Research (Germany).