nystrom-ncut 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

nystrom_ncut/__init__.py

@@ -2,13 +2,16 @@ from .nystrom import (
     NCut,
     axis_align,
 )
-from .propagation_utils import (
+from .distance_utils import (
     distance_from_features,
     affinity_from_features,
-    extrapolate_knn_with_subsampling,
-    extrapolate_knn,
+)
+from .sampling_utils import (
+    SampleConfig,
 )
 from .visualize_utils import (
+    extrapolate_knn,
+    extrapolate_knn_with_subsampling,
     rgb_from_tsne_3d,
     rgb_from_umap_sphere,
     rgb_from_tsne_2d,

nystrom_ncut/common.py

@@ -1,14 +1,10 @@
-from typing import Any, Literal
+from typing import Any
 
 import numpy as np
 import torch
 import torch.nn.functional as Fn
 
 
-DistanceOptions = Literal["cosine", "euclidean", "rbf"]
-SampleOptions = Literal["farthest", "random"]
-
-
 def ceildiv(a: int, b: int) -> int:
     return -(-a // b)
 
@@ -24,16 +20,7 @@ def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> t
         return Fn.normalize(x, **normalize_kwargs)
 
 
-def to_euclidean(x: torch.Tensor, disttype: DistanceOptions) -> torch.Tensor:
-    if disttype == "cosine":
-        return lazy_normalize(x, p=2, dim=-1)
-    elif disttype == "rbf":
-        return x
-    else:
-        raise ValueError(f"to_euclidean not implemented for disttype {disttype}.")
-
-
-def quantile_min_max(x, q1=0.01, q2=0.99, n_sample=10000):
+def quantile_min_max(x: torch.Tensor, q1: float, q2: float, n_sample: int = 10000):
     if x.shape[0] > n_sample:
         np.random.seed(0)
         random_idx = np.random.choice(x.shape[0], n_sample, replace=False)
@@ -43,7 +30,7 @@ def quantile_min_max(x, q1=0.01, q2=0.99, n_sample=10000):
     return vmin, vmax
 
 
-def quantile_normalize(x, q=0.95):
+def quantile_normalize(x: torch.Tensor, q: float = 0.95):
     """normalize each dimension of x to [0, 1], take 95-th percentage, this robust to outliers
         </br> 1. sort x
         </br> 2. take q-th quantile
@@ -68,3 +55,7 @@ def quantile_normalize(x, q=0.95):
     x = (x - vmin) / (vmax - vmin)
     x = x.clamp(0, 1)
     return x
+
+
+def profile(name: str, t: torch.Tensor) -> None:
+    print(f"{name} --- nan: {t.isnan().any()}, inf: {t.isinf().any()}, max: {t.abs().max()}, min: {t.abs().min()}")

nystrom_ncut/distance_utils.py

@@ -0,0 +1,81 @@
+from typing import Literal
+
+import torch
+
+from .common import lazy_normalize
+
+
+DistanceOptions = Literal["cosine", "euclidean", "rbf"]
+
+
+def to_euclidean(x: torch.Tensor, disttype: DistanceOptions) -> torch.Tensor:
+    if disttype == "cosine":
+        return lazy_normalize(x, p=2, dim=-1)
+    elif disttype == "rbf":
+        return x
+    else:
+        raise ValueError(f"to_euclidean not implemented for disttype {disttype}.")
+
+
+def distance_from_features(
+    features: torch.Tensor,
+    features_B: torch.Tensor,
+    distance: DistanceOptions,
+):
+    """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
+        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
+    Returns:
+        (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
+    """
+    # compute distance matrix from input features
+    if distance == "cosine":
+        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)
+    elif distance == "rbf":
+        D = torch.cdist(features, features_B, p=2) ** 2
+
+        # Outlier-robust scale invariance using quantiles to estimate standard deviation
+        stds = torch.quantile(features, q=torch.tensor((0.158655, 0.841345), device=features.device), dim=0)
+        stds = (stds[1] - stds[0]) / 2
+        D = D / (2 * torch.linalg.norm(stds) ** 2)
+    else:
+        raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
+    return D
+
+
+def affinity_from_features(
+    features: torch.Tensor,
+    features_B: torch.Tensor = None,
+    affinity_focal_gamma: float = 1.0,
+    distance: DistanceOptions = "cosine",
+):
+    """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'.
+    Returns:
+        (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
+    """
+    # compute affinity matrix from input 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
+    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)
+
+    # torch.exp make affinity matrix positive definite,
+    # lower affinity_focal_gamma reduce the weak edge weights
+    A = torch.exp(-D / affinity_focal_gamma)
+    return A

nystrom_ncut/nystrom/distance_realization.py

@@ -1,20 +1,18 @@
-from typing import Tuple
-
 import torch
 
-from .nystrom import (
+from .nystrom_utils import (
     EigSolverOptions,
     OnlineKernel,
     OnlineNystromSubsampleFit,
     solve_eig,
 )
-from ..common import (
+from ..distance_utils import (
     DistanceOptions,
-    SampleOptions,
-)
-from ..propagation_utils import (
     distance_from_features,
 )
+from ..sampling_utils import (
+    SampleConfig,
+)
 
 
 class GramKernel(OnlineKernel):
@@ -100,8 +98,7 @@ class DistanceRealization(OnlineNystromSubsampleFit):
     def __init__(
         self,
         n_components: int = 100,
-        num_sample: int = 10000,
-        sample_method: SampleOptions = "farthest",
+        sample_config: SampleConfig = SampleConfig(),
         distance: DistanceOptions = "cosine",
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
@@ -109,9 +106,7 @@ class DistanceRealization(OnlineNystromSubsampleFit):
         """
         Args:
             n_components (int): number of top eigenvectors to return
-            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'].
+            sample_config (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'].
             eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
@@ -121,9 +116,8 @@ class DistanceRealization(OnlineNystromSubsampleFit):
             self,
             n_components=n_components,
             kernel=GramKernel(distance, eig_solver),
-            num_sample=num_sample,
             distance=distance,
-            sample_method=sample_method,
+            sample_config=sample_config,
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )
@@ -138,5 +132,5 @@ class DistanceRealization(OnlineNystromSubsampleFit):
         return V * (L ** 0.5)
 
     def transform(self, features: torch.Tensor = None) -> torch.Tensor:
-        V, L = OnlineNystromSubsampleFit.transform(features)
+        V, L = OnlineNystromSubsampleFit.transform(self, features)
         return V * (L ** 0.5)

nystrom_ncut/nystrom/normalized_cut.py

@@ -1,19 +1,19 @@
 import torch
 import torch.nn.functional as Fn
 
-from .nystrom import (
+from .nystrom_utils import (
     EigSolverOptions,
     OnlineKernel,
     OnlineNystromSubsampleFit,
     solve_eig,
 )
-from ..common import (
+from ..distance_utils import (
     DistanceOptions,
-    SampleOptions,
-)
-from ..propagation_utils import (
     affinity_from_features,
 )
+from ..sampling_utils import (
+    SampleConfig,
+)
 
 
 class LaplacianKernel(OnlineKernel):
@@ -94,8 +94,7 @@ class NCut(OnlineNystromSubsampleFit):
         self,
         n_components: int = 100,
         affinity_focal_gamma: float = 1.0,
-        num_sample: int = 10000,
-        sample_method: SampleOptions = "farthest",
+        sample_config: SampleConfig = SampleConfig(),
         distance: DistanceOptions = "cosine",
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
@@ -105,9 +104,7 @@ class NCut(OnlineNystromSubsampleFit):
             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,
-                reduce only if memory is not enough, increase for better approximation
-            sample_method (str): subgraph sampling, ['farthest', 'random'].
+            sample_config (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'].
             eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
@@ -117,9 +114,8 @@ class NCut(OnlineNystromSubsampleFit):
             self,
             n_components=n_components,
             kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
-            num_sample=num_sample,
             distance=distance,
-            sample_method=sample_method,
+            sample_config=sample_config,
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )

nystrom_ncut/nystrom/{nystrom.py → nystrom_utils.py}

@@ -1,14 +1,17 @@
+import copy
 import logging
 from typing import Literal, Tuple
 
 import torch
 
 from ..common import (
-    DistanceOptions,
-    SampleOptions,
     ceildiv,
 )
-from ..propagation_utils import (
+from ..distance_utils import (
+    DistanceOptions,
+)
+from ..sampling_utils import (
+    SampleConfig,
     run_subgraph_sampling,
 )
 
@@ -145,9 +148,8 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         self,
         n_components: int,
         kernel: OnlineKernel,
-        num_sample: int,
         distance: DistanceOptions,
-        sample_method: SampleOptions,
+        sample_config: SampleConfig,
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
     ):
@@ -158,9 +160,9 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )
-        self.num_sample: int = num_sample
         self.distance: DistanceOptions = distance
-        self.sample_method: SampleOptions = sample_method
+        self.sample_config: SampleConfig = sample_config
+        self.sample_config._ncut_obj = copy.deepcopy(self)
         self.anchor_indices: torch.Tensor = None
 
     def _fit_helper(
@@ -169,7 +171,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         precomputed_sampled_indices: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         _n = features.shape[0]
-        if self.num_sample >= _n:
+        if self.sample_config.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}"
             )
@@ -180,9 +182,8 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         else:
             self.anchor_indices = run_subgraph_sampling(
                 features=features,
-                num_sample=self.num_sample,
                 disttype=self.distance,
-                sample_method=self.sample_method,
+                config=self.sample_config,
             )
         sampled_features = features[self.anchor_indices]
         OnlineNystrom.fit(self, sampled_features)
@@ -243,6 +244,7 @@ def solve_eig(
     A: torch.Tensor,
     num_eig: int,
     eig_solver: EigSolverOptions,
+    eig_value_buffer: float = 0.0,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """PyTorch implementation of Eigensolver cut without Nystrom-like approximation.
 
@@ -250,11 +252,13 @@ def solve_eig(
         A (torch.Tensor): input 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']
-
+        eig_value_buffer (float): value added to diagonal to buffer symmetric but non-PSD matrices
     Returns:
         (torch.Tensor): eigenvectors corresponding to the eigenvalues, shape (n_samples, num_eig)
         (torch.Tensor): eigenvalues of the eigenvectors, sorted in descending order
     """
+    A = A + eig_value_buffer * torch.eye(A.shape[0], device=A.device)
+
     # compute eigenvectors
     if eig_solver == "svd_lowrank":  # default
         # only top q eigenvectors, fastest
@@ -272,15 +276,14 @@ def solve_eig(
         raise ValueError(
             "eigen_solver should be 'lobpcg', 'svd_lowrank', 'svd' or 'eigh'"
         )
+    eigen_value = eigen_value - eig_value_buffer
 
     # sort eigenvectors by eigenvalues, take top (descending order)
-    eigen_value = eigen_value.real
-    eigen_vector = eigen_vector.real
-    eigen_value, indices = torch.topk(eigen_value, k=num_eig, dim=0)
-    eigen_vector = eigen_vector[:, indices]
+    indices = torch.topk(eigen_value.abs(), k=num_eig, dim=0).indices
+    eigen_value, eigen_vector = eigen_value[indices], eigen_vector[:, indices]
 
     # correct the random rotation (flipping sign) of eigenvectors
-    sign = torch.sum(eigen_vector, dim=0).sign()
+    sign = torch.sum(eigen_vector.real, dim=0).sign()
     sign[sign == 0] = 1.0
     eigen_vector = eigen_vector * sign
     return eigen_vector, eigen_value

nystrom_ncut/sampling_utils.py

@@ -0,0 +1,96 @@
+import logging
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+from dgl.geometry import farthest_point_sampler
+
+from .distance_utils import (
+    DistanceOptions,
+    affinity_from_features,
+    to_euclidean,
+)
+
+
+SampleOptions = Literal["random", "fps", "fps_recursive"]
+
+
+@dataclass
+class SampleConfig:
+    method: SampleOptions = "fps"
+    num_sample: int = 10000
+    fps_dim: int = 12
+    n_iter: int = None
+    _ncut_obj: object = None
+
+
+@torch.no_grad()
+def run_subgraph_sampling(
+    features: torch.Tensor,
+    disttype: DistanceOptions,
+    config: SampleConfig,
+    max_draw: int = 1000000,
+):
+    features = features.detach()
+    if config.num_sample >= features.shape[0]:
+        # if too many samples, use all samples and bypass Nystrom-like approximation
+        logging.info(
+            "num_sample is larger than total, bypass Nystrom-like approximation"
+        )
+        sampled_indices = torch.arange(features.shape[0])
+    else:
+        # sample subgraph
+        if config.method == "fps":  # default
+            features = to_euclidean(features, disttype)
+            if config.num_sample > max_draw:
+                logging.warning(
+                    f"num_sample is larger than max_draw, apply farthest point sampling on random sampled {max_draw} samples"
+                )
+                draw_indices = torch.randperm(features.shape[0])[:max_draw]
+                sampled_indices = fpsample(features[draw_indices], config)
+                sampled_indices = draw_indices[sampled_indices]
+            else:
+                sampled_indices = fpsample(features, config)
+
+        elif config.method == "random":  # not recommended
+            sampled_indices = torch.randperm(features.shape[0])[:config.num_sample]
+
+        elif config.method == "fps_recursive":
+            features = to_euclidean(features, disttype)
+            sampled_indices = run_subgraph_sampling(
+                features=features,
+                disttype=disttype,
+                config=SampleConfig(method="fps", num_sample=config.num_sample, fps_dim=config.fps_dim)
+            )
+
+            nc = config._ncut_obj
+
+            A = affinity_from_features(features, affinity_focal_gamma=nc.kernel.affinity_focal_gamma, distance=nc.kernel.distance)
+            R = torch.diag(torch.sum(A, dim=-1) ** -0.5)
+            L = R @ A @ R
+
+            for _ in range(config.n_iter):
+                fps_features, eigenvalues = nc.fit_transform(features, precomputed_sampled_indices=sampled_indices)
+
+                _L = fps_features @ torch.diag(eigenvalues) @ fps_features.mT
+                RE = torch.abs(_L / L - 1)
+
+                print(f"Iteration {_} --- max: {RE.max().item()}, mean: {RE.mean().item()}, min: {RE.min().item()}")
+                fps_features = to_euclidean(fps_features[:, :config.fps_dim], "cosine")
+                sampled_indices = torch.sort(fpsample(fps_features, config)).values
+        else:
+            raise ValueError("sample_method should be 'farthest' or 'random'")
+        sampled_indices = torch.sort(sampled_indices).values
+    return sampled_indices.to(features.device)
+
+
+def fpsample(
+    features: torch.Tensor,
+    config: SampleConfig,
+):
+    # PCA to reduce the dimension
+    if features.shape[1] > config.fps_dim:
+        U, S, V = torch.pca_lowrank(features, q=config.fps_dim)
+        features = U * S
+
+    return farthest_point_sampler(features[None], config.num_sample)[0]

nystrom_ncut/visualize_utils.py

@@ -1,26 +1,150 @@
 import logging
-from typing import Any, Callable, Dict, Literal, Tuple
+from typing import Any, Callable, Dict, Literal
 
 import numpy as np
 import torch
-import torch.nn.functional as F
+import torch.nn.functional as Fn
 from sklearn.base import BaseEstimator
 
 from .common import (
+    ceildiv,
     lazy_normalize,
-    to_euclidean,
     quantile_min_max,
     quantile_normalize,
 )
-from .nystrom import (
-    DistanceRealization,
+from .distance_utils import (
+    DistanceOptions,
+    to_euclidean,
+    affinity_from_features,
 )
-from .propagation_utils import (
+from .sampling_utils import (
+    SampleConfig,
     run_subgraph_sampling,
-    extrapolate_knn,
 )
 
 
+def extrapolate_knn(
+    anchor_features: torch.Tensor,          # [n x d]
+    anchor_output: torch.Tensor,            # [n x d']
+    extrapolation_features: torch.Tensor,   # [m x d]
+    distance: DistanceOptions,
+    knn: int = 10,                          # k
+    affinity_focal_gamma: float = 1.0,
+    chunk_size: int = 8192,
+    device: str = None,
+    move_output_to_cpu: bool = False,
+) -> torch.Tensor:                          # [m x d']
+    """A generic function to propagate new nodes using KNN.
+
+    Args:
+        anchor_features (torch.Tensor): features from subgraph, shape (num_sample, n_features)
+        anchor_output (torch.Tensor): output from subgraph, shape (num_sample, D)
+        extrapolation_features (torch.Tensor): features from existing nodes, shape (new_num_samples, n_features)
+        knn (int): number of KNN to propagate eige nvectors
+        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
+    Returns:
+        torch.Tensor: propagated eigenvectors, shape (new_num_samples, D)
+
+    Examples:
+        >>> old_eigenvectors = torch.randn(3000, 20)
+        >>> old_features = torch.randn(3000, 100)
+        >>> new_features = torch.randn(200, 100)
+        >>> new_eigenvectors = extrapolate_knn(old_features, old_eigenvectors, new_features, knn=3)
+        >>> # new_eigenvectors.shape = (200, 20)
+
+    """
+    device = anchor_output.device if device is None else device
+
+    # used in nystrom_ncut
+    # propagate eigen_vector from subgraph to full graph
+    anchor_output = anchor_output.to(device)
+
+    n_chunks = ceildiv(extrapolation_features.shape[0], chunk_size)
+    V_list = []
+    for _v in torch.chunk(extrapolation_features, n_chunks, dim=0):
+        _v = _v.to(device)                                                                              # [_m x d]
+
+        _A = affinity_from_features(anchor_features, _v, affinity_focal_gamma, distance).mT             # [_m x n]
+        if knn is not None:
+            _A, indices = _A.topk(k=knn, dim=-1, largest=True)                                          # [_m x k], [_m x k]
+            _anchor_output = anchor_output[indices]                                                     # [_m x k x d]
+        else:
+            _anchor_output = anchor_output[None]                                                        # [1 x n x d]
+
+        _A = Fn.normalize(_A, p=1, dim=-1)                                                              # [_m x k]
+        _V = (_A[:, None, :] @ _anchor_output).squeeze(1)                                               # [_m x d]
+
+        if move_output_to_cpu:
+            _V = _V.cpu()
+        V_list.append(_V)
+
+    extrapolation_output = torch.cat(V_list, dim=0)
+    return extrapolation_output
+
+
+# wrapper functions for adding new nodes to existing graph
+def extrapolate_knn_with_subsampling(
+    full_features: torch.Tensor,            # [n x d]
+    full_output: torch.Tensor,              # [n x d']
+    extrapolation_features: torch.Tensor,   # [m x d]
+    sample_config: SampleConfig,
+    distance: DistanceOptions,
+    knn: int = 10,                          # k
+    affinity_focal_gamma: float = 1.0,
+    chunk_size: int = 8192,
+    device: str = None,
+    move_output_to_cpu: bool = False,
+) -> torch.Tensor:                          # [m x d']
+    """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:
+        full_output (torch.Tensor): eigenvectors from existing nodes, shape (num_sample, num_eig)
+        full_features (torch.Tensor): features from existing nodes, shape (n_samples, n_features)
+        extrapolation_features (torch.Tensor): features from new nodes, shape (n_new_samples, n_features)
+        knn (int): number of KNN to propagate eigenvectors, default 3
+        sample_config (str): sample method, 'farthest' (default) or 'random'
+        chunk_size (int): chunk size for matrix multiplication, default 8192
+        device (str): device to use for computation, if None, will not change device
+    Returns:
+        torch.Tensor: propagated eigenvectors, shape (n_new_samples, num_eig)
+
+    Examples:
+        >>> old_eigenvectors = torch.randn(3000, 20)
+        >>> old_features = torch.randn(3000, 100)
+        >>> new_features = torch.randn(200, 100)
+        >>> new_eigenvectors = extrapolate_knn_with_subsampling(extrapolation_features,old_eigenvectors,old_features,knn=3,num_sample=,sample_method=,chunk_size=,device=)
+        >>> # new_eigenvectors.shape = (200, 20)
+    """
+
+    device = full_output.device if device is None else device
+
+    # sample subgraph
+    anchor_indices = run_subgraph_sampling(
+        features=full_features,
+        disttype=distance,
+        config=sample_config,
+    )
+
+    anchor_output = full_output[anchor_indices].to(device)
+    anchor_features = full_features[anchor_indices].to(device)
+    extrapolation_features = extrapolation_features.to(device)
+
+    # propagate eigenvectors from subgraph to new nodes
+    extrapolation_output = extrapolate_knn(
+        anchor_features,
+        anchor_output,
+        extrapolation_features,
+        distance,
+        knn=knn,
+        affinity_focal_gamma=affinity_focal_gamma,
+        chunk_size=chunk_size,
+        device=device,
+        move_output_to_cpu=move_output_to_cpu,
+    )
+    return extrapolation_output
+
+
 def _rgb_with_dimensionality_reduction(
     features: torch.Tensor,
     num_sample: int,
@@ -38,9 +162,8 @@ def _rgb_with_dimensionality_reduction(
     if True:
         _subgraph_indices = run_subgraph_sampling(
             features=features,
-            num_sample=10000,
             disttype=disttype,
-            sample_method="farthest",
+            config=SampleConfig(method="fps"),
         )
         features = extrapolate_knn(
             anchor_features=features[_subgraph_indices],
@@ -51,9 +174,8 @@ def _rgb_with_dimensionality_reduction(
 
     subgraph_indices = run_subgraph_sampling(
         features=features,
-        num_sample=num_sample,
         disttype=disttype,
-        sample_method="farthest",
+        config=SampleConfig(method="fps", num_sample=num_sample),
     )
 
     _inp = features[subgraph_indices].numpy(force=True)
@@ -334,14 +456,14 @@ def rgb_from_umap_3d(
     return rgb
 
 
-def flatten_sphere(X_3d):
-    x = np.arctan2(X_3d[:, 0], X_3d[:, 1])
-    y = -np.arccos(X_3d[:, 2])
-    X_2d = np.stack([x, y], axis=1)
+def flatten_sphere(X_3d: torch.Tensor) -> torch.Tensor:
+    x = torch.atan2(X_3d[:, 0], X_3d[:, 1])
+    y = -torch.acos(X_3d[:, 2])
+    X_2d = torch.stack((x, y), dim=1)
     return X_2d
 
 
-def rotate_rgb_cube(rgb, position=1):
+def rotate_rgb_cube(rgb: torch.Tensor, position: int = 1) -> torch.Tensor:
     """rotate RGB cube to different position
 
     Args:
@@ -365,7 +487,7 @@ def rotate_rgb_cube(rgb, position=1):
     return rgb
 
 
-def rgb_from_3d_rgb_cube(X_3d, q=0.95):
+def rgb_from_3d_rgb_cube(X_3d: torch.Tensor, q: float = 0.95) -> torch.Tensor:
     """convert 3D t-SNE to RGB color space
     Args:
         X_3d (torch.Tensor): 3D t-SNE embedding, shape (n_samples, 3)
@@ -383,6 +505,26 @@ def rgb_from_3d_rgb_cube(X_3d, q=0.95):
     return rgb
 
 
+def rgb_from_3d_lab_cube(X_3d: torch.Tensor, q: float = 0.95, full_range: bool = True) -> torch.Tensor:
+    from skimage import color
+    X_3d = X_3d - torch.mean(X_3d, dim=0)
+    U, S, VT = torch.linalg.svd(X_3d)
+    X_3d = torch.flip(U[:, :3] * S, dims=(1,))
+
+    AB_scale = 128.0 / torch.quantile(torch.linalg.norm(X_3d[:, 1:], dim=1), q=q, dim=0)
+    L_min, L_max = torch.quantile(X_3d[:, 0], q=torch.tensor(((1 - q) / 2, (1 + q) / 2)), dim=0)
+    L_scale = 100.0 / (L_max - L_min)
+
+    X_3d[:, 0] = X_3d[:, 0] - L_min
+    if full_range:
+        lab = X_3d * torch.tensor((L_scale, AB_scale, AB_scale))
+    else:
+        lab = X_3d * L_scale
+
+    rgb = torch.tensor(color.lab2rgb(lab.numpy(force=True)))
+    return rgb
+
+
 def convert_to_lab_color(rgb, full_range=True):
     from skimage import color
     import copy
@@ -401,7 +543,7 @@ def convert_to_lab_color(rgb, full_range=True):
     return lab_rgb
 
 
-def rgb_from_2d_colormap(X_2d, q=0.95):
+def rgb_from_2d_colormap(X_2d: torch.Tensor, q: float = 0.95):
     xy = X_2d.clone()
     for i in range(2):
         xy[:, i] = quantile_normalize(xy[:, i], q=q)
@@ -446,7 +588,7 @@ def _transform_heatmap(heatmap, gamma=1.0):
     # large gamma means more focus on the high values, hence smaller mask
     heatmap = 1 / heatmap ** gamma
     # min-max normalization [0, 1]
-    vmin, vmax = quantile_min_max(heatmap.flatten())
+    vmin, vmax = quantile_min_max(heatmap.flatten(), 0.01, 0.99)
     heatmap = (heatmap - vmin) / (vmax - vmin)
     return heatmap
 
@@ -514,7 +656,7 @@ def get_mask(
 
     # normalize the eigenvectors to unit norm, to compute cosine similarity
     all_eigvecs = lazy_normalize(all_eigvecs, p=2, dim=-1)
-    prompt_eigvec = F.normalize(prompt_eigvec, p=2, dim=-1)
+    prompt_eigvec = Fn.normalize(prompt_eigvec, p=2, dim=-1)
 
     # compute the cosine similarity
     cos_sim = all_eigvecs @ prompt_eigvec.unsqueeze(-1)  # (B, H, W, 1)

{nystrom_ncut-0.1.0.dist-info → nystrom_ncut-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.1.0
+Version: 0.1.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.1.2.dist-info/RECORD

@@ -0,0 +1,15 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nystrom_ncut/__init__.py,sha256=Wgud0tGaNkK2m_qVU47rXQqKdoR-4ztDXXD9UKzc4c8,488
+nystrom_ncut/common.py,sha256=_PGJoImSk_Fb_5Ri-e_IsFoCcSfbGS8CxYUUHVoNM50,2036
+nystrom_ncut/distance_utils.py,sha256=U1223ri8OuIzj0wjhAUhHWcsEvREDitgz8i1rRlCfj8,3069
+nystrom_ncut/sampling_utils.py,sha256=uoWWSyfttv5fnOSq8KFXomWiNO-THiPPbLXfupnVar0,3444
+nystrom_ncut/visualize_utils.py,sha256=xDlkE5sMXehK5hNz9U1twqgHZVzmV5tf5O9bL96AiaM,22982
+nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
+nystrom_ncut/nystrom/distance_realization.py,sha256=9GX_XSISTvsEWUu8bG5AxtlkYYNItFspcH5wXiwSOKY,5789
+nystrom_ncut/nystrom/normalized_cut.py,sha256=ZxFV8Sckp6wtpNyoA15DS7Vfu9QLvzNpwrwY0n9_GNs,6953
+nystrom_ncut/nystrom/nystrom_utils.py,sha256=MEmW5xgOu8u2HCwjFapHAOFFXhoVslBbLG4Cn-mYMDU,12995
+nystrom_ncut-0.1.2.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.1.2.dist-info/METADATA,sha256=0wsHYtW3cY4Bzq-lH_y_Blazt6YMrwWxsmR7SOHMyzs,6058
+nystrom_ncut-0.1.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.1.2.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.1.2.dist-info/RECORD,,

nystrom_ncut/propagation_utils.py

@@ -1,268 +0,0 @@
-import logging
-
-import numpy as np
-import torch
-import torch.nn.functional as Fn
-
-from .common import (
-    DistanceOptions,
-    SampleOptions,
-    ceildiv,
-    lazy_normalize,
-    to_euclidean,
-)
-
-
-# @torch.no_grad()
-def run_subgraph_sampling(
-    features: torch.Tensor,
-    num_sample: int,
-    disttype: DistanceOptions,
-    sample_method: SampleOptions,
-    max_draw: int = 1000000,
-):
-    if num_sample >= features.shape[0]:
-        # if too many samples, use all samples and bypass Nystrom-like approximation
-        logging.info(
-            "num_sample is larger than total, bypass Nystrom-like approximation"
-        )
-        sampled_indices = torch.arange(features.shape[0])
-    else:
-        # sample subgraph
-        if sample_method == "farthest":  # default
-            features = to_euclidean(features, disttype)
-            if num_sample > max_draw:
-                logging.warning(
-                    f"num_sample is larger than max_draw, apply farthest point sampling on random sampled {max_draw} samples"
-                )
-                draw_indices = torch.randperm(features.shape[0])[:max_draw]
-                sampled_indices = farthest_point_sampling(
-                    features[draw_indices].detach(),
-                    num_sample=num_sample,
-                )
-                sampled_indices = draw_indices[sampled_indices]
-            else:
-                sampled_indices = farthest_point_sampling(
-                    features.detach(),
-                    num_sample=num_sample,
-                )
-        elif sample_method == "random":  # not recommended
-            sampled_indices = torch.randperm(features.shape[0])[:num_sample]
-        else:
-            raise ValueError("sample_method should be 'farthest' or 'random'")
-        sampled_indices = torch.sort(sampled_indices).values
-    return sampled_indices.to(features.device)
-
-
-def farthest_point_sampling(
-    features: torch.Tensor,
-    num_sample: int = 300,
-    h: int = 9,
-):
-    try:
-        import fpsample
-    except ImportError:
-        raise ImportError(
-            "fpsample import failed, please install `pip install fpsample`"
-        )
-
-    # PCA to reduce the dimension
-    if features.shape[1] > 8:
-        u, s, v = torch.pca_lowrank(features, q=8)
-        features = u @ torch.diag(s)
-
-    h = min(h, int(np.log2(features.shape[0])))
-
-    kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(
-        features.numpy(force=True), num_sample, h
-    ).astype(np.int64)
-    return torch.from_numpy(kdline_fps_samples_idx)
-
-
-def distance_from_features(
-    features: torch.Tensor,
-    features_B: torch.Tensor,
-    distance: DistanceOptions,
-):
-    """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
-        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
-    Returns:
-        (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
-    """
-    # compute distance matrix from input features
-    if distance == "cosine":
-        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)
-    elif distance == "rbf":
-        D = torch.cdist(features, features_B, p=2) ** 2
-
-        # Outlier-robust scale invariance using quantiles to estimate standard deviation
-        stds = torch.quantile(features, q=torch.tensor((0.158655, 0.841345), device=features.device), dim=0)
-        stds = (stds[1] - stds[0]) / 2
-        D = D / (2 * torch.linalg.norm(stds) ** 2)
-    else:
-        raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
-    return D
-
-
-def affinity_from_features(
-    features: torch.Tensor,
-    features_B: torch.Tensor = None,
-    affinity_focal_gamma: float = 1.0,
-    distance: DistanceOptions = "cosine",
-):
-    """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'.
-    Returns:
-        (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
-    """
-    # compute affinity matrix from input 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
-    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)
-
-    # torch.exp make affinity matrix positive definite,
-    # lower affinity_focal_gamma reduce the weak edge weights
-    A = torch.exp(-D / affinity_focal_gamma)
-    return A
-
-
-def extrapolate_knn(
-    anchor_features: torch.Tensor,          # [n x d]
-    anchor_output: torch.Tensor,            # [n x d']
-    extrapolation_features: torch.Tensor,   # [m x d]
-    distance: DistanceOptions,
-    knn: int = 10,                          # k
-    affinity_focal_gamma: float = 1.0,
-    chunk_size: int = 8192,
-    device: str = None,
-    move_output_to_cpu: bool = False,
-) -> torch.Tensor:                          # [m x d']
-    """A generic function to propagate new nodes using KNN.
-
-    Args:
-        anchor_features (torch.Tensor): features from subgraph, shape (num_sample, n_features)
-        anchor_output (torch.Tensor): output from subgraph, shape (num_sample, D)
-        extrapolation_features (torch.Tensor): features from existing nodes, shape (new_num_samples, n_features)
-        knn (int): number of KNN to propagate eige nvectors
-        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
-    Returns:
-        torch.Tensor: propagated eigenvectors, shape (new_num_samples, D)
-
-    Examples:
-        >>> old_eigenvectors = torch.randn(3000, 20)
-        >>> old_features = torch.randn(3000, 100)
-        >>> new_features = torch.randn(200, 100)
-        >>> new_eigenvectors = extrapolate_knn(old_features, old_eigenvectors, new_features, knn=3)
-        >>> # new_eigenvectors.shape = (200, 20)
-
-    """
-    device = anchor_output.device if device is None else device
-
-    # used in nystrom_ncut
-    # propagate eigen_vector from subgraph to full graph
-    anchor_output = anchor_output.to(device)
-
-    n_chunks = ceildiv(extrapolation_features.shape[0], chunk_size)
-    V_list = []
-    for _v in torch.chunk(extrapolation_features, n_chunks, dim=0):
-        _v = _v.to(device)                                                                              # [_m x d]
-
-        _A = affinity_from_features(anchor_features, _v, affinity_focal_gamma, distance).mT             # [_m x n]
-        if knn is not None:
-            _A, indices = _A.topk(k=knn, dim=-1, largest=True)                                          # [_m x k], [_m x k]
-            _anchor_output = anchor_output[indices]                                                     # [_m x k x d]
-        else:
-            _anchor_output = anchor_output[None]                                                        # [1 x n x d]
-
-        _A = Fn.normalize(_A, p=1, dim=-1)                                                              # [_m x k]
-        _V = (_A[:, None, :] @ _anchor_output).squeeze(1)                                               # [_m x d]
-
-        if move_output_to_cpu:
-            _V = _V.cpu()
-        V_list.append(_V)
-
-    extrapolation_output = torch.cat(V_list, dim=0)
-    return extrapolation_output
-
-
-# wrapper functions for adding new nodes to existing graph
-def extrapolate_knn_with_subsampling(
-    full_features: torch.Tensor,            # [n x d]
-    full_output: torch.Tensor,              # [n x d']
-    extrapolation_features: torch.Tensor,   # [m x d]
-    num_sample: int,                        # n'
-    sample_method: SampleOptions,
-    distance: DistanceOptions,
-    knn: int = 10,                          # k
-    affinity_focal_gamma: float = 1.0,
-    chunk_size: int = 8192,
-    device: str = None,
-    move_output_to_cpu: bool = False,
-) -> torch.Tensor:                          # [m x d']
-    """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:
-        full_output (torch.Tensor): eigenvectors from existing nodes, shape (num_sample, num_eig)
-        full_features (torch.Tensor): features from existing nodes, shape (n_samples, n_features)
-        extrapolation_features (torch.Tensor): features from new nodes, shape (n_new_samples, n_features)
-        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 8192
-        device (str): device to use for computation, if None, will not change device
-    Returns:
-        torch.Tensor: propagated eigenvectors, shape (n_new_samples, num_eig)
-
-    Examples:
-        >>> old_eigenvectors = torch.randn(3000, 20)
-        >>> old_features = torch.randn(3000, 100)
-        >>> new_features = torch.randn(200, 100)
-        >>> new_eigenvectors = extrapolate_knn_with_subsampling(extrapolation_features,old_eigenvectors,old_features,knn=3,num_sample=,sample_method=,chunk_size=,device=)
-        >>> # new_eigenvectors.shape = (200, 20)
-    """
-
-    device = full_output.device if device is None else device
-
-    # sample subgraph
-    anchor_indices = run_subgraph_sampling(
-        features=full_features,
-        num_sample=num_sample,
-        disttype=distance,
-        sample_method=sample_method,
-    )
-
-    anchor_output = full_output[anchor_indices].to(device)
-    anchor_features = full_features[anchor_indices].to(device)
-    extrapolation_features = extrapolation_features.to(device)
-
-    # propagate eigenvectors from subgraph to new nodes
-    extrapolation_output = extrapolate_knn(
-        anchor_features,
-        anchor_output,
-        extrapolation_features,
-        distance,
-        knn=knn,
-        affinity_focal_gamma=affinity_focal_gamma,
-        chunk_size=chunk_size,
-        device=device,
-        move_output_to_cpu=move_output_to_cpu,
-    )
-    return extrapolation_output

nystrom_ncut-0.1.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nystrom_ncut/__init__.py,sha256=ffExLdGTaPsUweHcYc61Ose6a5A5Tfo9hm48zjEl6ho,441
-nystrom_ncut/common.py,sha256=l5kjF6neukdtbMrt5bE69pN0JM9r-93JYltLlYJuBik,2227
-nystrom_ncut/propagation_utils.py,sha256=79M61iJfp_RWj_xLOn51PHiextWcEWTQ7NWl2T51-3Y,10907
-nystrom_ncut/visualize_utils.py,sha256=uGfBBkETQ8uRJ-UXPWpiXT8KEYMLmEHXPKJMLvN0c34,16543
-nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
-nystrom_ncut/nystrom/distance_realization.py,sha256=FGH7VjbtRrSROH0d8OPuCUxLQy5j7Z8BuE4hrSGGZG4,6031
-nystrom_ncut/nystrom/normalized_cut.py,sha256=s9ZS3-tQbWnxAlPc01v9l7fqBhl28lvOalaCO2y-Gd8,7175
-nystrom_ncut/nystrom/nystrom.py,sha256=OV5o9UL9fkrz9HdsD6rXh7MTsenPKrtCNRIczMuDS_4,12779
-nystrom_ncut-0.1.0.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.1.0.dist-info/METADATA,sha256=wAhGONU0ZM1VWoLqEwtfAs3_GUAt4CItkHq3ISuFyVE,6058
-nystrom_ncut-0.1.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.1.0.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.1.0.dist-info/RECORD,,