BFS layout implementation (#5179)

* first BFS layout implementation

* use of multipartite layout in bfs layout

* Add some nominal test cases for bfs_layout.

* Apply suggestions from code review

Co-authored-by: Ross Barnowski <rossbar@berkeley.edu>

* fix style

* add option of dict to multipartite_layout, use it from bfs_layout

* connect bsf_layout to docs

* Add test case for non-connected graphs.

* Apply latest pre-commit.

---------

Co-authored-by: Dan Schult <dschult@colgate.edu>
Co-authored-by: Ross Barnowski <rossbar@berkeley.edu>

doc/reference/drawing.rst

@@ -84,6 +84,7 @@ Graph Layout
    :toctree: generated/
 
    bipartite_layout
+   bfs_layout
    circular_layout
    kamada_kawai_layout
    planar_layout

networkx/drawing/layout.py

@@ -32,6 +32,7 @@
     "fruchterman_reingold_layout",
     "spiral_layout",
     "multipartite_layout",
+    "bfs_layout",
     "arf_layout",
 ]
 
@@ -1030,8 +1031,9 @@ def multipartite_layout(G, subset_key="subset", align="vertical", scale=1, cente
     G : NetworkX graph or list of nodes
         A position will be assigned to every node in G.
 
-    subset_key : string (default='subset')
-        Key of node data to be used as layer subset.
+    subset_key : string or dict (default='subset')
+        If a string, the key of node data in G that holds the node subset.
+        If a dict, keyed by layer number to the nodes in that layer/subset.
 
     align : string (default='vertical')
         The alignment of nodes. Vertical or horizontal.
@@ -1052,6 +1054,12 @@ def multipartite_layout(G, subset_key="subset", align="vertical", scale=1, cente
     >>> G = nx.complete_multipartite_graph(28, 16, 10)
     >>> pos = nx.multipartite_layout(G)
 
+    or use a dict to provide the layers of the layout
+
+    >>> G = nx.Graph([(0, 1), (1, 2), (1, 3), (3, 4)])
+    >>> layers = {"a": [0], "b": [1], "c": [2, 3], "d": [4]}
+    >>> pos = nx.multipartite_layout(G, subset_key=layers)
+
     Notes
     -----
     This algorithm currently only works in two dimensions and does not
@@ -1071,25 +1079,31 @@ def multipartite_layout(G, subset_key="subset", align="vertical", scale=1, cente
     if len(G) == 0:
         return {}
 
-    layers = {}
-    for v, data in G.nodes(data=True):
-        try:
-            layer = data[subset_key]
-        except KeyError:
-            msg = "all nodes must have subset_key (default='subset') as data"
-            raise ValueError(msg)
-        layers[layer] = [v] + layers.get(layer, [])
+    try:
+        # check if subset_key is dict-like
+        if len(G) != sum(len(nodes) for nodes in subset_key.values()):
+            raise nx.NetworkXError(
+                "all nodes must be in one subset of `subset_key` dict"
+            )
+    except AttributeError:
+        # subset_key is not a dict, hence a string
+        node_to_subset = nx.get_node_attributes(G, subset_key)
+        if len(node_to_subset) != len(G):
+            raise nx.NetworkXError(
+                f"all nodes need a subset_key attribute: {subset_key}"
+            )
+        subset_key = nx.utils.groups(node_to_subset)
 
     # Sort by layer, if possible
     try:
-        layers = sorted(layers.items())
+        layers = dict(sorted(subset_key.items()))
     except TypeError:
-        layers = list(layers.items())
+        layers = subset_key
 
     pos = None
     nodes = []
     width = len(layers)
-    for i, (_, layer) in enumerate(layers):
+    for i, layer in enumerate(layers.values()):
         height = len(layer)
         xs = np.repeat(i, height)
         ys = np.arange(0, height, dtype=float)
@@ -1295,3 +1309,50 @@ def rescale_layout_dict(pos, scale=1):
     pos_v = np.array(list(pos.values()))
     pos_v = rescale_layout(pos_v, scale=scale)
     return dict(zip(pos, pos_v))
+
+
+def bfs_layout(G, start, *, align="vertical", scale=1, center=None):
+    """Position nodes according to breadth-first search algorithm.
+
+    Parameters
+    ----------
+    G : NetworkX graph
+        A position will be assigned to every node in G.
+
+    start : node in `G`
+        Starting node for bfs
+
+    center : array-like or None
+        Coordinate pair around which to center the layout.
+
+    Returns
+    -------
+    pos : dict
+        A dictionary of positions keyed by node.
+
+    Examples
+    --------
+    >>> G = nx.path_graph(4)
+    >>> pos = nx.bfs_layout(G, 0)
+
+    Notes
+    -----
+    This algorithm currently only works in two dimensions and does not
+    try to minimize edge crossings.
+
+    """
+    G, center = _process_params(G, center, 2)
+
+    # Compute layers with BFS
+    layers = dict(enumerate(nx.bfs_layers(G, start)))
+
+    if len(G) != sum(len(nodes) for nodes in layers.values()):
+        raise nx.NetworkXError(
+            "bfs_layout didn't include all nodes. Perhaps use input graph:\n"
+            "        G.subgraph(nx.node_connected_component(G, start))"
+        )
+
+    # Compute node positions with multipartite_layout
+    return multipartite_layout(
+        G, subset_key=layers, align=align, scale=scale, center=center
+    )

networkx/drawing/tests/test_layout.py

@@ -452,12 +452,18 @@ def test_multipartite_layout_layer_order():
     """Return the layers in sorted order if the layers of the multipartite
     graph are sortable. See gh-5691"""
     G = nx.Graph()
-    for node, layer in zip(("a", "b", "c", "d", "e"), (2, 3, 1, 2, 4)):
+    node_group = dict(zip(("a", "b", "c", "d", "e"), (2, 3, 1, 2, 4)))
+    for node, layer in node_group.items():
         G.add_node(node, subset=layer)
 
     # Horizontal alignment, therefore y-coord determines layers
     pos = nx.multipartite_layout(G, align="horizontal")
 
+    layers = nx.utils.groups(node_group)
+    pos_from_layers = nx.multipartite_layout(G, align="horizontal", subset_key=layers)
+    for (n1, p1), (n2, p2) in zip(pos.items(), pos_from_layers.items()):
+        assert n1 == n2 and (p1 == p2).all()
+
     # Nodes "a" and "d" are in the same layer
     assert pos["a"][-1] == pos["d"][-1]
     # positions should be sorted according to layer
@@ -467,3 +473,43 @@ def test_multipartite_layout_layer_order():
     G.nodes["a"]["subset"] = "layer_0"  # Can't sort mixed strs/ints
     pos_nosort = nx.multipartite_layout(G)  # smoke test: this should not raise
     assert pos_nosort.keys() == pos.keys()
+
+
+def _num_nodes_per_bfs_layer(pos):
+    """Helper function to extract the number of nodes in each layer of bfs_layout"""
+    x = np.array(list(pos.values()))[:, 0]  # node positions in layered dimension
+    _, layer_count = np.unique(x, return_counts=True)
+    return layer_count
+
+
+@pytest.mark.parametrize("n", range(2, 7))
+def test_bfs_layout_complete_graph(n):
+    """The complete graph should result in two layers: the starting node and
+    a second layer containing all neighbors."""
+    G = nx.complete_graph(n)
+    pos = nx.bfs_layout(G, start=0)
+    assert np.array_equal(_num_nodes_per_bfs_layer(pos), [1, n - 1])
+
+
+def test_bfs_layout_barbell():
+    G = nx.barbell_graph(5, 3)
+    # Start in one of the "bells"
+    pos = nx.bfs_layout(G, start=0)
+    # start, bell-1, [1] * len(bar)+1, bell-1
+    expected_nodes_per_layer = [1, 4, 1, 1, 1, 1, 4]
+    assert np.array_equal(_num_nodes_per_bfs_layer(pos), expected_nodes_per_layer)
+    # Start in the other "bell" - expect same layer pattern
+    pos = nx.bfs_layout(G, start=12)
+    assert np.array_equal(_num_nodes_per_bfs_layer(pos), expected_nodes_per_layer)
+    # Starting in the center of the bar, expect layers to be symmetric
+    pos = nx.bfs_layout(G, start=6)
+    # Expected layers: {6 (start)}, {5, 7}, {4, 8}, {8 nodes from remainder of bells}
+    expected_nodes_per_layer = [1, 2, 2, 8]
+    assert np.array_equal(_num_nodes_per_bfs_layer(pos), expected_nodes_per_layer)
+
+
+def test_bfs_layout_disconnected():
+    G = nx.complete_graph(5)
+    G.add_edges_from([(10, 11), (11, 12)])
+    with pytest.raises(nx.NetworkXError, match="bfs_layout didn't include all nodes"):
+        nx.bfs_layout(G, start=0)