Generative Modeling of Atmospheric Convection

Boxed in by computational limits, many of the details of our atmosphere remain too minute to explicitly resolve in climate models [1, 2, 3]. Key physics driving convection and cloud formation occur on the scale of meters to a few kilometers, while typical modern climate models have a resolution of $100-200{km}^2$ horizontally - meaning important sub-grid processes are parameterized. Computational capabilities are advancing, and climate models are increasingly common, in particular those with three-dimensional explicit resolution of clouds systems. However, the capability to run these models for the $\sim$100-year timescales needed is often impractical [4, 5, 6] and the information content they generate about the details of cloud and storm organization are frequently overwhelming to analyze at its native scale. This has left significant gaps in knowledge about many of the details of cloud-climate feedbacks and the relationship between storm organization and its thermodynamic environment [1, 6]. However, deep learning, and in particular generative models, may provide a path to a better understanding of these phenomena and their role driving the weather and climate of our world.

The application of machine learning in the physical sciences has increased exponentially in recent years but with important avenues still largely unexplored. In climate modeling, deep neural networks have been re-purposed to emulate the large-scale consequences of storm-level heating and moistening over the atmospheric column to replicate mean climate and expected precipitation patterns and extremes [7, 8, 6, 9, 10]. However, much of this work has been confined to deterministic neural networks that ignore the interesting stochastic details of eddy and storm organization. The recent application of Generative Adversarial Networks (GANs, [11]) to the Lorenz ’96 Model suggests a potential, under-explored role for generative models in atmospheric sciences – particularly towards stochastic parameterizations [12, 13]. There have also been initial successes using various types of GAN architectures to generate plausible Rayleigh-Bernard convection. In particular, adding informed physical constraints to GAN loss functions seem to improve the generation of these non-linear fluid flow systems [14, 15, 16, 17]. While promising, such techniques have thus far been restricted to idealized turbulent flows of reduced dimension and complexity; there is ample room to explore generative modeling methods for representing convective details amidst settings of realistic geographic complexity. Meanwhile, generative modeling besides GANs have not been as thoroughly considered for turbulent flow emulation and could potentially power climate models down the line.

VAEs may prove more appropriate than GANs for these climate applications given their design containing both a generative and representational model, their often superior log-likelihoods and reconstruction simulations, and practical advantages including stabler training results, easier performance benchmarking, and more interpretable latent manifold representations [18, 19, 20]. Modified VAEs can reconstruct plausible two-dimensional laminar flow with computational efficiency beyond what is common when numerically solving linear differential equations [21]. There has been preliminary work using VAEs for the clustering of atmospheric dynamics – a gain again relying on simplified Lorenz ’96 model data as well as potential vorticity fields and geopotential heights [22, 23]. This application of representation learning across a variety of simplified simulations suggests VAEs offer great potential as both an engineering tool to help escape computational limits on the generative side and may provide the ability to learn and extract hidden organizational details in atmospheric dynamics on the representation side. However, to the best of our knowledge, this is the first study to use a VAE for representational learning on the details of convective organization and associated gravity wave radiation¹ as revealed by spatial snapshots of vertical velocity – an inherently chaotic and bimodal variable [24] – across a dataset large enough to nonetheless encompass the spatiotemporal diversity of turbulence regimes in the atmosphere. As far as we know, this is also the first study to constrain a VAE’s output statistics by adding a covariance constraint term to its loss function to improve representation and capture variance details at small spatial scales in the turbulent atmospheric boundary layer, which can be considered one of the most difficult locations for climate models. Our results demonstrate the power of VAEs to accurately reconstruct high-resolution climate data, as well as their ability to leverage dimensionality reduction for high level feature learning and anomaly detection. This casts VAEs as promising tools for both dynamical analysis and stochastic parameterization of fine-scale atmospheric processes from cloud-resolving data.

In this Section, we discuss the architecture of the three machine-learning models used here, the design of our covariance constrained VAE loss function, and the generation and preprocessing of the atmospheric simulation data.

