nystrom-ncut 0.0.1__tar.gz → 0.0.2__tar.gz

{nystrom_ncut-0.0.1/src/nystrom_ncut.egg-info → nystrom_ncut-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nystrom_ncut
-Version: 0.0.1
+Version: 0.0.2
 Summary: Normalized Cut and Nyström Approximation
 Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
 Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nystrom_ncut"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
     { name = "Huzheng Yang", email = "huze.yann@gmail.com" },
     { name = "Wentinn Liao", email = "wentinn.liao@gmail.com" },

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/requirements.txt

@@ -3,4 +3,5 @@ scikit-learn
 umap-learn
 fpsample>=0.3.2
 pycolormap-2d
-tqdm
+tqdm
+torch

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/src/nystrom_ncut/__init__.py

@@ -1,4 +1,7 @@
-from .ncut_pytorch import NCUT
+from .ncut_pytorch import (
+    NCUT,
+    axis_align,
+)
 from .propagation_utils import (
     affinity_from_features,
     propagate_eigenvectors,
@@ -6,7 +9,6 @@ from .propagation_utils import (
     quantile_normalize,
 )
 from .visualize_utils import (
-    eigenvector_to_rgb,
     rgb_from_tsne_3d,
     rgb_from_umap_sphere,
     rgb_from_tsne_2d,
@@ -18,5 +20,3 @@ from .visualize_utils import (
     propagate_rgb_color,
     get_mask,
 )
-from .ncut_pytorch import nystrom_ncut, ncut
-from .ncut_pytorch import kway_ncut, axis_align

nystrom_ncut-0.0.2/src/nystrom_ncut/common.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as Fn
+
+
+def ceildiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> torch.Tensor:
+    numel = np.prod(x.shape[:-1])
+    n = min(n, numel)
+    random_indices = torch.randperm(numel)[:n]
+    _x = x.flatten(0, -2)[random_indices]
+    if torch.allclose(torch.norm(_x, **normalize_kwargs), torch.ones(n, device=x.device)):
+        return x
+    else:
+        return Fn.normalize(x, **normalize_kwargs)

nystrom_ncut-0.0.1/src/nystrom_ncut/new_ncut_pytorch.py → nystrom_ncut-0.0.2/src/nystrom_ncut/ncut_pytorch.py

@@ -2,6 +2,7 @@ import logging
 from typing import Literal, Tuple
 
 import torch
+import torch.nn.functional as Fn
 
 from .nystrom import (
     EigSolverOptions,
@@ -44,7 +45,6 @@ class LaplacianKernel(OnlineKernel):
             self.anchor_features,                               # [n x d]
             affinity_focal_gamma=self.affinity_focal_gamma,
             distance=self.distance,
-            fill_diagonal=False,
         )                                                       # [n x n]
         U, L = solve_eig(
             self.A,
@@ -61,7 +61,6 @@ class LaplacianKernel(OnlineKernel):
             features,                                           # [m x d]
             affinity_focal_gamma=self.affinity_focal_gamma,
             distance=self.distance,
-            fill_diagonal=False,
         )                                                       # [n x m]
         b_r = torch.sum(B, dim=-1)                              # [n]
         b_c = torch.sum(B, dim=-2)                              # [m]
@@ -83,7 +82,6 @@ class LaplacianKernel(OnlineKernel):
                 features,                                       # [m x d]
                 affinity_focal_gamma=self.affinity_focal_gamma,
                 distance=self.distance,
-                fill_diagonal=False,
             )                                                   # [n x m]
             b_c = torch.sum(B, dim=-2)                          # [m]
             colscale = b_c + B.mT @ self.Ainv @ self.b_r        # [m]
@@ -91,7 +89,7 @@ class LaplacianKernel(OnlineKernel):
         return (B * scale).mT                                   # [m x n]
 
 
-class NewNCUT(OnlineNystrom):
+class NCUT(OnlineNystrom):
     """Nystrom Normalized Cut for large scale graph."""
 
     def __init__(
@@ -211,7 +209,7 @@ class NewNCUT(OnlineNystrom):
         Returns:
             (NCUT): self
         """
-        NewNCUT._fit_helper(self, features, precomputed_sampled_indices)
+        NCUT._fit_helper(self, features, precomputed_sampled_indices)
         return self
 
     def fit_transform(
@@ -229,7 +227,7 @@ class NewNCUT(OnlineNystrom):
             (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
             (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
         """
-        unsampled_indices, V_unsampled = NewNCUT._fit_helper(self, features, precomputed_sampled_indices)
+        unsampled_indices, V_unsampled = NCUT._fit_helper(self, features, precomputed_sampled_indices)
         V_sampled, L = OnlineNystrom.transform(self)
 
         if unsampled_indices is not None:
@@ -239,3 +237,50 @@ class NewNCUT(OnlineNystrom):
         else:
             V = V_sampled
         return V, L
+
+
+def axis_align(eigen_vectors: torch.Tensor, max_iter=300):
+    """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
+    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.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 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
+    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)
+
+        # Compute the NCut value
+        objective = torch.norm(M)
+
+        # Check for convergence
+        if torch.abs(objective - prev_objective) < eps:
+            break
+        prev_objective = objective
+
+        # SVD decomposition
+        U, S, Vh = torch.linalg.svd(M, full_matrices=False)
+        R = U @ Vh
+
+    return Fn.one_hot(idx, num_classes=k).to(torch.float), R

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/src/nystrom_ncut/nystrom.py

@@ -2,6 +2,8 @@ from typing import Literal, Tuple
 
 import torch
 
+from .common import ceildiv
+
 
 EigSolverOptions = Literal["svd_lowrank", "lobpcg", "svd", "eigh"]
 
@@ -75,7 +77,7 @@ class OnlineNystrom:
         return U[:, :self.n_components], L[:self.n_components]                                      # [n x n_components], [n_components]
 
     def update(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        n_chunks = -(-len(features) // self.chunk_size)
+        n_chunks = ceildiv(len(features), self.chunk_size)
         if n_chunks > 1:
             """ Chunked version """
             chunks = torch.chunk(features, n_chunks, dim=0)
@@ -111,7 +113,7 @@ class OnlineNystrom:
         if features is None:
             VS = self.A @ self.transform_matrix                                                     # [n x n_components]
         else:
-            n_chunks = -(-len(features) // self.chunk_size)
+            n_chunks = ceildiv(len(features), self.chunk_size)
             if n_chunks > 1:
                 """ Chunked version """
                 chunks = torch.chunk(features, n_chunks, dim=0)

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/src/nystrom_ncut/propagation_utils.py

@@ -1,11 +1,12 @@
 import logging
-import math
 from typing import Literal
 
 import numpy as np
 import torch
 import torch.nn.functional as F
 
+from .common import ceildiv, lazy_normalize
+
 
 @torch.no_grad()
 def run_subgraph_sampling(
@@ -60,14 +61,12 @@ def farthest_point_sampling(
     # PCA to reduce the dimension
     if features.shape[1] > 8:
         u, s, v = torch.pca_lowrank(features, q=8)
-        _n = features.shape[0]
-        s /= math.sqrt(_n)
         features = u @ torch.diag(s)
 
     h = min(h, int(np.log2(features.shape[0])))
 
     kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(
-        features.cpu().numpy(), num_sample, h
+        features.numpy(force=True), num_sample, h
     ).astype(np.int64)
     return torch.from_numpy(kdline_fps_samples_idx)
 
@@ -76,26 +75,19 @@ def distance_from_features(
     features: torch.Tensor,
     features_B: torch.Tensor,
     distance: Literal["cosine", "euclidean", "rbf"],
-    fill_diagonal: bool,
 ):
     """Compute 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)
     """
     # compute distance matrix from input features
     if distance == "cosine":
-        if not check_if_normalized(features):
-            features = F.normalize(features, dim=-1)
-        if not check_if_normalized(features_B):
-            features_B = F.normalize(features_B, dim=-1)
+        features = lazy_normalize(features, dim=-1)
+        features_B = lazy_normalize(features_B, dim=-1)
         D = 1 - features @ features_B.T
     elif distance == "euclidean":
         D = torch.cdist(features, features_B, p=2)
@@ -105,8 +97,6 @@ def distance_from_features(
     else:
         raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
 
-    if fill_diagonal:
-        D[torch.arange(D.shape[0]), torch.arange(D.shape[0])] = 0
     return D
 
 
@@ -115,7 +105,6 @@ def affinity_from_features(
     features_B: torch.Tensor = None,
     affinity_focal_gamma: float = 1.0,
     distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
-    fill_diagonal: bool = True,
 ):
     """Compute affinity matrix from input features.
 
@@ -125,8 +114,6 @@ def affinity_from_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)
     """
@@ -134,12 +121,10 @@ def affinity_from_features(
 
     # if feature_B is not provided, compute affinity matrix on features x features
     # if feature_B is provided, compute affinity matrix on features x feature_B
-    if features_B is not None:
-        assert not fill_diagonal, "fill_diagonal should be False when feature_B is None"
     features_B = features if features_B is None else features_B
 
     # compute distance matrix from input features
-    D = distance_from_features(features, features_B, distance, fill_diagonal)
+    D = distance_from_features(features, features_B, distance)
 
     # torch.exp make affinity matrix positive definite,
     # lower affinity_focal_gamma reduce the weak edge weights
@@ -156,7 +141,6 @@ def propagate_knn(
     affinity_focal_gamma: float = 1.0,
     chunk_size: int = 8096,
     device: str = None,
-    use_tqdm: bool = False,
     move_output_to_cpu: bool = False,
 ):
     """A generic function to propagate new nodes using KNN.
@@ -169,8 +153,6 @@ def propagate_knn(
         distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'
         chunk_size (int): chunk size for matrix multiplication
         device (str): device to use for computation, if None, will not change device
-        use_tqdm (bool): show progress bar when propagating eigenvectors from subgraph to full graph
-
     Returns:
         torch.Tensor: propagated eigenvectors, shape (new_num_samples, D)
 
@@ -197,24 +179,16 @@ def propagate_knn(
     # used in nystrom_ncut
     # propagate eigen_vector from subgraph to full graph
     subgraph_output = subgraph_output.to(device)
-    V_list = []
-    iterator = range(0, inp_features.shape[0], chunk_size)
-    try:
-        assert use_tqdm
-        from tqdm import tqdm
-        iterator = tqdm(iterator, "propagate by KNN")
-    except (AssertionError, ImportError):
-        pass
 
-    subgraph_features = subgraph_features.to(device)
-    for i in iterator:
-        end = min(i + chunk_size, inp_features.shape[0])
-        _v = inp_features[i:end].to(device)
-        _A = affinity_from_features(subgraph_features, _v, affinity_focal_gamma, distance, False).mT
+    n_chunks = ceildiv(inp_features.shape[0], chunk_size)
+    V_list = []
+    for _v in torch.chunk(inp_features, n_chunks, dim=0):
+        _v = _v.to(device)
+        _A = affinity_from_features(subgraph_features, _v, affinity_focal_gamma, distance).mT
 
         if knn is not None:
             mask = torch.full_like(_A, True, dtype=torch.bool)
-            mask[torch.arange(end - i)[:, None], _A.topk(knn, dim=-1, largest=True).indices] = False
+            mask[torch.arange(len(_v))[:, None], _A.topk(knn, dim=-1, largest=True).indices] = False
             _A[mask] = 0.0
         _A = F.normalize(_A, p=1, dim=-1)
 
@@ -238,10 +212,8 @@ def propagate_nearest(
 ):
     device = subgraph_output.device if device is None else device
     if distance == 'cosine':
-        if not check_if_normalized(inp_features):
-            inp_features = F.normalize(inp_features, dim=-1)
-        if not check_if_normalized(subgraph_features):
-            subgraph_features = F.normalize(subgraph_features, dim=-1)
+        inp_features = lazy_normalize(inp_features, dim=-1)
+        subgraph_features = lazy_normalize(subgraph_features, dim=-1)
 
     # used in nystrom_tsne, equivalent to propagate_by_knn with knn=1
     # propagate tSNE from subgraph to full graph
@@ -250,7 +222,7 @@ def propagate_nearest(
     for i in range(0, inp_features.shape[0], chunk_size):
         end = min(i + chunk_size, inp_features.shape[0])
         _v = inp_features[i:end].to(device)
-        _A = -distance_from_features(subgraph_features, _v, distance, False).mT
+        _A = -distance_from_features(subgraph_features, _v, distance).mT
 
         # keep top1 for each row
         top_idx = _A.argmax(dim=-1).cpu()
@@ -273,7 +245,6 @@ def propagate_eigenvectors(
     sample_method: Literal["farthest", "random"],
     chunk_size: int,
     device: str,
-    use_tqdm: bool,
 ):
     """Propagate eigenvectors to new nodes using KNN. Note: this is equivalent to the class API `NCUT.tranform(new_features)`, expect for the sampling is re-done in this function.
     Args:
@@ -285,8 +256,6 @@ def propagate_eigenvectors(
         sample_method (str): sample method, 'farthest' (default) or 'random'
         chunk_size (int): chunk size for matrix multiplication, default 8096
         device (str): device to use for computation, if None, will not change device
-        use_tqdm (bool): show progress bar when propagating eigenvectors from subgraph to full graph
-
     Returns:
         torch.Tensor: propagated eigenvectors, shape (n_new_samples, num_eig)
 
@@ -319,21 +288,10 @@ def propagate_eigenvectors(
         knn=knn,
         chunk_size=chunk_size,
         device=device,
-        use_tqdm=use_tqdm,
     )
-
     return new_eigenvectors
 
 
-def check_if_normalized(x, n=1000):
-    """check if the input tensor is normalized (unit norm)"""
-    n = min(n, x.shape[0])
-    random_indices = torch.randperm(x.shape[0])[:n]
-    _x = x[random_indices]
-    flag = torch.allclose(torch.norm(_x, dim=-1), torch.ones(n, device=x.device))
-    return flag
-
-
 def quantile_min_max(x, q1=0.01, q2=0.99, n_sample=10000):
     if x.shape[0] > n_sample:
         np.random.seed(0)

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/src/nystrom_ncut/visualize_utils.py

@@ -6,11 +6,11 @@ import torch
 import torch.nn.functional as F
 from sklearn.base import BaseEstimator
 
+from .common import lazy_normalize
 from .propagation_utils import (
     run_subgraph_sampling,
     propagate_knn,
     propagate_eigenvectors,
-    check_if_normalized,
     quantile_min_max,
     quantile_normalize
 )
@@ -20,75 +20,6 @@ def _identity(X: torch.Tensor) -> torch.Tensor:
     return X
 
 
-def eigenvector_to_rgb(
-    eigen_vector: torch.Tensor,
-    method: Literal["tsne_2d", "tsne_3d", "umap_sphere", "umap_2d", "umap_3d"] = "tsne_3d",
-    num_sample: int = 1000,
-    perplexity: int = 150,
-    n_neighbors: int = 150,
-    min_distance: float = 0.1,
-    metric: Literal["cosine", "euclidean"] = "cosine",
-    device: str = None,
-    q: float = 0.95,
-    knn: int = 10,
-    seed: int = 0,
-):
-    """Use t-SNE or UMAP to convert eigenvectors (more than 3) to RGB color (3D RGB CUBE).
-
-    Args:
-        eigen_vector (torch.Tensor): eigenvectors, shape (n_samples, num_eig)
-        method (str): method to convert eigenvectors to RGB,
-            choices are: ['tsne_2d', 'tsne_3d', 'umap_sphere', 'umap_2d', 'umap_3d']
-        num_sample (int): number of samples for Nystrom-like approximation, increase for better approximation
-        perplexity (int): perplexity for t-SNE, increase for more global structure
-        n_neighbors (int): number of neighbors for UMAP, increase for more global structure
-        min_distance (float): minimum distance for UMAP
-        metric (str): distance metric, default 'cosine'
-        device (str): device to use for computation, if None, will not change device
-        q (float): quantile for RGB normalization, default 0.95. lower q results in more sharp colors
-        knn (int): number of KNN for propagating eigenvectors from subgraph to full graph,
-            smaller knn result in more sharp colors, default 1. knn>1 will smooth-out the embedding
-            in the t-SNE or UMAP space.
-        seed (int): random seed for t-SNE or UMAP
-
-    Examples:
-        >>> from ncut_pytorch import eigenvector_to_rgb
-        >>> X_3d, rgb = eigenvector_to_rgb(eigenvectors, method='tsne_3d')
-        >>> print(X_3d.shape, rgb.shape)
-        >>> # (10000, 3) (10000, 3)
-
-    Returns:
-        (torch.Tensor): t-SNE or UMAP embedding, shape (n_samples, 2) or (n_samples, 3)
-        (torch.Tensor): RGB color for each data sample, shape (n_samples, 3)
-    """
-    kwargs = {
-        "num_sample": num_sample,
-        "perplexity": perplexity,
-        "n_neighbors": n_neighbors,
-        "min_distance": min_distance,
-        "metric": metric,
-        "device": device,
-        "q": q,
-        "knn": knn,
-        "seed": seed,
-    }
-
-    if method == "tsne_2d":
-        embed, rgb = rgb_from_tsne_2d(eigen_vector, **kwargs)
-    elif method == "tsne_3d":
-        embed, rgb = rgb_from_tsne_3d(eigen_vector, **kwargs)
-    elif method == "umap_sphere":
-        embed, rgb = rgb_from_umap_sphere(eigen_vector, **kwargs)
-    elif method == "umap_2d":
-        embed, rgb = rgb_from_umap_2d(eigen_vector, **kwargs)
-    elif method == "umap_3d":
-        embed, rgb = rgb_from_umap_3d(eigen_vector, **kwargs)
-    else:
-        raise ValueError("method should be 'tsne_2d', 'tsne_3d' or 'umap_sphere'")
-
-    return embed, rgb
-
-
 def _rgb_with_dimensionality_reduction(
     features: torch.Tensor,
     num_sample: int,
@@ -126,7 +57,7 @@ def _rgb_with_dimensionality_reduction(
         move_output_to_cpu=True,
     ))
     rgb = rgb_func(X_nd, q)
-    return X_nd.numpy(force=True), rgb
+    return X_nd, rgb
 
 
 def rgb_from_tsne_2d(
@@ -138,7 +69,6 @@ def rgb_from_tsne_2d(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -169,7 +99,6 @@ def rgb_from_tsne_2d(
             "perplexity": perplexity,
         },
     )
-    
     return x2d, rgb
 
 
@@ -182,7 +111,6 @@ def rgb_from_tsne_3d(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -213,7 +141,6 @@ def rgb_from_tsne_3d(
             "perplexity": perplexity,
         },
     )
-
     return x3d, rgb
 
 
@@ -225,7 +152,6 @@ def rgb_from_cosine_tsne_3d(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -272,7 +198,6 @@ def rgb_from_cosine_tsne_3d(
             "perplexity": perplexity,
         },
     )
-    
     return x3d, rgb
 
 
@@ -286,7 +211,6 @@ def rgb_from_umap_2d(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -310,7 +234,6 @@ def rgb_from_umap_2d(
             "min_dist": min_dist,
         },
     )
-    
     return x2d, rgb
 
 
@@ -324,7 +247,6 @@ def rgb_from_umap_sphere(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -357,7 +279,6 @@ def rgb_from_umap_sphere(
         },
         transform_func=transform_func
     )
-    
     return x3d, rgb
 
 
@@ -371,7 +292,6 @@ def rgb_from_umap_3d(
     seed: int = 0,
     q: float = 0.95,
     knn: int = 10,
-    **kwargs: Any,
 ):
     """
     Returns:
@@ -395,7 +315,6 @@ def rgb_from_umap_3d(
             "min_dist": min_dist,
         },
     )
-    
     return x3d, rgb
 
 
@@ -417,13 +336,11 @@ def rotate_rgb_cube(rgb, position=1):
         torch.Tensor: RGB color space, shape (n_samples, 3)
     """
     assert position in range(0, 7), "position should be 0, 1, 2, 3, 4, 5, 6"
-    rotation_matrix = torch.tensor(
-        [
-            [0, 1, 0],
-            [0, 0, 1],
-            [1, 0, 0],
-        ]
-    ).float()
+    rotation_matrix = torch.tensor((
+        (0., 1., 0.),
+        (0., 0., 1.),
+        (1., 0., 0.),
+    ))
     n_mul = position % 3
     rotation_matrix = torch.matrix_power(rotation_matrix, n_mul)
     rgb = rgb @ rotation_matrix
@@ -505,7 +422,6 @@ def propagate_rgb_color(
     sample_method: Literal["farthest", "random"] = "farthest",
     chunk_size: int = 8096,
     device: str = None,
-    use_tqdm: bool = False,
 ):
     """Propagate RGB color to new nodes using KNN.
     Args:
@@ -517,8 +433,6 @@ def propagate_rgb_color(
         sample_method (str): sample method, 'farthest' (default) or 'random'
         chunk_size (int): chunk size for matrix multiplication, default 8096
         device (str): device to use for computation, if None, will not change device
-        use_tqdm (bool): show progress bar when propagating RGB color from subgraph to full graph
-
     Returns:
         torch.Tensor: propagated RGB color for each data sample, shape (n_new_samples, 3)
 
@@ -538,7 +452,6 @@ def propagate_rgb_color(
         sample_method=sample_method,
         chunk_size=chunk_size,
         device=device,
-        use_tqdm=use_tqdm,
     )
 
 
@@ -627,9 +540,7 @@ def get_mask(
     """
 
     # normalize the eigenvectors to unit norm, to compute cosine similarity
-    if not check_if_normalized(all_eigvecs.reshape(-1, all_eigvecs.shape[-1])):
-        all_eigvecs = F.normalize(all_eigvecs, p=2, dim=-1)
-
+    all_eigvecs = lazy_normalize(all_eigvecs, p=2, dim=-1)
     prompt_eigvec = F.normalize(prompt_eigvec, p=2, dim=-1)
 
     # compute the cosine similarity
@@ -642,7 +553,7 @@ def get_mask(
     heatmap = _transform_heatmap(heatmap, gamma=gamma)
 
     masks = heatmap > threshold
-    masks = masks.cpu().numpy().astype(np.uint8)
+    masks = masks.numpy(force=True).astype(np.uint8)
 
     if denoise:
         cleaned_masks = []

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2/src/nystrom_ncut.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nystrom_ncut
-Version: 0.0.1
+Version: 0.0.2
 Summary: Normalized Cut and Nyström Approximation
 Author-email: Huzheng Yang <huze.yann@gmail.com>, Wentinn Liao <wentinn.liao@gmail.com>
 Project-URL: Documentation, https://github.com/JophiArcana/Nystrom-NCUT/

{nystrom_ncut-0.0.1 → nystrom_ncut-0.0.2}/src/nystrom_ncut.egg-info/SOURCES.txt

@@ -4,12 +4,13 @@ README.md
 pyproject.toml
 requirements.txt
 src/nystrom_ncut/__init__.py
+src/nystrom_ncut/common.py
 src/nystrom_ncut/ncut_pytorch.py
-src/nystrom_ncut/new_ncut_pytorch.py
 src/nystrom_ncut/nystrom.py
 src/nystrom_ncut/propagation_utils.py
 src/nystrom_ncut/visualize_utils.py
 src/nystrom_ncut.egg-info/PKG-INFO
 src/nystrom_ncut.egg-info/SOURCES.txt
 src/nystrom_ncut.egg-info/dependency_links.txt
-src/nystrom_ncut.egg-info/top_level.txt
+src/nystrom_ncut.egg-info/top_level.txt
+tests/test.py