nystrom-ncut 0.3.4__tar.gz → 0.3.6__tar.gz

{nystrom_ncut-0.3.4/src/nystrom_ncut.egg-info → nystrom_ncut-0.3.6}/PKG-INFO

@@ -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 → nystrom_ncut-0.3.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nystrom_ncut"
-version = "0.3.4"
+version = "0.3.6"
 authors = [
     { name = "Huzheng Yang", email = "huze.yann@gmail.com" },
     { name = "Wentinn Liao", email = "wentinn.liao@gmail.com" },

{nystrom_ncut-0.3.4 → nystrom_ncut-0.3.6}/src/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-0.3.4 → nystrom_ncut-0.3.6}/src/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 → nystrom_ncut-0.3.6/src/nystrom_ncut.egg-info}/PKG-INFO

@@ -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/