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

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/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(
@@ -42,7 +43,7 @@ def run_subgraph_sampling(
             sampled_indices = torch.randperm(features.shape[0])[:num_sample]
         else:
             raise ValueError("sample_method should be 'farthest' or 'random'")
-    return sampled_indices
+    return sampled_indices.to(features.device)
 
 
 def farthest_point_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
@@ -154,9 +139,8 @@ def propagate_knn(
     knn: int = 10,
     distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
     affinity_focal_gamma: float = 1.0,
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     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)
 
@@ -232,16 +206,14 @@ def propagate_nearest(
     inp_features: torch.Tensor,
     subgraph_features: torch.Tensor,
     distance: Literal["cosine", "euclidean", "rbf"] = "cosine",
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     device: str = None,
     move_output_to_cpu: bool = False,
 ):
     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:
@@ -283,10 +254,8 @@ def propagate_eigenvectors(
         knn (int): number of KNN to propagate eigenvectors, default 3
         num_sample (int): number of samples for subgraph sampling, default 50000
         sample_method (str): sample method, 'farthest' (default) or 'random'
-        chunk_size (int): chunk size for matrix multiplication, default 8096
+        chunk_size (int): chunk size for matrix multiplication, default 8192
         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/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
@@ -503,9 +420,8 @@ def propagate_rgb_color(
     knn: int = 10,
     num_sample: int = 1000,
     sample_method: Literal["farthest", "random"] = "farthest",
-    chunk_size: int = 8096,
+    chunk_size: int = 8192,
     device: str = None,
-    use_tqdm: bool = False,
 ):
     """Propagate RGB color to new nodes using KNN.
     Args:
@@ -515,10 +431,8 @@ def propagate_rgb_color(
         knn (int): number of KNN to propagate RGB color, default 1
         num_sample (int): number of samples for subgraph sampling, default 50000
         sample_method (str): sample method, 'farthest' (default) or 'random'
-        chunk_size (int): chunk size for matrix multiplication, default 8096
+        chunk_size (int): chunk size for matrix multiplication, default 8192
         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.dist-info → nystrom_ncut-0.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nystrom_ncut
-Version: 0.0.1
+Version: 0.0.3
 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.3.dist-info/RECORD

@@ -0,0 +1,11 @@
+nystrom_ncut/__init__.py,sha256=Cww-_OsyQHLKpgw_Wh28_tUOvIMMr7Ey8w-tH7v99xQ,452
+nystrom_ncut/common.py,sha256=qdR_JwknT9H1Cv5LopwdwZfORFx-O8MLiRI6ZF1Qohc,558
+nystrom_ncut/ncut_pytorch.py,sha256=8LfznDwhq-WL_vQxbFBFLSzymg9SEDti_zzf9QQLnrA,11651
+nystrom_ncut/nystrom.py,sha256=Jo-P-2vnYk8yEZinGZnN3jHMiiB5AueoaLAYK4OmRqE,8604
+nystrom_ncut/propagation_utils.py,sha256=pigecB0rAmlbCoMNb8zhCyyNwh3QzkxXEnaBsDRE_ns,12136
+nystrom_ncut/visualize_utils.py,sha256=oNaDz_Xn12g3knEZZTb-QWVN-wTrnCNE5gn9cu8Xl_U,18569
+nystrom_ncut-0.0.3.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.0.3.dist-info/METADATA,sha256=yh1pDFHUL2Z4WPVnuQyeHjgSEokUtpB6OPzPMxEClsM,6058
+nystrom_ncut-0.0.3.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+nystrom_ncut-0.0.3.dist-info/top_level.txt,sha256=j7g_j0S048EvguFFnGgD5Ewd3r2H6klsxd5A4dd-wHw,13
+nystrom_ncut-0.0.3.dist-info/RECORD,,