Weather Forecasting with Graph Neural Networks
Project description
Graph Weather
Implementation of the Graph Weather paper (https://arxiv.org/pdf/2202.07575.pdf) in PyTorch. Additionally, an implementation of a modified model that assimilates raw or processed observations into analysis files.
Installation
This library can be installed through
pip install graph-weather
Example Usage
The models generate the graphs internally, so the only thing that needs to be passed to the model is the node features
in the same order as the lat_lons
.
import torch
from graph_weather import GraphWeatherForecaster
from graph_weather.models.losses import NormalizedMSELoss
lat_lons = []
for lat in range(-90, 90, 1):
for lon in range(0, 360, 1):
lat_lons.append((lat, lon))
model = GraphWeatherForecaster(lat_lons)
# Generate 78 random features + 24 non-NWP features (i.e. landsea mask)
features = torch.randn((2, len(lat_lons), 102))
target = torch.randn((2, len(lat_lons), 78))
out = model(features)
criterion = NormalizedMSELoss(lat_lons=lat_lons, feature_variance=torch.randn((78,)))
loss = criterion(out, target)
loss.backward()
And for the assimilation model, which assumes each lat/lon point also has a height above ground, and each observation is a single value + the relative time. The assimlation model also assumes the desired output grid is given to it as well.
import torch
import numpy as np
from graph_weather import GraphWeatherAssimilator
from graph_weather.models.losses import NormalizedMSELoss
obs_lat_lons = []
for lat in range(-90, 90, 7):
for lon in range(0, 180, 6):
obs_lat_lons.append((lat, lon, np.random.random(1)))
for lon in 360 * np.random.random(100):
obs_lat_lons.append((lat, lon, np.random.random(1)))
output_lat_lons = []
for lat in range(-90, 90, 5):
for lon in range(0, 360, 5):
output_lat_lons.append((lat, lon))
model = GraphWeatherAssimilator(output_lat_lons=output_lat_lons, analysis_dim=24)
features = torch.randn((1, len(obs_lat_lons), 2))
lat_lon_heights = torch.tensor(obs_lat_lons)
out = model(features, lat_lon_heights)
assert not torch.isnan(out).all()
assert out.size() == (1, len(output_lat_lons), 24)
criterion = torch.nn.MSELoss()
loss = criterion(out, torch.randn((1, len(output_lat_lons), 24)))
loss.backward()
Pretrained Weights
Coming soon! We plan to train a model on GFS 0.25 degree operational forecasts, as well as MetOffice NWP forecasts. We also plan trying out adaptive meshes, and predicting future satellite imagery as well.
Training Data
Training data will be available through HuggingFace Datasets for the GFS forecasts. The initial set of data is available for GFSv16 forecasts, raw observations, and FNL Analysis files from 2016 to 2022, and for ERA5 Reanlaysis. MetOffice NWP forecasts we cannot redistribute, but can be accessed through CEDA.
Contributors ✨
Thanks goes to these wonderful people (emoji key):
Jacob Bieker 💻 |
Jack Kelly 🤔 |
byphilipp 🤔 |
Markus Kaukonen 💬 |
MoHawastaken 🐛 |
Mihai 💬 |
Vitus Benson 🐛 |
dongZheX 💬 |
sabbir2331 💬 |
Lorenzo Breschi 💻 |
gbruno16 💻 |
This project follows the all-contributors specification. Contributions of any kind welcome!
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for graph_weather-1.0.69-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 322ae928e6a6978aed08e0d1cd1937ed630844211575bc958bfea794f868366c |
|
MD5 | c9ca1382a32f5b611f75282a2633d0ba |
|
BLAKE2b-256 | dd51f39d1c0fa8b701b11bb14163d0e72906a7b696a6fb78cff980aa7d25a177 |