nystrom-ncut 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

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/transformer/axis_align.py

@@ -3,6 +3,9 @@ from typing import Literal
 import torch
 import torch.nn.functional as Fn
 
+from ..common import (
+    default_device,
+)
 from .transformer_mixin import (
     TorchTransformerMixin,
 )
@@ -27,50 +30,59 @@ class AxisAlign(TorchTransformerMixin):
 
     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
+        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:
         """
@@ -82,9 +94,9 @@ class AxisAlign(TorchTransformerMixin):
             torch.Tensor: Discretized eigenvectors, shape (n, k), each row is a one-hot vector.
         """
         if normalize:
-            X = Fn.normalize(X, p=2, dim=1)
+            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
+        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.dist-info → nystrom_ncut-0.3.2.dist-info}/METADATA

@@ -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.dist-info → nystrom_ncut-0.3.2.dist-info}/RECORD

@@ -6,13 +6,13 @@ nystrom_ncut/sampling_utils.py,sha256=6lP8F6gftl4mgkavPsD7Vuk4erj4RtgILPhcj3YqLX
 nystrom_ncut/visualize_utils.py,sha256=Sfi_kKpvFFzBFoJnbo-pQpH2jhs-A6tH64SV_WGoq58,22740
 nystrom_ncut/nystrom/__init__.py,sha256=1aUXK87g4cXRXqNt6XkZsfyauw1-yv3sv0NmdmkWo-8,42
 nystrom_ncut/nystrom/distance_realization.py,sha256=RTI1_Q8fCUGAPSbXaVuNA-2B-11CEAfy2CwKWPJj6xQ,5830
-nystrom_ncut/nystrom/normalized_cut.py,sha256=jB_QALMY3l5CFfZPsrOFpEaquTrJP17muTrDZXxzUA8,7177
+nystrom_ncut/nystrom/normalized_cut.py,sha256=cjkG8JeDmTPDK8KwfkAIqF9f1dI-D9s1muJ9WWZlUoc,7237
 nystrom_ncut/nystrom/nystrom_utils.py,sha256=hksDO8uuAb9xKoA1ZafGwXDlQN_gZJn_qHscaSoO8JE,14120
 nystrom_ncut/transformer/__init__.py,sha256=jjXjcNp3LrxeF6mqG9VY5k3asrqaY6bXzJz6wTpH78Q,105
-nystrom_ncut/transformer/axis_align.py,sha256=6LTR-syJ-f4pcbnMexFmFNn1QADDhH5gka6979YBRrI,3549
+nystrom_ncut/transformer/axis_align.py,sha256=j3LlAPrp8O_jQAlwZz-gu3D7n_wICEJranye-YK5wvA,4880
 nystrom_ncut/transformer/transformer_mixin.py,sha256=YAjrDWTL5Hjnk9J2OsoxvtwT2N0u8IdgMSx0rRFmZzE,1653
-nystrom_ncut-0.3.0.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.3.0.dist-info/METADATA,sha256=lhxicufu5Eo9HQsUiS_K-CzocemOeNravAaIXeCtriM,6058
-nystrom_ncut-0.3.0.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-nystrom_ncut-0.3.0.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.3.0.dist-info/RECORD,,
+nystrom_ncut-0.3.2.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.3.2.dist-info/METADATA,sha256=4E42fHnXLNvbErWrwxE5K_oeOdpi5Bfabed9VF-YkV0,6058
+nystrom_ncut-0.3.2.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+nystrom_ncut-0.3.2.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.3.2.dist-info/RECORD,,