openclimatefix / graph_weather Goto Github PK

PyTorch implementation of Ryan Keisler's 2022 "Forecasting Global Weather with Graph Neural Networks" paper (https://arxiv.org/abs/2202.07575)

License: MIT License

Python 99.65% Dockerfile 0.35%

pytorch weather graph-neural-networks forecasting-models

graph_weather's Introduction

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!

graph_weather's People

Contributors

Stargazers

Watchers

graph_weather's Issues

Add support for multiple outputs/inputs

Detailed Description

The code as it is currently does take multiple graphs at once, and can outputs multiple graphs, along the batch dimension though. Having multiple inputs, such as for the NWP forecast grid + raw observations to generate the next analysis grid to more closely match how analysis files are actually generated could be helpful. Additionally, multiple outputs would be helpful, for example for #20 to output both the next NWP state and satellite image, potentially at different resolutions (i.e. HRV vs non-HRV channels for EUMETSAT), or for having more dense local outputs, it could be the global graph, and a local graph output in the decoder.

Context

Possible Implementation

Update the latent graph processor to take multiple encoders. Multple outputs can be achieved by taking the output from the graph processor and giving it as inputs into multiple decoders. There would just need to be a way to give multiple decoder configs to the overall model.

Error after Installation

Hey,
I just installed your package and all requirements, tried to run your first example usage and got the following error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/Users/me/Documents/gnn_code/graph_weather.py in <cell line: 15>()
     [12](file:///Users/me/Documents/gnn_code/graph_weather.py?line=11) model = GraphWeatherForecaster(lat_lons)
     [14](file:///Users/me/Documents/gnn_code/graph_weather.py?line=13) features = torch.randn((2, len(lat_lons), 78))
---> [16](file:///Users/me/Documents/gnn_code/graph_weather.py?line=15) out = model(features)
     [17](file:///Users/me/Documents/gnn_code/graph_weather.py?line=16) criterion = NormalizedMSELoss(lat_lons=lat_lons, feature_variance=torch.randn((78,)))
     [18](file:///Users/me/Documents/gnn_code/graph_weather.py?line=17) loss = criterion(out, features)

File ~/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   [1126](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1125) # If we don't have any hooks, we want to skip the rest of the logic in
   [1127](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1126) # this function, and just call forward.
   [1128](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1127) if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   [1129](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1128)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1130](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1129)     return forward_call(*input, **kwargs)
   [1131](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1130) # Do not call functions when jit is used
   [1132](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1131) full_backward_hooks, non_full_backward_hooks = [], []

File ~/Documents/gnn_code/graph_weather/models/forecast.py:97, in GraphWeatherForecaster.forward(self, features)
     [87](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=86) def forward(self, features: torch.Tensor) -> torch.Tensor:
     [88](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=87)     """
     [89](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=88)     Compute the new state of the forecast
     [90](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=89) 
   (...)
     [95](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=94)         The next state in the forecast
     [96](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=95)     """
---> [97](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=96)     x, edge_idx, edge_attr = self.encoder(features)
     [98](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=97)     x = self.processor(x, edge_idx, edge_attr)
     [99](file:///Users/me/Documents/gnn_code/graph_weather/models/forecast.py?line=98)     x = self.decoder(x, features.shape[0])

File ~/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   [1126](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1125) # If we don't have any hooks, we want to skip the rest of the logic in
   [1127](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1126) # this function, and just call forward.
   [1128](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1127) if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   [1129](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1128)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1130](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1129)     return forward_call(*input, **kwargs)
   [1131](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1130) # Do not call functions when jit is used
   [1132](file:///Users/me/Documents/gnn_code/graphenv/lib/python3.9/site-packages/torch/nn/modules/module.py?line=1131) full_backward_hooks, non_full_backward_hooks = [], []

File ~/Documents/gnn_code/graph_weather/models/layers/encoder.py:158, in Encoder.forward(self, features)
    [156](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=155) self.graph = self.graph.to(features.device)
    [157](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=156) self.latent_graph = self.latent_graph.to(features.device)
--> [158](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=157) features = torch.cat(
    [159](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=158)     [features, einops.repeat(self.h3_nodes, "n f -> b n f", b=batch_size)], dim=1
    [160](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=159) )
    [161](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=160) # Cat with the h3 nodes to have correct amount of nodes, and in right order
    [162](file:///Users/me/Documents/gnn_code/graph_weather/models/layers/encoder.py?line=161) features = einops.rearrange(features, "b n f -> (b n) f")

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 78 but got size 102 for tensor number 1 in the list.

Is this a mistake on my side?

Best,
Moritz

[Paper] Generating physically-consistent high-resolution climate data with hard-constrained neural networks

https://www.climatechange.ai/papers/aaaifss2022/5

Interesting work on making sure the outputs of high-resolution models are physically consistent. Might be good to include some of their work in this

How to get landsea.zarr file?

Thanks for the code.

I have tried to execute train/run.py, but the landsea.zarr is missed? How I get this file?

self.landsea = xr.open_zarr("/home/bieker/Downloads/landsea.zarr", consolidated=True).load()
(If it's because I'm blind, please don't mind -.-)
Thanks.

NWP up sampling predictions paper

https://arxiv.org/abs/2203.12297

Detailed Description

The paper shows making a Gan that can up sample global forecasts for precipitation to high resolution local forecasts, which is similar to something we want to try here, of training on various resolution graph inputs so that we can make high res forecasts over arbitrary areas.

Context

Possible Implementation

Update to take advantage of PyG 2.2.0

Detailed Description

https://github.com/pyg-team/pytorch_geometric/releases/tag/2.2.0

Context

New version has improvements for working on larger-than-memory graphs, etc. which might be helpful here, and also the speed ups that should be included.

Possible Implementation

[Question] Where are the constants coming from?

If I'm blind I'm sorry, but currently I can not find a source (in paper or Readme) regarding the constants in graph_weather/graph_weather/data/const.py , where are they coming from / calculated from exactly?

Great work btw

Keisler code?

Is the original Keisler code for available somewhere (the code that produced https://arxiv.org/abs/2202.07575 results)?
Has Keisler anyhing to do with this project? Is he willing to help?

Map NWP outputs to Satellite Imagery

Detailed Description

Courtesy @JackKelly but satellite imagery seems very useful for predicting PV output, and to go out more than a few hours, such as day ahead, etc. the actual satellite imagery isn't very useful. Instead, we could try to map from the NWP outputs to Satellite imagery.

In this repo, it could be similar to the data assimilation model, in that it takes the NWP output on one graph, and then maps that to the new graph that is the native satellite resolution over areas of interest, and maps down to the 12 satellite channels for EUMETSAT for example.

Context

There has been other work, I think detailed in other issues, that maps from NWPs to sub-grid outputs, like precipitation rate, so it seems like it should be possible.

Possible Implementation

latent graph structure - Encoder.forward()

@jacobbieker - Hi, I've been looking at the GNN encoder / decoder code to try to understand the overall logic and structure (I'm new to pytorch-geometric). I'm not quite able to unpack the following snippet https://github.com/openclimatefix/graph_weather/blob/main/graph_weather/models/layers/encoder.py#L179:

out = einops.rearrange(out, "b n f -> (b n) f")
return (
            out,
            torch.cat(
                [self.latent_graph.edge_index + i * torch.max(self.latent_graph.edge_index) + i for i in range(batch_size)],
                dim=1,
            ),
            self.latent_edge_encoder(einops.repeat(self.latent_graph.edge_attr, "e f -> (repeat e) f", repeat=batch_size)),
        )  # New graph

Judging by the comment, this builds the structure of the "new" latent graph (with updated node and edge info). Why are the edge indexes and edge attribute tensors concatenated / repeated this way, though? Thanks!

Add TPU support

Detailed Description

Currently, the graph neural network library dependencies don't support TPUs with pytorch geometric, or don't seem to at least because of custom kernels. We could add a Jax version for TPU support? The original model was implemented in Jax apparently.

Context

Being able to use TPUs could speed up training quite a bit.

Possible Implementation

deepspeed_graph.py run error

Thank you for your sharing code. I try to run the deepspeed_graph.py, but have some errors. I made a little modification to the code and it is as follows:

import lightning.pytorch as pl
import torch
from lightning.pytorch import Trainer

from graph_weather import GraphWeatherForecaster
from lightning.pytorch.strategies import DeepSpeedStrategy
from deepspeed.ops.adam import DeepSpeedCPUAdam

lat_lons = []
for lat in range(-90, 90, 1):
    for lon in range(0, 360, 1):
        lat_lons.append((lat, lon))


class LitModel(pl.LightningModule):
    def __init__(self, lat_lons, feature_dim, aux_dim):
        super().__init__()
        self.model = GraphWeatherForecaster(
            lat_lons=lat_lons, feature_dim=feature_dim, aux_dim=aux_dim
        )

    def training_step(self, batch):
        x, y = batch
        x = x.half()
        y = y.half()
             
        out = self.model(x)

        criterion = torch.nn.MSELoss()
        loss = criterion(out, y)
        return loss

    def configure_optimizers(self):

        optimizer = DeepSpeedCPUAdam(self.parameters())
        # optimizer = torch.optim.AdamW(self.parameters())

        return optimizer

    # def forward(self, x):
    #     return self.model(x)


# Fake data
from torch.utils.data import DataLoader, Dataset


class FakeDataset(Dataset):
    def __init__(self):
        super(FakeDataset, self).__init__()

    def __len__(self):
        return 64000

    def __getitem__(self, item):
        return torch.randn((64800, 78 + 0)), torch.randn((64800, 78))


model = LitModel(lat_lons=lat_lons, feature_dim=78, aux_dim=0)
trainer = Trainer(
    accelerator="gpu",
    devices=1,
    strategy="deepspeed_stage_2_offload", 
    precision=16,
    max_epochs=10,
    limit_train_batches=2000,
)
dataset = FakeDataset()
train_dataloader = DataLoader(
    dataset, batch_size=1, num_workers=1, pin_memory=True, prefetch_factor=1
)
trainer.fit(model=model, train_dataloaders=train_dataloader)

The error is :

  | Name  | Type                   | Params
-------------------------------------------------
0 | model | GraphWeatherForecaster | 7.6 M 
-------------------------------------------------
7.6 M     Trainable params
0         Non-trainable params
7.6 M     Total params
30.342    Total estimated model params size (MB)
/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:430: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 40 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Epoch 0:   0%|                                                                                                                    | 0/2000 [00:00<?, ?it/s]Traceback (most recent call last):
  File "/media/lk/lksgcc/lk_git/21_RenewablePower/WeatherForecast/graph_weather/train/deepspeed_graph.py", line 73, in <module>
    trainer.fit(model=model, train_dataloaders=train_dataloader)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 520, in fit
    call._call_and_handle_interrupt(
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 92, in launch
    return function(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 559, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 935, in _run
    results = self._run_stage()
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 978, in _run_stage
    self.fit_loop.run()
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 201, in run
    self.advance()
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 354, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 133, in run
    self.advance(data_fetcher)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 218, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 185, in run
    self._optimizer_step(kwargs.get("batch_idx", 0), closure)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 261, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 142, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/core/module.py", line 1266, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/core/optimizer.py", line 158, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/strategies/ddp.py", line 257, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/strategies/strategy.py", line 224, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/plugins/precision/deepspeed.py", line 92, in optimizer_step
    closure_result = closure()
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 140, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 126, in closure
    step_output = self._step_fn()
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 308, in _training_step
    training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 288, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/strategies/ddp.py", line 329, in training_step
    return self.model(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1846, in forward
    loss = self.module(*inputs, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/lightning/pytorch/overrides/base.py", line 90, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "/media/lk/lksgcc/lk_git/21_RenewablePower/WeatherForecast/graph_weather/train/deepspeed_graph.py", line 28, in training_step
    out = self.model(x)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/graph_weather/models/forecast.py", line 109, in forward
    x, edge_idx, edge_attr = self.encoder(features)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/graph_weather/models/layers/encoder.py", line 172, in forward
    edge_attr = self.edge_encoder(self.graph.edge_attr)  # Update attributes based on distance
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/graph_weather/models/layers/graph_net_block.py", line 75, in forward
    out = self.model(x)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/container.py", line 204, in forward
    input = module(input)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lk/anaconda3/envs/rlai/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 115, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 must have the same dtype

Following the running process, I found the linear input become torch.float32 dtype, however, the weights and bias are all torch.float16. So how could I fix this problem? Thanks.

[Paper] ClimatX

https://arxiv.org/abs/2301.10343

Detailed Description

Transformer based foundation model for climate forecasting, only very very lightly skimmed it for now.

Context

Possible Implementation

[Paper] FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators

Another interesting-looking paper!

https://arxiv.org/abs/2202.11214

To quote the abstract:

FourCastNet, short for Fourier Forecasting Neural Network, is a global data-driven weather forecasting model that provides accurate short to medium-range global predictions at 0.25∘ resolution. FourCastNet accurately forecasts high-resolution, fast-timescale variables such as the surface wind speed, precipitation, and atmospheric water vapor. It has important implications for planning wind energy resources, predicting extreme weather events such as tropical cyclones, extra-tropical cyclones, and atmospheric rivers. FourCastNet matches the forecasting accuracy of the ECMWF Integrated Forecasting System (IFS), a state-of-the-art Numerical Weather Prediction (NWP) model, at short lead times for large-scale variables, while outperforming IFS for variables with complex fine-scale structure, including precipitation. FourCastNet generates a week-long forecast in less than 2 seconds, orders of magnitude faster than IFS. The speed of FourCastNet enables the creation of rapid and inexpensive large-ensemble forecasts with thousands of ensemble-members for improving probabilistic forecasting. We discuss how data-driven deep learning models such as FourCastNet are a valuable addition to the meteorology toolkit to aid and augment NWP models.

[Paper & Dataset] Verification against in-situ observations forData-Driven Weather Prediction

Arxiv/Blog/Paper Link

https://arxiv.org/pdf/2305.00048.pdf

Detailed Description

Dataset is available through Huggingface here: https://huggingface.co/datasets/excarta/madis2020
This paper goes more into verifying how well data driven model perform, specifically against real observations as well, compared to the common benchmark of ERA5 reanalysis. The data in the dataset is from MADIS.

Context

It would be good to compare what these models can do against observations.

Missing dependencies in setup.py ?

Dear Developers,
I installed graph_weather from source (from the main branch) using
pip install -e . in the main directory (directory graph_weather, where the setup.py lives).

I wanted to do some training, but before that I needed to install additional dependencies:

pip install pytorch-lightning
pip install pysolar
pip install torch-geometric
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.12.0+cu102.html

In the last one check correct pytorch + cuda/cpu combination from https://pypi.org/project/torch-sparse/
(I had pytorch 1.12 and cu102 (I had gpu with CUDA 10.2) )

With these additions I got python3 graph_weather/train/pl_graph_weather.py going.

Terveisin, Markus

Implement Model

Detailed Description

Implement the basic MeshGraphNet model of an encoder-processor-decoder

Context

The point of this repo is that model.

Possible Implementation

how to run the code

Can someone please tell me how to run the code? Because, When I tried to run the program after installing necessary dependencies in colab says out of memories and crashes.

[Paper] Earthformer

Not exactly graph neural network, but does focus on forecasting NWP and climate data

Detailed Description

https://arxiv.org/pdf/2207.05833.pdf

Context

Possible Implementation

Sin/Cos of Lat/Lon in degrees?

Hi,
I stumbled upon this line:

graph_weather/train/pl_graph_weather.py

Line 178 in 6524ec4

sin_lat_lons = np.sin(lat_lons)

If I am not mistaken, the lat & lon given as input to the sine & cose functions are in degrees.
Should it then not be?

sin_lat_lons = np.sin(lat_lons * np.pi / 180.)
cos_lat_lons = np.cos(lat_lons * np.pi / 180.)

Cheers!:)

Paper: Graph Transformer

https://andrew.cmu.edu/~yangli1/Forecaster_ECAI2020.pdf

Detailed Description

Only skimmed it but a transformer graph model, but for more general forecasting.

Context

Possible Implementation

Model performance plots

Hi @jacobbieker ,

Happy new year!:)
I was wondering if you ever measured the performance of your models with this code. Like is it similar to Keisler etc?
I saw there are some pretrained weights on Huggingface, but I am a bit puzzled how to use them (otherwise would just create the plots myself). Is there some tutorial or similar? Thought I'd ask you directly before trying to reverse-engineer what you did.

Thanks in advance,
Vitus

Add support for training on Google's ARCO ERA5

Detailed Description

The Google ERA5 archive is available: https://github.com/google-research/ARCO-ERA5 and goes back quite far. Its currently being expanded, but already covers years of data at the native model resolution. This is probably quicker and faster to pull data from than HuggingFace, and is already on Google Public Datasets, so should be easy to integrate.

Context

Its a lot of data in an accessible and publicly available place that we can try to leverage for training the graph models easily.

Possible Implementation

A HuggingFace dataset script like for the EUMETSAT data? Or can just add Xarray opening it in this repo directly as well.

Forecaster doesn't work with H3 resolutions > 2

Describe the bug
The Encoder's scatter_sum results in the input graph not matching the new graph, and being off by a little bite in the first dimension. This results in the model failing at the encoder forward pass

To Reproduce
Steps to reproduce the behavior:

Go to '...'
Click on '....'
Scroll down to '....'
See error

Expected behavior
A clear and concise description of what you expected to happen.

Additional context
Add any other context about the problem here.

Paper: Transformer UNet for data assimilation and forecasting

https://gmd.copernicus.org/articles/15/2221/2022/gmd-15-2221-2022.pdf

More data is better, and keeping data in native resolution is probably better.

Possible Implementation

Add another graph with the invariant features, and then a second input to the forward pass that can resuse the invariant features in multiple forward passes during rollout.

The attribute 'dayofyear' in line 428 of train/run.py

Original Traceback (most recent call last):
File "/home/data/yh/miniconda/envs/nudt/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
data = fetcher.fetch(index)
File "/home/data/yh/miniconda/envs/nudt/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 32, in fetch
data.append(next(self.dataset_iter))
File "/home/data/yh/nudt/graph_weather-main/train/run.py", line 428, in iter
day_of_year = data["timestamps"][0].dayofyear / 366.0
AttributeError: 'datetime.datetime' object has no attribute 'dayofyear'

I think the attribute used here should be 'day' instead.

NREL's Sup3r

Arxiv/Blog/Paper Link

https://github.com/NREL/sup3r uses GANs to create high-res solar and wind data from other NWP data. Not graph related, but weather related here.

Context

It is an interesting way of using GANs for increasing the resolution of wind and solar data both temporally and spatially.

[Paper] Learning the Evolutionary and Multi-scale Graph Structure for Multivariate Time Series Forecasting

Arxiv/Blog/Paper Link

https://arxiv.org/pdf/2206.13816.pdf

Detailed Description

Uses GNN for time series forecasting, including of solar power data. Not directly weather, but putting here as GNN based.

Context

Could be good benchmark/interesting to try out. @thomasarmstrong98 found it!

Seems like might be interesting for the new national model, or benchmarking against GSP forecasts? @JackKelly @dfulu

model(features) -> RuntimeError: expected scalar type Float but found Half

I just changed aux_dim as 0
But I got a error on out = model(features)

Error msg:
RuntimeError: expected scalar type Float but found Half

So I tried to change the tensor dtype on torch.randn(... dtype=torch.float64 )
But it is same..

This is my code:
_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, aux_dim=0)

features = torch.randn((2, len(lat_lons), 78), dtype=torch.float64)

out = model(features)
criterion = NormalizedMSELoss(lat_lons=lat_lons, feature_variance=torch.randn((78,)))
loss = criterion(out, features)
loss.backward()_

Forecasting API call

I'm currently running a residential solar system and using online forecasting predictions to set battery charge levels as to not over charge during off peak times.

My current provider has been pretty hit and miss (by a large percentage, today it's predicted less than 1kw at 1am this morning, but we ended up with 18.4kw)

You can review my public GitHub project for more details on what my application does. But I track it's accuracy over time

Detailed Description

Context

I wish to reduce my co2 footprint by reducing demand on my electrical provider by using more accurate predictions to set minimum charge on my battery to reach the point at which the solar panels can take over the runner or the house hold

Possible Implementation

A secure call to a exposed public API that gives the caller the total expected solar generation based on pv size, efficiency, orientation, geo location and pitch.

Model(features) Error

I succeeded to install graph-weather by pip

But I got an error message for model(features)

Message:
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 78 but got size 102 for tensor number 1 in the list.

I checked feature shape -> torch.Size([2, 64800, 78])

I don't know why the input feature got 102 tensor.. anyone who help me?

Add adaptive meshing

Detailed Description

Being able to get more detailed output around areas of interest is quite important, such as around the UK or some other country of interest.

Context

Possible Implementation

Uses the adaptive mesh from the original DeepMind paper, and train on 0.25 degree GFS data downsampled in some areas to 1 degree. Or something like that

Add HuggingFace integration

Detailed Description

Like our other models MetNet, and DGMR, have support for saving/loading model and data from HuggingFace

Context

Makes it a lot easier for people to use the models

Possible Implementation

Hydrostatic balance

Your forecasting system used 13 levels instead <=5 in others similar systems

The hydrostatic balance equation can greatly improve your system.

You could add to the loss function the hydrostatic balance discrepancy in each of the layers:
( (Tv_i+Tv_{i+1})/2 + g/R * (H_{i+1}-H_{i}) / log(p_{i+1}/p_i) )^2 -> min
where
g = 9.8065 m/s^2 is the gravitational acceleration
R = 287.04 is gas constant for air
Tv [K] = T [K] *(1+Q [kg/kg]/(Q + R/Rv)) is virtual temperature
Rv = 461.51 is gas constant for vapor

Undersurface values

Some values on p-levels coresponds to the virtual values under Earth surface

May be we dont takes into account this values in loss function ?

era5 training to reproduce Keisler?

Dear Developers,

Is there any code for training the network with era5 data to reproduce Keisler paper?

There seem to be scripts in graph_weather/train but all for gfs data?

Terveisin, Markus

Perform Data Assimilation with model

Detailed Description

One of the most compute intensive steps in NWPs is doing the data assimilation to create the starting conditions for running the simulation. This type of model could theoretically speed that up if it can integrate lots of observations and essentially interpolate them to the regular grid for NWP initialization.

Context

Possible Implementation

NOAA makes the raw observations public as well as the analysis files which are what the NWP's are started from.

Running pl_graph_weather.py in graph_weather/train causes "Killed"

Dear Developers,
Wanted to test how training goes 'out of the box' with graph_weather.
First problem was that I ran out of disk space, but thats ok now.

Other problem (not fixed) is that I get 'Killed' and program stops (see below).
Probably one has run out of memory.
Is there any way to fix this? Typically one sets smaller batch size, but I think it is already smallelst possible (1?).

My machine is
nvidia tesla v100
product: GV100GL [Tesla V100 SXM2 16GB]

Terveisin, Markus

ubuntu@gpumok:~/git/graph_weather/train$ python3 pl_graph_weather.py 
/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 0.1.43ubuntu1 is an invalid version and will not be supported in a future release
  warnings.warn(
/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release
  warnings.warn(
Using custom data configuration openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744
Reusing dataset parquet (/data10t/huggingface/datasets/openclimatefix___parquet/openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)
Using custom data configuration openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744
Reusing dataset parquet (/data10t/huggingface/datasets/openclimatefix___parquet/openclimatefix--gfs-surface-pressure-2deg-ca6875a1d925d744/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)
#1: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:22<00:00, 282.63s/ex]
#7: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:23<00:00, 282.65s/ex]
#2: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:29<00:00, 282.77s/ex]
#9: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:35<00:00, 282.89s/ex]
#0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:42<00:00, 283.02s/ex]
#4: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:50<00:00, 283.15s/ex]
#14: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:14:59<00:00, 283.33s/ex]
#11: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:01<00:00, 283.36s/ex]
#8: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:08<00:00, 283.49s/ex]
#3: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:09<00:00, 283.51s/ex]
#10: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:10<00:00, 283.54s/ex]
#5: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:15<00:00, 283.61s/ex]
#12: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:23<00:00, 283.76s/ex]
#6: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:24<00:00, 283.78s/ex]
#13: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 283.92s/ex]
#15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 283.92s/ex]
  0%|                                                                                                                                                            | 0/1 [00:00<?, ?ba/s]Killed
Killed
ubuntu@gpumok:~/git/graph_weather/train$ ███████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:01<00:00, 279.11s/ex]
#6: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:24<00:00, 272.73s/ex]
#10: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:10<00:00, 279.13s/ex]
#12: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:23<00:00, 272.14s/ex]
#13: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 264.74s/ex]
#15: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 54/54 [4:15:31<00:00, 270.98s/ex]

Create Ensembles

Detailed Description

If the model is fast enough at inference, it could be quite easy and cheap to generate a lot of forecasts at once. The original implementation took 0.8 seconds for a 5 day forecast, so a 100 member ensemble could be created in 80 seconds, which is much larger than most of the ensembles generated by GFS at least.

Context

Ensembles can be quite important for helpiung with probablistic forecasts.

Possible Implementation

One way would be the same as the skillful nowcasting model, where we have a latent vector that adds some noise to the inputs to generate ensembles. But even easier, and more similar to real NWPs, is to just perturb the initial state and run again, i.e. change the 2m temperature by a bit, or something.

[Paper] Positional Encoder Graph Neural Networks for Geographic Data

This paper might be relevant: https://arxiv.org/abs/2111.10144

To quote the abstract:

Graph neural networks (GNNs) provide a powerful and scalable solution for modeling continuous spatial data. However, in the absence of further context on the geometric structure of the data, they often rely on Euclidean distances to construct the input graphs. This assumption can be improbable in many real-world settings, where the spatial structure is more complex and explicitly non-Euclidean (e.g., road networks). In this paper, we propose PE-GNN, a new framework that incorporates spatial context and correlation explicitly into the models. Building on recent advances in geospatial auxiliary task learning and semantic spatial embeddings, our proposed method (1) learns a context-aware vector encoding of the geographic coordinates and (2) predicts spatial autocorrelation in the data in parallel with the main task. On spatial regression tasks, we show the effectiveness of our approach, improving performance over different state-of-the-art GNN approaches. We also test our approach for spatial interpolation, i.e., spatial regression without node features, a task that GNNs are currently not competitive at. We observe that our approach not only vastly improves over the GNN baselines, but can match Gaussian processes, the most commonly utilized method for spatial interpolation problems. The code for this study can be accessed via: this https URL

Memory optimization in graph data

Not a bug, but seemed like the best issue template to use :)

Passing x is unnecessary here as it is never accessed. The nodes tensor does not need to be created. Could save a bit of memory here that's allocated to it.

graph_weather/graph_weather/models/layers/encoder.py

Line 107 in ef13589

self.graph = Data(x=nodes, edge_index=edge_index, edge_attr=self.h3_distances)

Here also.

graph_weather/graph_weather/models/layers/assimilator_decoder.py

Line 104 in ef13589

 self.graph = Data(x=nodes, edge_index=edge_index, edge_attr=self.h3_to_lat_distances) 

And here.

graph_weather/graph_weather/models/layers/assimilator_encoder.py

Line 216 in ef13589

return Data(x=nodes, edge_index=edge_index, edge_attr=h3_distances)

Fix NormalizedMSELoss and tests

Describe the bug
No current tests for just NormalizedMSELoss, and a test fails because loss got deleted I guess

To Reproduce
Steps to reproduce the behavior:

Go to '...'
Click on '....'
Scroll down to '....'
See error

Expected behavior
A clear and concise description of what you expected to happen.

Additional context
Add any other context about the problem here.

[Paper] GraphCast: Learning skillful medium-range global weather forecasting

Detailed Description

https://arxiv.org/abs/2212.12794

openclimatefix / graph_weather Goto Github PK

graph_weather's Introduction

Graph Weather

Installation

Example Usage

Pretrained Weights

Training Data

Contributors ✨

graph_weather's People

Contributors

Stargazers

Watchers

Forkers

graph_weather's Issues

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Arxiv/Blog/Paper Link

Detailed Description

Context

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Arxiv/Blog/Paper Link

Context

Arxiv/Blog/Paper Link

Detailed Description

Context

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description

Context

Possible Implementation

Detailed Description