spectral bisection for graphs using fiedler vector (#6404)

Add spectral_bisection function to bisect node sets based on the
Fiedler vector

Co-authored-by: Benjamin Edwards <bedwards@cs.unm.edu>
Co-authored-by: Ross Barnowski <rossbar@berkeley.edu>

2 files changed, 88 insertions, 1 deletions

diff --git a/networkx/linalg/algebraicconnectivity.py b/networkx/linalg/algebraicconnectivity.py
index a309b1c4..2e1aca60 100644
--- a/networkx/linalg/algebraicconnectivity.py
+++ b/networkx/linalg/algebraicconnectivity.py
@@ -10,7 +10,12 @@ from networkx.utils import (
     reverse_cuthill_mckee_ordering,
 )
 
-__all__ = ["algebraic_connectivity", "fiedler_vector", "spectral_ordering"]
+__all__ = [
+    "algebraic_connectivity",
+    "fiedler_vector",
+    "spectral_ordering",
+    "spectral_bisection",
+]
 
 
 class _PCGSolver:
@@ -585,3 +590,70 @@ def spectral_ordering(
             order.extend(component)
 
     return order
+
+
+def spectral_bisection(
+    G, weight="weight", normalized=False, tol=1e-8, method="tracemin_pcg", seed=None
+):
+    """Bisect the graph using the Fiedler vector.
+
+    This method uses the Fiedler vector to bisect a graph.
+    The partition is defined by the nodes which are associated with
+    either positive or negative values in the vector.
+
+    Parameters
+    ----------
+    G : NetworkX Graph
+
+    weight : str, optional (default: weight)
+        The data key used to determine the weight of each edge. If None, then
+        each edge has unit weight.
+
+    normalized : bool, optional (default: False)
+        Whether the normalized Laplacian matrix is used.
+
+    tol : float, optional (default: 1e-8)
+        Tolerance of relative residual in eigenvalue computation.
+
+    method : string, optional (default: 'tracemin_pcg')
+        Method of eigenvalue computation. It must be one of the tracemin
+        options shown below (TraceMIN), 'lanczos' (Lanczos iteration)
+        or 'lobpcg' (LOBPCG).
+
+        The TraceMIN algorithm uses a linear system solver. The following
+        values allow specifying the solver to be used.
+
+        =============== ========================================
+        Value           Solver
+        =============== ========================================
+        'tracemin_pcg'  Preconditioned conjugate gradient method
+        'tracemin_lu'   LU factorization
+        =============== ========================================
+
+    seed : integer, random_state, or None (default)
+        Indicator of random number generation state.
+        See :ref:`Randomness<randomness>`.
+
+    Returns
+    --------
+    bisection : tuple of sets
+        Sets with the bisection of nodes
+
+    Examples
+    --------
+    >>> G = nx.barbell_graph(3, 0)
+    >>> nx.spectral_bisection(G)
+    ({0, 1, 2}, {3, 4, 5})
+
+    References
+    ----------
+    .. [1] M. E. J Newman 'Networks: An Introduction', pages 364-370
+       Oxford University Press 2011.
+    """
+    import numpy as np
+
+    v = nx.fiedler_vector(G, weight, normalized, tol, method, seed)
+    nodes = np.array(list(G))
+    pos_vals = v >= 0
+
+    return set(nodes[~pos_vals]), set(nodes[pos_vals])
diff --git a/networkx/linalg/tests/test_algebraic_connectivity.py b/networkx/linalg/tests/test_algebraic_connectivity.py
index 7a86bd0a..16ee09c1 100644
--- a/networkx/linalg/tests/test_algebraic_connectivity.py
+++ b/networkx/linalg/tests/test_algebraic_connectivity.py
@@ -46,6 +46,21 @@ def test_fiedler_vector_tracemin_unknown():
         )
 
 
+def test_spectral_bisection():
+    pytest.importorskip("scipy")
+    G = nx.barbell_graph(3, 0)
+    C = nx.spectral_bisection(G)
+    assert C == ({0, 1, 2}, {3, 4, 5})
+
+    mapping = dict(enumerate("badfec"))
+    G = nx.relabel_nodes(G, mapping)
+    C = nx.spectral_bisection(G)
+    assert C == (
+        {mapping[0], mapping[1], mapping[2]},
+        {mapping[3], mapping[4], mapping[5]},
+    )
+
+
 def check_eigenvector(A, l, x):
     nx = np.linalg.norm(x)
     # Check zeroness.