From bb368b65cdb81ed72ebd06d845dbf1c00b5f5a46 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Fri, 10 May 2024 11:34:50 +0200 Subject: [PATCH] fix build path and update readme (#11) * fix build path * add badge and remove docs from README --- .github/workflows/deploy-docs-book.yml | 3 +- README.md | 173 +------------------------ 2 files changed, 6 insertions(+), 170 deletions(-) diff --git a/.github/workflows/deploy-docs-book.yml b/.github/workflows/deploy-docs-book.yml index 292ec42..24776ad 100644 --- a/.github/workflows/deploy-docs-book.yml +++ b/.github/workflows/deploy-docs-book.yml @@ -6,7 +6,6 @@ on: branches: - master - main - - jupyterbook-docs* # This job installs dependencies, builds the book, and pushes it to `gh-pages` jobs: @@ -38,7 +37,7 @@ jobs: - name: Upload artifact uses: actions/upload-pages-artifact@v2 with: - path: "_build/html" + path: "docs/_build/html" # Deploy the book's HTML to GitHub Pages - name: Deploy to GitHub Pages diff --git a/README.md b/README.md index d7a510c..7542076 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # weather-model-graphs -[![linting](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml/badge.svg)](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml) +[![linting](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml/badge.svg)](https://github.com/mllam/weather-model-graphs/actions/workflows/pre-commit.yml) [![Jupyter Book Badge](https://jupyterbook.org/badge.svg)](https://mllam.github.io/weather-model-graphs) `weather-model-graphs` is a package for creating, visualising and storing message-passing graphs for data-driven weather models. @@ -47,176 +47,13 @@ xy_grid = np.meshgrid(np.linspace(0, 1, 32), np.linspace(0, 1, 32), indexing='ij graph = wmg.create.archetype.create_keisler_graph(xy=xy_grid) # split the graph by component -graph_components = wmg.split_graph_by_edge_attribute(graph=graph, attribute='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) + wmg.save.to_pyg(graph=graph, name=component, output_directory=".") ``` -# Background and design +# Documentation -The only input the graph generation in `weather-model-graphs` requires is the static `(x,y)` *grid* coordinates of the atmospheric state as the state changes over time. These coordinates are used to create the **grid nodes** nodes of the graph, with a node for each `(x,y)` coordinate. -In addition to grid nodes the graph also contains **mesh nodes** that represent the latent space of the model at a set of `(x,y)` coordinates (this is in general a different set of coordinates to the **grid nodes** coordinates). - -In summary, there are two sets of nodes in the graph: - -- **grid nodes**: representing the physical variables of the atmospheric state at a specific `(x,y)` coordinate - -- **mesh nodes**: representing the latent space of the model at specific `(x,y)` coordinate - -With these two sets of nodes, the graph is constructed by connecting the **grid nodes** to the **mesh nodes** and the **mesh nodes** to each other. -The edges between the **grid nodes** and the **mesh nodes** represent the encoding of the physical variables into the latent space of the model, while the edges between the **mesh nodes** represent the processing of the latent space through the time evolution of the atmospheric state. - -In summary, the complete message-passing graph consists of three components: - -- **grid-to-mesh** (`g2m`): the encoding compenent, where edges represent the encoding of physical variables into the latent space of the model - -- **mesh-to-mesh** (`m2m`): the processing component, where edges represent information flow between nodes updating the latent presentation at mesh nodes through the time evolution of the atmospheric state - -- **mesh-to-grid** (`m2g`): the decoding component, where edges represent the decoding of the latent space back into physical variables - -Practically, the **mesh-to-grid** and **grid-to-mesh** updates can probably also encode some of the time evolution processing, in addition to the latent space encoding/decoding, unless the GNN is trained specifically as an auto-encoder using the same graph as input and output. - -## Design principles - -The key design principle of `weather-model-graphs` is to work with `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. -Using only `networkx.DiGraph` objects as the intermediate representations makes it possible to - -1) easily modularise the whole generation process, with every step outputting a `networkx.DiGraph` object, -2) easily visualise the graph resulting from any step and -3) easily connect graph nodes across graph components, combine graphs and split graphs based on node and edge attributes. - -The graph generation in `weather-model-graphs` is split into to the following steps: - -1. Create the three graph components of the message-passing graph that constitute the auto-regressive atmospheric flow model, all represented by `networkx.DiGraph` objects: - - - **grid-to-mesh** (`g2m`): the encoding compenent, where edges represent the encoding of physical variables into the latent space of the model - - - **mesh-to-mesh** (`m2m`): the processing component, where edges represent information flow between nodes through the time evolution of the atmospheric state - - - **mesh-to-grid** (`m2g`): the decoding component, where edges represent the decoding of the latent space back into physical variables - -2. Combine all three graph components into a single `networkx.DiGraph` object and create a unique node identifier for each node in the combined graph. -3. Split the combined graph into the three *output* graph components again (or more if the specific graph architecture requires it). -4. Store each of the *output* graph components in the desired format, for example: - - - **networkx** `.pickle` file: save `networkx.DiGraph` objects using `pickle` to disk (`weather_model_graphs.save.to_pickle(...)`) - - - [pytorch-geometric](https://github.com/pyg-team/pytorch_geometric) for [neural-lam](https://github.com/mllam/neural-lam): edges indexes and features are stored in separate `torch.Tensor` objects serialised to disk that can then be loaded into `torch_geometric.data.Data` objects (`weather_model_graphs.save.to_pyg(...)`) - -### Diagram of the graph generation process: - -Below, the graph generation process is visualised in `weather-model-graphs` for the example given above: - - -```mermaid -graph TB - -xy_grid["(x,y) grid coordinates [numpy.ndarray]"] - -subgraph weather_model_graphs["weather-model-graphs"] - subgraph graph_components["Graph components"] - G_g2m["G_g2m[networkx.DiGraph]"] - G_m2m["G_m2m[networkx.DiGraph]"] - G_m2g["G_m2g[networkx.DiGraph]"] - end - - G_g2m --merged into --> G_full - G_m2m --merged into --> G_full - G_m2g --merged into --> G_full - - G_full["complete graph\nG_full[networkx.Digraph]"] - - G_full --replace node labels with unique integer id --> G_full_int["G_int[networkx.Digraph]"] - - G_full_int --split and converted into --> pyg_g2m["pyg_g2m[pyg.data.Data]"] - G_full_int --split and converted into --> pyg_m2m["pyg_m2m[pyg.data.Data]"] - G_full_int --split and converted into --> pyg_m2g["pyg_m2g[pyg.data.Data]"] -end - -xy_grid --used to create node coordinates in --> G_g2m -xy_grid --used to create node coordinates in --> G_m2g -xy_grid --used to create node coordinates in --> G_m2m - -subgraph stored_on_disk["stored on disk"] - subgraph hidden_graph1[ ] - pyg_g2m_file["g2m_edge_index.pt\ng2m_features.pt"] - pyg_m2m_file["m2m_edge_index.pt\nm2m_features.pt"] - pyg_m2g_file["m2g_edge_index.pt\nm2g_features.pt"] - end -end - -pyg_g2m --stored in--> pyg_g2m_file -pyg_m2m --stored in--> pyg_m2m_file -pyg_m2g --stored in--> pyg_m2g_file - -subgraph pyg_loaded["Loaded into model"] - subgraph hidden_graph2["(e.g. torch.nn.Module)"] - pyg_g2m_loaded["pyg_g2m[pyg.data.Data]"] - pyg_m2m_loaded["pyg_m2m[pyg.data.Data]"] - pyg_m2g_loaded["pyg_m2g[pyg.data.Data]"] - end -end - -pyg_g2m_file --loaded into --> pyg_g2m_loaded -pyg_m2m_file --loaded into --> pyg_m2m_loaded -pyg_m2g_file --loaded into --> pyg_m2g_loaded - -classDef title_left margin-right:250px,display:block -classDef hidden stroke-width:0px,margin-right:250px,display:block - -class hidden_graph1 hidden -class hidden_graph2 hidden -class weather_model_graphs title_left -class graph_components title_left -class stored_on_disk title_left -class pyg_loaded title_left -``` - -### Node and edge attributes - -There are a number of node and edge attributes with special meanings in `weather-model-graphs` which enable the splitting and visualisation of the graph components. - -#### Node attributes - -- `pos`: the `(x,y)` coordinates of the node in the grid -- `type`: the type of node, either `grid` or `mesh` - - -#### Edge attributes - -- `component`: the component of the graph the edge belongs to, either `g2m`, `m2m` or `m2g` -- `level`: for multi-range mesh graphs this denotes the refinement level of mesh connection. For hierarchical graphs the different ranges of connections are split into different levels and so here `level` also denotes the level in the hierarchy that the edge belongs to. -- `len`: the length of the edge in the (x,y) coordinate space of the grid nodes, i.e. the distance between the two nodes in the grid -- `vdiff`: the vector spanning between the (x,y) coordinates of the two nodes -- `direction`: for hierarchical graphs this denotes the direction of the edge, either `up`, `down` and `same` - - -### Splitting graphs - -The splitting of the graph is done with by utilising the edge attributes, and thus it is easy to split the complete graph by either which component the edge belongs to, or by the level of the edge in the graph. This is done using the `weather_model_graphs.split_graph_by_edge_attribute(...)` function. - -## Code layout - -The code layout of `weather-model-graphs` is organised into submodules by the functionality they provide. The main submodules are: - -``` -weather_model_graphs - .create - .archetype: - for creating specific archetype graph - architectures (e.g. Keisler 2021, Lam et al 2023, - Oscarsson et al 2023) - .base - general interface for creating graph architectures - (here you define the g2m, m2m and m2g connectivity directly) - .mesh - for creating the mesh nodes and edges - .grid - for creating the grid nodes - .visualise - for plotting graphs, allowing for easy visualisation using any - edge or node attribute for colouring - .save - for saving the graph to specific formats (e.g. pytorch-geometric) +The documentation is built using [Jupyter Book](https://jupyterbook.org/intro.html) and can be found at [https://mllam.github.io/weather-model-graphs](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.