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

nystrom_ncut/kernel/kernel_ncut.py

@@ -1,5 +1,10 @@
+from typing import Dict
+
 import torch
 
+from ..common import (
+    lazy_normalize,
+)
 from ..distance_utils import (
     AffinityOptions,
     AFFINITY_TO_DISTANCE,
@@ -28,44 +33,61 @@ class KernelNCutBaseTransformer(OnlineTorchTransformerMixin):
         self.affinity_focal_gamma = affinity_focal_gamma
 
         # Anchor matrices
-        self.W: torch.Tensor = None                     # [... x d x kernel_dim]
+        self.anchor_count: int = None                   # n
         self.kernelized_anchor: torch.Tensor = None     # [... x n x (2 * kernel_dim)]
+        self.store: Dict[str, torch.Tensor] = {}
 
         # Updated matrices
+        self.total_count: int = None                    # m
         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 _kernelize_features(self, features: torch.Tensor) -> torch.Tensor:
+        match self.affinity_type:
+            case "cosine" | "rbf":
+                if self.affinity_type == "cosine":
+                    features = lazy_normalize(features)
+                W_features = features @ self.store["W"] # [... x m x kernel_dim]
+                return torch.cat((
+                    torch.cos(W_features),
+                    torch.sin(W_features),
+                ), dim=-1) / (self.kernel_dim ** 0.5)   # [... x m x (2 * kernel_dim)]
+
+            case _:
+                raise ValueError(self.affinity_type)
+
     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]
+        _, 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]
+        S = S * (self.total_count / self.anchor_count) ** 0.5
+        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.anchor_count = self.total_count = features.shape[-2]
+        shape, d = features.shape[:-2], features.shape[-1]
+
+        match self.affinity_type:
+            case "cosine" | "rbf":
+                scale = self.affinity_focal_gamma ** 0.5
+                if self.affinity_type == "rbf":
+                    scale = get_normalization_factor(features)[..., None, None] * scale                     # [... x 1 x 1]
+                self.store["W"] = torch.randn((*shape, d, self.kernel_dim), device=features.device) / scale # [... x d x kernel_dim]
+
+            case _:
+                raise ValueError(self.affinity_type)
+
+        self.kernelized_anchor = self._kernelize_features(features)                     # [... 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.total_count += features.shape[-2]
+        kernelized_features = self._kernelize_features(features)                        # [... 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()
 
@@ -77,11 +99,8 @@ class KernelNCutBaseTransformer(OnlineTorchTransformerMixin):
         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)]
+            kernelized_features = self._kernelize_features(features)                # [... 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]

nystrom_ncut/sampling_utils.py

@@ -193,3 +193,7 @@ class OnlineTransformerSubsampleFit(TorchTransformerMixin, OnlineTorchTransforme
 
     def transform(self, features: torch.Tensor = None, **transform_kwargs) -> torch.Tensor:
         return self.base_transformer.transform(features)
+
+    @property
+    def eigenvalues_(self) -> torch.Tensor:
+        return getattr(self.base_transformer, "eigenvalues_", None)

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

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

@@ -3,10 +3,10 @@ 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/sampling_utils.py,sha256=Njfge6E0DzwWgS0hk1Os3VEhw896Z4eLWKhX_FTQado,9164
 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/kernel/kernel_ncut.py,sha256=QoQ-wRBZmdL9WcuzaHr2A3aPN-iuFL3UIUKpqFmX1_k,5809
 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
@@ -14,8 +14,8 @@ nystrom_ncut/nystrom/nystrom_utils.py,sha256=5cMoF8UFgi_N-nEzbSQqVGhuep_eOn-FzEr
 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.6.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.3.6.dist-info/METADATA,sha256=T-Skh2IsSTV7xwJ7QO9FFTqe0yatgv_kw-7StlQj568,6058
+nystrom_ncut-0.3.6.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
+nystrom_ncut-0.3.6.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.3.6.dist-info/RECORD,,

{nystrom_ncut-0.3.4.dist-info → nystrom_ncut-0.3.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.0.0)
+Generator: setuptools (76.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any