nystrom-ncut 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

nystrom_ncut/transformer/axis_align.py

@@ -1,4 +1,5 @@
 import random
+from typing import Literal
 
 import torch
 import torch.nn.functional as Fn
@@ -13,31 +14,39 @@ class AxisAlign(TorchTransformerMixin):
     Args:
         max_iter (int, optional): Maximum number of iterations.
     """
-    def __init__(self, max_iter: int = 100):
-        self.max_iter = max_iter
+    SortOptions = Literal["count", "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
-        X = Fn.normalize(X, p=2, dim=-1)
+        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] = X[random.randint(0, n - 1)]
+        self.R[0] = normalized_X[random.randint(0, n - 1)]
 
         # Loop to populate R with k orthogonal directions
         c = torch.zeros((n,), device=X.device)
         for i in range(1, d):
-            c += torch.abs(X @ self.R[i - 1])
-            self.R[i] = X[torch.argmin(c, dim=0)]
+            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(X @ self.R.mT, dim=-1)
-            M = torch.zeros((d, d)).index_add_(0, idx, X)
+            idx = torch.argmax(normalized_X @ self.R.mT, dim=-1)
+            M = torch.zeros((d, d)).index_add_(0, idx, normalized_X)
 
             # Check for convergence
             objective = torch.norm(M)
@@ -49,8 +58,15 @@ class AxisAlign(TorchTransformerMixin):
             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 counts
-        self.R = self.R[torch.argsort(torch.bincount(idx, minlength=d), dim=0, descending=True)]
+        # 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":
+            sort_metric = torch.linalg.norm(X @ self.R.mT, p=2, dim=0)
+        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, hard: bool = False) -> torch.Tensor:

{nystrom_ncut-0.1.8.dist-info → nystrom_ncut-0.1.9.dist-info}/METADATA

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

@@ -9,10 +9,10 @@ nystrom_ncut/nystrom/distance_realization.py,sha256=InajllGtRVnLVlZoipZNbHFTGHaT
 nystrom_ncut/nystrom/normalized_cut.py,sha256=5aR-CbRAWQVOA1FlQCuxSKEik9tR9sNLsJVBA7_LXyE,5905
 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=bea_5QVm5gY0o_rZB9jOr5fbXWCL_IWB9FI-wSygx00,2513
+nystrom_ncut/transformer/axis_align.py,sha256=pX7wk4O6fj-CwRv1TYyPYXsTmmXUtQ5q0c5fDQBVE6Q,3068
 nystrom_ncut/transformer/transformer_mixin.py,sha256=fTNtDFYPw2Fc8mjvK2xNHOw5mCkbO0usUpOnnJdyr5M,1743
-nystrom_ncut-0.1.8.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.1.8.dist-info/METADATA,sha256=9hsyObQahN0v45sjnBh9-tJuwY-62T2iU_7vEl3MqkY,6058
-nystrom_ncut-0.1.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.1.8.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.1.8.dist-info/RECORD,,
+nystrom_ncut-0.1.9.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.1.9.dist-info/METADATA,sha256=8ez3ayc8UcBR8R8Ds7nRAKbrEa3766WNDrQwXToQ9ZM,6058
+nystrom_ncut-0.1.9.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.1.9.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.1.9.dist-info/RECORD,,