Neural Compression of Atmospheric States

TL;DR

This work presents a neural compression framework for atmospheric states that leverages a HEALPix projection to enable efficient, transform-based coding of spherical data. By comparing VQ-VAE/VQ-GAN and hyperprior models, it demonstrates compression ratios exceeding $10^3$ with low mean errors and strong preservation of extreme events and spectral properties, enabling desktop-scale access to multi-decadal ERA5-like data. The hyperprior approach emerges as the most robust across variables and levels, while analyses reveal trade-offs related to error bounds, artefacts, and cross-variable interactions. The study lays a practical foundation for deploying neural compression in weather/climate workflows, with implications for data accessibility and dissemination, and points to future work on error-bounded schemes and artefact-aware training.

Abstract

Atmospheric states derived from reanalysis comprise a substantial portion of weather and climate simulation outputs. Many stakeholders -- such as researchers, policy makers, and insurers -- use this data to better understand the earth system and guide policy decisions. Atmospheric states have also received increased interest as machine learning approaches to weather prediction have shown promising results. A key issue for all audiences is that dense time series of these high-dimensional states comprise an enormous amount of data, precluding all but the most well resourced groups from accessing and using historical data and future projections. To address this problem, we propose a method for compressing atmospheric states using methods from the neural network literature, adapting spherical data to processing by conventional neural architectures through the use of the area-preserving HEALPix projection. We investigate two model classes for building neural compressors: the hyperprior model from the neural image compression literature and recent vector-quantised models. We show that both families of models satisfy the desiderata of small average error, a small number of high-error reconstructed pixels, faithful reproduction of extreme events such as hurricanes and heatwaves, preservation of the spectral power distribution across spatial scales. We demonstrate compression ratios in excess of 1000x, with compression and decompression at a rate of approximately one second per global atmospheric state.

Figures (25)

• Figure 1: Hyperprior reconstructions of a global frame on 2016/10/5 at 0 UTC, compressed $1000\times$. From top to bottom: temperature at 850 hPa, specific humidity at 1000 hPa, the zonal component of wind at 1000 hPa (and a zoom over the Caribbean), and geopotential at 500 hPa over the Caribbean. Last row shows reconstruction using huang2022compressing compressed $1150\times$ as a baseline. Columns show the ground truth, the reconstruction and the residual along mean absolute error (MAE). Note the improvement in Hurricane Matthew's reconstruction over the baseline.
• Figure 2: Overview of our neural data compression system. Global surface and atmospheric data are projected onto the 12 HEALPix square base pixels. On the left, the figure illustrates 3 such squares, numbered 7 (in red, over the Indian Ocean), 10 (in green, over continental Europe and Arabic peninsula) and 11 (in blue, over East Asia), for $C$ variables (which may include different physical variables at one or multiple pressure levels), collected from a single atmospheric state at a given timestamp. We jointly compress $C$ separate channels for every $256 \times 256$ square into a quantized representation (here, we illustrate a $32 \times 32$ map of discrete VQ indices corresponding to a 4-downsampling-block VQ-VAE or VQ-GAN). During training and evaluation, we reconstruct the HEALPix squares back from compressed, quantized representations. After reconstruction, spherical harmonics coefficients are calculated via forward transform of the HEALPix reconstructions. These coefficients are used to synthesize re-projected images via inverse transform in latitude/longitude coordinates. On the right, the figure illustrates two re-projected reconstructions for variables 1 and $C$, with areas corresponding to different HEALPix squares coded in different colours. Metrics are computed during reconstruction (comparison of reconstruction vs. ground truth in HEALPix projection), after spherical harmonics forward transform (power spectrum) and after re-projection back onto latitude/longitude representations (comparison of re-projected reconstruction vs. ground truth in latitude/longitude coordinates).
• Figure 3: The VQ-GAN network, which contains a VQ-VAE network and an auxiliary discriminator. The encoder and decoder are residual convnets with attention, which reduce the input data into an 18x18 grid of vectors of size 1024. These are then quantized to their nearest match in the learned codebook of N vectors. See Appendix \ref{['sx:appendix:models']} for details.
• Figure 4: Hyperprior network. This augments the factorized prior network (dashed rectangle) by adding an additional encoder-decoder pathway that encodes the FP network's latents, and decodes them to a grid of variances.
• Figure 5: Distribution of signed errors, with the error in the variable units on the X axis and the fraction of pixels in the evaluation set (2010-2023) at that error levels. Top row shows hyperprior results, bottom row shows 3-block VQ-VAE results, with columns corresponding to temperature, geopotential and zonal wind speed. On each plot we show coloured curves for each of the 13 levels in the vertical dataset.