nystrom-ncut 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

nystrom_ncut/distance_utils.py

@@ -38,14 +38,14 @@ def distance_from_features(
     elif distance == "euclidean":
         D = torch.cdist(features, features_B, p=2)
     elif distance == "rbf":
-        D = torch.cdist(features, features_B, p=2) ** 2
+        D = 0.5 * torch.cdist(features, features_B, p=2) ** 2
 
         # 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.quantile(features, q=(p + 1) / 2, dim=0)
-        stds = (stds[1] - stds[0]) / 2
-        D = D / (2 * torch.linalg.norm(stds) ** 2)
+        stds = (stds[1] - stds[0]) / (2 * c)
+        D = D / (torch.linalg.norm(stds) ** 2)
     else:
         raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
     return D

nystrom_ncut/nystrom/distance_realization.py

@@ -98,8 +98,8 @@ class DistanceRealization(OnlineNystromSubsampleFit):
     def __init__(
         self,
         n_components: int = 100,
-        sample_config: SampleConfig = SampleConfig(),
         distance: DistanceOptions = "cosine",
+        sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
     ):

nystrom_ncut/nystrom/normalized_cut.py

@@ -1,3 +1,4 @@
+import einops
 import torch
 import torch.nn.functional as Fn
 
@@ -21,10 +22,12 @@ class LaplacianKernel(OnlineKernel):
         self,
         affinity_focal_gamma: float,
         distance: DistanceOptions,
+        adaptive_scaling: bool,
         eig_solver: EigSolverOptions,
     ):
         self.affinity_focal_gamma = affinity_focal_gamma
         self.distance: DistanceOptions = distance
+        self.adaptive_scaling: bool = adaptive_scaling
         self.eig_solver: EigSolverOptions = eig_solver
 
         # Anchor matrices
@@ -53,13 +56,25 @@ class LaplacianKernel(OnlineKernel):
         self.a_r = torch.sum(self.A, dim=-1)                    # [n]
         self.b_r = torch.zeros_like(self.a_r)                   # [n]
 
-    def update(self, features: torch.Tensor) -> torch.Tensor:
+    def _affinity(self, features: torch.Tensor) -> torch.Tensor:
         B = affinity_from_features(
             self.anchor_features,                               # [n x d]
             features,                                           # [m x d]
             affinity_focal_gamma=self.affinity_focal_gamma,
             distance=self.distance,
         )                                                       # [n x m]
+        if self.adaptive_scaling:
+            diagonal = (
+                einops.rearrange(B, "n m -> m 1 n")             # [m x 1 x n]
+                @ self.Ainv                                     # [n x n]
+                @ einops.rearrange(B, "n m -> m n 1")           # [m x n x 1]
+            ).squeeze(1, 2)                                     # [m]
+            adaptive_scale = diagonal ** -0.5                   # [m]
+            B = B * adaptive_scale
+        return B                                                # [n x m]
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        B = self._affinity(features)                            # [n x m]
         b_r = torch.sum(B, dim=-1)                              # [n]
         b_c = torch.sum(B, dim=-2)                              # [m]
         self.b_r = self.b_r + b_r                               # [n]
@@ -75,12 +90,7 @@ class LaplacianKernel(OnlineKernel):
             B = self.A                                          # [n x n]
             col_sum = row_sum                                   # [n]
         else:
-            B = affinity_from_features(
-                self.anchor_features,                           # [n x d]
-                features,                                       # [m x d]
-                affinity_focal_gamma=self.affinity_focal_gamma,
-                distance=self.distance,
-            )                                                   # [n x m]
+            B = self._affinity(features)
             b_c = torch.sum(B, dim=-2)                          # [m]
             col_sum = b_c + B.mT @ self.Ainv @ self.b_r         # [m]
         scale = (row_sum[:, None] * col_sum) ** -0.5            # [n x m]
@@ -94,8 +104,9 @@ class NCut(OnlineNystromSubsampleFit):
         self,
         n_components: int = 100,
         affinity_focal_gamma: float = 1.0,
-        sample_config: SampleConfig = SampleConfig(),
         distance: DistanceOptions = "cosine",