Our VAE trained on cloud-resolving climate data produces accurate vertical velocity field reconstructions. When we provide the high resolution training dataset and appropriate convolutional architecture, our VAE learns remarkably accurate representations of any type of convection found within the test dataset. Our VAE captures the magnitude, proper height, and structure across deep convective regimes, shallow convective regimes, and non-convecting regimes (Figure 4). When the “Covariance Constraining” term is added to create a physically informed loss, the CC VAE performance improves enough to match a linear baseline (Table III). But unlike many other image recognition tasks generative models perform, reconstructing the mean of the convection is necessary but not sufficient – we must capture the variance and correlation in the vertical velocity fields. The CC VAE reconstructs variance better than a traditional convolutional VAE and at least on par with the linear baseline (Table III, Figure 2). Our CC VAE is the most versatile of our models with an accurate reconstruction performance overall at different levels of the atmospheric column and different convective spatial scales based on the power spectra of the three models (Figure 2). This precision across both small and large spatial scales revealing our CC VAEs ability to emulate both the overall large pattern of convective plumes and the details within convective composition. Our CC VAEs results replicate disparate structures of convection in both areas of high stochasticity near the atmospheric boundary layer, characteristic of shallow convection, as well as in the upper troposphere, where deep convective regimes dominate. At this stage CC VAEs match the performance of our linear baseline but do not exceed it.

However, unlike the linear baseline, our VAE and CC VAE discover the details of convective organization by representation learning via dimensionality reduction and feature extraction. A 2D, deterministic PCA projection of our CC VAE latent space clusters and separates different convective types (Figure 1). In particular, the distinction between deep and shallow convective regimes and non-convective regimes is encouraging (Figure 1, please visit this link for a complete animation of the 2D Projection of the latent space). The physical knowledge represented in our CC VAEs latent space stands alone from other forms of dimenionality reduction (PCA and t-SNE on the preprocessed data) where there is no evidence of distinction based on convective type. Furthermore, CC VAE predictions of convective type based solely on latent space location map back to a physically sensible pattern over the tropics with deep convection concentrated on land over the Amazon and African Rainforests as well as over the Pacific Warmpool (Figure 3). These predictions from latent space location not only map convection type in a spatially coherent pattern, but also capture the change in convection type with the diurnal cycle over moist, tropical continents (Figure 3, please visit this link for a complete animation of the tropical diurnal cycle). When we exclusively restrict the test dataset to an Amazon Diurnal Composite, the known coherent transitions from shallow to deep convection that occur over tropical rain-forest in response to solar heating of the diurnal cycle correspond to monotonic trajectories in the latent space projection, verified using both t-SNE and PCA (Figure 5). Further tests are required on more complex convective transitions to understand the extent of the physical meaning of the CC VAE latent space, but these initial positive results suggest great potential for physically constrained VAEs as a tool in atmospheric dynamics to uncover information about convective transitions, storm morphology and propagation.

We also evaluate ELBO (Equation 3) for each sample of our test data to find unusual storm development and activity in the dense CRM data.

ELBO allows us to determine the degree to which a vertical velocity field, drawn from our models latent variables is an aberration in the data. Our VAEs inherent ability to detect anomalies in the vertical velocity data proves to be an elegant way to identify deep convection in a more thorough manner than traditional vertical velocity thresholding. An example of one such anomaly we identify is Figure 6 – in this case an instance of two moderate storms developing in one CRM array. This phenomena would be less straightforward to locate through conventional methods, particularly given the size and density of data involved. Our VAEs attribute of anomaly detection learns characteristics of the data instead of naively thresholding based on priors experiences that may not reflect the composition of the dataset. This feature provides the potential to help identify interesting and unexpected weather phenomena from noise – artifacts that might otherwise never be studied in overwhelmingly large and rich datasets.

We develop a VAE to reconstruct immaculate convection images from a high-resolution, cloud-resolving dataset. Our VAE, particularly once a statistically constrained loss function is added, captures the variance and magnitude of distinct convective regimes. The latent space of the VAE proves to be a potent tool for making physically sensible predictions of convection type that accurately reflect the tropical atmosphere and capture the effects of solar heating through the diurnal cycle. The unique VAE loss function allows us to use ELBO to find anomalous storm development in a dense, high resolution dataset that traditional methods might miss. But there is much work to be done before a VAE could be implemented to power stochastic parameterizations for a climate model, likely requiring to condition the VAE on large-scale thermodynamics via expansion of the input vector. If successful, the ability to quickly and efficiently generate synthetic, detailed vertical velocity fields to help run climate models would be a valuable resource for the atmospheric sciences and meteorology communities. But improvements in the generative capabilities would likely come at the expense of the representation learning and the VAEs diagnosis of the physics of convection. We believe these preliminary physical intuitions achieved via representation learning represent a promising avenue for the broader application of generative modeling for advancing the field of atmospheric dynamics [27, 26] and warrant further investigation to understand their full potential.