nystrom-ncut 0.1.4__py3-none-any.whl → 0.1.6__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 * c)
-        D = D / (2 * torch.linalg.norm(stds) ** 2)
+        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,
     ):
@@ -128,9 +128,9 @@ class DistanceRealization(OnlineNystromSubsampleFit):
         features: torch.Tensor,
         precomputed_sampled_indices: torch.Tensor = None,
     ) -> torch.Tensor:
-        V, L = OnlineNystromSubsampleFit.fit_transform(self, features, precomputed_sampled_indices)
-        return V * (L ** 0.5)
+        V = OnlineNystromSubsampleFit.fit_transform(self, features, precomputed_sampled_indices)
+        return V * (self.eigenvalues_ ** 0.5)
 
     def transform(self, features: torch.Tensor = None) -> torch.Tensor:
-        V, L = OnlineNystromSubsampleFit.transform(self, features)
-        return V * (L ** 0.5)
+        V = OnlineNystromSubsampleFit.transform(self, features)
+        return V * (self.eigenvalues_ ** 0.5)

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/nystrom/nystrom_utils.py

@@ -20,7 +20,7 @@ EigSolverOptions = Literal["svd_lowrank", "lobpcg", "svd", "eigh"]
 
 
 class OnlineKernel:
-    def fit(self, features: torch.Tensor) -> None:                          # [n x d]
+    def fit(self, features: torch.Tensor) -> "OnlineKernel":                # [n x d]
         raise NotImplementedError()
 
     def update(self, features: torch.Tensor) -> torch.Tensor:               # [m x d] -> [m x n]
@@ -47,7 +47,6 @@ class OnlineNystrom:
         self.n_components: int = n_components
         self.kernel: OnlineKernel = kernel
         self.eig_solver: EigSolverOptions = eig_solver
-        self.inverse_approximation_dim: int = None
 
         self.chunk_size = chunk_size
 
@@ -61,13 +60,13 @@ class OnlineNystrom:
         # Updated matrices
         self.S: torch.Tensor = None                 # [n x n]
         self.transform_matrix: torch.Tensor = None  # [n x n_components]
-        self.LS: torch.Tensor = None                # [n]
+        self.eigenvalues_: torch.Tensor = None      # [n]
 
-    def _update_to_kernel(self) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _update_to_kernel(self, d: int) -> Tuple[torch.Tensor, torch.Tensor]:
         self.A = self.S = self.kernel.transform()
         U, L = solve_eig(
             self.A,
-            num_eig=self.inverse_approximation_dim,
+            num_eig=d + 1,  # d * (d + 3) // 2 + 1,
             eig_solver=self.eig_solver,
         )                                                                                           # [n x (? + 1)], [? + 1]
         self.Ahinv_UL = U * (L ** -0.5)                                                             # [n x (? + 1)]
@@ -75,57 +74,57 @@ class OnlineNystrom:
         self.Ahinv = self.Ahinv_UL @ self.Ahinv_VT                                                  # [n x n]
         return U, L
 
-    def fit(self, features: torch.Tensor):
+    def fit(self, features: torch.Tensor) -> "OnlineNystrom":
         OnlineNystrom.fit_transform(self, features)
         return self
 
