nystrom-ncut 0.0.1__py3-none-any.whl

nystrom_ncut/new_ncut_pytorch.py

@@ -0,0 +1,241 @@
+import logging
+from typing import Literal, Tuple
+
+import torch
+
+from .nystrom import (
+    EigSolverOptions,
+    OnlineKernel,
+    OnlineNystrom,
+    solve_eig,
+)
+from .propagation_utils import (
+    affinity_from_features,
+    run_subgraph_sampling,
+)
+
+
+DistanceOptions = Literal["cosine", "euclidean", "rbf"]
+
+
+class LaplacianKernel(OnlineKernel):
+    def __init__(
+        self,
+        affinity_focal_gamma: float,
+        distance: DistanceOptions,
+        eig_solver: EigSolverOptions,
+    ):
+        self.affinity_focal_gamma = affinity_focal_gamma
+        self.distance: DistanceOptions = distance
+        self.eig_solver: EigSolverOptions = eig_solver
+
+        # Anchor matrices
+        self.anchor_features: torch.Tensor = None               # [n x d]
+        self.A: torch.Tensor = None                             # [n x n]
+        self.Ainv: torch.Tensor = None                          # [n x n]
+
+        # Updated matrices
+        self.a_r: torch.Tensor = None                           # [n]
+        self.b_r: torch.Tensor = None                           # [n]
+
+    def fit(self, features: torch.Tensor) -> None:
+        self.anchor_features = features                         # [n x d]
+        self.A = affinity_from_features(
+            self.anchor_features,                               # [n x d]
+            affinity_focal_gamma=self.affinity_focal_gamma,
+            distance=self.distance,
+            fill_diagonal=False,
+        )                                                       # [n x n]
+        U, L = solve_eig(
+            self.A,
+            num_eig=features.shape[-1] + 1,
+            eig_solver=self.eig_solver,
+        )                                                       # [n x (d + 1)], [d + 1]
+        self.Ainv = U @ torch.diag(1 / L) @ U.mT                # [n x n]
+        self.a_r = torch.sum(self.A, dim=-1)                    # [n]
+        self.b_r = torch.zeros_like(self.a_r)                   # [n]
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:
+        B = affinity_from_features(
+            self.anchor_features,                               # [n x d]
+            features,                                           # [m x d]
+            affinity_focal_gamma=self.affinity_focal_gamma,
+            distance=self.distance,
+            fill_diagonal=False,
+        )                                                       # [n x m]
+        b_r = torch.sum(B, dim=-1)                              # [n]
+        b_c = torch.sum(B, dim=-2)                              # [m]
+        self.b_r = self.b_r + b_r                               # [n]
+
+        rowscale = self.a_r + self.b_r                          # [n]
+        colscale = b_c + B.mT @ self.Ainv @ self.b_r            # [m]
+        scale = (rowscale[:, None] * colscale) ** -0.5          # [n x m]
+        return (B * scale).mT                                   # [m x n]
+
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:
+        rowscale = self.a_r + self.b_r                          # [n]
+        if features is None:
+            B = self.A                                          # [n x n]
+            colscale = rowscale                                 # [n]
+        else:
+            B = affinity_from_features(
+                self.anchor_features,                           # [n x d]
+                features,                                       # [m x d]
+                affinity_focal_gamma=self.affinity_focal_gamma,
+                distance=self.distance,
+                fill_diagonal=False,
+            )                                                   # [n x m]
+            b_c = torch.sum(B, dim=-2)                          # [m]
+            colscale = b_c + B.mT @ self.Ainv @ self.b_r        # [m]
+        scale = (rowscale[:, None] * colscale) ** -0.5          # [n x m]
+        return (B * scale).mT                                   # [m x n]
+
+
+class NewNCUT(OnlineNystrom):
+    """Nystrom Normalized Cut for large scale graph."""
+
+    def __init__(
+        self,
+        num_eig: int = 100,
+        affinity_focal_gamma: float = 1.0,
+        num_sample: int = 10000,
+        sample_method: Literal["farthest", "random"] = "farthest",
+        distance: DistanceOptions = "cosine",
+        eig_solver: EigSolverOptions = "svd_lowrank",
+        normalize_features: bool = None,
+        device: str = None,
+        move_output_to_cpu: bool = False,
+        matmul_chunk_size: int = 8096,
+    ):
+        """
+        Args:
+            num_eig (int): number of top eigenvectors to return
+            affinity_focal_gamma (float): affinity matrix temperature, lower t reduce the not-so-connected edge weights,
+                smaller t result in more sharp eigenvectors.
+            num_sample (int): number of samples for Nystrom-like approximation,
+                reduce only if memory is not enough, increase for better approximation
+            sample_method (str): subgraph sampling, ['farthest', 'random'].
+                farthest point sampling is recommended for better Nystrom-approximation accuracy
+            distance (str): distance metric for affinity matrix, ['cosine', 'euclidean', 'rbf'].
+            eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
+            normalize_features (bool): normalize input features before computing affinity matrix,
+                default 'None' is True for cosine distance, False for euclidean distance and rbf
+            device (str): device to use for eigen computation,
+                move to GPU to speeds up a bit (~5x faster)
+            move_output_to_cpu (bool): move output to CPU, set to True if you have memory issue
+            matmul_chunk_size (int): chunk size for large-scale matrix multiplication
+        """
+        OnlineNystrom.__init__(
+            self,
+            n_components=num_eig,
+            kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
+            eig_solver=eig_solver,
+            chunk_size=matmul_chunk_size,
+        )
+        self.num_sample = num_sample
+        self.sample_method = sample_method
+        self.distance = distance
+        self.normalize_features = normalize_features
+        if self.normalize_features is None:
+            if distance in ["cosine"]:
+                self.normalize_features = True
+            if distance in ["euclidean", "rbf"]:
+                self.normalize_features = False
+
+        self.device = device
+        self.move_output_to_cpu = move_output_to_cpu
+        self.matmul_chunk_size = matmul_chunk_size
+
+    def _fit_helper(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # move subgraph gpu to speed up
+        original_device = features.device
+        device = original_device if self.device is None else self.device
+
+        _n = features.shape[0]
+        if self.num_sample >= _n:
+            logging.info(
+                f"NCUT nystrom num_sample is larger than number of input samples, nyström approximation is not needed, setting num_sample={_n}"
+            )
+            self.num_sample = _n
+
+        # check if features dimension greater than num_eig
+        if self.eig_solver in ["svd_lowrank", "lobpcg"]:
+            assert (
+                _n >= self.n_components * 2
+            ), "number of nodes should be greater than 2*num_eig"
+        elif self.eig_solver in ["svd", "eigh"]:
+            assert (
+                _n >= self.n_components
+            ), "number of nodes should be greater than num_eig"
+
+        assert self.distance in ["cosine", "euclidean", "rbf"], "distance should be 'cosine', 'euclidean', 'rbf'"
+
+        if self.normalize_features:
+            # features need to be normalized for affinity matrix computation (cosine distance)
+            features = torch.nn.functional.normalize(features, dim=-1)
+
+        if precomputed_sampled_indices is not None:
+            sampled_indices = precomputed_sampled_indices
+        else:
+            sampled_indices = run_subgraph_sampling(
+                features,
+                num_sample=self.num_sample,
+                sample_method=self.sample_method,
+            )
+        sampled_features = features[sampled_indices].to(device)
+        OnlineNystrom.fit(self, sampled_features)
+
+        _n_not_sampled = _n - len(sampled_features)
+        if _n_not_sampled > 0:
+            unsampled_indices = torch.full((_n,), True).scatter(0, sampled_indices, False)
+            unsampled_features = features[unsampled_indices].to(device)
+            V_unsampled, _ = OnlineNystrom.update(self, unsampled_features)
+        else:
+            unsampled_indices = V_unsampled = None
+        return unsampled_indices, V_unsampled
+
+    def fit(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor = None,
+    ):
+        """Fit Nystrom Normalized Cut on the input features.
+        Args:
+            features (torch.Tensor): input features, shape (n_samples, n_features)
+            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+                override the sample_method, if not None
+        Returns:
+            (NCUT): self
+        """
+        NewNCUT._fit_helper(self, features, precomputed_sampled_indices)
+        return self
+
+    def fit_transform(
+        self,
+        features: torch.Tensor,
+        precomputed_sampled_indices: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            features (torch.Tensor): input features, shape (n_samples, n_features)
+            precomputed_sampled_indices (torch.Tensor): precomputed sampled indices, shape (num_sample,)
+                override the sample_method, if not None
+
+        Returns:
+            (torch.Tensor): eigen_vectors, shape (n_samples, num_eig)
+            (torch.Tensor): eigen_values, sorted in descending order, shape (num_eig,)
+        """
+        unsampled_indices, V_unsampled = NewNCUT._fit_helper(self, features, precomputed_sampled_indices)
+        V_sampled, L = OnlineNystrom.transform(self)
+
+        if unsampled_indices is not None:
+            V = torch.zeros((len(unsampled_indices), self.n_components))
+            V[~unsampled_indices] = V_sampled
+            V[unsampled_indices] = V_unsampled
+        else:
+            V = V_sampled
+        return V, L

nystrom_ncut/nystrom.py

@@ -0,0 +1,170 @@
+from typing import Literal, Tuple
+
+import torch
+
+
+EigSolverOptions = Literal["svd_lowrank", "lobpcg", "svd", "eigh"]
+
+
+class OnlineKernel:
+    def fit(self, features: torch.Tensor) -> None:                          # [n x d]
+        raise NotImplementedError()
+
+    def update(self, features: torch.Tensor) -> torch.Tensor:               # [m x d] -> [m x n]
+        raise NotImplementedError()
+
+    def transform(self, features: torch.Tensor = None) -> torch.Tensor:     # [m x d] -> [m x n]
+        raise NotImplementedError()
+
+
+class OnlineNystrom:
+    def __init__(
+        self,
+        n_components: int,
+        kernel: OnlineKernel,
+        eig_solver: EigSolverOptions,
+        chunk_size: int = 8192,
+    ):
+        """
+        Args:
+            n_components (int): number of top eigenvectors to return
+            kernel (OnlineKernel): Online kernel that computes pairwise matrix entries from input features and allows updates
+            eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh'].
+        """
+        self.n_components: int = n_components
+        self.kernel: OnlineKernel = kernel
+        self.eig_solver: EigSolverOptions = eig_solver
+        self.inverse_approximation_dim: int = None
+
+        self.chunk_size = chunk_size
+
+        # Anchor matrices
+        self.anchor_features: torch.Tensor = None   # [n x d]
+        self.A: torch.Tensor = None                 # [n x n]
+        self.Ahinv: torch.Tensor = None             # [n x n]
+        self.Ahinv_UL: torch.Tensor = None          # [n x indirect_pca_dim]
+        self.Ahinv_VT: torch.Tensor = None          # [indirect_pca_dim x n]
+
+        # Updated matrices
+        self.S: torch.Tensor = None                 # [n x n]
+        self.transform_matrix: torch.Tensor = None  # [n x n_components]
+        self.LS: torch.Tensor = None                # [n]
+
+    def fit(self, features: torch.Tensor):
+        OnlineNystrom.fit_transform(self, features)
+        return self
+
+    def fit_transform(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self.anchor_features = features
+
+        self.kernel.fit(self.anchor_features)
+        self.A = self.S = self.kernel.transform()                                                   # [n x n]
+
+        self.inverse_approximation_dim = max(self.n_components, features.shape[-1]) + 1
+        U, L = solve_eig(
+            self.A,
+            num_eig=self.inverse_approximation_dim,
+            eig_solver=self.eig_solver,
+        )                                                                                           # [n x (? + 1)], [? + 1]
+        self.Ahinv_UL = U * (L ** -0.5)                                                             # [n x (? + 1)]
+        self.Ahinv_VT = U.mT                                                                        # [(? + 1) x n]
+        self.Ahinv = self.Ahinv_UL @ self.Ahinv_VT                                                  # [n x n]
+
+        self.transform_matrix = (U / L)[:, :self.n_components]                                      # [n x n_components]
+        self.LS = L[:self.n_components]                                                             # [n_components]
+        return U[:, :self.n_components], L[:self.n_components]                                      # [n x n_components], [n_components]
+
+    def update(self, features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        n_chunks = -(-len(features) // self.chunk_size)
+        if n_chunks > 1:
+            """ Chunked version """
+            chunks = torch.chunk(features, n_chunks, dim=0)
+            for chunk in chunks:
+                self.kernel.update(chunk)
+
+            compressed_BBT = torch.zeros((self.inverse_approximation_dim, self.inverse_approximation_dim))  # [(? + 1) x (? + 1))]
+            for i, chunk in enumerate(chunks):
+                _B = self.kernel.transform(chunk).mT                                                    # [n x _m]
+                _compressed_B = self.Ahinv_VT @ _B                                                      # [(? + 1) x _m]
+                compressed_BBT = compressed_BBT + _compressed_B @ _compressed_B.mT                      # [(? + 1) x (? + 1)]
+            self.S = self.S + self.Ahinv_UL @ compressed_BBT @ self.Ahinv_UL.mT                         # [n x n]
+            US, self.LS = solve_eig(self.S, self.n_components, self.eig_solver)                         # [n x n_components], [n_components]
+            self.transform_matrix = self.Ahinv @ US * (self.LS ** -0.5)                                 # [n x n_components]
+
+            VS = []
+            for chunk in chunks:
+                VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                         # [_m x n_components]
+            VS = torch.cat(VS, dim=0)
+            return VS, self.LS                                                                          # [m x n_components], [n_components]
+        else:
+            """ Unchunked version """
+            B = self.kernel.update(features).mT                                                         # [n x m]
+            compressed_B = self.Ahinv_VT @ B                                                            # [indirect_pca_dim x m]
+
+            self.S = self.S + self.Ahinv_UL @ (compressed_B @ compressed_B.mT) @ self.Ahinv_UL.mT       # [n x n]
+            US, self.LS = solve_eig(self.S, self.n_components, self.eig_solver)                         # [n x n_components], [n_components]
+            self.transform_matrix = self.Ahinv @ US * (self.LS ** -0.5)                                 # [n x n_components]
+
+            return B.mT @ self.transform_matrix, self.LS                                                # [m x n_components], [n_components]
+
+    def transform(self, features: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        if features is None:
+            VS = self.A @ self.transform_matrix                                                     # [n x n_components]
+        else:
+            n_chunks = -(-len(features) // self.chunk_size)
+            if n_chunks > 1:
+                """ Chunked version """
+                chunks = torch.chunk(features, n_chunks, dim=0)
+                VS = []
+                for chunk in chunks:
+                    VS.append(self.kernel.transform(chunk) @ self.transform_matrix)                     # [_m x n_components]
+                VS = torch.cat(VS, dim=0)
+            else:
+                """ Unchunked version """
+                VS = self.kernel.transform(features) @ self.transform_matrix                            # [m x n_components]
+        return VS, self.LS                                                                          # [m x n_components], [n_components]
+
+
+def solve_eig(
+    A: torch.Tensor,
+    num_eig: int,
+    eig_solver: Literal["svd_lowrank", "lobpcg", "svd", "eigh"],
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """PyTorch implementation of Eigensolver cut without Nystrom-like approximation.
+
+    Args:
+        A (torch.Tensor): input matrix, shape (n_samples, n_samples)
+        num_eig (int): number of eigenvectors to return
+        eig_solver (str): eigen decompose solver, ['svd_lowrank', 'lobpcg', 'svd', 'eigh']
+
+    Returns:
+        (torch.Tensor): eigenvectors corresponding to the eigenvalues, shape (n_samples, num_eig)
+        (torch.Tensor): eigenvalues of the eigenvectors, sorted in descending order
+    """
+    # compute eigenvectors
+    if eig_solver == "svd_lowrank":  # default
+        # only top q eigenvectors, fastest
+        eigen_vector, eigen_value, _ = torch.svd_lowrank(A, q=num_eig)
+    elif eig_solver == "lobpcg":
+        # only top k eigenvectors, fast
+        eigen_value, eigen_vector = torch.lobpcg(A, k=num_eig)
+    elif eig_solver == "svd":
+        # all eigenvectors, slow
+        eigen_vector, eigen_value, _ = torch.svd(A)
+    elif eig_solver == "eigh":
+        # all eigenvectors, slow
+        eigen_value, eigen_vector = torch.linalg.eigh(A)
+    else:
+        raise ValueError(
+            "eigen_solver should be 'lobpcg', 'svd_lowrank', 'svd' or 'eigh'"
+        )
+
+    # sort eigenvectors by eigenvalues, take top (descending order)
+    eigen_value = eigen_value.real
+    eigen_vector = eigen_vector.real
+    eigen_value, indices = torch.topk(eigen_value, k=num_eig, dim=0)
+    eigen_vector = eigen_vector[:, indices]
+
+    # correct the random rotation (flipping sign) of eigenvectors
+    eigen_vector = eigen_vector * torch.sum(eigen_vector, dim=0).sign()
+    return eigen_vector, eigen_value