Fix IndexError when visualizung subgraphs with zerod edges or nodes

CHANGELOG.md

@@ -18,6 +18,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `__version__` attribute to the package init
   [\#56](https://github.com/mllam/weather-model-graphs/pull/56) @AdMub
 
+### Fixed
+
+- Fix an `IndexError` crash when plotting a graph component with zero edges or zero nodes using `wmg.visualise.nx_draw_with_pos_and_attr`. [\#131](https://github.com/mllam/weather-model-graphs/pull/131) @sudhansu-24
+
 ## [v0.3.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.3.0)
 
 

src/weather_model_graphs/visualise/plot_2d.py

@@ -21,9 +21,21 @@ def nx_draw_with_pos(g, with_labels=False, **kwargs):
 
 def _get_graph_attr_values(g, attr_name, component="edges"):
     if component == "edges":
-        features = list(g.edges(data=True))[0][2].keys()
+        edge_data = list(g.edges(data=True))
+        if len(edge_data) == 0:
+            raise ValueError(
+                f"Cannot colour by edge attribute '{attr_name}' because the "
+                "graph has no edges."
+            )
+        features = edge_data[0][2].keys()
     elif component == "nodes":
-        features = list(g.nodes(data=True))[0][1].keys()
+        node_data = list(g.nodes(data=True))
+        if len(node_data) == 0:
+            raise ValueError(
+                f"Cannot colour by node attribute '{attr_name}' because the "
+                "graph has no nodes."
+            )
+        features = node_data[0][1].keys()
     else:
         raise ValueError(
             f"`component` should be either 'edges' or 'nodes', but got '{component}'"
@@ -141,7 +153,7 @@ def nx_draw_with_pos_and_attr(
     if node_zorder_attr is not None:
         graph = nx_utils.sort_nodes_internally(graph, node_attr=node_zorder_attr)
 
-    if edge_color_attr is not None:
+    if edge_color_attr is not None and graph.number_of_edges() > 0:
         edge_attr_vals = _get_graph_attr_values(
             graph, edge_color_attr, component="edges"
         )
@@ -154,8 +166,10 @@ def nx_draw_with_pos_and_attr(
         kwargs["edge_color"] = edge_attr_vals["values"]
         kwargs["edge_vmin"] = min(edge_attr_vals["values"])
         kwargs["edge_vmax"] = max(edge_attr_vals["values"])
+    elif edge_color_attr is not None:
+        edge_color_attr = None  # nothing to colour
 
-    if node_color_attr is not None:
+    if node_color_attr is not None and graph.number_of_nodes() > 0:
         node_attr_vals = _get_graph_attr_values(
             graph, node_color_attr, component="nodes"
         )
@@ -167,6 +181,8 @@ def nx_draw_with_pos_and_attr(
         kwargs["node_color"] = node_attr_vals["values"]
         kwargs["vmin"] = min(node_attr_vals["values"])
         kwargs["vmax"] = max(node_attr_vals["values"])
+    elif node_color_attr is not None:
+        node_color_attr = None  # nothing to colour
 
     ax = nx_draw_with_pos(
         graph,

tests/test_graph_plots.py

@@ -64,3 +64,54 @@ def fn():
 
     with tempfile.NamedTemporaryFile(suffix=".png") as f:
         fig.savefig(f.name)
+
+
+def test_plot_nodes_only_graph_with_edge_color_attr():
+    """Plotting a graph with nodes but no edges and edge_color_attr should
+    not crash (previously raised IndexError)."""
+    import networkx as nx
+
+    G = nx.DiGraph()
+    G.add_node(0, pos=np.array([1.0, 2.0]), type="mesh")
+    G.add_node(1, pos=np.array([3.0, 4.0]), type="grid")
+
+    # Should silently skip edge colouring, not raise IndexError
+    ax = wmg.visualise.nx_draw_with_pos_and_attr(G, edge_color_attr="len")
+    assert ax is not None
+    plt.close("all")
+
+
+def test_plot_empty_graph_with_node_color_attr():
+    """Plotting a completely empty graph with node_color_attr should not crash."""
+    import networkx as nx
+
+    G = nx.DiGraph()
+
+    ax = wmg.visualise.nx_draw_with_pos_and_attr(G, node_color_attr="type")
+    assert ax is not None
+    plt.close("all")
+
+
+def test_get_graph_attr_values_raises_on_empty_edges():
+    """_get_graph_attr_values should raise a clear ValueError for empty edges."""
+    import networkx as nx
+
+    from weather_model_graphs.visualise.plot_2d import _get_graph_attr_values
+
+    G = nx.DiGraph()
+    G.add_node(0, pos=np.array([1.0, 2.0]))
+
+    with pytest.raises(ValueError, match="no edges"):
+        _get_graph_attr_values(G, "len", component="edges")
+
+
+def test_get_graph_attr_values_raises_on_empty_nodes():
+    """_get_graph_attr_values should raise a clear ValueError for empty nodes."""
+    import networkx as nx
+
+    from weather_model_graphs.visualise.plot_2d import _get_graph_attr_values
+
+    G = nx.DiGraph()
+
+    with pytest.raises(ValueError, match="no nodes"):
+        _get_graph_attr_values(G, "type", component="nodes")