nystrom-ncut 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

nystrom_ncut/__init__.py

@@ -1,5 +1,8 @@
+from .kernel import (
+    KernelNCut,
+)
 from .nystrom import (
-    NCut,
+    NystromNCut,
 )
 from .transformer import (
     AxisAlign,

nystrom_ncut/distance_utils.py

@@ -6,13 +6,25 @@ import torch
 from .common import lazy_normalize
 
 
-DistanceOptions = Literal["cosine", "euclidean"]
-AffinityOptions = Literal["cosine", "rbf", "laplacian"]
+DistanceOptions = Literal[
+    "cosine",
+    "euclidean",
+]
+AffinityOptions = Literal[
+    "cosine",
+    "rbf",
+    # "laplacian",
+]
 
 # noinspection PyTypeChecker
 DISTANCE_TO_AFFINITY: OrderedDict[DistanceOptions, List[AffinityOptions]] = collections.OrderedDict([
-    ("cosine", ["cosine"]),
-    ("euclidean", ["rbf", "laplacian"]),
+    ("cosine", [
+        "cosine",
+    ]),
+    ("euclidean", [
+        "rbf",
+        # "laplacian",
+    ]),
 ])
 # noinspection PyTypeChecker
 AFFINITY_TO_DISTANCE: OrderedDict[AffinityOptions, DistanceOptions] = collections.OrderedDict(sum([
@@ -32,45 +44,51 @@ def to_euclidean(x: torch.Tensor, distance_type: DistanceOptions) -> torch.Tenso
 
 
 def distance_from_features(
-    features: torch.Tensor,
+    features_A: torch.Tensor,
     features_B: torch.Tensor,
     distance_type: DistanceOptions,
 ):
     """Compute distance matrix from input features.
     Args:
-        features (torch.Tensor): input features, shape (n_samples, n_features)
+        features_A (torch.Tensor): input features, shape (n_samples, n_features)
         features_B (torch.Tensor, optional): optional, if not None, compute affinity between two features
         distance_type (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
     Returns:
         (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
     """
     # compute distance matrix from input features
-    shape: torch.Size = features.shape[:-2]
-    features = features.view((-1, *features.shape[-2:]))
+    shape: torch.Size = features_A.shape[:-2]
+    features_A = features_A.view((-1, *features_A.shape[-2:]))
     features_B = features_B.view((-1, *features_B.shape[-2:]))
 
     match distance_type:
         case "cosine":
-            features = lazy_normalize(features, dim=-1)
+            features_A = lazy_normalize(features_A, dim=-1)
             features_B = lazy_normalize(features_B, dim=-1)
-            D = 1 - features @ features_B.mT
+            D = 1 - features_A @ features_B.mT
         case "euclidean":
-            D = torch.cdist(features, features_B, p=2)
+            D = torch.cdist(features_A, features_B, p=2)
         case _:
             raise ValueError("Distance should be 'cosine' or 'euclidean'")
     return D.view((*shape, *D.shape[-2:]))
 
 
+def get_normalization_factor(features: torch.Tensor, c: float = 2.0) -> torch.Tensor:
+    p = torch.erf(torch.tensor((-c, c), device=features.device) * (2 ** -0.5))
+    lo, hi = torch.nanquantile(features, q=(p + 1) / 2, dim=-2)     # [... x d], [... x d]
+    return torch.norm(hi - lo, dim=-1) / (2 * c)                    # [...]
+
+
 def affinity_from_features(
-    features: torch.Tensor,
-    features_B: torch.Tensor = None,
-    affinity_focal_gamma: float = 1.0,
-    affinity_type: AffinityOptions = "cosine",
+    features_A: torch.Tensor,
+    features_B: torch.Tensor,
+    affinity_type: AffinityOptions,
+    affinity_focal_gamma: float,
 ):
     """Compute affinity matrix from input features.
 
     Args:
-        features (torch.Tensor): input features, shape (n_samples, n_features)
+        features_A (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
@@ -82,24 +100,17 @@ 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
-    features_B = features if features_B is None else features_B
-
-    # compute distance matrix from input features
-    D = distance_from_features(features, features_B, AFFINITY_TO_DISTANCE[affinity_type])
+    D = distance_from_features(features_A, features_B, AFFINITY_TO_DISTANCE[affinity_type])
 
     # lower affinity_focal_gamma reduce the weak edge weights
     match affinity_type:
-        case "cosine" | "laplacian":
-            A = torch.exp(-D / affinity_focal_gamma)                                        # [... x n x n]
+        case "cosine":
+            pass
+        # case "laplacian":
+        #     D = D / get_normalization_factor(features_A)[..., None, None]
         case "rbf":
-            # Outlier-robust scale invariance using quantiles to estimate standard deviation
-            c = 2.0
-            p = torch.erf(torch.tensor((-c, c), device=features.device) * (2 ** -0.5))
-            stds = torch.nanquantile(features, q=(p + 1) / 2, dim=-2)                       # [2 x ... x d]
-            stds = (stds[1] - stds[0]) / (2 * c)                                            # [... x d]
-            D = 0.5 * (D / torch.norm(stds, dim=-1)[..., None, None]) ** 2
-            A = torch.exp(-D / affinity_focal_gamma)
+            D = 0.5 * (D / get_normalization_factor(features_A)[..., None, None]) ** 2
         case _:
             raise ValueError("Affinity should be 'cosine', 'rbf', or 'laplacian'")
-
+    A = torch.exp(-D / affinity_focal_gamma)    # [... x n x n]
     return A

nystrom_ncut/global_settings.py

@@ -0,0 +1,2 @@
+
+CHUNK_SIZE: int = 8192

nystrom_ncut/kernel/__init__.py

@@ -0,0 +1,3 @@
+from .kernel_ncut import (
+    KernelNCut,
+)

nystrom_ncut/kernel/kernel_ncut.py

@@ -0,0 +1,113 @@
+import torch
+
+from ..distance_utils import (
+    AffinityOptions,
+    AFFINITY_TO_DISTANCE,
+    get_normalization_factor,
+)
+from ..sampling_utils import (
+    SampleConfig,
+    OnlineTransformerSubsampleFit,
+)
+from ..transformer import (
+    OnlineTorchTransformerMixin,
+)
+
+
+class KernelNCutBaseTransformer(OnlineTorchTransformerMixin):
+    def __init__(
+        self,
+        n_components: int,
+        kernel_dim: int,
+        affinity_type: AffinityOptions,
+        affinity_focal_gamma: float,
+    ):
+        self.n_components: int = n_components
+        self.kernel_dim: int = kernel_dim
+        self.affinity_type: AffinityOptions = affinity_type
+        self.affinity_focal_gamma = affinity_focal_gamma
+
+        # Anchor matrices
+        self.W: torch.Tensor = None                     # [... x d x kernel_dim]
+        self.kernelized_anchor: torch.Tensor = None     # [... x n x (2 * kernel_dim)]
+
+        # Updated matrices
+        self.r: torch.Tensor = None                     # [... x (2 * kernel_dim)]
+        self.transform_matrix: torch.Tensor = None      # [... x (2 * kernel_dim) x n_components]
+        self.eigenvalues_: torch.Tensor = None          # [... x n_components]
+
+    def _update(self) -> None:
+        row_sum = self.kernelized_anchor @ self.r[..., None]                        # [... x n x 1]
+        normalized_kernelized_anchor = self.kernelized_anchor / (row_sum ** 0.5)    # [... x n x (2 * kernel_dim)]
+        _, S, V = torch.svd_lowrank(torch.nan_to_num(
+            normalized_kernelized_anchor, nan=0.0,
+        ), q=self.n_components)                                                     # [... x n_components], [... x (2 * kernel_dim) x n_components]
+        self.transform_matrix = V * torch.nan_to_num(1 / S, posinf=0.0, neginf=0.0)[..., None, :]       # [... x (2 * kernel_dim) x n_components]
+        self.eigenvalues_ = S ** 2
+
+    def fit(self, features: torch.Tensor) -> "KernelNCutBaseTransformer":
+        d = features.shape[-1]
+        scale = get_normalization_factor(features) * (self.affinity_focal_gamma ** 0.5)             # [...]
+        self.W = torch.randn((*features.shape[:-2], d, self.kernel_dim)) / scale[..., None, None]   # [... x d x kernel_dim]
+
+        W_anchor = features @ self.W                                                                # [... x n x kernel_dim]
+        self.kernelized_anchor = torch.cat((
+            torch.cos(W_anchor),
+            torch.sin(W_anchor),
+        ), dim=-1) / (self.kernel_dim ** 0.5)                                                       # [... x n * (2 * kernel_dim)]
+        self.r = torch.sum(torch.nan_to_num(self.kernelized_anchor, nan=0.0), dim=-2)               # [... x (2 * kernel_dim)]
+        self._update()
+        return self
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        W_features = features @ self.W                                              # [... x m x kernel_dim]
+        kernelized_features = torch.cat((
+            torch.cos(W_features),
+            torch.sin(W_features),
+        ), dim=-1) / (self.kernel_dim ** 0.5)                                       # [... x m x (2 * kernel_dim)]
+        b_r = torch.sum(torch.nan_to_num(kernelized_features, nan=0.0), dim=-2)     # [... x (2 * kernel_dim)]
+        self.r = self.r + b_r
+        self._update()
+
+        row_sum = kernelized_features @ self.r[..., None]                           # [... x m x 1]
+        normalized_kernelized_features = kernelized_features / (row_sum ** 0.5)     # [... x m x (2 * kernel_dim)]
+        return normalized_kernelized_features @ self.transform_matrix               # [... x m x n_components]
+
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:
+        if features is None:
+            kernelized_features = self.kernelized_anchor                            # [... x n x (2 * kernel_dim)]
+        else:
+            W_features = features @ self.W
+            kernelized_features = torch.cat((
+                torch.cos(W_features),
+                torch.sin(W_features),
+            ), dim=-1) / (self.kernel_dim ** 0.5)                                   # [... x m x (2 * kernel_dim)]
+        row_sum = kernelized_features @ self.r[..., None]                           # [... x m x 1]
+        normalized_kernelized_features = kernelized_features / (row_sum ** 0.5)     # [... x m x (2 * kernel_dim)]
+        return normalized_kernelized_features @ self.transform_matrix               # [... x m x n_components]
+
+
+class KernelNCut(OnlineTransformerSubsampleFit):
+    """Kernelized Normalized Cut for large scale graph."""
+
+    def __init__(
+        self,
+        n_components: int,
+        kernel_dim: int = 1024,
+        affinity_type: AffinityOptions = "cosine",
+        affinity_focal_gamma: float = 1.0,
+        sample_config: SampleConfig = SampleConfig(),
+    ):
+        OnlineTransformerSubsampleFit.__init__(
+            self,
+            base_transformer=KernelNCutBaseTransformer(
+                n_components=n_components,
+                kernel_dim=kernel_dim,
+                affinity_type=affinity_type,
+                affinity_focal_gamma=affinity_focal_gamma,
+            ),
+            distance_type=AFFINITY_TO_DISTANCE[affinity_type],
+            sample_config=sample_config,
+        )
+
+

nystrom_ncut/nystrom/__init__.py

@@ -1,3 +1,3 @@
 from .normalized_cut import (
-    NCut,
+    NystromNCut,
 )

nystrom_ncut/nystrom/distance_realization.py

@@ -101,7 +101,6 @@ class DistanceRealization(OnlineNystromSubsampleFit):
         distance_type: DistanceOptions = "cosine",
         sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
-        chunk_size: int = 8192,
     ):
         """
         Args:
@@ -110,7 +109,6 @@ class DistanceRealization(OnlineNystromSubsampleFit):
                 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'].
-            chunk_size (int): chunk size for large-scale matrix multiplication
         """
         OnlineNystromSubsampleFit.__init__(
             self,
@@ -119,7 +117,6 @@ class DistanceRealization(OnlineNystromSubsampleFit):
             distance_type=distance_type,
             sample_config=sample_config,
             eig_solver=eig_solver,
-            chunk_size=chunk_size,
         )
 
     def fit_transform(

nystrom_ncut/nystrom/normalized_cut.py

@@ -4,7 +4,7 @@ import torch
 from .nystrom_utils import (
     EigSolverOptions,
     OnlineKernel,
-    OnlineNystromSubsampleFit,
+    OnlineNystrom,
     solve_eig,
 )
 from ..distance_utils import (
@@ -14,19 +14,20 @@ from ..distance_utils import (
 )
 from ..sampling_utils import (
     SampleConfig,
+    OnlineTransformerSubsampleFit,
 )
 
 
 class LaplacianKernel(OnlineKernel):
     def __init__(
         self,
-        affinity_focal_gamma: float,
         affinity_type: AffinityOptions,
+        affinity_focal_gamma: float,
         adaptive_scaling: bool,
         eig_solver: EigSolverOptions,
     ):
-        self.affinity_focal_gamma = affinity_focal_gamma
         self.affinity_type: AffinityOptions = affinity_type
+        self.affinity_focal_gamma = affinity_focal_gamma
         self.adaptive_scaling: bool = adaptive_scaling
         self.eig_solver: EigSolverOptions = eig_solver
 
@@ -44,27 +45,29 @@ class LaplacianKernel(OnlineKernel):
         self.anchor_features = features                                             # [... x n x d]
         self.anchor_mask = torch.all(torch.isnan(self.anchor_features), dim=-1)     # [... x n]
 
-        self.A = torch.nan_to_num(affinity_from_features(
-            self.anchor_features,                                                   # [... x n x d]
-            affinity_focal_gamma=self.affinity_focal_gamma,
+
+        self.A = torch.where(self.anchor_mask[..., None], 0.0, affinity_from_features(
+            features_A=self.anchor_features,                                        # [... x n x d]
+            features_B=self.anchor_features,                                        # [... x n x d]
             affinity_type=self.affinity_type,
-        ), nan=0.0)                                                                 # [... x n x n]
+            affinity_focal_gamma=self.affinity_focal_gamma,
+        ))                                                                          # [... x n x n]
         d = features.shape[-1]
         U, L = solve_eig(
-            self.A,
+            torch.nan_to_num(self.A, nan=0.0),
             num_eig=d + 1,  # d * (d + 3) // 2 + 1,
             eig_solver=self.eig_solver,
         )                                                                                           # [... x n x (d + 1)], [... x (d + 1)]
         self.Ainv = U @ torch.nan_to_num(torch.diag_embed(1 / L), posinf=0.0, neginf=0.0) @ U.mT    # [... x n x n]
-        self.a_r = torch.where(self.anchor_mask, torch.inf, torch.sum(self.A, dim=-1))              # [... x n]
+        self.a_r = torch.where(self.anchor_mask, torch.inf, torch.sum(self.A.mT, dim=-1))           # [... x n]
         self.b_r = torch.zeros_like(self.a_r)                                                       # [... x n]
 
     def _affinity(self, features: torch.Tensor) -> torch.Tensor:
         B = torch.where(self.anchor_mask[..., None], 0.0, affinity_from_features(
-            self.anchor_features,                                                   # [... x n x d]
-            features,                                                               # [... x m x d]
-            affinity_focal_gamma=self.affinity_focal_gamma,
+            features_A=self.anchor_features,                                        # [... x n x d]
+            features_B=features,                                                    # [... x m x d]
             affinity_type=self.affinity_type,
+            affinity_focal_gamma=self.affinity_focal_gamma,
         ))                                                                          # [... x n x m]
         if self.adaptive_scaling:
             diagonal = (
@@ -93,44 +96,43 @@ class LaplacianKernel(OnlineKernel):
             B = self.A                                                              # [... x n x n]
             col_sum = row_sum                                                       # [... x n]
         else:
-            B = self._affinity(features)
+            B = self._affinity(features)                                            # [... x n x m]
             b_c = torch.sum(B, dim=-2)                                              # [... x m]
             col_sum = b_c + (B.mT @ (self.Ainv @ self.b_r[..., None]))[..., 0]      # [... x m]
         scale = (row_sum[..., :, None] * col_sum[..., None, :]) ** -0.5             # [... x n x m]
         return (B * scale).mT                                                       # [... x m x n]
 
 
-class NCut(OnlineNystromSubsampleFit):
+class NystromNCut(OnlineTransformerSubsampleFit):
     """Nystrom Normalized Cut for large scale graph."""
 
     def __init__(
         self,
-        n_components: int = 100,
-        affinity_focal_gamma: float = 1.0,
+        n_components: int,
         affinity_type: AffinityOptions = "cosine",
+        affinity_focal_gamma: float = 1.0,
         adaptive_scaling: bool = False,
         sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
-        chunk_size: int = 8192,
     ):
         """
         Args:
             n_components (int): number of top eigenvectors to return
+            affinity_type (str): distance metric for affinity matrix, ['cosine', 'euclidean', 'rbf'].
             affinity_focal_gamma (float): affinity matrix temperature, lower t reduce the not-so-connected edge weights,
                 smaller t result in more sharp eigenvectors.
-            distance (str): distance metric for affinity matrix, ['cosine', 'euclidean', 'rbf'].
             adaptive_scaling (bool): whether to scale off-diagonal affinity vectors so extended diagonal equals 1
             sample_config (str): subgraph sampling, ['farthest', 'random'].
                 farthest point sampling is recommended for better Nystrom-approximation accuracy
             eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
-            chunk_size (int): chunk size for large-scale matrix multiplication
         """
-        OnlineNystromSubsampleFit.__init__(
+        OnlineTransformerSubsampleFit.__init__(
             self,
-            n_components=n_components,
-            kernel=LaplacianKernel(affinity_focal_gamma, affinity_type, adaptive_scaling, eig_solver),
+            base_transformer=OnlineNystrom(
+                n_components=n_components,
+                kernel=LaplacianKernel(affinity_type, affinity_focal_gamma, adaptive_scaling, eig_solver),
+                eig_solver=eig_solver,
+            ),
             distance_type=AFFINITY_TO_DISTANCE[affinity_type],
             sample_config=sample_config,
-            eig_solver=eig_solver,
-            chunk_size=chunk_size,
         )

nystrom_ncut/nystrom/nystrom_utils.py

@@ -1,5 +1,3 @@
-import copy
-import logging
 from abc import abstractmethod
 from typing import Literal, Tuple
 
@@ -8,15 +6,11 @@ import torch
 from ..common import (
     ceildiv,
 )
-from ..distance_utils import (
-    DistanceOptions,
-)
-from ..sampling_utils import (
-    SampleConfig,
-    subsample_features,
+from ..global_settings import (
+    CHUNK_SIZE,
 )
 from ..transformer import (
-    TorchTransformerMixin,
+    OnlineTorchTransformerMixin,
 )
 
 
@@ -37,13 +31,12 @@ class OnlineKernel:
         """"""
 
 
-class OnlineNystrom(TorchTransformerMixin):
+class OnlineNystrom(OnlineTorchTransformerMixin):
     def __init__(
         self,
         n_components: int,
         kernel: OnlineKernel,
         eig_solver: EigSolverOptions,
-        chunk_size: int = 8192,
     ):
         """
         Args:
@@ -56,8 +49,6 @@ class OnlineNystrom(TorchTransformerMixin):
         self.eig_solver: EigSolverOptions = eig_solver
         self.shape: torch.Size = None               # ...
 
-        self.chunk_size = chunk_size
-
         # Anchor matrices
         self.anchor_features: torch.Tensor = None   # [... x n x d]
         self.A: torch.Tensor = None                 # [... x n x n]
@@ -68,12 +59,13 @@ class OnlineNystrom(TorchTransformerMixin):
         # Updated matrices
         self.S: torch.Tensor = None                 # [... x n x n]
         self.transform_matrix: torch.Tensor = None  # [... x n x n_components]
-        self.eigenvalues_: torch.Tensor = None      # [... x n]
+        self.eigenvalues_: torch.Tensor = None      # [... x n_components]
 
     def _update_to_kernel(self, d: int) -> Tuple[torch.Tensor, torch.Tensor]:
-        self.A = self.S = self.kernel.transform()
+        self.A = self.kernel.transform()
+        self.S = torch.nan_to_num(self.A, nan=0.0)
         U, L = solve_eig(
-            self.A,
+            self.S,
             num_eig=d + 1,  # d * (d + 3) // 2 + 1,
             eig_solver=self.eig_solver,
         )                                                                                           # [... x n x (? + 1)], [... x (? + 1)]
@@ -83,10 +75,6 @@ class OnlineNystrom(TorchTransformerMixin):
         return U, L
 
     def fit(self, features: torch.Tensor) -> "OnlineNystrom":
-        OnlineNystrom.fit_transform(self, features)
-        return self
-
-    def fit_transform(self, features: torch.Tensor) -> torch.Tensor:
         self.anchor_features = features
 
         self.kernel.fit(self.anchor_features)
@@ -94,11 +82,11 @@ class OnlineNystrom(TorchTransformerMixin):
 
         self.transform_matrix = (U / L[..., None, :])[..., :, :self.n_components]                   # [... x n x n_components]
         self.eigenvalues_ = L[..., :self.n_components]                                              # [... x n_components]
-        return U[..., :, :self.n_components]                                                        # [... x n x n_components]
+        return self
 
     def update(self, features: torch.Tensor) -> torch.Tensor:
         d = features.shape[-1]
-        n_chunks = ceildiv(features.shape[-2], self.chunk_size)
+        n_chunks = ceildiv(features.shape[-2], CHUNK_SIZE)
         if n_chunks > 1:
             """ Chunked version """
             chunks = torch.chunk(features, n_chunks, dim=-2)
@@ -134,119 +122,24 @@ class OnlineNystrom(TorchTransformerMixin):
 
             return B.mT @ self.transform_matrix                                                     # [... x m x n_components]
 
-    def transform(self, features: torch.Tensor) -> torch.Tensor:
-        n_chunks = ceildiv(features.shape[-2], self.chunk_size)
-        if n_chunks > 1:
-            """ Chunked version """
-            chunks = torch.chunk(features, n_chunks, dim=-2)
-            VS = []
-            for chunk in chunks:
-                VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                     # [... x _m x n_components]
-            VS = torch.cat(VS, dim=-2)
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:
+        if features is None:
+            VS = self.A @ self.transform_matrix                                                     # [... x n x n_components]
         else:
-            """ Unchunked version """
-            VS = self.kernel.transform(features) @ self.transform_matrix                            # [... x m x n_components]
+            n_chunks = ceildiv(features.shape[-2], CHUNK_SIZE)
+            if n_chunks > 1:
+                """ Chunked version """
+                chunks = torch.chunk(features, n_chunks, dim=-2)
+                VS = []
+                for chunk in chunks:
+                    VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                 # [... x _m x n_components]
+                VS = torch.cat(VS, dim=-2)
+            else:
+                """ Unchunked version """
+                VS = self.kernel.transform(features) @ self.transform_matrix                        # [... x m x n_components]
         return VS                                                                                   # [... x m x n_components]
 
 
-class OnlineNystromSubsampleFit(OnlineNystrom):
-    def __init__(
-        self,
-        n_components: int,
-        kernel: OnlineKernel,
-        distance_type: DistanceOptions,
-        sample_config: SampleConfig,
-        eig_solver: EigSolverOptions = "svd_lowrank",
-        chunk_size: int = 8192,
-    ):
-        OnlineNystrom.__init__(
-            self,
-            n_components=n_components,
-            kernel=kernel,
-            eig_solver=eig_solver,
-            chunk_size=chunk_size,
-        )
-        self.distance_type: DistanceOptions = distance_type
-        self.sample_config: SampleConfig = sample_config
-        self.sample_config._ncut_obj = copy.deepcopy(self)
-        self.anchor_indices: torch.Tensor = None
-
-    def _fit_helper(
-        self,
-        features: torch.Tensor,
-        precomputed_sampled_indices: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        _n = features.shape[-2]
-        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}"
-            )
-            self.num_sample = _n
-
-        if precomputed_sampled_indices is not None:
-            self.anchor_indices = precomputed_sampled_indices
-        else:
-            self.anchor_indices = subsample_features(
-                features=features,
-                distance_type=self.distance_type,
-                config=self.sample_config,
-            )
-        sampled_features = torch.gather(features, -2, self.anchor_indices[..., None].expand([-1] * self.anchor_indices.ndim + [features.shape[-1]]))
-        OnlineNystrom.fit(self, sampled_features)
-
-        _n_not_sampled = _n - self.anchor_indices.shape[-1]
-        if _n_not_sampled > 0:
-            unsampled_mask = torch.full(features.shape[:-1], True, device=features.device).scatter_(-1, self.anchor_indices, False)
-            unsampled_indices = torch.where(unsampled_mask)[-1].view((*features.shape[:-2], -1))
-            unsampled_features = torch.gather(features, -2, unsampled_indices[..., None].expand([-1] * unsampled_indices.ndim + [features.shape[-1]]))
-            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,
-    ) -> "OnlineNystromSubsampleFit":
-        """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
-        """
-        OnlineNystromSubsampleFit._fit_helper(self, features, precomputed_sampled_indices)
-        return self
-
-    def fit_transform(
-        self,
-        features: torch.Tensor,
-        precomputed_sampled_indices: torch.Tensor = None,
-    ) -> torch.Tensor:
-        """
-        Args:
-            features (torch.Tensor): input features, shape (n_samples, n_features)
-            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
-                override the sample_method, if not None
-
-        Returns:
-            (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
-            (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
-        """
-        unsampled_indices, V_unsampled = OnlineNystromSubsampleFit._fit_helper(self, features, precomputed_sampled_indices)
-        V_sampled = OnlineNystrom.transform(self, self.anchor_features)
-
-        if unsampled_indices is not None:
-            V = torch.zeros((*features.shape[:-1], self.n_components), device=features.device)
-            for (indices, _V) in [(self.anchor_indices, V_sampled), (unsampled_indices, V_unsampled)]:
-                V.scatter_(-2, indices[..., None].expand([-1] * indices.ndim + [self.n_components]), _V)
-        else:
-            V = V_sampled
-        return V
-
-
 def solve_eig(
     A: torch.Tensor,
     num_eig: int,
@@ -269,7 +162,7 @@ def solve_eig(
     bsz: int = A.shape[0]
 
     A = A + eig_value_buffer * torch.eye(A.shape[-1], device=A.device)
-
+    num_eig = min(A.shape[-1], num_eig)
     # compute eigenvectors
     if eig_solver == "svd_lowrank":  # default
         # only top q eigenvectors, fastest

nystrom_ncut/sampling_utils.py

@@ -1,5 +1,6 @@
+import copy
 from dataclasses import dataclass
-from typing import Literal
+from typing import Any, Literal, Tuple
 
 import torch
 from pytorch3d.ops import sample_farthest_points
@@ -13,6 +14,7 @@ from .distance_utils import (
 )
 from .transformer import (
     TorchTransformerMixin,
+    OnlineTorchTransformerMixin,
 )
 
 
@@ -21,11 +23,11 @@ SampleOptions = Literal["full", "random", "fps", "fps_recursive"]
 
 @dataclass
 class SampleConfig:
-    method: SampleOptions = "fps"
+    method: SampleOptions = "full"
     num_sample: int = 10000
     fps_dim: int = 12
     n_iter: int = None
-    _ncut_obj: TorchTransformerMixin = None
+    _recursive_obj: TorchTransformerMixin = None
 
 
 @torch.no_grad()
@@ -56,7 +58,7 @@ def subsample_features(
                         distance_type=distance_type,
                         config=SampleConfig(method="fps", num_sample=config.num_sample, fps_dim=config.fps_dim)
                     )                                                                                   # int: [... x num_sample]
-                    nc = config._ncut_obj
+                    nc = config._recursive_obj
                     for _ in range(config.n_iter):
                         fps_features, eigenvalues = nc.fit_transform(features, precomputed_sampled_indices=sampled_indices)
 
@@ -99,3 +101,95 @@ def fpsample(
     sample_indices = torch.gather(order, 1, sample_indices)                             # int: [(...) x num_sample]
 
     return sample_indices.view((*shape, *sample_indices.shape[-1:]))                    # int: [... x num_sample]
+
+
+class OnlineTransformerSubsampleFit(TorchTransformerMixin, OnlineTorchTransformerMixin):
+    def __init__(
+        self,
+        base_transformer: OnlineTorchTransformerMixin,
+        distance_type: DistanceOptions,
+        sample_config: SampleConfig,
+    ):
+        OnlineTorchTransformerMixin.__init__(self)
+        self.base_transformer: OnlineTorchTransformerMixin = base_transformer
+        self.distance_type: DistanceOptions = distance_type
+        self.sample_config: SampleConfig = sample_config
+        self.sample_config._recursive_obj = copy.deepcopy(self)
+        self.anchor_indices: torch.Tensor = None
+
+    def _fit_helper(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        _n = features.shape[-2]
+        self.sample_config.num_sample = min(self.sample_config.num_sample, _n)
+
+        if precomputed_sampled_indices is not None:
+            self.anchor_indices = precomputed_sampled_indices
+        else:
+            self.anchor_indices = subsample_features(
+                features=features,
+                distance_type=self.distance_type,
+                config=self.sample_config,
+            )
+        sampled_features = torch.gather(features, -2, self.anchor_indices[..., None].expand([-1] * self.anchor_indices.ndim + [features.shape[-1]]))
+        self.base_transformer.fit(sampled_features)
+
+        _n_not_sampled = _n - self.anchor_indices.shape[-1]
+        if _n_not_sampled > 0:
+            unsampled_mask = torch.full(features.shape[:-1], True, device=features.device).scatter_(-1, self.anchor_indices, False)
+            unsampled_indices = torch.where(unsampled_mask)[-1].view((*features.shape[:-2], -1))
+            unsampled_features = torch.gather(features, -2, unsampled_indices[..., None].expand([-1] * unsampled_indices.ndim + [features.shape[-1]]))
+            V_unsampled = self.base_transformer.update(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,
+    ) -> "OnlineTransformerSubsampleFit":
+        """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._fit_helper(features, precomputed_sampled_indices)
+        return self
+
+    def fit_transform(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            features (torch.Tensor): input features, shape (n_samples, n_features)
+            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+                override the sample_method, if not None
+
+        Returns:
+            (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
+            (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
+        """
+        unsampled_indices, V_unsampled = self._fit_helper(features, precomputed_sampled_indices)
+        V_sampled = self.base_transformer.transform()
+
+        if unsampled_indices is not None:
+            V = torch.zeros((*features.shape[:-1], V_sampled.shape[-1]), device=features.device)
+            for (indices, _V) in [(self.anchor_indices, V_sampled), (unsampled_indices, V_unsampled)]:
+                V.scatter_(-2, indices[..., None].expand([-1] * indices.ndim + [V_sampled.shape[-1]]), _V)
+        else:
+            V = V_sampled
+        return V
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        return self.base_transformer.update(features)
+
+    def transform(self, features: torch.Tensor = None, **transform_kwargs) -> torch.Tensor:
+        return self.base_transformer.transform(features)

nystrom_ncut/transformer/__init__.py

@@ -1,5 +1,6 @@
 from .transformer_mixin import (
     TorchTransformerMixin,
+    OnlineTorchTransformerMixin,
 )
 from .axis_align import (
     AxisAlign,

nystrom_ncut/transformer/transformer_mixin.py

@@ -36,15 +36,28 @@ class TorchTransformerMixin:
     >>> transformer.fit_transform(X)
     array([1, 1, 1])
     """
+    @abstractmethod
+    def fit(self, X: torch.Tensor, **fit_kwargs: Any) -> "OnlineTorchTransformerMixin":
+        """"""
+
+    @abstractmethod
+    def transform(self, X: torch.Tensor = None, **transform_kwargs: Any) -> torch.Tensor:
+        """"""
+
+    @abstractmethod
+    def fit_transform(self, X: torch.Tensor, **fit_transform_kwargs: Any) -> torch.Tensor:
+        """"""
+
 
+class OnlineTorchTransformerMixin:
     @abstractmethod
-    def fit(self, X: torch.Tensor, **fit_kwargs: Any) -> "TorchTransformerMixin":
+    def fit(self, X: torch.Tensor) -> "OnlineTorchTransformerMixin":
         """"""
 
     @abstractmethod
-    def transform(self, X: torch.Tensor, **transform_kwargs: Any) -> torch.Tensor:
+    def transform(self, X: torch.Tensor = None) -> torch.Tensor:
         """"""
 
     @abstractmethod
-    def fit_transform(self, X: torch.Tensor, **kwargs: Any) -> torch.Tensor:
+    def update(self, X: torch.Tensor) -> torch.Tensor:
         """"""

nystrom_ncut/visualize_utils.py

@@ -17,6 +17,9 @@ from .distance_utils import (
     to_euclidean,
     affinity_from_features,
 )
+from .global_settings import (
+    CHUNK_SIZE,
+)
 from .sampling_utils import (
     SampleConfig,
     subsample_features,
@@ -30,7 +33,6 @@ def extrapolate_knn(
     affinity_type: AffinityOptions,
     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']
@@ -42,7 +44,6 @@ def extrapolate_knn(
         extrapolation_features (torch.Tensor): features from existing nodes, shape (new_num_samples, n_features)
         knn (int): number of KNN to propagate eige nvectors
         affinity_type (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)
@@ -61,12 +62,17 @@ def extrapolate_knn(
     # propagate eigen_vector from subgraph to full graph
     anchor_output = anchor_output.to(device)
 
-    n_chunks = ceildiv(extrapolation_features.shape[0], chunk_size)
+    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, affinity_type).mT        # [_m x n]
+        _A = affinity_from_features(
+            features_A=anchor_features,
+            features_B=_v,
+            affinity_type=affinity_type,
+            affinity_focal_gamma=affinity_focal_gamma,
+        ).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]
@@ -93,7 +99,6 @@ def extrapolate_knn_with_subsampling(
     affinity_type: AffinityOptions,
     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']
@@ -104,7 +109,6 @@ def extrapolate_knn_with_subsampling(
         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)
@@ -138,7 +142,6 @@ def extrapolate_knn_with_subsampling(
         affinity_type,
         knn=knn,
         affinity_focal_gamma=affinity_focal_gamma,
-        chunk_size=chunk_size,
         device=device,
         move_output_to_cpu=move_output_to_cpu,
     )

{nystrom_ncut-0.3.2.dist-info → nystrom_ncut-0.3.4.dist-info}/METADATA

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

@@ -0,0 +1,21 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nystrom_ncut/__init__.py,sha256=4qNyWD5s1Uvd9OpfiMV4mF-3yFCi_K2QVRJIcAOXh70,587
+nystrom_ncut/common.py,sha256=eie19AHTMk6AGTxNnYq1UcFkHJVimeywAUYryXwaiHk,2428
+nystrom_ncut/distance_utils.py,sha256=zMI651RlIbd6ygvIxRp6jY5Ilfu7j9WQ5FD4I1wmmeo,4198
+nystrom_ncut/global_settings.py,sha256=TckHuF8geWM2ofd99jBupHD3TQdeEB583_3pdIVRU34,24
+nystrom_ncut/sampling_utils.py,sha256=lPWWNcDMBIvK9MmtNkVUeokeCOH9P1Fm8TZZ4sjlCpM,9037
+nystrom_ncut/visualize_utils.py,sha256=A1qmL8eNZjtvOOlyl9KIeObnpPVUGZVsCB9QBRV7n9I,22762
+nystrom_ncut/kernel/__init__.py,sha256=pvJ3tFukmlNZqw8VUB_iKPY9gbchUAlNkmnRCZcp0uU,44
+nystrom_ncut/kernel/kernel_ncut.py,sha256=u2oIu-SmNGn9tg-SXBJY4G9WGJ0SeGSnyZcvX8DGAE0,5276
+nystrom_ncut/nystrom/__init__.py,sha256=NHse-dW4nTo9wJUEJ6G4_Gw8uAi2vP6Hxm9aeBWxmqc,49
+nystrom_ncut/nystrom/distance_realization.py,sha256=kvPS-jGUn85MRJx-Dh2IZJ3IwvavRDCbXq6wh_aEBxc,5684
+nystrom_ncut/nystrom/normalized_cut.py,sha256=BD1F9Wz1BXbTGC-AVwT4IGmsPp334z-7jE9CuwhpNjY,7415
+nystrom_ncut/nystrom/nystrom_utils.py,sha256=5cMoF8UFgi_N-nEzbSQqVGhuep_eOn-FzErCFMqT7VM,9784
+nystrom_ncut/transformer/__init__.py,sha256=2FJEG9CXavfDDdDk1i9OOGkd8uSOHMkP8LBH49nnPnM,138
+nystrom_ncut/transformer/axis_align.py,sha256=j3LlAPrp8O_jQAlwZz-gu3D7n_wICEJranye-YK5wvA,4880
+nystrom_ncut/transformer/transformer_mixin.py,sha256=9wYdWknnJP7jijz70lRAyDn_kpC9aWwYN7Pbt5Mf6gQ,2012
+nystrom_ncut-0.3.4.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.3.4.dist-info/METADATA,sha256=I3V55oiDNJl4hw9v7TdEb1ibmJWtU5Ife9KwJ03rJPQ,6058
+nystrom_ncut-0.3.4.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+nystrom_ncut-0.3.4.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.3.4.dist-info/RECORD,,

nystrom_ncut-0.3.2.dist-info/RECORD

@@ -1,18 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nystrom_ncut/__init__.py,sha256=tKq9-2QRNFetckHY77qAaKEMjMCYTYcorS2f74aNtvk,540
-nystrom_ncut/common.py,sha256=eie19AHTMk6AGTxNnYq1UcFkHJVimeywAUYryXwaiHk,2428
-nystrom_ncut/distance_utils.py,sha256=pJA8NcIKyS7-YDpRGOkc7mwBQQEsYVemdkHiTjyU4n8,4300
-nystrom_ncut/sampling_utils.py,sha256=6lP8F6gftl4mgkavPsD7Vuk4erj4RtgILPhcj3YqLXk,4840
-nystrom_ncut/visualize_utils.py,sha256=Sfi_kKpvFFzBFoJnbo-pQpH2jhs-A6tH64SV_WGoq58,22740
-nystrom_ncut/nystrom/__init__.py,sha256=1aUXK87g4cXRXqNt6XkZsfyauw1-yv3sv0NmdmkWo-8,42
-nystrom_ncut/nystrom/distance_realization.py,sha256=RTI1_Q8fCUGAPSbXaVuNA-2B-11CEAfy2CwKWPJj6xQ,5830
-nystrom_ncut/nystrom/normalized_cut.py,sha256=cjkG8JeDmTPDK8KwfkAIqF9f1dI-D9s1muJ9WWZlUoc,7237
-nystrom_ncut/nystrom/nystrom_utils.py,sha256=hksDO8uuAb9xKoA1ZafGwXDlQN_gZJn_qHscaSoO8JE,14120
-nystrom_ncut/transformer/__init__.py,sha256=jjXjcNp3LrxeF6mqG9VY5k3asrqaY6bXzJz6wTpH78Q,105
-nystrom_ncut/transformer/axis_align.py,sha256=j3LlAPrp8O_jQAlwZz-gu3D7n_wICEJranye-YK5wvA,4880
-nystrom_ncut/transformer/transformer_mixin.py,sha256=YAjrDWTL5Hjnk9J2OsoxvtwT2N0u8IdgMSx0rRFmZzE,1653
-nystrom_ncut-0.3.2.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.3.2.dist-info/METADATA,sha256=4E42fHnXLNvbErWrwxE5K_oeOdpi5Bfabed9VF-YkV0,6058
-nystrom_ncut-0.3.2.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-nystrom_ncut-0.3.2.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.3.2.dist-info/RECORD,,