Add a decoding mask option to only include subset of grid nodes in m2g

src/weather_model_graphs/create/archetype.py

@@ -1,12 +1,7 @@
 from .base import create_all_graph_components
 
 
-def create_keisler_graph(
-    coords,
-    mesh_node_distance=3,
-    coords_crs=None,
-    graph_crs=None,
-):
+def create_keisler_graph(coords, mesh_node_distance=3, **kwargs):
     """
     Create a flat LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
     This graph setup is inspired by the global graph used by Keisler (2022, https://arxiv.org/abs/2202.07575).
@@ -28,11 +23,8 @@ def create_keisler_graph(
     mesh_node_distance: float
         Distance (in x- and y-direction) between created mesh nodes,
         in coordinate system of coords
-    coords_crs: pyproj.crs.CRS or None
-        CRS of the given coordinates
-    graph_crs:
-        CRS to build graph in. If given, coords will be transformed from
-        coords_crs to graph_crs before graph construction
+    **kwargs:
+        Additional keyword arguments passed on to create_all_graph_components.
 
     Returns
     -------
@@ -51,8 +43,7 @@ def create_keisler_graph(
         m2g_connectivity_kwargs=dict(
             max_num_neighbours=4,
         ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        **kwargs,
     )
 
 
@@ -61,8 +52,7 @@ def create_graphcast_graph(
     mesh_node_distance=3,
     level_refinement_factor=3,
     max_num_levels=None,
-    coords_crs=None,
-    graph_crs=None,
+    **kwargs,
 ):
     """
     Create a multiscale LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
