weather-model-graphs
is a package for creating, visualising and storing graphs used in message-passing graph-based data-driven weather models.
The package is designed to use networkx.DiGraph
objects as the primary data structure for the graph representation right until the graph is to be stored on disk into a specific format.
This makes the graph generation process modular (every step outputs a networkx.DiGraph
), easy to debug (visualise the graph at any step) and allows output to different file-formats and file-structures to be easily implemented. More details are given in the background and design section of the online documentation.
If you simply want to install and use weather-model-graphs
as-is you can install the most recent release directly from pypi with pip
python -m pip install weather-model-graphs
If you want to be able to save to pytorch-geometric data-structure used in
neural-lam then you will need to install
pytorch and pytorch-geometric too. This can be done by with the pytorch
optional extra in weather-model-graphs
:
python -m pip install weather-model-graphs[pytorch]
This will install the CPU version of pytorch by default. If you want to install
a GPU variant you should install that
first before installing
weather-model-graphs
.
The easiest way to work on developing weather-model-graphs
is to fork the main repo under your github account, clone this repo locally, install pdm, create a venv with pdm and then install weather-model-graphs
(and all development dependencies):
git clone https://github.com/<your-github-username>/weather-model-graphs
cd weather-model-graphs
pdm venv create
pdm use --venv in-project
pdm install --dev
All linting is handeled with pre-commit which you can ensure automatically executes on all commits by installing the git hook:
pdm run pre-commit install
Then branch, commit, push and create a pull-request!
cpu only:
PIP_INDEX_URL=https://download.pytorch.org/whl/cpu pdm install --group pytorch
gpu support (see https://pytorch.org/get-started/locally/#linux-pip for older versions of CUDA):
pdm install --group pytorch
The best way to understand how to use weather-model-graphs
is to look at the documentation (which are executable Jupyter notebooks!), to have look at the tests in tests/ or simply to read through the source code.
import numpy as np
import weather_model_graphs as wmg
# define your (x,y) grid coodinates
xy_grid = np.meshgrid(np.linspace(0, 1, 32), np.linspace(0, 1, 32), indexing='ij')
xy_grid = np.stack(xy_grid, axis=0)
# create the full graph
graph = wmg.create.archetype.create_keisler_graph(xy_grid=xy_grid)
# split the graph by component
graph_components = wmg.split_graph_by_edge_attribute(graph=graph, attr='component')
# save the graph components to disk in pytorch-geometric format
for component, graph in graph_components.items():
wmg.save.to_pyg(graph=graph, name=component, output_directory=".")
The documentation is built using Jupyter Book and can be found at https://mllam.github.io/weather-model-graphs. This includes background on graph-based weather models, the design principles of weather-model-graphs
and how to use it to create your own graph architectures.