Representation learning of near-surface atmospheric fields using SHViT modules

With improving resolution of NWP products and observation data, efficiency and scalability of processing methods are becoming increasingly relevant. Substituting the inputs with their lower-dimensional representations in downstream computations is a common technique to address this issue.

This study examines the Encoder, Decoder, and the Encoder, Decoder, and Discriminator setup, utilizing Single-Head Vision Transformer (SHViT) as a backbone architecture. Scalability with a further increase in input size and/or resolution was considered in the selection. SHViT architecture demonstrates good scaling with larger input sizes, making it a suitable candidate.

We optimize for representation learning tasks on fields of selected near-surface variables from the CERRA dataset from 01/2010 to 12/2019. The CERRA dataset was chosen due to its availability and diversity within the spatial extent. We prepared a set of smaller regions, focusing on areas with prominent orographical features. The test set consists of regions not used during optimization, while the validation set is a temporal split from the training set regions. During the evaluation, the scaling constants are assumed to be known. Orography and land-sea mask are provided as input channels for the Encoder at inference and for the Discriminator during optimization for each region.

The Discriminator is enhanced with an explicit Haar wavelet transform and spectral power transform of each variable field as additional input channels.  The latent representation is a one-dimensional vector with 256 features. We evaluate the spatial distribution of mean absolute error (MAE) and associated standard deviation.

This study is based on work carried out at IBL Software Engineering, Innovation Department as part of research and development of NWP products post-processing framework.