SUPR
Global weather forecasting with equivariant neural networks
Dnr:

NAISS 2023/5-452

Type:

NAISS Medium Compute

Principal Investigator:

Hampus Linander

Affiliation:

Chalmers tekniska högskola

Start Date:

2023-11-30

End Date:

2024-06-01

Primary Classification:

10799: Other Natural Sciences not elsewhere specified

Allocation

Abstract

Accurate global weather forecasting is important for many levels of society. Recently there has been an increase in the availability of high quality global climate data that enables a data-driven approach to predictive models using machine learning. This data has been used to train large neural networks that outperform [1a,1b,1c] classical simulation methods for medium range forecasting. This project will apply methods from geometric deep learning to construct and train models that use the geometry and symmetries of the data to achieve more efficient models. Our group have previously [2-5] developed neural networks that operate on spherical data that we have shown to be more efficient in terms of training data, training time and efficacy. The global weather forecasting problem consists of predicting different quantities such as the temperature and wind velocity on the surface and in the atmosphere of the earth at some time in the future. The input and output data are thus maps of values on the sphere, which are perfectly suited for our models. In contrast to other recent work, we will leverage the symmetries available on the sphere, for example predicting a certain weather system at a given longitude should generalize to predicting the same weather system if present on another longitude. We will use the ERA5 dataset [7], a reanalysis of the global weather at a resolution of 0.25 degrees with hourly samples from 1940 to present time, produced by the European Centre for Medium-Range Weather Forecasts. To train and evaluate our models we will initially restrict the dataset to a subset of an 11-year span that has been shown [6] to generate models with good performance at a fraction of the compute and storage required for the full dataset. [1a] "Accurate medium-range global weather forecasting with 3D neural networks" , Kaifeng Bi et al. https://www.nature.com/articles/s41586-023-06185-3 [1b] Kurth, Thorsten et al. “FourCastNet: Accelerating Global High-Resolution Weather Forecasting Using Adaptive Fourier Neural Operators.” Proceedings of the Platform for Advanced Scientific Computing Conference (2022) [1c] Nguyen, Tung et al. “ClimaX: A foundation model for weather and climate.” International Conference on Machine Learning (2023). [2] Gerken, Jan E., et al. “Geometric Deep Learning and Equivariant Neural Networks.” Artificial Intelligence Review, vol. 56, no. 12, June 2023, pp. 14605–62. Crossref, https://doi.org/10.1007/s10462-023-10502-7. [3] Aronsson, Jimmy. “Homogeneous Vector Bundles and G-Equivariant Convolutional Neural Networks.” Sampling Theory, Signal Processing, and Data Analysis, vol. 20, no. 2, July 2022. Crossref, https://doi.org/10.1007/s43670-022-00029-3. [4] Gerken, Jan, et al. "Equivariance versus augmentation for spherical images." International Conference on Machine Learning. PMLR, 2022. [5] Carlsson, Oscar, et al. "HEAL-SWIN: A Vision Transformer On The Sphere." arXiv preprint arXiv:2307.07313 (2023). [6] https://github.com/198808xc/Pangu-Weather [7] Hersbach, H, Bell, B, Berrisford, P, et al. The ERA5 global reanalysis. Q J R Meteorol Soc. 2020; 146: 1999–2049.