-    def fit_transform(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def fit_transform(self, features: torch.Tensor) -> torch.Tensor:
         self.anchor_features = features
 
         self.kernel.fit(self.anchor_features)
-        self.inverse_approximation_dim = max(self.n_components, features.shape[-1] + 1)
-        U, L = self._update_to_kernel()                                                             # [n x (? + 1)], [? + 1]
+        U, L = self._update_to_kernel(features.shape[-1])                                           # [n x (d + 1)], [d + 1]
 
         self.transform_matrix = (U / L)[:, :self.n_components]                                      # [n x n_components]
-        self.LS = L[:self.n_components]                                                             # [n_components]
-        return U[:, :self.n_components], L[:self.n_components]                                      # [n x n_components], [n_components]
+        self.eigenvalues_ = L[:self.n_components]                                                   # [n_components]
+        return U[:, :self.n_components]                                                             # [n x n_components]
 
-    def update(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        d = features.shape[-1]
         n_chunks = ceildiv(len(features), self.chunk_size)
         if n_chunks > 1:
             """ Chunked version """
             chunks = torch.chunk(features, n_chunks, dim=0)
             for chunk in chunks:
                 self.kernel.update(chunk)
-            self._update_to_kernel()
+            self._update_to_kernel(d)
 
-            compressed_BBT = torch.zeros((self.inverse_approximation_dim, self.inverse_approximation_dim))  # [(? + 1) x (? + 1))]
-            for i, chunk in enumerate(chunks):
-                _B = self.kernel.transform(chunk).mT                                                    # [n x _m]
-                _compressed_B = self.Ahinv_VT @ _B                                                      # [(? + 1) x _m]
-                compressed_BBT = compressed_BBT + _compressed_B @ _compressed_B.mT                      # [(? + 1) x (? + 1)]
-            self.S = self.S + self.Ahinv_UL @ compressed_BBT @ self.Ahinv_UL.mT                         # [n x n]
-            US, self.LS = solve_eig(self.S, self.n_components, self.eig_solver)                         # [n x n_components], [n_components]
-            self.transform_matrix = self.Ahinv @ US * (self.LS ** -0.5)                                 # [n x n_components]
+            compressed_BBT = 0.0                                                                    # [(? + 1) x (? + 1))]
+            for chunk in chunks:
+                _B = self.kernel.transform(chunk).mT                                                # [n x _m]
+                _compressed_B = self.Ahinv_VT @ _B                                                  # [(? + 1) x _m]
+                compressed_BBT = compressed_BBT + _compressed_B @ _compressed_B.mT                  # [(? + 1) x (? + 1)]
+            self.S = self.S + self.Ahinv_UL @ compressed_BBT @ self.Ahinv_UL.mT                     # [n x n]
+            US, self.eigenvalues_ = solve_eig(self.S, self.n_components, self.eig_solver)           # [n x n_components], [n_components]
+            self.transform_matrix = self.Ahinv @ US * (self.eigenvalues_ ** -0.5)                   # [n x n_components]
 
             VS = []
             for chunk in chunks:
-                VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                         # [_m x n_components]
+                VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                     # [_m x n_components]
             VS = torch.cat(VS, dim=0)
-            return VS, self.LS                                                                          # [m x n_components], [n_components]
+            return VS                                                                               # [m x n_components]
         else:
             """ Unchunked version """
-            B = self.kernel.update(features).mT                                                         # [n x m]
-            self._update_to_kernel()
-            compressed_B = self.Ahinv_VT @ B                                                            # [indirect_pca_dim x m]
+            B = self.kernel.update(features).mT                                                     # [n x m]
+            self._update_to_kernel(d)
+            compressed_B = self.Ahinv_VT @ B                                                        # [indirect_pca_dim x m]
 
-            self.S = self.S + self.Ahinv_UL @ (compressed_B @ compressed_B.mT) @ self.Ahinv_UL.mT       # [n x n]
-            US, self.LS = solve_eig(self.S, self.n_components, self.eig_solver)                         # [n x n_components], [n_components]
-            self.transform_matrix = self.Ahinv @ US * (self.LS ** -0.5)                                 # [n x n_components]
+            self.S = self.S + self.Ahinv_UL @ (compressed_B @ compressed_B.mT) @ self.Ahinv_UL.mT   # [n x n]
+            US, self.eigenvalues_ = solve_eig(self.S, self.n_components, self.eig_solver)           # [n x n_components], [n_components]
+            self.transform_matrix = self.Ahinv @ US * (self.eigenvalues_ ** -0.5)                   # [n x n_components]
 
-            return B.mT @ self.transform_matrix, self.LS                                                # [m x n_components], [n_components]
+            return B.mT @ self.transform_matrix                                                     # [m x n_components]
 
-    def transform(self, features: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:
         if features is None:
             VS = self.A @ self.transform_matrix                                                     # [n x n_components]
         else:
@@ -135,12 +134,12 @@ class OnlineNystrom:
                 chunks = torch.chunk(features, n_chunks, dim=0)
                 VS = []
                 for chunk in chunks:
-                    VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                     # [_m x n_components]
+                    VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                 # [_m x n_components]
                 VS = torch.cat(VS, dim=0)
             else:
                 """ Unchunked version """
-                VS = self.kernel.transform(features) @ self.transform_matrix                            # [m x n_components]
-        return VS, self.LS                                                                          # [m x n_components], [n_components]
+                VS = self.kernel.transform(features) @ self.transform_matrix                        # [m x n_components]
+        return VS                                                                                   # [m x n_components]
 
 
 class OnlineNystromSubsampleFit(OnlineNystrom):
@@ -192,7 +191,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         if _n_not_sampled > 0:
             unsampled_indices = torch.full((_n,), True, device=features.device).scatter_(0, self.anchor_indices, False)
             unsampled_features = features[unsampled_indices]
-            V_unsampled, _ = OnlineNystrom.update(self, unsampled_features)
+            V_unsampled = OnlineNystrom.update(self, unsampled_features)
         else:
             unsampled_indices = V_unsampled = None
         return unsampled_indices, V_unsampled
@@ -201,7 +200,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         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)
@@ -217,7 +216,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         self,
         features: torch.Tensor,
         precomputed_sampled_indices: torch.Tensor = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
         """
         Args:
             features (torch.Tensor): input features, shape (n_samples, n_features)
@@ -229,7 +228,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
             (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, L = OnlineNystrom.transform(self)
+        V_sampled = OnlineNystrom.transform(self)
 
         if unsampled_indices is not None:
             V = torch.zeros((len(unsampled_indices), self.n_components), device=features.device)
@@ -237,7 +236,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
             V[unsampled_indices] = V_unsampled
         else:
             V = V_sampled
-        return V, L
+        return V
 
 
 def solve_eig(

{nystrom_ncut-0.1.4.dist-info → nystrom_ncut-0.1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.1.4
+Version: 0.1.6
 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.6.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=InajllGtRVnLVlZoipZNbHFTGHaTs3zxizKe3kI2Los,5815
+nystrom_ncut/nystrom/normalized_cut.py,sha256=N-M5wkTo59vpbBfIx8evkSQBxlo4j80qCtuoifxQa_A,7578
+nystrom_ncut/nystrom/nystrom_utils.py,sha256=UVs1tC7vnVq2mWSTpcrP4C19x9wDJ77ACht0EltOO2E,12698
+nystrom_ncut-0.1.6.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.1.6.dist-info/METADATA,sha256=DfH6Ylr79RYc0JiYkDpz0kv08RxSP73dBhK8TVJSUfE,6058
+nystrom_ncut-0.1.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.1.6.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.1.6.dist-info/RECORD,,

nystrom_ncut-0.1.4.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=wEXyKxkLv9VvF5xn42ypiwUAAq9xTh3RmMvdHhFuPms,3127
-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=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.4.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.1.4.dist-info/METADATA,sha256=lVURVQRK6krrpu6iUvDdX0-I5FD7mHFIK_uQjzHlxhU,6058
-nystrom_ncut-0.1.4.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.1.4.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.1.4.dist-info/RECORD,,