nystrom-ncut 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

nystrom_ncut/ncut_pytorch.py

@@ -1,46 +1,112 @@
-# %%
 import logging
-import math
-from typing import Literal
+from typing import Literal, Tuple
 
 import torch
+import torch.nn.functional as Fn
 
 from .nystrom import (
-    solve_eig
+    EigSolverOptions,
+    OnlineKernel,
+    OnlineNystrom,
+    solve_eig,
 )
 from .propagation_utils import (
-    run_subgraph_sampling,
-    propagate_knn,
     affinity_from_features,
+    run_subgraph_sampling,
 )
 
 
-class NCUT:
+DistanceOptions = Literal["cosine", "euclidean", "rbf"]
+
+
+class LaplacianKernel(OnlineKernel):
+    def __init__(
+        self,
+        affinity_focal_gamma: float,
+        distance: DistanceOptions,
+        eig_solver: EigSolverOptions,
+    ):
+        self.affinity_focal_gamma = affinity_focal_gamma
+        self.distance: DistanceOptions = distance
+        self.eig_solver: EigSolverOptions = eig_solver
+
+        # Anchor matrices
+        self.anchor_features: torch.Tensor = None               # [n x d]
+        self.A: torch.Tensor = None                             # [n x n]
+        self.Ainv: torch.Tensor = None                          # [n x n]
+
+        # Updated matrices
+        self.a_r: torch.Tensor = None                           # [n]
+        self.b_r: torch.Tensor = None                           # [n]
+
+    def fit(self, features: torch.Tensor) -> None:
+        self.anchor_features = features                         # [n x d]
+        self.A = affinity_from_features(
+            self.anchor_features,                               # [n x d]
+            affinity_focal_gamma=self.affinity_focal_gamma,
+            distance=self.distance,
+        )                                                       # [n x n]
+        U, L = solve_eig(
+            self.A,
+            num_eig=features.shape[-1] + 1,
+            eig_solver=self.eig_solver,
+        )                                                       # [n x (d + 1)], [d + 1]
+        self.Ainv = U @ torch.diag(1 / L) @ U.mT                # [n x n]
+        self.a_r = torch.sum(self.A, dim=-1)                    # [n]
+        self.b_r = torch.zeros_like(self.a_r)                   # [n]
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        B = affinity_from_features(
+            self.anchor_features,                               # [n x d]
+            features,                                           # [m x d]
+            affinity_focal_gamma=self.affinity_focal_gamma,
+            distance=self.distance,
+        )                                                       # [n x m]
+        b_r = torch.sum(B, dim=-1)                              # [n]
+        b_c = torch.sum(B, dim=-2)                              # [m]
+        self.b_r = self.b_r + b_r                               # [n]
+
+        rowscale = self.a_r + self.b_r                          # [n]
+        colscale = b_c + B.mT @ self.Ainv @ self.b_r            # [m]
+        scale = (rowscale[:, None] * colscale) ** -0.5          # [n x m]
+        return (B * scale).mT                                   # [m x n]
+
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:
+        rowscale = self.a_r + self.b_r                          # [n]
+        if features is None:
+            B = self.A                                          # [n x n]
+            colscale = rowscale                                 # [n]
+        else:
+            B = affinity_from_features(
+                self.anchor_features,                           # [n x d]
+                features,                                       # [m x d]
+                affinity_focal_gamma=self.affinity_focal_gamma,
+                distance=self.distance,
+            )                                                   # [n x m]
+            b_c = torch.sum(B, dim=-2)                          # [m]
+            colscale = b_c + B.mT @ self.Ainv @ self.b_r        # [m]
+        scale = (rowscale[:, None] * colscale) ** -0.5          # [n x m]
+        return (B * scale).mT                                   # [m x n]
+
+
+class NCUT(OnlineNystrom):
     """Nystrom Normalized Cut for large scale graph."""
 
     def __init__(
         self,
-        num_eig: int = 100,
-        knn: int = 10,
+        n_components: int = 100,
         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",
+        distance: DistanceOptions = "cosine",
+        eig_solver: EigSolverOptions = "svd_lowrank",
         normalize_features: bool = None,
-        matmul_chunk_size: int = 8096,
-        make_orthogonal: bool = False,
-        verbose: bool = False,
+        move_output_to_cpu: bool = False,
+        chunk_size: int = 8192,
     ):
         """
         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.
+            n_components (int): number of top eigenvectors to return
             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,
@@ -48,140 +114,101 @@ class NCUT:
             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)
+            move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
+            chunk_size (int): chunk size for large-scale matrix multiplication
         """
-        self.num_eig = num_eig
+        OnlineNystrom.__init__(
+            self,
+            n_components=n_components,
+            kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
+            eig_solver=eig_solver,
+            chunk_size=chunk_size,
+        )
         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
