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

nystrom_ncut/__init__.py

@@ -1,6 +1,8 @@
 from .nystrom import (
     NCut,
-    axis_align,
+)
+from .transformer import (
+    AxisAlign,
 )
 from .distance_utils import (
     distance_from_features,

nystrom_ncut/nystrom/__init__.py

@@ -3,5 +3,4 @@ from .distance_realization import (
 )
 from .normalized_cut import (
     NCut,
-    axis_align,
 )

nystrom_ncut/nystrom/normalized_cut.py

@@ -1,6 +1,5 @@
 import einops
 import torch
-import torch.nn.functional as Fn
 
 from .nystrom_utils import (
     EigSolverOptions,
@@ -131,50 +130,3 @@ class NCut(OnlineNystromSubsampleFit):
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )
-
-
-def axis_align(eigen_vectors: torch.Tensor, max_iter=300):
-    """Multiclass Spectral Clustering, SX Yu, J Shi, 2003
-
-    Args:
-        eigen_vectors (torch.Tensor): continuous eigenvectors from NCUT, shape (n, k)
-        max_iter (int, optional): Maximum number of iterations.
-
-    Returns:
-        torch.Tensor: Discretized eigenvectors, shape (n, k), each row is a one-hot vector.
-    """
-    # Normalize eigenvectors
-    n, k = eigen_vectors.shape
-    eigen_vectors = Fn.normalize(eigen_vectors, p=2, dim=-1)
-
-    # Initialize R matrix with the first column from a random row of EigenVectors
-    R = torch.empty((k, k), device=eigen_vectors.device)
-    R[0] = eigen_vectors[torch.randint(0, n, (1,))].squeeze()
-
-    # Loop to populate R with k orthogonal directions
-    c = torch.zeros(n, device=eigen_vectors.device)
-    for i in range(1, k):
-        c += torch.abs(eigen_vectors @ R[i - 1])
-        R[i] = eigen_vectors[torch.argmin(c, dim=0)]
-
-    # Iterative optimization loop
-    eps = torch.finfo(torch.float32).eps
-    prev_objective = torch.inf
-    for _ in range(max_iter):
-        # Discretize the projected eigenvectors
-        idx = torch.argmax(eigen_vectors @ R.mT, dim=-1)
-        M = torch.zeros((k, k)).index_add_(0, idx, eigen_vectors)
-
-        # Compute the NCut value
-        objective = torch.norm(M)
-
-        # Check for convergence
-        if torch.abs(objective - prev_objective) < eps:
-            break
-        prev_objective = objective
-
-        # SVD decomposition
-        U, S, Vh = torch.linalg.svd(M, full_matrices=False)
-        R = U @ Vh
-
-    return Fn.one_hot(idx, num_classes=k).to(torch.float), R

nystrom_ncut/nystrom/nystrom_utils.py

@@ -1,5 +1,6 @@
 import copy
 import logging
+from abc import abstractmethod
 from typing import Literal, Tuple
 
 import torch
@@ -14,23 +15,29 @@ from ..sampling_utils import (
     SampleConfig,
     subsample_features,
 )
+from ..transformer import (
+    TorchTransformerMixin,
+)
 
 
 EigSolverOptions = Literal["svd_lowrank", "lobpcg", "svd", "eigh"]
 
 
 class OnlineKernel:
+    @abstractmethod
     def fit(self, features: torch.Tensor) -> "OnlineKernel":                # [n x d]
-        raise NotImplementedError()
+        """"""
 
+    @abstractmethod
     def update(self, features: torch.Tensor) -> torch.Tensor:               # [m x d] -> [m x n]
-        raise NotImplementedError()
+        """"""
 
+    @abstractmethod
     def transform(self, features: torch.Tensor = None) -> torch.Tensor:     # [m x d] -> [m x n]
-        raise NotImplementedError()
+        """"""
 
 
-class OnlineNystrom:
+class OnlineNystrom(TorchTransformerMixin):
     def __init__(
         self,
         n_components: int,

nystrom_ncut/transformer/__init__.py

@@ -0,0 +1,6 @@
+from .transformer_mixin import (
+    TorchTransformerMixin,
+)
+from .axis_align import (
+    AxisAlign,
+)

nystrom_ncut/transformer/axis_align.py

@@ -0,0 +1,84 @@
+import random
+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"]
+
+    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[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(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)).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":
+            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:
+        """
+        Args:
+            X (torch.Tensor): continuous eigenvectors from NCUT, shape (n, k)
+            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.
+        """
+        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, hard: bool = False) -> torch.Tensor:
+        return self.fit(X).transform(X, hard=hard)

nystrom_ncut/transformer/transformer_mixin.py

@@ -0,0 +1,51 @@
+from abc import abstractmethod
+from typing import Any
+
+import torch
+from sklearn.base import TransformerMixin, BaseEstimator
+
+
+class TorchTransformerMixin(TransformerMixin, BaseEstimator):
+    """Mixin class for all transformers in scikit-learn.
+
+    This mixin defines the following functionality:
+
+    - a `fit_transform` method that delegates to `fit` and `transform`;
+    - a `set_output` method to output `X` as a specific container type.
+
+    If :term:`get_feature_names_out` is defined, then :class:`BaseEstimator` will
+    automatically wrap `transform` and `fit_transform` to follow the `set_output`
+    API. See the :ref:`developer_api_set_output` for details.
+
+    :class:`OneToOneFeatureMixin` and
+    :class:`ClassNamePrefixFeaturesOutMixin` are helpful mixins for
+    defining :term:`get_feature_names_out`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from sklearn.base import BaseEstimator, TransformerMixin
+    >>> class MyTransformer(TransformerMixin, BaseEstimator):
+    ...     def __init__(self, *, param=1):
+    ...         self.param = param
+    ...     def fit(self, X, y=None):
+    ...         return self
+    ...     def transform(self, X):
+    ...         return np.full(shape=len(X), fill_value=self.param)
+    >>> transformer = MyTransformer()
+    >>> X = [[1, 2], [2, 3], [3, 4]]
+    >>> transformer.fit_transform(X)
+    array([1, 1, 1])
+    """
+
+    @abstractmethod
+    def fit(self, X: torch.Tensor, **fit_kwargs: Any) -> "TorchTransformerMixin":
+        """"""
+
+    @abstractmethod
+    def transform(self, X: torch.Tensor, **transform_kwargs: Any) -> torch.Tensor:
+        """"""
+
+    @abstractmethod
+    def fit_transform(self, X: torch.Tensor, **kwargs: Any) -> torch.Tensor:
+        """"""

nystrom_ncut/visualize_utils.py

@@ -1,10 +1,10 @@
 import logging
-from typing import Any, Callable, Dict, Literal
+from typing import Any, Callable, Dict, Literal, Union
 
 import numpy as np
 import torch
 import torch.nn.functional as Fn
-from sklearn.base import BaseEstimator
+from sklearn.base import TransformerMixin, BaseEstimator
 
 from .common import (
     ceildiv,
@@ -152,7 +152,7 @@ def _rgb_with_dimensionality_reduction(
     rgb_func: Callable[[torch.Tensor, float], torch.Tensor],
     q: float,
     knn: int,
-    reduction: Callable[..., BaseEstimator],
+    reduction: Callable[..., Union[TransformerMixin, BaseEstimator]],
     reduction_dim: int,
     reduction_kwargs: Dict[str, Any],
     seed: int,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nystrom_ncut
-Version: 0.1.7
+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.9.dist-info/RECORD

@@ -0,0 +1,18 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nystrom_ncut/__init__.py,sha256=tKq9-2QRNFetckHY77qAaKEMjMCYTYcorS2f74aNtvk,540
+nystrom_ncut/common.py,sha256=_PGJoImSk_Fb_5Ri-e_IsFoCcSfbGS8CxYUUHVoNM50,2036
+nystrom_ncut/distance_utils.py,sha256=p-pYdpRrJsIhzxM_IxUqja7N8okngx52WGXD9pu_Aec,3129
+nystrom_ncut/sampling_utils.py,sha256=oMmhFcd_N_D15Ht7F0rCGPSgLeitJszAKMD3ICKwHNU,3105
+nystrom_ncut/visualize_utils.py,sha256=d3VXjzJPZPPyUMg_b8hKLQoBaRWvutu6u7l36S2gmIM,23007
+nystrom_ncut/nystrom/__init__.py,sha256=lAoO00i4FG5xqGKDO_OYcSvO4qPK64x_X_hDNBvuLUc,105
+nystrom_ncut/nystrom/distance_realization.py,sha256=InajllGtRVnLVlZoipZNbHFTGHaTs3zxizKe3kI2Los,5815
+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=pX7wk4O6fj-CwRv1TYyPYXsTmmXUtQ5q0c5fDQBVE6Q,3068
+nystrom_ncut/transformer/transformer_mixin.py,sha256=fTNtDFYPw2Fc8mjvK2xNHOw5mCkbO0usUpOnnJdyr5M,1743
+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,,

nystrom_ncut-0.1.7.dist-info/RECORD

@@ -1,15 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nystrom_ncut/__init__.py,sha256=HifrTcqX2-hYjBDe6xIThHvuIBYMPBA3EzjR8-qPMUM,512
-nystrom_ncut/common.py,sha256=_PGJoImSk_Fb_5Ri-e_IsFoCcSfbGS8CxYUUHVoNM50,2036
-nystrom_ncut/distance_utils.py,sha256=p-pYdpRrJsIhzxM_IxUqja7N8okngx52WGXD9pu_Aec,3129
-nystrom_ncut/sampling_utils.py,sha256=oMmhFcd_N_D15Ht7F0rCGPSgLeitJszAKMD3ICKwHNU,3105
-nystrom_ncut/visualize_utils.py,sha256=_J6YjWUsBe0VqW6KXsQx_iPmRCcO-ie0g6t5mD289UI,22957
-nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
-nystrom_ncut/nystrom/distance_realization.py,sha256=InajllGtRVnLVlZoipZNbHFTGHaTs3zxizKe3kI2Los,5815
-nystrom_ncut/nystrom/normalized_cut.py,sha256=N-M5wkTo59vpbBfIx8evkSQBxlo4j80qCtuoifxQa_A,7578
-nystrom_ncut/nystrom/nystrom_utils.py,sha256=UVs1tC7vnVq2mWSTpcrP4C19x9wDJ77ACht0EltOO2E,12698
-nystrom_ncut-0.1.7.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.1.7.dist-info/METADATA,sha256=eb0Q6bwCKC4c5bcuJI_PnIaPW5qFYGSgwbPeuWI7EUk,6058
-nystrom_ncut-0.1.7.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.1.7.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.1.7.dist-info/RECORD,,