Document node-size drawing and cover randwalk ids

Author: FraLotito

hypergraphx/viz/draw_hypergraph.py

@@ -80,7 +80,7 @@ def draw_hypergraph(
     hyperedge_facecolor_by_order: Optional[dict] = None,
     edge_width: float = 1.2,
     hyperedge_alpha: Union[float, np.array] = 0.8,
-    node_size: Union[int, np.array] = 150,
+    node_size: Union[int, float, np.array, dict] = 150,
     node_color: Union[str, np.array] = "#E2E0DD",
     node_facecolor: Union[str, np.array] = "black",
     node_shape: str = "o",
@@ -97,6 +97,9 @@ def draw_hypergraph(
 
     Parameters
     ----------
+    node_size : int, float, numpy.ndarray, or dict
+        Node marker size. If a dict is provided, keys must be the hypergraph nodes and
+        values are the corresponding marker sizes.
     show : bool
         If True, call plt.show().
 
@@ -168,6 +171,10 @@ def _stable_color(order_value):
         node_size = {n: node_size[i] for i, n in enumerate(nodes)}
     elif not isinstance(node_size, dict):
         node_size = {n: node_size for n in nodes}
+    if isinstance(node_size, dict):
+        missing_sizes = set(nodes) - set(node_size.keys())
+        if missing_sizes:
+            raise ValueError("node_size is missing entries for some nodes.")
     if isinstance(node_color, np.ndarray):
         if len(node_color) != len(nodes):
             raise ValueError("node_color length must match number of nodes.")

tests/dynamics/test_randwalk.py

@@ -22,6 +22,23 @@ def test_transition_matrix_rows_sum_to_one():
     assert np.allclose(T.sum(axis=1), 1.0)
 
 
+def test_transition_matrix_handles_non_contiguous_node_ids():
+    hg = Hypergraph(edge_list=[(492, 938), (938, 1200), (1200, 5000)])
+
+    T = transition_matrix(hg).toarray()
+    _, mapping = hg.binary_incidence_matrix(return_mapping=True)
+    node_to_idx = {node: idx for idx, node in mapping.items()}
+
+    assert T.shape == (4, 4)
+    assert np.allclose(T.sum(axis=1), 1.0)
+    assert T[node_to_idx[492], node_to_idx[938]] == 1.0
+    assert T[node_to_idx[938], node_to_idx[492]] == 0.5
+    assert T[node_to_idx[938], node_to_idx[1200]] == 0.5
+    assert T[node_to_idx[1200], node_to_idx[938]] == 0.5
+    assert T[node_to_idx[1200], node_to_idx[5000]] == 0.5
+    assert T[node_to_idx[5000], node_to_idx[1200]] == 1.0
+
+
 def test_random_walk_length():
     """Test random walk length equals time + 1."""
     np.random.seed(0)
@@ -31,6 +48,15 @@ def test_random_walk_length():
     assert len(path) == 4
 
 
+def test_random_walk_returns_non_contiguous_node_ids():
+    hg = Hypergraph(edge_list=[(492, 938), (938, 1200), (1200, 5000)])
+
+    path = random_walk(hg, s=492, time=5, seed=0)
+
+    assert path[0] == 492
+    assert set(path).issubset({492, 938, 1200, 5000})
+
+
 def test_stationary_state_properties():
     """Test stationary state is a valid distribution."""
     hg = _make_connected_hypergraph()

tests/viz/test_draws.py

@@ -3,6 +3,7 @@
 matplotlib.use("Agg")
 
 import matplotlib.pyplot as plt
+import pytest
 
 from hypergraphx import Hypergraph
 from hypergraphx.viz.draw_hypergraph import draw_hypergraph
@@ -24,6 +25,23 @@ def test_draw_hypergraph_smoke(monkeypatch):
     draw_hypergraph(hg)
 
 
+def test_draw_hypergraph_accepts_node_size_dict(monkeypatch):
+    monkeypatch.setattr(plt, "show", lambda: None)
+    hg = _make_hypergraph()
+
+    ax = draw_hypergraph(hg, node_size={0: 100, 1: 200, 2: 300, 3: 400})
+
+    assert ax is not None
+
+
+def test_draw_hypergraph_node_size_dict_requires_all_nodes(monkeypatch):
+    monkeypatch.setattr(plt, "show", lambda: None)
+    hg = _make_hypergraph()
+
+    with pytest.raises(ValueError, match="node_size is missing entries"):
+        draw_hypergraph(hg, node_size={0: 100, 1: 200})
+
+
 def test_draw_projections(monkeypatch):
     """Test draw_bipartite and draw_clique return axes."""
     monkeypatch.setattr(plt, "show", lambda: None)

tutorials/basics.ipynb

@@ -270,6 +270,23 @@
     "draw_hypergraph(H)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Node sizes can also be set per node with a dictionary keyed by node id.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "node_sizes = {node: 80 + 40 * H.degree(node) for node in H.get_nodes()}\n",
+    "draw_hypergraph(H, node_size=node_sizes, with_node_labels=True)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 12,