-        """
+
+        self.move_output_to_cpu = move_output_to_cpu
+        self.chunk_size = chunk_size
+
+    def _fit_helper(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         _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"
+                f"NCUT nystrom num_sample is larger than number of input samples, nyström approximation is not needed, setting num_sample={_n}"
             )
             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.
+        # check if features dimension greater than num_eig
+        if self.eig_solver in ["svd_lowrank", "lobpcg"]:
+            assert (
+                _n >= self.n_components * 2
+            ), "number of nodes should be greater than 2*num_eig"
+        elif self.eig_solver in ["svd", "eigh"]:
+            assert (
+                _n >= self.n_components
+            ), "number of nodes should be greater than num_eig"
+
+        assert self.distance in ["cosine", "euclidean", "rbf"], "distance should be 'cosine', 'euclidean', 'rbf'"
+
+        if self.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=self.num_sample,
+                sample_method=self.sample_method,
+            )
+        sampled_features = features[sampled_indices]
+        OnlineNystrom.fit(self, sampled_features)
+
+        _n_not_sampled = _n - len(sampled_features)
+        if _n_not_sampled > 0:
+            unsampled_indices = torch.full((_n,), True, device=features.device).scatter_(0, sampled_indices, False)
+            unsampled_features = features[unsampled_indices]
+            V_unsampled, _ = OnlineNystrom.update(self, unsampled_features)
+        else:
+            unsampled_indices = V_unsampled = None
+        return unsampled_indices, V_unsampled
+
+    def fit(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor = None,
+    ):
+        """Fit Nystrom Normalized Cut on the input features.
         Args:
-            features (torch.Tensor): new features, shape (n_samples, n_features)
-            knn (int): number of KNN for propagating eigenvectors from subgraph to full graph,
+            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,)
+            (NCUT): self
         """
+        NCUT._fit_helper(self, features, precomputed_sampled_indices)
+        return self
 
-        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
-                      ):
+    def fit_transform(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Args:
             features (torch.Tensor): input features, shape (n_samples, n_features)
@@ -192,285 +219,19 @@ class NCUT:
             (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)
+        unsampled_indices, V_unsampled = NCUT._fit_helper(self, features, precomputed_sampled_indices)
+        V_sampled, L = OnlineNystrom.transform(self)
 
-    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
+        if unsampled_indices is not None:
+            V = torch.zeros((len(unsampled_indices), self.n_components), device=features.device)
+            V[~unsampled_indices] = V_sampled
+            V[unsampled_indices] = V_unsampled
+        else:
+            V = V_sampled
+        return V, L
 
 
-def kway_ncut(eigen_vectors: torch.Tensor, max_iter=300, return_rotation=False):
+def axis_align(eigen_vectors: torch.Tensor, max_iter=300):
     """Multiclass Spectral Clustering, SX Yu, J Shi, 2003
 
     Args:
@@ -482,80 +243,36 @@ def kway_ncut(eigen_vectors: torch.Tensor, max_iter=300, return_rotation=False):
     """
     # Normalize eigenvectors
     n, k = eigen_vectors.shape
-    vm = torch.sqrt(torch.sum(eigen_vectors ** 2, dim=1))
-    eigen_vectors = eigen_vectors / vm.unsqueeze(1)
+    eigen_vectors = Fn.normalize(eigen_vectors, p=2, dim=-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()
+    R = torch.empty((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]
+    for i in range(1, k):
+        c += torch.abs(eigen_vectors @ R[i - 1])
+        R[i] = eigen_vectors[torch.argmin(c, dim=0)]
 
     # Iterative optimization loop
-    last_objective_value = 0
-    exit_loop = False
-    nb_iterations_discretisation = 0
+    eps = torch.finfo(torch.float32).eps
+    prev_objective = torch.inf
+    for _ in range(max_iter):
+        # Discretize the projected eigenvectors
+        idx = torch.argmax(eigen_vectors @ R.mT, dim=-1)
+        M = torch.zeros((k, k)).index_add_(0, idx, eigen_vectors)
 
-    while not exit_loop:
-        nb_iterations_discretisation += 1
+        # Compute the NCut value
+        objective = torch.norm(M)
 
-        # Discretize the projected eigenvectors
-        eigenvectors_discrete = _discretisation_eigenvector(eigen_vectors @ R)
+        # Check for convergence
+        if torch.abs(objective - prev_objective) < eps:
+            break
+        prev_objective = objective
 
         # SVD decomposition
-        U, S, Vh = torch.linalg.svd(eigenvectors_discrete.T @ eigen_vectors, full_matrices=False)
-        V = Vh.T
+        U, S, Vh = torch.linalg.svd(M, full_matrices=False)
+        R = U @ Vh
 
-        # 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")
+    return Fn.one_hot(idx, num_classes=k).to(torch.float), R