nystrom-ncut 0.3.0__tar.gz → 0.3.2__tar.gz

{nystrom_ncut-0.3.0/src/nystrom_ncut.egg-info → nystrom_ncut-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.3.0
+Version: 0.3.2
 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.0 → nystrom_ncut-0.3.2}/pyproject.toml

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

{nystrom_ncut-0.3.0 → nystrom_ncut-0.3.2}/src/nystrom_ncut/nystrom/normalized_cut.py

@@ -54,10 +54,10 @@ class LaplacianKernel(OnlineKernel):
             self.A,
             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.diag_embed(1 / L) @ 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.b_r = torch.zeros_like(self.a_r)                                       # [... x n]
+        )                                                                                           # [... 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.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(

nystrom_ncut-0.3.2/src/nystrom_ncut/transformer/axis_align.py

@@ -0,0 +1,102 @@
+from typing import Literal
+
+import torch
+import torch.nn.functional as Fn
+
+from ..common import (
+    default_device,
+)
+from .transformer_mixin import (
+    TorchTransformerMixin,
+)
+
+
+class AxisAlign(TorchTransformerMixin):
+    """Multiclass Spectral Clustering, SX Yu, J Shi, 2003
+    Args:
+        max_iter (int, optional): Maximum number of iterations.
+    """
+    SortOptions = Literal["count", "norm", "marginal_norm"]
+
+    def __init__(
+        self,
+        sort_method: SortOptions = "norm",
+        max_iter: int = 100,
+    ):
+        self.sort_method: AxisAlign.SortOptions = sort_method
+        self.max_iter: int = max_iter
+
+        self.R: torch.Tensor = None
+
+    def fit(self, X: torch.Tensor) -> "AxisAlign":
+        # Normalize eigenvectors
+        with default_device(X.device):
+            d = X.shape[-1]
+            normalized_X = Fn.normalize(X, p=2, dim=-1)                                                         # float: [... x n x d]
+
+            # Initialize R matrix with the first column from a random row of EigenVectors
+            def get_idx(idx: torch.Tensor) -> torch.Tensor:
+                return torch.gather(normalized_X, -2, idx[..., None, None].expand([-1] * (X.ndim - 2) + [1, d]))[..., 0, :]
+
+            self.R = torch.empty((*X.shape[:-2], d, d))                                                         # float: [... x d x d]
+            mask = torch.all(torch.isfinite(normalized_X), dim=-1)                                              # bool: [... x n]
+            start_idx = torch.argmax(mask.to(torch.float) + torch.rand(mask.shape), dim=-1)                     # int: [...]
+            self.R[..., 0, :] = get_idx(start_idx)
+
+            # Loop to populate R with k orthogonal directions
+            c = torch.zeros(X.shape[:-1])                                                                       # float: [... x n]
+            for i in range(1, d):
+                c += torch.abs(normalized_X @ self.R[..., i - 1, :, None])[..., 0]
+                self.R[..., i, :] = get_idx(torch.argmin(c.nan_to_num(nan=torch.inf), dim=-1))
+
+            # Iterative optimization loop
+            normalized_X = torch.nan_to_num(normalized_X, nan=0.0)
+            idx, prev_objective = None, torch.inf
+            for _ in range(self.max_iter):
+                # Discretize the projected eigenvectors
+                idx = torch.argmax(normalized_X @ self.R.mT, dim=-1)                                                    # int: [... x n]
+                M = torch.sum((idx[..., None] == torch.arange(d))[..., None] * normalized_X[..., :, None, :], dim=-3)   # float: [... x d x d]
+
+                # Check for convergence
+                objective = torch.norm(M)
+                if torch.abs(objective - prev_objective) < torch.finfo(torch.float32).eps:
+                    break
+                prev_objective = objective
+
+                # SVD decomposition to compute the next R
+                U, S, Vh = torch.linalg.svd(M, full_matrices=False)
+                self.R = U @ Vh
+
+            # Permute the rotation matrix so the dimensions are sorted in descending cluster significance
+            match self.sort_method:
+                case "count":
+                    sort_metric = torch.sum((idx[..., None] == torch.arange(d)), dim=-2)
+                case "norm":
+                    rotated_X = torch.nan_to_num(X @ self.R.mT, nan=0.0)
+                    sort_metric = torch.linalg.norm(rotated_X, dim=-2)
+                case "marginal_norm":
+                    rotated_X = torch.nan_to_num(X @ self.R.mT, nan=0.0)
+                    sort_metric = torch.sum((idx[..., None] == torch.arange(d)) * (torch.gather(rotated_X, -1, idx[..., None]) ** 2), dim=-2)
+                case _:
+                    raise ValueError(f"Invalid sort method {self.sort_method}.")
+
+            order = torch.argsort(sort_metric, dim=-1, descending=True)
+            self.R = torch.gather(self.R, -2, order[..., None].expand([-1] * order.ndim + [d]))
+            return self
+
+    def transform(self, X: torch.Tensor, normalize: bool = True, hard: bool = False) -> torch.Tensor:
+        """
+        Args:
+            X (torch.Tensor): continuous eigenvectors from NCUT, shape (n, k)
+            normalize (bool): whether to normalize input features before rotating
+            hard (bool): whether to return cluster indices of input features or just the rotated features
+        Returns:
+            torch.Tensor: Discretized eigenvectors, shape (n, k), each row is a one-hot vector.
+        """
+        if normalize:
+            X = Fn.normalize(X, p=2, dim=-1)
+        rotated_X = X @ self.R.mT
+        return torch.argmax(rotated_X, dim=-1) if hard else rotated_X
+
+    def fit_transform(self, X: torch.Tensor, normalize: bool = True, hard: bool = False) -> torch.Tensor:
+        return self.fit(X).transform(X, normalize=normalize, hard=hard)

