Spectral layout improvements

src/pixelator/common/graph/backends/implementations/_networkx.py

@@ -1095,13 +1095,16 @@ def _spectral_layout_eigen(g, nodes, dim, normalize, seed):
 
     # Exclude trivial eigenvector
     order = np.argsort(vals)
-    raw_coordinates = vecs[:, order[1:k]]
+    selection = order[1:k]
+
+    eigenvalues = vals[selection]
+    raw_coordinates = vecs[:, selection]
 
     if normalize:
         # De-normalize
         raw_coordinates = diag_inv_sqrt @ raw_coordinates
 
-    return raw_coordinates
+    return eigenvalues, raw_coordinates
 
 
 def _spectral_layout_psvd(g, nodes, dim, seed):
@@ -1160,10 +1163,11 @@ def _spectral_layout_psvd(g, nodes, dim, seed):
 
     # Do partial singular value decomposition
     # (k=dim + 1 to account for trivial 0-eigenvalue)
-    vecs_u, s, vecs_v_t = sp.sparse.linalg.svds(B_norm, k=k, v0=v0, tol=0)
+    vecs_u, s, vecs_v_t = sp.sparse.linalg.svds(B_norm, k=k, v0=v0, tol=0, which="LM")
 
     # Sort by descending and exclude trivial value, transpose
     selection = np.argsort(s)[::-1][1:k]
+    eigenvalues = 1.0 - s[selection]
 
     vecs_v = vecs_v_t.T
 
@@ -1180,7 +1184,7 @@ def _spectral_layout_psvd(g, nodes, dim, seed):
 
     raw_coordinates = raw_coords_unord[order]
 
-    return raw_coordinates
+    return eigenvalues, raw_coordinates
 
 
 def spectral_layout(
@@ -1189,8 +1193,48 @@ def spectral_layout(
     normalize: bool = True,
     seed: Optional[int] = None,
     method: Literal["eigen", "psvd"] = "psvd",
+    diffusion_time: Optional[float] = None,
 ):
-    """Use a spectral layout algorithm to compute coordinates from a graph."""
+    """Use a spectral layout algorithm to compute coordinates from a graph.
+
+    The spectral layout algorithm computes ``dim`` eigenvectors of the graph Laplacian
+    corresponding to the smallest non-zero eigenvalues. The layout coordinates
+    are derived from the vectors by using each eigenvector as an embedding axis.
+    Resulting coordinates are normalized (centered at origin and scaled by median
+    radius) before return.
+
+    Args:
+        g: A connected NetworkX graph.
+        dim: Number of output dimension (3 by default).
+        normalize: If the symmetrically normalized Laplacian should be used in the
+        computations. Note that ``normalize=False`` is incompatible with
+        ``method='psvd'``. Defaults to ``True``.
+        seed: Optional random seed.
+        method: Which computational method to use. ``'eigen'`` uses an eigensolver to
+        compute the eigenvectors, while ``'psvd'`` solves an equivalent partial singular
+        value decomposition problem. Solutions generated by the different approaches are
+        equivalent, but ``'psvd'``is generally faster.
+        diffusion_time: Optional 'time' parameter for diffusion map coordinates.
+
+    Returns:
+        Dictionary mapping each node in ``g`` to its normalized layout coordinates as a
+        one-dimensional array with length ``dim``.
+
+    Raises:
+        ValueError: If the graph (``g``) does not have more than 1 node.
+        ValueError: If the graph is directed.
+        ValueError: If the graph is not connected.
+        ValueError: If ``dim`` is not 2 or 3.
+        ValueError: If ``method`` is not 'eigen' or 'psvd'.
+        ValueError: If ``method`` is ``'psvd'`` and ``normalize`` is ``False``.
+        ValueError: If using ``'eigen'`` and the number ``dim + 1`` is not less than the
+        number of nodes in the input graph.
+        ValueError: If using the ``'psvd'`` method and ``g`` isn't a bipartite graph whose
+        nodes are partitioned exactly into 'A' and 'B' via the 'pixel_type' attribute.
+        ValueError: If using ``'psvd'`` and ``g`` is not a bipartite graph whose partition
+        sizes ``M``, ``N`` satisfy ``dim + 1 <= min(M, N) - 1``.
+
+    """
     # Validate inputs
     if g.number_of_nodes() in (0, 1):
         raise ValueError("g must have more than 1 node.")
@@ -1208,27 +1252,34 @@ def spectral_layout(
             "Use normalize=True or switch to method='eigen' "
             "and keep normalize=False for non-normalized."
         )
+    if diffusion_time is not None and diffusion_time < 0:
+        raise ValueError("diffusion_time must be non-negative")
 
     # Get nodes
     nodes = list(g.nodes)
 
     # Determine method and compute raw coordinates
     if method == "psvd":
-        raw_coordinates = _spectral_layout_psvd(
+        eigenvalues, raw_coordinates = _spectral_layout_psvd(
             g,
             nodes=nodes,
             dim=dim,
             seed=seed,
         )
     elif method == "eigen":
-        raw_coordinates = _spectral_layout_eigen(
+        eigenvalues, raw_coordinates = _spectral_layout_eigen(
             g,
             nodes=nodes,
             dim=dim,
             normalize=normalize,
             seed=seed,
         )
 
+    if diffusion_time is not None:
+        # Diffusion map
+        diffusion_weights = np.exp(-eigenvalues * diffusion_time)
+        raw_coordinates = raw_coordinates @ np.diag(diffusion_weights)
+
     # Return normalized coordinates
     coordinates = normalize_layout_coordinates(raw_coordinates)