Advancing Parsimonious Deep Learning Weather Prediction using the HEALPix Mesh

Abstract

We present a parsimonious deep learning weather prediction model to forecast seven atmospheric variables with 3-h time resolution for up to one-year lead times on a 110-km global mesh using the Hierarchical Equal Area isoLatitude Pixelization (HEALPix). In comparison to state-of-the-art (SOTA) machine learning (ML) weather forecast models, such as Pangu-Weather and GraphCast, our DLWP-HPX model uses coarser resolution and far fewer prognostic variables. Yet, at one-week lead times, its skill is only about one day behind both SOTA ML forecast models and the SOTA numerical weather prediction model from the European Centre for Medium-Range Weather Forecasts. We report several improvements in model design, including switching from the cubed sphere to the HEALPix mesh, inverting the channel depth of the U-Net, and introducing gated recurrent units (GRU) on each level of the U-Net hierarchy. The consistent east-west orientation of all cells on the HEALPix mesh facilitates the development of location-invariant convolution kernels that successfully propagate weather patterns across the globe without requiring separate kernels for the polar and equatorial faces of the cube sphere. Without any loss of spectral power after the first two days, the model can be unrolled autoregressively for hundreds of steps into the future to generate realistic states of the atmosphere that respect seasonal trends, as showcased in one-year simulations.

Figures (12)

• Figure 1: Division of the sphere into twelve faces according to the HEALPix. Four faces to represent either the northern (blue) and southern extratropics, while four more faces arrange around the equator to represent the tropics (yellow). Each face can be subdivided into patches with divisions along the side of each face given by powers of two. The sphere in (a) has a pixel-count of one per face side; we call it hpx1. The sphere in (b) counts two pixels per side (hpx2), whereas the two spheres in (c) and (d) have eight pixels per side, i.e., hpx8. Several latitude lines in red emphasize the iso-latitudinal arrangement of the patches. The saturated blue area depicts a $3\times3$ stencil, as applied by a standard convolution. To apply the $3\times3$ stencil at the top corner of the equatorial faces, i.e., stencil position in (d), we fill in the missing corner patch with the average of the values in the two adjacent patches on the extratropical faces.
• Figure 2: Schematic representation of our DLWP-HPX architecture as a sequence of operations on layers (see legend). Individual layers are labeled by their channel depth, with $D_1=136$, $D_2=68$, and $D_3=34$ being associated with the first convolutions in each of the three U-Net levels. Each ConvNeXt block (blue) is replaced by the layers and operations shown in the inset labeled CNB, with generic depths $D$ and $I$ determined by the channel depth of the input and the labeled value of $D_n$. The purple blocks labeled GRU denote convolutional Gated Recurrent Unit layers, which are augmented with $1\times 1$ spatial convolutions. Other layers evaluated by the encoder are shown as dark green, while those evaluated by the decoder are shown as light green.
• Figure 3: Two time-level input-output scheme with GRU for training and inference assuming 6h time resolution. The output from the preliminary initialization step (in orange) is discarded, but the hidden state $h_1$ is generated and used in the first model step. The hidden state $h_3$ (in orange) at the end of the 24h forecast is discarded as the GRU will be re-initialized for the next recursive inference step (lowest row). For training (top right), the loss function is computed from the four forecast times spanning a 24h period at 6h resolution, as indicated in red.
• Figure 4: Comparison of the performance of the DLWP-HPX, Weyn et al. (2021), ECMWF IFS S2S, and GraphCast models. GraphCast is averaged over 104 forecasts for 2018, while other forecasts are averaged over 204 forecasts from 2017 through 2018. RMSE for (a) $Z_{500}$, (b) $T_{2m}$, and (c) $T_{850}$; climatology is indicated by the gray dashed line. ACC for or (d) $Z_{500}$, (e) $T_{2m}$ and (f) $T_{850}$.
• Figure 5: Impact of successive model improvements on the accuracy of $Z_{500}$ RMSE. Each successive change builds on top of the previous architecture, adding the modification indicated in the legend: (a) RMSE, (b) ACC. Inset in (a) provides a magnified view of the error growth between 5 and 6 forecast days.