diff --git a/src/pixelator/common/graph/backends/implementations/_networkx.py b/src/pixelator/common/graph/backends/implementations/_networkx.py index 4e3028145..5bde52dec 100644 --- a/src/pixelator/common/graph/backends/implementations/_networkx.py +++ b/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,20 +1252,22 @@ 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, @@ -1229,6 +1275,11 @@ def spectral_layout( 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)