Large-eddy simulations on TPUs

In this work, we focus on stratocumulus clouds, which cover ~20% of the tropical oceans and are the most prevalent cloud type on earth. Current climate models are not yet able to reproduce stratocumulus cloud behavior correctly, which has been one of the largest sources of errors in these models. Our work will provide a much more accurate ground truth for large-scale climate models.

Our simulations of clouds on TPUs exhibit unprecedented computational throughput and scaling, making it possible, for example, to simulate stratocumulus clouds with 10× speedup over real-time evolution across areas up to about 35 × 54 km². Such domain sizes are close to the cross-sectional area of typical global climate model grid boxes. Our results open up new avenues for computational experiments, and for substantially enlarging the sample of LES available to train parameterizations of clouds for global climate models.

The LES code is written in TensorFlow, an open-source software platform developed by Google for ML applications. The code takes advantage of TensorFlow’s graph computation and Accelerated Linear Algebra (XLA) optimizations, which enable the full exploitation of TPU hardware, including the high-speed, low-latency inter-chip interconnects (ICI) that helped us achieve this unprecedented performance. At the same time, the TensorFlow code makes it easy to incorporate ML components directly within the physics-based fluid solver.

We validated the code by simulating canonical test cases for atmospheric flow solvers, such as a buoyant bubble that rises in neutral stratification, and a negatively buoyant bubble that sinks and impinges on the surface. These test cases show that the TPU-based code faithfully simulates the flows, with increasingly fine turbulent details emerging as the resolution increases. The validation tests culminate in simulations of the conditions during the DYCOMS field campaign. The TPU-based code reliably reproduces the cloud fields and turbulence characteristics observed by aircraft during a field campaign — a feat that is notoriously difficult to achieve for LES because of the rapid changes in temperature and other thermodynamic properties at the top of the stratocumulus decks.

One of the test cases used to validate our TPU Cloud simulator. The fine structures from the density current generated by the negatively buoyant bubble impinging on the surface are much better resolved with a high resolution grid (10m, bottom row) compared to a low resolution grid (200 m, top row).

Outlook

With this foundation established, our next goal is to substantially enlarge existing databases of high-resolution cloud simulations that researchers building climate models can use to develop better cloud parameterizations — whether these are for physics-based models, ML models, or hybrids of the two. This requires additional physical processes beyond that described in the paper; for example, the need to integrate radiative transfer processes into the code. Our goal is to generate data across a variety of cloud types, e.g., thunderstorm clouds.

Rendering of a thunderstorm simulation using the same simulator as the stratocumulus simulation work. Rainfall can also be observed near the ground.

This work illustrates how advances in hardware for ML can be surprisingly effective when repurposed in other research areas — in this case, climate modeling. These simulations provide detailed training data for processes such as in-cloud turbulence, which are not directly observable, yet are crucially important for climate modeling and prediction.

Acknowledgements

We would like to thank the co-authors of the paper: Sheide Chammas, Qing Wang, Matthias Ihme, and John Anderson. We’d also like to thank Carla Bromberg, Rob Carver, Fei Sha, and Tyler Russell for their insights and contributions to the work.