nystrom-ncut 0.0.1__py3-none-any.whl

nystrom_ncut/__init__.py

@@ -0,0 +1,22 @@
+from .ncut_pytorch import NCUT
+from .propagation_utils import (
+    affinity_from_features,
+    propagate_eigenvectors,
+    propagate_knn,
+    quantile_normalize,
+)
+from .visualize_utils import (
+    eigenvector_to_rgb,
+    rgb_from_tsne_3d,
+    rgb_from_umap_sphere,
+    rgb_from_tsne_2d,
+    rgb_from_umap_3d,
+    rgb_from_umap_2d,
+    rgb_from_cosine_tsne_3d,
+    rotate_rgb_cube,
+    convert_to_lab_color,
+    propagate_rgb_color,
+    get_mask,
+)
+from .ncut_pytorch import nystrom_ncut, ncut
+from .ncut_pytorch import kway_ncut, axis_align

nystrom_ncut/ncut_pytorch.py

@@ -0,0 +1,561 @@
+# %%
+import logging
+import math
+from typing import Literal
+
+import torch
+
+from .nystrom import (
+    solve_eig
+)
+from .propagation_utils import (
+    run_subgraph_sampling,
+    propagate_knn,
+    affinity_from_features,
+)
+
+
+class NCUT:
+    """Nystrom Normalized Cut for large scale graph."""
+
+    def __init__(
+        self,
+        num_eig: int = 100,
+        knn: int = 10,
+        affinity_focal_gamma: float = 1.0,
+        num_sample: int = 10000,
+        sample_method: Literal["farthest", "random"] = "farthest",
+        distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
+        indirect_connection: bool = False,
+        indirect_pca_dim: int = 100,
+        device: str = None,
+        move_output_to_cpu: bool = False,
+        eig_solver: Literal["svd_lowrank", "lobpcg", "svd", "eigh"] = "svd_lowrank",
+        normalize_features: bool = None,
+        matmul_chunk_size: int = 8096,
+        make_orthogonal: bool = False,
+        verbose: bool = False,
+    ):
+        """
+        Args:
+            num_eig (int): number of top eigenvectors to return
+            knn (int): number of KNN for propagating eigenvectors from subgraph to full graph,
+                smaller knn result in more sharp eigenvectors.
+            affinity_focal_gamma (float): affinity matrix temperature, lower t reduce the not-so-connected edge weights,
+                smaller t result in more sharp eigenvectors.
+            num_sample (int): number of samples for Nystrom-like approximation,
+                reduce only if memory is not enough, increase for better approximation
+            sample_method (str): subgraph sampling, ['farthest', 'random'].
+                farthest point sampling is recommended for better Nystrom-approximation accuracy
+            distance (str): distance metric for affinity matrix, ['cosine', 'euclidean', 'rbf'].
+            indirect_connection (bool): include indirect connection in the Nystrom-like approximation
+            indirect_pca_dim (int): when compute indirect connection, PCA to reduce the node dimension,
+            device (str): device to use for eigen computation,
+                move to GPU to speeds up a bit (~5x faster)
+            move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
+            eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
+            normalize_features (bool): normalize input features before computing affinity matrix,
+                default 'None' is True for cosine distance, False for euclidean distance and rbf
+            matmul_chunk_size (int): chunk size for large-scale matrix multiplication
+            make_orthogonal (bool): make eigenvectors orthogonal post-hoc
+            verbose (bool): progress bar
+
+        Examples:
+            >>> from ncut_pytorch import NCUT
+            >>> import torch
+            >>> features = torch.rand(10000, 100)
+            >>> ncut = NCUT(num_eig=20)
+            >>> ncut.fit(features)
+            >>> eigenvectors, eigenvalues = ncut.transform(features)
+            >>> print(eigenvectors.shape, eigenvalues.shape)
+            >>> # (10000, 20) (20,)
+
+            >>> from ncut_pytorch import eigenvector_to_rgb
+            >>> # use t-SNE or UMAP to convert eigenvectors to RGB
+            >>> X_3d, rgb = eigenvector_to_rgb(eigenvectors, method='tsne_3d')
+            >>> print(X_3d.shape, rgb.shape)
+            >>> # (10000, 3) (10000, 3)
+
+            >>> # transform new features
+            >>> new_features = torch.rand(500, 100)
+            >>> new_eigenvectors, _ = ncut.transform(new_features)
+            >>> print(new_eigenvectors.shape)
+            >>> # (500, 20)
+        """
+        self.num_eig = num_eig
+        self.num_sample = num_sample
+        self.knn = knn
+        self.sample_method = sample_method
+        self.distance = distance
+        self.affinity_focal_gamma = affinity_focal_gamma
+        self.indirect_connection = indirect_connection
+        self.indirect_pca_dim = indirect_pca_dim
+        self.device = device
+        self.move_output_to_cpu = move_output_to_cpu
+        self.eig_solver = eig_solver
+        self.normalize_features = normalize_features
+        if self.normalize_features is None:
+            if distance in ["cosine"]:
+                self.normalize_features = True
+            if distance in ["euclidean", "rbf"]:
+                self.normalize_features = False
+        self.matmul_chunk_size = matmul_chunk_size
+        self.make_orthogonal = make_orthogonal
+        self.verbose = verbose
+
+        self.subgraph_eigen_vector = None
+        self.eigen_value = None
+        self.subgraph_indices = None
+        self.subgraph_features = None
+
+    def fit(self,
+            features: torch.Tensor,
+            precomputed_sampled_indices: torch.Tensor = None
+            ):
+        """Fit Nystrom Normalized Cut on the input features.
+        Args:
+            features (torch.Tensor): input features, shape (n_samples, n_features)
+            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+                override the sample_method, if not None
+        Returns:
+            (NCUT): self
+        """
+        _n = features.shape[0]
+        if self.num_sample >= _n:
+            logging.info(
+                f"NCUT nystrom num_sample is larger than number of input samples, nyström approximation is not needed, setting num_sample={_n} and knn=1"
+            )
+            self.num_sample = _n
+            self.knn = 1
+
+        # save the eigenvectors solution on the sub-sampled graph, do not propagate to full graph yet
+        self.subgraph_eigen_vector, self.eigen_value, self.subgraph_indices = nystrom_ncut(
+            features,
+            num_eig=self.num_eig,
+            num_sample=self.num_sample,
+            sample_method=self.sample_method,
+            precomputed_sampled_indices=precomputed_sampled_indices,
+            distance=self.distance,
+            affinity_focal_gamma=self.affinity_focal_gamma,
+            indirect_connection=self.indirect_connection,
+            indirect_pca_dim=self.indirect_pca_dim,
+            device=self.device,
+            eig_solver=self.eig_solver,
+            normalize_features=self.normalize_features,
+            matmul_chunk_size=self.matmul_chunk_size,
+            verbose=self.verbose,
+            no_propagation=True,
+            move_output_to_cpu=self.move_output_to_cpu,
+        )
+        self.subgraph_features = features[self.subgraph_indices]
+        return self
+
+    def transform(self, features: torch.Tensor, knn: int = None):
+        """Transform new features using the fitted Nystrom Normalized Cut.
+        Args:
+            features (torch.Tensor): new features, shape (n_samples, n_features)
+            knn (int): number of KNN for propagating eigenvectors from subgraph to full graph,
+        Returns:
+            (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
+            (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
+        """
+
+        knn = self.knn if knn is None else knn
+
+        # propagate eigenvectors from subgraph to full graph
+        eigen_vector = propagate_knn(
+            self.subgraph_eigen_vector,
+            features,
+            self.subgraph_features,
+            knn,
+            distance=self.distance,
+            chunk_size=self.matmul_chunk_size,
+            device=self.device,
+            use_tqdm=self.verbose,
+            move_output_to_cpu=self.move_output_to_cpu,
+        )
+        if self.make_orthogonal:
+            eigen_vector = gram_schmidt(eigen_vector)
+        return eigen_vector, self.eigen_value
+
+    def fit_transform(self,
+                      features: torch.Tensor,
+                      precomputed_sampled_indices: torch.Tensor = None
+                      ):
+        """
+        Args:
+            features (torch.Tensor): input features, shape (n_samples, n_features)
+            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+                override the sample_method, if not None
+
+        Returns:
+            (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
+            (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
+        """
+        return self.fit(features, precomputed_sampled_indices=precomputed_sampled_indices).transform(features)
+
+
+def nystrom_ncut(
+    features: torch.Tensor,
+    num_eig: int = 100,
+    num_sample: int = 10000,
+    knn: int = 10,
+    sample_method: Literal["farthest", "random"] = "farthest",
+    precomputed_sampled_indices: torch.Tensor = None,
+    distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
+    affinity_focal_gamma: float = 1.0,
+    indirect_connection: bool = True,
+    indirect_pca_dim: int = 100,
+    device: str = None,
+    eig_solver: Literal["svd_lowrank", "lobpcg", "svd", "eigh"] = "svd_lowrank",
+    normalize_features: bool = None,
+    matmul_chunk_size: int = 8096,
+    make_orthogonal: bool = True,
+    verbose: bool = False,
+    no_propagation: bool = False,
+    move_output_to_cpu: bool = False,
+):
+    """PyTorch implementation of Faster Nystrom Normalized cut.
+    Args:
+        features (torch.Tensor): feature matrix, shape (n_samples, n_features)
+        num_eig (int): default 100, number of top eigenvectors to return
+        num_sample (int): default 10000, number of samples for Nystrom-like approximation
+        knn (int): default 10, number of KNN for propagating eigenvectors from subgraph to full graph,
+            smaller knn will result in more sharp eigenvectors,
+        sample_method (str): sample method, 'farthest' (default) or 'random'
+            'farthest' is recommended for better approximation
+        precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+            override the sample_method, if not None
+        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'
+        affinity_focal_gamma (float): affinity matrix parameter, lower t reduce the weak edge weights,
+            resulting in more sharp eigenvectors, default 1.0
+        indirect_connection (bool): include indirect connection in the subgraph, default True
+        indirect_pca_dim (int): default 100, PCA dimension to reduce the node dimension, only applied to
+            the not sampled nodes, not applied to the sampled nodes
+        device (str): device to use for computation, if None, will not change device
+            a good practice is to pass features by CPU since it's usually large,
+            and move subgraph affinity to GPU to speed up eigenvector computation
+        eig_solver (str): eigen decompose solver, 'svd_lowrank' (default), 'lobpcg', 'svd', 'eigh'
+            'svd_lowrank' is recommended for large scale graph, it's the fastest
+            they correspond to torch.svd_lowrank, torch.lobpcg, torch.svd, torch.linalg.eigh
+        normalize_features (bool): normalize input features before computing affinity matrix,
+            default 'None' is True for cosine distance, False for euclidean distance and rbf
+        matmul_chunk_size (int): chunk size for matrix multiplication
+            large matrix multiplication is chunked to reduce memory usage,
+            smaller chunk size will reduce memory usage but slower computation, default 8096
+        make_orthogonal (bool): make eigenvectors orthogonal after propagation, default True
+        verbose (bool): show progress bar when propagating eigenvectors from subgraph to full graph
+        no_propagation (bool): if True, skip the eigenvector propagation step, only return the subgraph eigenvectors
+        move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
+    Returns:
+        (torch.Tensor): eigenvectors, shape (n_samples, num_eig)
+        (torch.Tensor): eigenvalues, sorted in descending order, shape (num_eig,)
+        (torch.Tensor): sampled_indices used by Nystrom-like approximation subgraph, shape (num_sample,)
+    """
+
+    # check if features dimension greater than num_eig
+    if eig_solver in ["svd_lowrank", "lobpcg"]:
+        assert features.shape[0] > (
+            num_eig * 2
+        ), "number of nodes should be greater than 2*num_eig"
+    if eig_solver in ["svd", "eigh"]:
+        assert (
+            features.shape[0] > num_eig
+        ), "number of nodes should be greater than num_eig"
+
+    assert distance in ["cosine", "euclidean", "rbf"], "distance should be 'cosine', 'euclidean', 'rbf'"
+
+    if normalize_features:
+        # features need to be normalized for affinity matrix computation (cosine distance)
+        features = torch.nn.functional.normalize(features, dim=-1)
+
+    if precomputed_sampled_indices is not None:
+        sampled_indices = precomputed_sampled_indices
+    else:
+        sampled_indices = run_subgraph_sampling(
+            features,
+            num_sample=num_sample,
+            sample_method=sample_method,
+        )
+
+    sampled_features = features[sampled_indices]
+    # move subgraph gpu to speed up
+    original_device = sampled_features.device
+    device = original_device if device is None else device
+    sampled_features = sampled_features.to(device)
+
+    # compute affinity matrix on subgraph
+    A = affinity_from_features(
+        sampled_features,
+        affinity_focal_gamma=affinity_focal_gamma,
+        distance=distance,
+    )
+
+    # check if all nodes are sampled, if so, no need for Nystrom approximation
+    not_sampled = torch.full((features.shape[0],), True)
+    not_sampled[sampled_indices] = False
+    _n_not_sampled = not_sampled.sum()
+
+    if _n_not_sampled == 0:
+        # if sampled all nodes, no need for nyström approximation
+        eigen_vector, eigen_value = ncut(A, num_eig, eig_solver=eig_solver)
+        return eigen_vector, eigen_value, sampled_indices
+
+    # 1) PCA to reduce the node dimension for the not sampled nodes
+    # 2) compute indirect connection on the PC nodes
+    if _n_not_sampled > 0 and indirect_connection:
+        indirect_pca_dim = min(indirect_pca_dim, *features.shape)
+        U, S, V = torch.pca_lowrank(features[not_sampled].T, q=indirect_pca_dim)
+        S = S / math.sqrt(_n_not_sampled)
+        feature_B_T = U @ torch.diag(S)
+        feature_B = feature_B_T.T
+        feature_B = feature_B.to(device)
+
+        B = affinity_from_features(
+            sampled_features,
+            feature_B,
+            affinity_focal_gamma=affinity_focal_gamma,
+            distance=distance,
+            fill_diagonal=False,
+        )
+        # P is 1-hop random walk matrix
+        B_row = B / B.sum(dim=1, keepdim=True)
+        B_col = B / B.sum(dim=0, keepdim=True)
+        P = B_row @ B_col.T
+        P = (P + P.T) / 2
+        # fill diagonal with 0
+        P[torch.arange(P.shape[0]), torch.arange(P.shape[0])] = 0
+        A = A + P
+
+    # compute normalized cut on the subgraph
+    eigen_vector, eigen_value = ncut(A, num_eig, eig_solver=eig_solver)
+    eigen_vector = eigen_vector.to(dtype=features.dtype, device=original_device)
+    eigen_value = eigen_value.to(dtype=features.dtype, device=original_device)
+
+    if no_propagation:
+        return eigen_vector, eigen_value, sampled_indices
+
+    # propagate eigenvectors from subgraph to full graph
+    eigen_vector = propagate_knn(
+        eigen_vector,
+        features,
+        sampled_features,
+        knn,
+        distance=distance,
+        chunk_size=matmul_chunk_size,
+        device=device,
+        use_tqdm=verbose,
+        move_output_to_cpu=move_output_to_cpu,
+    )
+
+    # post-hoc orthogonalization
+    if make_orthogonal:
+        eigen_vector = gram_schmidt(eigen_vector)
+
+    return eigen_vector, eigen_value, sampled_indices
+
+
+def normalized_affinity_transform(D: torch.Tensor, affinity_focal_gamma: float):
+    """Compute Laplacian-normalized affinity matrix from input features.
+
+    Args:
+        features (torch.Tensor): input features, shape (n_samples, n_features)
+        features_B (torch.Tensor, optional): optional, if not None, compute affinity between two features
+        affinity_focal_gamma (float): affinity matrix parameter, lower t reduce the edge weights
+            on weak connections, default 1.0
+        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
+        normalize_features (bool): normalize input features before computing affinity matrix
+
+    Returns:
+        (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
+    """
+    # make sure D is symmetric
+    D = (D + D.T) / 2
+    A = torch.exp(-D / affinity_focal_gamma)
+
+    # symmetrical normalization; A = D^(-1/2) A D^(-1/2)
+    D = A.sum(dim=-1).detach().clone()
+    A /= torch.sqrt(D)[:, None]
+    A /= torch.sqrt(D)[None, :]
+    return A
+
+
+def ncut(
+    A: torch.Tensor,
+    num_eig: int = 100,
+    eig_solver: Literal["svd_lowrank", "lobpcg", "svd", "eigh"] = "svd_lowrank",
+):
+    """PyTorch implementation of Normalized cut without Nystrom-like approximation.
+
+    Args:
+        A (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
+        num_eig (int): number of eigenvectors to return
+        eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh']
+
+    Returns:
+        (torch.Tensor): eigenvectors corresponding to the eigenvalues, shape (n_samples, num_eig)
+        (torch.Tensor): eigenvalues of the eigenvectors, sorted in descending order
+    """
+    # make sure A is symmetric
+    A = (A + A.T) / 2
+
+    # symmetrical normalization; A = D^(-1/2) A D^(-1/2)
+    D = A.sum(dim=-1).detach().clone()
+    A /= torch.sqrt(D)[:, None]
+    A /= torch.sqrt(D)[None, :]
+
+    # compute eigenvectors
+    eigen_vector, eigen_value = solve_eig(A, num_eig, eig_solver)
+
+    if eigen_value.min() < 0:
+        logging.warning(
+            "negative eigenvalues detected, please make sure the affinity matrix is positive definite"
+        )
+
+    return eigen_vector, eigen_value
+
+
+def gram_schmidt(matrix):
+    """Orthogonalize a matrix column-wise using the Gram-Schmidt process.
+
+    Args:
+        matrix (torch.Tensor): A matrix to be orthogonalized (m x n).
+            the second dimension is orthogonalized
+    Returns:
+        torch.Tensor: Orthogonalized matrix (m x n).
+    """
+
+    # Get the number of rows (m) and columns (n) of the input matrix
+    m, n = matrix.shape
+
+    # Create an empty matrix to store the orthogonalized columns
+    orthogonal_matrix = torch.zeros((m, n), dtype=matrix.dtype)
+
+    for i in range(n):
+        # Start with the i-th column of the input matrix
+        vec = matrix[:, i]
+
+        for j in range(i):
+            # Subtract the projection of vec onto the j-th orthogonal column
+            proj = torch.dot(orthogonal_matrix[:, j], matrix[:, i]) / torch.dot(
+                orthogonal_matrix[:, j], orthogonal_matrix[:, j]
+            )
+            vec = vec - proj * orthogonal_matrix[:, j]
+
+        # Store the orthogonalized vector
+        orthogonal_matrix[:, i] = vec / torch.norm(vec)
+
+    return orthogonal_matrix
+
+
+def correct_rotation(eigen_vector):
+    # correct the random rotation (flipping sign) of eigenvectors
+    rand_w = torch.ones(
+        eigen_vector.shape[0], device=eigen_vector.device, dtype=eigen_vector.dtype
+    )
+    s = rand_w[None, :] @ eigen_vector
+    s = s.sign()
+    return eigen_vector * s
+
+
+# Multiclass Spectral Clustering, SX Yu, J Shi, 2003
+def _discretisation_eigenvector(eigen_vector):
+    # Function that discretizes rotated eigenvectors
+    n, k = eigen_vector.shape
+
+    # Find the maximum index along each row
+    _, J = torch.max(eigen_vector, dim=1)
+    Y = torch.zeros(n, k, device=eigen_vector.device).scatter_(1, J.unsqueeze(1), 1)
+
+    return Y
+
+
+def kway_ncut(eigen_vectors: torch.Tensor, max_iter=300, return_rotation=False):
+    """Multiclass Spectral Clustering, SX Yu, J Shi, 2003
+
+    Args:
+        eigen_vectors (torch.Tensor): continuous eigenvectors from NCUT, shape (n, k)
+        max_iter (int, optional): Maximum number of iterations.
+
+    Returns:
+        torch.Tensor: Discretized eigenvectors, shape (n, k), each row is a one-hot vector.
+    """
+    # Normalize eigenvectors
+    n, k = eigen_vectors.shape
+    vm = torch.sqrt(torch.sum(eigen_vectors ** 2, dim=1))
+    eigen_vectors = eigen_vectors / vm.unsqueeze(1)
+
+    # Initialize R matrix with the first column from a random row of EigenVectors
+    R = torch.zeros(k, k, device=eigen_vectors.device)
+    R[:, 0] = eigen_vectors[torch.randint(0, n, (1,))].squeeze()
+
+    # Loop to populate R with k orthogonal directions
+    c = torch.zeros(n, device=eigen_vectors.device)
+    for j in range(1, k):
+        c += torch.abs(eigen_vectors @ R[:, j - 1])
+        _, i = torch.min(c, dim=0)
+        R[:, j] = eigen_vectors[i]
+
+    # Iterative optimization loop
+    last_objective_value = 0
+    exit_loop = False
+    nb_iterations_discretisation = 0
+
+    while not exit_loop:
+        nb_iterations_discretisation += 1
+
+        # Discretize the projected eigenvectors
+        eigenvectors_discrete = _discretisation_eigenvector(eigen_vectors @ R)
+
+        # SVD decomposition
+        U, S, Vh = torch.linalg.svd(eigenvectors_discrete.T @ eigen_vectors, full_matrices=False)
+        V = Vh.T
+
+        # Compute the Ncut value
+        ncut_value = 2 * (n - torch.sum(S))
+
+        # Check for convergence
+        if torch.abs(ncut_value - last_objective_value) < torch.finfo(
+            torch.float32).eps or nb_iterations_discretisation > max_iter:
+            exit_loop = True
+        else:
+            last_objective_value = ncut_value
+            R = V @ U.T
+
+    if return_rotation:
+        return eigenvectors_discrete, R
+
+    return eigenvectors_discrete
+
+
+def axis_align(eigen_vectors, max_iter=300):
+    return kway_ncut(eigen_vectors, max_iter=max_iter, return_rotation=True)
+
+
+## for backward compatibility ##
+
+try:
+
+    from .propagation_utils import (
+        propagate_nearest,
+        propagate_eigenvectors,
+        quantile_normalize,
+        quantile_min_max,
+        farthest_point_sampling,
+    )
+    from .visualize_utils import (
+        eigenvector_to_rgb,
+        rgb_from_tsne_3d,
+        rgb_from_umap_sphere,
+        rgb_from_tsne_2d,
+        rgb_from_umap_3d,
+        rgb_from_umap_2d,
+        rotate_rgb_cube,
+        convert_to_lab_color,
+        _transform_heatmap,
+        _clean_mask,
+        get_mask,
+    )
+
+except ImportError:
+    print("some of viualization and nystrom_utils are not imported")