+        adaptive_scaling: bool = True,
+        sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
     ):
@@ -104,16 +115,17 @@ 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.
+            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
-            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,
             n_components=n_components,
-            kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
+            kernel=LaplacianKernel(affinity_focal_gamma, distance, adaptive_scaling, eig_solver),
             distance=distance,
             sample_config=sample_config,
             eig_solver=eig_solver,

nystrom_ncut/sampling_utils.py

@@ -82,4 +82,7 @@ def fpsample(
         U, S, V = torch.pca_lowrank(features, q=config.fps_dim)
         features = U * S
 
-    return sample_farthest_points(features[None], K=config.num_sample)[1][0]
+    try:
+        return sample_farthest_points(features[None], K=config.num_sample)[1][0]
+    except RuntimeError:
+        return sample_farthest_points(features[None].cpu(), K=config.num_sample)[1][0].to(features.device)

{nystrom_ncut-0.1.3.dist-info → nystrom_ncut-0.1.5.dist-info}/METADATA

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

@@ -0,0 +1,15 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nystrom_ncut/__init__.py,sha256=HifrTcqX2-hYjBDe6xIThHvuIBYMPBA3EzjR8-qPMUM,512
+nystrom_ncut/common.py,sha256=_PGJoImSk_Fb_5Ri-e_IsFoCcSfbGS8CxYUUHVoNM50,2036
+nystrom_ncut/distance_utils.py,sha256=p-pYdpRrJsIhzxM_IxUqja7N8okngx52WGXD9pu_Aec,3129
+nystrom_ncut/sampling_utils.py,sha256=oMmhFcd_N_D15Ht7F0rCGPSgLeitJszAKMD3ICKwHNU,3105
+nystrom_ncut/visualize_utils.py,sha256=RsQVjPhxoIdxDOQ2PI7ifFDuEL23YXpZBdJ0wjjafek,22970
+nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
+nystrom_ncut/nystrom/distance_realization.py,sha256=YaKEzuiiJfS6tdOTuJ4Ce1jqOWIfrVGb8ilvhzKeW9w,5789
+nystrom_ncut/nystrom/normalized_cut.py,sha256=N-M5wkTo59vpbBfIx8evkSQBxlo4j80qCtuoifxQa_A,7578
+nystrom_ncut/nystrom/nystrom_utils.py,sha256=Wq364xlxBhr74lqyCkPWLBxq5YSt2zr-DSfYUHpYfgE,12989
+nystrom_ncut-0.1.5.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.1.5.dist-info/METADATA,sha256=7gO_GxpjbitbEBB9juaha7W4e7SWs-hYqgkYfaGmuIg,6058
+nystrom_ncut-0.1.5.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.1.5.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.1.5.dist-info/RECORD,,

nystrom_ncut-0.1.3.dist-info/RECORD

@@ -1,15 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nystrom_ncut/__init__.py,sha256=HifrTcqX2-hYjBDe6xIThHvuIBYMPBA3EzjR8-qPMUM,512
-nystrom_ncut/common.py,sha256=_PGJoImSk_Fb_5Ri-e_IsFoCcSfbGS8CxYUUHVoNM50,2036
-nystrom_ncut/distance_utils.py,sha256=xSMKL-sFGrz0EL106vhVx0qSk3iSdSLRFhGL0KmAOnU,3121
-nystrom_ncut/sampling_utils.py,sha256=7zCneqmkaA_fUkaZcykFjHtn7pxdZdsjAbKxJplegc0,2960
-nystrom_ncut/visualize_utils.py,sha256=RsQVjPhxoIdxDOQ2PI7ifFDuEL23YXpZBdJ0wjjafek,22970
-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=Wq364xlxBhr74lqyCkPWLBxq5YSt2zr-DSfYUHpYfgE,12989
-nystrom_ncut-0.1.3.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.1.3.dist-info/METADATA,sha256=tK9gbt1b0c5mGNm2e0OdAKQm7rFUWCuhn9myh_Vf408,6058
-nystrom_ncut-0.1.3.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.1.3.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.1.3.dist-info/RECORD,,