nystrom-ncut 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

nystrom_ncut/transformer/axis_align.py

@@ -1,4 +1,3 @@
-import random
 from typing import Literal
 
 import torch
@@ -14,7 +13,7 @@ class AxisAlign(TorchTransformerMixin):
     Args:
         max_iter (int, optional): Maximum number of iterations.
     """
-    SortOptions = Literal["count", "norm"]
+    SortOptions = Literal["count", "norm", "marginal_norm"]
 
     def __init__(
         self,
@@ -33,7 +32,7 @@ class AxisAlign(TorchTransformerMixin):
 
         # 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[random.randint(0, n - 1)]
+        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)
@@ -62,7 +61,11 @@ class AxisAlign(TorchTransformerMixin):
         if self.sort_method == "count":
             sort_metric = torch.bincount(idx, minlength=d)
         elif self.sort_method == "norm":
-            sort_metric = torch.linalg.norm(X @ self.R.mT, dim=0)
+            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}.")
 

{nystrom_ncut-0.2.0.dist-info → nystrom_ncut-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.2.0
+Version: 0.2.1
 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.2.0.dist-info → nystrom_ncut-0.2.1.dist-info}/RECORD

@@ -9,10 +9,10 @@ nystrom_ncut/nystrom/distance_realization.py,sha256=InajllGtRVnLVlZoipZNbHFTGHaT
 nystrom_ncut/nystrom/normalized_cut.py,sha256=2ocwc4U3A6GGFs0cuL0DO1yNvt59SJ3uDtj00U0foPM,5906
 nystrom_ncut/nystrom/nystrom_utils.py,sha256=5w-2GAMb7b6ArZdPEnAnKPFFrsbHSfC-S78cvrR6O20,12806
 nystrom_ncut/transformer/__init__.py,sha256=jjXjcNp3LrxeF6mqG9VY5k3asrqaY6bXzJz6wTpH78Q,105
-nystrom_ncut/transformer/axis_align.py,sha256=UHPbZVs3XFDvxAQHJC2La8W534k3nJ776DEfQeJVKxg,3297
+nystrom_ncut/transformer/axis_align.py,sha256=6LTR-syJ-f4pcbnMexFmFNn1QADDhH5gka6979YBRrI,3549
 nystrom_ncut/transformer/transformer_mixin.py,sha256=YAjrDWTL5Hjnk9J2OsoxvtwT2N0u8IdgMSx0rRFmZzE,1653
-nystrom_ncut-0.2.0.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.2.0.dist-info/METADATA,sha256=rKIYYECCawfOcjVLx6OZJAo5ngi2uKriNLbgKVEs_IY,6058
-nystrom_ncut-0.2.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.2.0.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.2.0.dist-info/RECORD,,
+nystrom_ncut-0.2.1.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.2.1.dist-info/METADATA,sha256=l5t4vEFtPANsQY8PK0YHDJ1tw6dZUulU5daxX9T8QC0,6058
+nystrom_ncut-0.2.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.2.1.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.2.1.dist-info/RECORD,,