{nystrom_ncut-0.3.0 → nystrom_ncut-0.3.2/src/nystrom_ncut.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.3.0
+Version: 0.3.2
 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.0/src/nystrom_ncut/transformer/axis_align.py

@@ -1,90 +0,0 @@
-from typing import Literal
-
-import torch
-import torch.nn.functional as Fn
-
-from .transformer_mixin import (
-    TorchTransformerMixin,
-)
-
-
-class AxisAlign(TorchTransformerMixin):
-    """Multiclass Spectral Clustering, SX Yu, J Shi, 2003
-    Args:
-        max_iter (int, optional): Maximum number of iterations.
-    """
-    SortOptions = Literal["count", "norm", "marginal_norm"]
-
-    def __init__(
-        self,
-        sort_method: SortOptions = "norm",
-        max_iter: int = 100,
-    ):
-        self.sort_method: AxisAlign.SortOptions = sort_method
-        self.max_iter: int = max_iter
-
-        self.R: torch.Tensor = None
-
-    def fit(self, X: torch.Tensor) -> "AxisAlign":
-        # Normalize eigenvectors
-        n, d = X.shape
-        normalized_X = Fn.normalize(X, p=2, dim=-1)
-
-        # Initialize R matrix with the first column from a random row of EigenVectors
-        self.R = torch.empty((d, d), device=X.device)
-        self.R[0] = normalized_X[torch.randint(0, n, (), device=X.device)]
-
-        # Loop to populate R with k orthogonal directions
-        c = torch.zeros((n,), device=X.device)
-        for i in range(1, d):
-            c += torch.abs(normalized_X @ self.R[i - 1])
-            self.R[i] = normalized_X[torch.argmin(c, dim=0)]
-
-        # Iterative optimization loop
-        idx, prev_objective = None, torch.inf
-        for _ in range(self.max_iter):
-            # Discretize the projected eigenvectors
-            idx = torch.argmax(normalized_X @ self.R.mT, dim=-1)
-            M = torch.zeros((d, d), device=X.device).index_add_(0, idx, normalized_X)
-
-            # Check for convergence
-            objective = torch.norm(M)
-            if torch.abs(objective - prev_objective) < torch.finfo(torch.float32).eps:
-                break
-            prev_objective = objective
-
-            # SVD decomposition to compute the next R
-            U, S, Vh = torch.linalg.svd(M, full_matrices=False)
-            self.R = U @ Vh
-
-        # Permute the rotation matrix so the dimensions are sorted in descending cluster significance
-        if self.sort_method == "count":
-            sort_metric = torch.bincount(idx, minlength=d)
-        elif self.sort_method == "norm":
-            rotated_X = X @ self.R.mT
-            sort_metric = torch.linalg.norm(rotated_X, dim=0)
-        elif self.sort_method == "marginal_norm":
-            rotated_X = X @ self.R.mT
-            sort_metric = torch.zeros((d,), device=X.device).index_add_(0, idx, rotated_X[range(n), idx] ** 2)
-        else:
-            raise ValueError(f"Invalid sort method {self.sort_method}.")
-
-        self.R = self.R[torch.argsort(sort_metric, dim=0, descending=True)]
-        return self
-
-    def transform(self, X: torch.Tensor, normalize: bool = True, hard: bool = False) -> torch.Tensor:
-        """
-        Args:
-            X (torch.Tensor): continuous eigenvectors from NCUT, shape (n, k)
-            normalize (bool): whether to normalize input features before rotating
-            hard (bool): whether to return cluster indices of input features or just the rotated features
-        Returns:
-            torch.Tensor: Discretized eigenvectors, shape (n, k), each row is a one-hot vector.
-        """
-        if normalize:
-            X = Fn.normalize(X, p=2, dim=1)
-        rotated_X = X @ self.R.mT
-        return torch.argmax(rotated_X, dim=1) if hard else rotated_X
-
-    def fit_transform(self, X: torch.Tensor, normalize: bool = True, hard: bool = False) -> torch.Tensor:
-        return self.fit(X).transform(X, normalize=normalize, hard=hard)