Unify hierarchical mesh edge level metadata

CHANGELOG.md

@@ -27,7 +27,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   [\#32](https://github.com/mllam/weather-model-graphs/pull/32), @joeloskarsson
 
 ### Fixed
-
+- Fix issue #45 about inconsistent edge metadata and improves compatibility
+  [\#66](https://github.com/mllam/weather-model-graphs/pull/66), 
+  @siddhantcookie
+
 - Fix the bug with edgeless nodes being dropped
   [\#51](https://github.com/mllam/weather-model-graphs/pull/51), @pkhalaj, @wi-spang, @krikru
 

docs/creating_the_graph.ipynb

@@ -331,12 +331,14 @@
     "for (name, graph), ax in zip(graph_components.items(), axes.flatten()):\n",
     "    pl_kwargs = {}\n",
     "    if name == \"m2m_same\":\n",
-    "        pl_kwargs = dict(edge_color_attr=\"level\", node_color_attr=\"level\", node_size=10)\n",
+    "        pl_kwargs = dict(\n",
+    "            edge_color_attr=\"from_level\", node_color_attr=\"level\", node_size=10\n",
+    "        )\n",
     "    elif name == \"g2m\" or name == \"m2g\":\n",
     "        pl_kwargs = dict(edge_color_attr=\"len\", node_color_attr=\"type\", node_size=30)\n",
     "    elif name in [\"m2m_up\", \"m2m_down\"]:\n",
     "        pl_kwargs = dict(\n",
-    "            edge_color_attr=\"levels\", node_color_attr=\"level\", node_size=30\n",
+    "            edge_color_attr=\"connection\", node_color_attr=\"level\", node_size=30\n",
     "        )\n",
     "\n",
     "    wmg.visualise.nx_draw_with_pos_and_attr(graph, ax=ax, **pl_kwargs)\n",
@@ -363,7 +365,7 @@
     "m2m_same_graph_components = {\n",
     "    f\"m2m_same_level_{level}\": graph\n",
     "    for level, graph in wmg.split_graph_by_edge_attribute(\n",
-    "        graph=m2m_same_graph, attr=\"level\"\n",
+    "        graph=m2m_same_graph, attr=\"from_level\"\n",
     "    ).items()\n",
     "}\n",
     "m2m_same_graph_components"

docs/design.ipynb

@@ -122,7 +122,10 @@
     "#### Edge attributes\n",
     "\n",
     "- `component`: the component of the graph the edge belongs to, either `g2m`, `m2m` or `m2g`\n",
-    "- `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.\n",
+    "- `level`: for multi-range mesh graphs this denotes the refinement level of mesh connection.        \n",
+    "- `from_level`: for hierarchical mesh graphs, the source level of the edge\n",
+    "- `to_level`: for hierarchical mesh graphs, the target level of the edge\n",
+    "- `connection`: for hierarchical mesh graphs, a readable level-pair string in the form `\"from_level>to_level\"` (for example `\"1>0\"`)\n",
     "- `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\n",
     "- `vdiff`: the vector spanning between the (x,y) coordinates of the two nodes\n",
     "- `direction`: for hierarchical graphs this denotes the direction of the edge, either `up`, `down` and `same`"

src/weather_model_graphs/create/archetype.py

@@ -172,8 +172,8 @@ def create_oskarsson_hierarchical_graph(
     connections the mesh graph contains nearest neighbour connections between
     levels (up and down). To distinguish between these these three types of
     edge connections each edge has a `direction` attribute (with value "up",
-    "down", or "same"). In addition, the `levels` attribute indicates which two levels
-    are connected for cross-level edges (e.g. "1>2" for edges between level 1 and 2).
+    "down", or "same"). In addition, the `from_level` and `to_level` attributes
+    indicate which two levels are connected (e.g. from_level=1 and to_level=2).
 
     The grid to mesh connectivity connects each mesh node to grid nodes withing
     distance 0.51d, where d is the length of diagonal edges between neighbouring

src/weather_model_graphs/create/base.py

@@ -139,9 +139,16 @@ def create_all_graph_components(
             **m2m_connectivity_kwargs,
         )
         # Only connect grid to bottom level of hierarchy
-        grid_connect_graph = split_graph_by_edge_attribute(
-            graph_components["m2m"], "level"
-        )[0]
+        bottom_same_level_edges = [
+            (u, v)
+            for u, v, edge_data in graph_components["m2m"].edges(data=True)
+            if edge_data.get("direction") == "same"
+            and edge_data.get("from_level") == 0
+            and edge_data.get("to_level") == 0
+        ]
+        grid_connect_graph = graph_components["m2m"].edge_subgraph(
+            bottom_same_level_edges
+        )
     elif m2m_connectivity == "flat_multiscale":
         graph_components["m2m"] = create_flat_multiscale_mesh_graph(
             xy=xy,
@@ -365,20 +372,39 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target):
             for edge_check_graph in (G_source, G_target):
                 # Check if graph has edges
                 if len(edge_check_graph.edges) > 0:
-                    (
-                        level_subgraph,
-                        no_level_subgraph,
-                    ) = split_on_edge_attribute_existance(edge_check_graph, "level")
-
-                    # Check if graph has levels (hierarchical or multi-scale edges)
-                    if nx.is_empty(level_subgraph):
-                        # Consider edges in whole graph (whole graph is level 1)
-                        first_level_graph = edge_check_graph  # == no_level_subgraph
+                    # Hierarchical mesh graphs expose a direction + level-range
+                    # edge schema. In that case use only same-level edges on
+                    # bottom level to estimate local characteristic edge length.
+                    if any(
+                        "from_level" in edge_data and "to_level" in edge_data
+                        for _, _, edge_data in edge_check_graph.edges(data=True)
+                    ):
+                        first_level_graph = edge_check_graph.edge_subgraph(
+                            [
+                                (u, v)
+                                for u, v, edge_data in edge_check_graph.edges(data=True)
+                                if edge_data.get("direction") == "same"
+                                and edge_data.get("from_level") == 0
+                                and edge_data.get("to_level") == 0
+                            ]
+                        )
                     else:
-                        # Has levels, only consider edges in level 1 graph
-                        first_level_graph = split_graph_by_edge_attribute(
-                            level_subgraph, "level"
-                        )[0]
+                        (
+                            level_subgraph,
+                            _no_level_subgraph,
+                        ) = split_on_edge_attribute_existance(edge_check_graph, "level")
+
+                        # Check if graph has levels (hierarchical or
+                        # multi-scale edges)
+                        if nx.is_empty(level_subgraph):
+                            # Consider edges in whole graph (whole graph is
+                            # level 1)
+                            first_level_graph = edge_check_graph  # == no_level_subgraph
+                        else:
+                            # Has levels, only consider edges in level 1 graph
+                            first_level_graph = split_graph_by_edge_attribute(
+                                level_subgraph, "level"
+                            )[0]
                     longest_graph_edge = max(
                         first_level_graph.edges(data=True),
                         key=lambda x: x[2].get("len", 0),

src/weather_model_graphs/create/mesh/kinds/hierarchical.py

@@ -61,11 +61,16 @@ def create_hierarchical_multiscale_mesh_graph(
         for level_i, graph in enumerate(Gs_all_levels)
     ]
 
-    # add `direction` attribute to all edges with value `same``
+    # Add edge metadata using a consistent schema across all directions.
     for i, G in enumerate(Gs_all_levels):
         for u, v in G.edges:
+            # Remove legacy keys to keep one consistent edge schema.
+            G.edges[u, v].pop("level", None)
+            G.edges[u, v].pop("levels", None)
             G.edges[u, v]["direction"] = "same"
-            G.edges[u, v]["level"] = i
+            G.edges[u, v]["from_level"] = i
+            G.edges[u, v]["to_level"] = i
+            G.edges[u, v]["connection"] = f"{i}>{i}"
 
     # Create inter-level mesh edges
     up_graphs = []
@@ -105,15 +110,22 @@ def create_hierarchical_multiscale_mesh_graph(
             G_down.edges[u, v]["vdiff"] = (
                 G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"]
             )
-            G_down.edges[u, v]["levels"] = f"{from_level}>{to_level}"
             G_down.edges[u, v]["direction"] = "down"
+            G_down.edges[u, v]["from_level"] = from_level
+            G_down.edges[u, v]["to_level"] = to_level
+            G_down.edges[u, v]["connection"] = f"{from_level}>{to_level}"
 
         G_up = networkx.DiGraph()
         G_up.add_nodes_from(G_down.nodes(data=True))
         for u, v, data in G_down.edges(data=True):
             data = data.copy()
-            data["levels"] = f"{to_level}>{from_level}"
+            # Remove legacy keys to keep one consistent edge schema.
+            data.pop("level", None)
+            data.pop("levels", None)
             data["direction"] = "up"
+            data["from_level"] = to_level
+            data["to_level"] = from_level
+            data["connection"] = f"{to_level}>{from_level}"
             G_up.add_edge(v, u, **data)
 
         up_graphs.append(G_up)

tests/test_graph_creation.py

@@ -239,6 +239,33 @@ def test_create_many_levels(kind):
     )
 
 
+def test_hierarchical_edges_have_consistent_level_metadata():
+    xy = test_utils.create_fake_xy(N=64)
+    graph = wmg.create.archetype.create_oskarsson_hierarchical_graph(coords=xy)
+
+    graph_components = wmg.split_graph_by_edge_attribute(graph=graph, attr="component")
+    m2m_graph = graph_components["m2m"]
+
+    for _, _, edge_data in m2m_graph.edges(data=True):
+        assert "direction" in edge_data
+        assert "from_level" in edge_data
+        assert "to_level" in edge_data
+        assert "connection" in edge_data
+
+        assert isinstance(edge_data["from_level"], int)
+        assert isinstance(edge_data["to_level"], int)
+        assert isinstance(edge_data["connection"], str)
+
+        if edge_data["direction"] == "same":
+            assert edge_data["from_level"] == edge_data["to_level"]
+
+        assert edge_data["connection"] == (
+            f"{edge_data['from_level']}>{edge_data['to_level']}"
+        )
+        assert "level" not in edge_data
+        assert "levels" not in edge_data
+
+
 @pytest.mark.parametrize("coordinates_offset", [0, 1, 10, 100])
 @pytest.mark.parametrize(
     ("method", "method_kwargs"),