Accelerating Large-Eddy Simulations of Clouds With Tensor Processing Units

Clouds, especially low clouds, are crucial for regulating Earth's energy balance and mediating the response of the climate system to changes in greenhouse gas concentrations. Despite their importance for climate, they remain relatively poorly understood and are inaccurately represented in climate models. A principal reason is that the high computational expense of simulating them with large-eddy simulations (LES) has inhibited broad and systematic numerical experimentation and the generation of large data sets for training parametrization schemes for climate models. Here we demonstrate LES of low clouds on tensor processing units (TPUs), application-specific integrated circuits that were originally developed for machine learning applications. We show that TPUs in conjunction with tailored software implementations can be used to simulate computationally challenging stratocumulus clouds in conditions observed during the Dynamics and Chemistry of Marine Stratocumulus (DYCOMS) field study. The TPU-based LES code successfully reproduces clouds during DYCOMS and opens up the large computational resources available on TPUs to cloud simulations. The code enables unprecedented weak and strong scaling of LES, making it possible, for example, to simulate stratocumulus with 10x speedup over real-time evolution in domains with a 34.7 km x 53.8 km horizontal cross section. The results open up new avenues for computational experiments and for substantially enlarging the sample of LES available to train parameterizations of low clouds. The study of clouds has been impeded by, among other factors, limitations in our ability to simulate them rapidly and on sufficiently large domains. In particular, computational limitations in simulating low clouds are among the reasons for the difficulties of representing them accurately in climate models; this is one of the dominant uncertainties in climate predictions. This paper demonstrates how the large computing power available on tensor processing units (TPUs) (integrated circuits originally designed for machine learning applications) can be harnessed for simulating low clouds. We demonstrate the largest simulations of low clouds to date, with hundreds of billions of variables, and we document their fidelity to aircraft observations. The results open up the large computational resources available on TPUs, hitherto primarily used for machine learning, to the study of clouds in the climate system. We introduce a large-eddy simulation (LES) framework that runs on tensor processing units (TPUs, accelerators designed for machine learning)The fidelity of the LES is established by reproducing aircraft observations of nocturnal stratocumulus clouds over the PacificThe LES exhibit unprecedented scalability on TPUs, enabling the large-scale generation of training data for cloud parameterizations