@@ -89,11 +79,8 @@ def create_graphcast_graph(
         NOTE: Must be an odd integer >1 to create proper multiscale graph
     max_num_levels: int
         The number of levels of longer-range connections in the mesh graph.
-    coords_crs: pyproj.crs.CRS or None
-        CRS of the given coordinates
-    graph_crs:
-        CRS to build graph in. If given, coords will be transformed from
-        coords_crs to graph_crs before graph construction
+    **kwargs:
+        Additional keyword arguments passed on to create_all_graph_components.
 
     Returns
     -------
@@ -116,8 +103,7 @@ def create_graphcast_graph(
         m2g_connectivity_kwargs=dict(
             max_num_neighbours=4,
         ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        **kwargs,
     )
 
 
@@ -126,8 +112,7 @@ def create_oskarsson_hierarchical_graph(
     mesh_node_distance=3,
     level_refinement_factor=3,
     max_num_levels=None,
-    coords_crs=None,
-    graph_crs=None,
+    **kwargs,
 ):
     """
     Create a LAM graph following Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
@@ -157,11 +142,10 @@ def create_oskarsson_hierarchical_graph(
         in coordinate system of coords
     level_refinement_factor: float
         Refinement factor between grid points and bottom level of mesh hierarchy
-    coords_crs: pyproj.crs.CRS or None
-        CRS of the given coordinates
-    graph_crs:
-        CRS to build graph in. If given, coords will be transformed from
-        coords_crs to graph_crs before graph construction
+    max_num_levels: int
+        The number of levels of longer-range connections in the mesh graph.
+    **kwargs:
+        Additional keyword arguments passed on to create_all_graph_components.
 
     Returns
     -------
@@ -184,6 +168,5 @@ def create_oskarsson_hierarchical_graph(
         m2g_connectivity_kwargs=dict(
             max_num_neighbours=4,
         ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        **kwargs,
     )

src/weather_model_graphs/create/base.py

@@ -9,6 +9,8 @@
 """
 
 
+from typing import Iterable
+
 import networkx
 import networkx as nx
 import numpy as np
@@ -39,6 +41,8 @@ def create_all_graph_components(
     g2m_connectivity_kwargs={},
     coords_crs: pyproj.crs.CRS | None = None,
     graph_crs: pyproj.crs.CRS | None = None,
+    decode_mask: Iterable | None = None,
+    return_components: bool = False,
 ):
     """
     Create all graph components used in creating the message-passing graph,
@@ -82,6 +86,11 @@ def create_all_graph_components(
     will be transformed from their original Coordinate Reference System (`coords_crs`)
     to the CRS where the graph creation should take place (`graph_crs`).
     If any one of them is None the graph creation is carried out using the original coords.
+
+    `decode_mask` should be an Iterable of booleans, masking which grid positions should be
+    decoded to (included in the m2g subgraph), i.e. which positions should be output. It should have the same length as the number of
+    grid position coordinates given in `coords`.  The mask being set to True means that corresponding
+    grid nodes should be included in g2m. If `decode_mask=None` (default), all grid nodes are included.
     """
     graph_components: dict[networkx.DiGraph] = {}
 
@@ -149,9 +158,19 @@ def create_all_graph_components(
     )
     graph_components["g2m"] = G_g2m
 
+    if decode_mask is None:
+        # decode to all grid nodes
+        decode_grid = G_grid
+    else:
+        # Select subset of grid nodes to decode to, where m2g should connect
+        filter_nodes = [
+            n for n, include in zip(G_grid.nodes, decode_mask, strict=True) if include
+        ]
+        decode_grid = G_grid.subgraph(filter_nodes)
+
     G_m2g = connect_nodes_across_graphs(
         G_source=grid_connect_graph,
-        G_target=G_grid,
+        G_target=decode_grid,
         method=m2g_connectivity,
         **m2g_connectivity_kwargs,
     )
@@ -162,6 +181,14 @@ def create_all_graph_components(
         for edge in graph.edges:
             graph.edges[edge]["component"] = name
 
+    if return_components:
+        # Give each component unique ids
+        graph_components = {
+            comp_name: replace_node_labels_with_unique_ids(subgraph)
+            for comp_name, subgraph in graph_components.items()
+        }
+        return graph_components
+
     # merge to single graph
     G_tot = networkx.compose_all(graph_components.values())
     # only keep graph attributes that are the same for all components

src/weather_model_graphs/networkx_utils.py

@@ -98,40 +98,29 @@ def split_graph_by_edge_attribute(graph, attr):
             f"No subgraphs were created. Check the edge attribute '{attr}'."
         )
 
-    # copy node attributes
-    for subgraph in subgraphs.values():
-        for node in subgraph.nodes:
-            subgraph.nodes[node].update(graph.nodes[node])
-
-    # check that at least one subgraph was created
-    if len(subgraphs) == 0:
-        raise ValueError(
-            f"No subgraphs were created. Check the edge attribute '{attr}'."
-        )
+    return subgraphs
 
-    # copy node attributes
-    for subgraph in subgraphs.values():
-        for node in subgraph.nodes:
-            subgraph.nodes[node].update(graph.nodes[node])
 
-    # check that at least one subgraph was created
-    if len(subgraphs) == 0:
-        raise ValueError(
-            f"No subgraphs were created. Check the edge attribute '{attr}'."
-        )
+def sort_nodes_in_graph(graph):
+    """
+    Creates a new networkx.DiGraph that is a copy of input, but with nodes
+    sorted according to their id
 
-    # copy node attributes
-    for subgraph in subgraphs.values():
-        for node in subgraph.nodes:
-            subgraph.nodes[node].update(graph.nodes[node])
+    Parameters
+    ----------
+    graph : networkx.DiGraph
+        Graph to sort nodes from
 
-    # check that at least one subgraph was created
-    if len(subgraphs) == 0:
-        raise ValueError(
-            f"No subgraphs were created. Check the edge attribute '{attr}'."
-        )
+    Returns
+    -------
+    networkx.DiGraph
+        Graph with sorted nodes
+    """
+    sorted_graph = networkx.DiGraph()
+    sorted_graph.add_nodes_from(sorted(graph.nodes(data=True)))
+    sorted_graph.add_edges_from(graph.edges(data=True))
 
-    return subgraphs
+    return sorted_graph
 
 
 def replace_node_labels_with_unique_ids(graph):
@@ -149,7 +138,7 @@ def replace_node_labels_with_unique_ids(graph):
         Graph with node labels renamed
     """
     return networkx.relabel_nodes(
-        graph, {node: i for i, node in enumerate(graph.nodes)}, copy=True
+        graph, {node: i for i, node in enumerate(sorted(graph.nodes))}, copy=True
     )
 
 

src/weather_model_graphs/save.py

@@ -5,7 +5,11 @@
 import networkx
 from loguru import logger
 
-from .networkx_utils import MissingEdgeAttributeError, split_graph_by_edge_attribute
+from .networkx_utils import (
+    MissingEdgeAttributeError,
+    sort_nodes_in_graph,
+    split_graph_by_edge_attribute,
+)
 
 try:
     import torch
@@ -106,18 +110,25 @@ def _concat_pyg_features(
     if list_from_attribute is not None:
         # create a list of graph objects by splitting the graph by the list_from_attribute
         try:
-            sub_graphs = list(
-                split_graph_by_edge_attribute(
-                    graph=graph, attr=list_from_attribute
-                ).values()
-            )
+            sub_graphs = [
+                value
+                for key, value in sorted(
+                    split_graph_by_edge_attribute(
+                        graph=graph, attr=list_from_attribute
+                    ).items()
+                )
+            ]
         except MissingEdgeAttributeError:
             # neural-lam still expects a list of graphs, so if the attribute is missing
             # we just return the original graph as a list
             sub_graphs = [graph]
-        pyg_graphs = [pyg_convert.from_networkx(g) for g in sub_graphs]
+        # Nodes must be sorted if we want to preserve any ordering
+        # when converted to pyg
+        pyg_graphs = [
+            pyg_convert.from_networkx(sort_nodes_in_graph(g)) for g in sub_graphs
+        ]
     else:
-        pyg_graphs = [pyg_convert.from_networkx(graph)]
+        pyg_graphs = [pyg_convert.from_networkx(sort_nodes_in_graph(graph))]
 
     edge_features_values = [
         _concat_pyg_features(pyg_g, features=edge_features) for pyg_g in pyg_graphs

src/weather_model_graphs/visualise/plot_2d.py

@@ -105,6 +105,7 @@ def nx_draw_with_pos_and_attr(
     node_zorder_attr=None,
     node_size=100,
     connectionstyle="arc3, rad=0.1",
+    with_labels=False,
     **kwargs,
 ):
     """
@@ -171,7 +172,7 @@ def nx_draw_with_pos_and_attr(
         graph,
         ax=ax,
         arrows=True,
-        with_labels=False,
+        with_labels=with_labels,
         node_size=node_size,
         connectionstyle=connectionstyle,
         **kwargs,

tests/test_graph_creation.py

@@ -192,3 +192,27 @@ def test_create_lat_lon(kind):
         coords_crs=coords_crs,
         graph_crs=graph_crs,
     )
+
+
+@pytest.mark.parametrize("kind", ["graphcast", "keisler", "oskarsson_hierarchical"])
+def test_create_decode_mask(kind):
+    """
+    Tests that the decode mask for m2g works, resulting in less edges than
+    no filtering.
+    """
+    xy = test_utils.create_fake_irregular_coords(100)
+    fn_name = f"create_{kind}_graph"
+    fn = getattr(wmg.create.archetype, fn_name)
+    # ~= 20 mesh nodes in bottom layer in each direction
+    mesh_node_distance = 0.05
+
+    unfiltered_graph = fn(coords=xy, mesh_node_distance=mesh_node_distance)
+
+    # Filter to only 20 / 100 grid nodes
+    decode_mask = np.concatenate((np.ones(20), np.zeros(80))).astype(bool)
+    filtered_graph = fn(
+        coords=xy, mesh_node_distance=mesh_node_distance, decode_mask=decode_mask
+    )
+
+    # Check that some filtering has been performed
+    assert len(filtered_graph.edges) < len(unfiltered_graph.edges)