nystrom-ncut 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

nystrom_ncut/common.py

@@ -12,8 +12,8 @@ def ceildiv(a: int, b: int) -> int:
 def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> torch.Tensor:
     numel = np.prod(x.shape[:-1])
     n = min(n, numel)
-    random_indices = torch.randperm(numel)[:n]
-    _x = x.flatten(0, -2)[random_indices]
+    random_indices = torch.randperm(numel, device=x.device)[:n]
+    _x = x.view((-1, x.shape[-1]))[random_indices]
     if torch.allclose(torch.norm(_x, **normalize_kwargs), torch.ones(n, device=x.device)):
         return x
     else:
@@ -21,13 +21,14 @@ def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> t
 
 
 def quantile_min_max(x: torch.Tensor, q1: float, q2: float, n_sample: int = 10000):
-    if x.shape[0] > n_sample:
+    x = x.flatten()
+    if len(x) > n_sample:
         np.random.seed(0)
-        random_idx = np.random.choice(x.shape[0], n_sample, replace=False)
+        random_idx = np.random.choice(len(x), n_sample, replace=False)
         vmin, vmax = x[random_idx].quantile(q1), x[random_idx].quantile(q2)
     else:
         vmin, vmax = x.quantile(q1), x.quantile(q2)
-    return vmin, vmax
+    return vmin.item(), vmax.item()
 
 
 def quantile_normalize(x: torch.Tensor, q: float = 0.95):
@@ -57,5 +58,17 @@ def quantile_normalize(x: torch.Tensor, q: float = 0.95):
     return x
 
 
+class default_device:
+    def __init__(self, device: torch.device):
+        self._device = device
+
+    def __enter__(self):
+        self._original_device = torch.get_default_device()
+        torch.set_default_device(self._device)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        torch.set_default_device(self._original_device)
+
+
 def profile(name: str, t: torch.Tensor) -> None:
     print(f"{name} --- nan: {t.isnan().any()}, inf: {t.isinf().any()}, max: {t.abs().max()}, min: {t.abs().min()}")

nystrom_ncut/distance_utils.py

@@ -1,61 +1,71 @@
-from typing import Literal
+import collections
+from typing import List, Literal, OrderedDict
 
 import torch
 
 from .common import lazy_normalize
 
 
-DistanceOptions = Literal["cosine", "euclidean", "rbf"]
+DistanceOptions = Literal["cosine", "euclidean"]
+AffinityOptions = Literal["cosine", "rbf", "laplacian"]
 
+# noinspection PyTypeChecker
+DISTANCE_TO_AFFINITY: OrderedDict[DistanceOptions, List[AffinityOptions]] = collections.OrderedDict([
+    ("cosine", ["cosine"]),
+    ("euclidean", ["rbf", "laplacian"]),
+])
+# noinspection PyTypeChecker
+AFFINITY_TO_DISTANCE: OrderedDict[AffinityOptions, DistanceOptions] = collections.OrderedDict(sum([
+    [(affinity_type, distance_type) for affinity_type in affinity_types]
+    for distance_type, affinity_types in DISTANCE_TO_AFFINITY.items()
+], start=[]))
 
-def to_euclidean(x: torch.Tensor, disttype: DistanceOptions) -> torch.Tensor:
-    if disttype == "cosine":
+
+
+def to_euclidean(x: torch.Tensor, distance_type: DistanceOptions) -> torch.Tensor:
+    if distance_type == "cosine":
         return lazy_normalize(x, p=2, dim=-1)
-    elif disttype == "rbf":
+    elif distance_type == "euclidean":
         return x
     else:
-        raise ValueError(f"to_euclidean not implemented for disttype {disttype}.")
+        raise ValueError(f"to_euclidean not implemented for distance_type {distance_type}.")
 
 
 def distance_from_features(
     features: torch.Tensor,
     features_B: torch.Tensor,
-    distance: DistanceOptions,
+    distance_type: DistanceOptions,
 ):
-    """Compute affinity matrix from input features.
+    """Compute distance matrix from input features.
     Args:
         features (torch.Tensor): input features, shape (n_samples, n_features)
         features_B (torch.Tensor, optional): optional, if not None, compute affinity between two features
-        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
+        distance_type (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
     Returns:
         (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
     """
     # compute distance matrix from input features
-    if distance == "cosine":
-        features = lazy_normalize(features, dim=-1)
-        features_B = lazy_normalize(features_B, dim=-1)
-        D = 1 - features @ features_B.T
-    elif distance == "euclidean":
-        D = torch.cdist(features, features_B, p=2)
-    elif distance == "rbf":
-        D = 0.5 * torch.cdist(features, features_B, p=2) ** 2
-
-        # Outlier-robust scale invariance using quantiles to estimate standard deviation
-        c = 2.0
-        p = torch.erf(torch.tensor((-c, c), device=features.device) * (2 ** -0.5))
-        stds = torch.quantile(features, q=(p + 1) / 2, dim=0)
-        stds = (stds[1] - stds[0]) / (2 * c)
-        D = D / (torch.linalg.norm(stds) ** 2)
-    else:
-        raise ValueError("distance should be 'cosine' or 'euclidean', 'rbf'")
-    return D
+    shape: torch.Size = features.shape[:-2]
+    features = features.view((-1, *features.shape[-2:]))
+    features_B = features_B.view((-1, *features_B.shape[-2:]))
+
+    match distance_type:
+        case "cosine":
+            features = lazy_normalize(features, dim=-1)
+            features_B = lazy_normalize(features_B, dim=-1)
+            D = 1 - features @ features_B.mT
+        case "euclidean":
+            D = torch.cdist(features, features_B, p=2)
+        case _:
+            raise ValueError("Distance should be 'cosine' or 'euclidean'")
+    return D.view((*shape, *D.shape[-2:]))
 
 
 def affinity_from_features(
     features: torch.Tensor,
     features_B: torch.Tensor = None,
     affinity_focal_gamma: float = 1.0,
-    distance: DistanceOptions = "cosine",
+    affinity_type: AffinityOptions = "cosine",
 ):
     """Compute affinity matrix from input features.
 
@@ -64,7 +74,7 @@ def affinity_from_features(
         features_B (torch.Tensor, optional): optional, if not None, compute affinity between two features
         affinity_focal_gamma (float): affinity matrix parameter, lower t reduce the edge weights
             on weak connections, default 1.0
-        distance (str): distance metric, 'cosine' (default) or 'euclidean', 'rbf'.
+        affinity_type (str): distance metric, 'cosine' (default) or 'euclidean'.
     Returns:
         (torch.Tensor): affinity matrix, shape (n_samples, n_samples)
     """
@@ -75,9 +85,21 @@ def affinity_from_features(
     features_B = features if features_B is None else features_B
 
     # compute distance matrix from input features
-    D = distance_from_features(features, features_B, distance)
+    D = distance_from_features(features, features_B, AFFINITY_TO_DISTANCE[affinity_type])
 
-    # torch.exp make affinity matrix positive definite,
     # lower affinity_focal_gamma reduce the weak edge weights
-    A = torch.exp(-D / affinity_focal_gamma)
+    match affinity_type:
+        case "cosine" | "laplacian":
+            A = torch.exp(-D / affinity_focal_gamma)                                        # [... x n x n]
+        case "rbf":
+            # Outlier-robust scale invariance using quantiles to estimate standard deviation
+            c = 2.0
+            p = torch.erf(torch.tensor((-c, c), device=features.device) * (2 ** -0.5))
+            stds = torch.nanquantile(features, q=(p + 1) / 2, dim=-2)                       # [2 x ... x d]
+            stds = (stds[1] - stds[0]) / (2 * c)                                            # [... x d]
+            D = 0.5 * (D / torch.norm(stds, dim=-1)[..., None, None]) ** 2
+            A = torch.exp(-D / affinity_focal_gamma)
+        case _:
+            raise ValueError("Affinity should be 'cosine', 'rbf', or 'laplacian'")
+
     return A

nystrom_ncut/nystrom/__init__.py

@@ -1,6 +1,3 @@
-from .distance_realization import (
-    DistanceRealization,
-)
 from .normalized_cut import (
     NCut,
 )

nystrom_ncut/nystrom/distance_realization.py

@@ -18,10 +18,10 @@ from ..sampling_utils import (
 class GramKernel(OnlineKernel):
     def __init__(
         self,
-        distance: DistanceOptions,
+        distance_type: DistanceOptions,
         eig_solver: EigSolverOptions,
     ):
-        self.distance: DistanceOptions = distance
+        self.distance_type: DistanceOptions = distance_type
         self.eig_solver: EigSolverOptions = eig_solver
 
         # Anchor matrices
@@ -40,7 +40,7 @@ class GramKernel(OnlineKernel):
         self.A = -0.5 * distance_from_features(
             self.anchor_features,                               # [n x d]
             self.anchor_features,
-            distance=self.distance,
+            distance_type=self.distance_type,
         )                                                       # [n x n]
         d = features.shape[-1]
         U, L = solve_eig(
@@ -58,7 +58,7 @@ class GramKernel(OnlineKernel):
         B = -0.5 * distance_from_features(
             self.anchor_features,                               # [n x d]
             features,                                           # [m x d]
-            distance=self.distance,
+            distance_type=self.distance_type,
         )                                                       # [n x m]
         b_r = torch.sum(B, dim=-1)                              # [n]
         b_c = torch.sum(B, dim=-2)                              # [m]
@@ -84,7 +84,7 @@ class GramKernel(OnlineKernel):
             B = -0.5 * distance_from_features(
                 self.anchor_features,
                 features,
-                distance=self.distance,
+                distance_type=self.distance_type,
             )
             b_c = torch.sum(B, dim=-2)                          # [m]
             col_sum = b_c + B.mT @ self.Ainv @ self.b_r         # [m]
@@ -98,7 +98,7 @@ class DistanceRealization(OnlineNystromSubsampleFit):
     def __init__(
         self,
         n_components: int = 100,
-        distance: DistanceOptions = "cosine",
+        distance_type: DistanceOptions = "cosine",
         sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
@@ -115,13 +115,12 @@ class DistanceRealization(OnlineNystromSubsampleFit):
         OnlineNystromSubsampleFit.__init__(
             self,
             n_components=n_components,
-            kernel=GramKernel(distance, eig_solver),
-            distance=distance,
+            kernel=GramKernel(distance_type, eig_solver),
+            distance_type=distance_type,
             sample_config=sample_config,
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )
-        self.distance: DistanceOptions = distance
 
     def fit_transform(
         self,

nystrom_ncut/nystrom/normalized_cut.py

@@ -8,7 +8,8 @@ from .nystrom_utils import (
     solve_eig,
 )
 from ..distance_utils import (
-    DistanceOptions,
+    AffinityOptions,
+    AFFINITY_TO_DISTANCE,
     affinity_from_features,
 )
 from ..sampling_utils import (
@@ -20,80 +21,83 @@ class LaplacianKernel(OnlineKernel):
     def __init__(
         self,
         affinity_focal_gamma: float,
-        distance: DistanceOptions,
+        affinity_type: AffinityOptions,
         adaptive_scaling: bool,
         eig_solver: EigSolverOptions,
     ):
         self.affinity_focal_gamma = affinity_focal_gamma
-        self.distance: DistanceOptions = distance
+        self.affinity_type: AffinityOptions = affinity_type
         self.adaptive_scaling: bool = adaptive_scaling
         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]
+        self.anchor_features: torch.Tensor = None                                   # [... x n x d]
+        self.anchor_mask: torch.Tensor = None
+        self.A: torch.Tensor = None                                                 # [... x n x n]
+        self.Ainv: torch.Tensor = None                                              # [... x n x n]
 
         # Updated matrices
-        self.a_r: torch.Tensor = None                           # [n]
-        self.b_r: torch.Tensor = None                           # [n]
+        self.a_r: torch.Tensor = None                                               # [... x n]
+        self.b_r: torch.Tensor = None                                               # [... x 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]
+        self.anchor_features = features                                             # [... x n x d]
+        self.anchor_mask = torch.all(torch.isnan(self.anchor_features), dim=-1)     # [... x n]
+
+        self.A = torch.nan_to_num(affinity_from_features(
+            self.anchor_features,                                                   # [... x n x d]
             affinity_focal_gamma=self.affinity_focal_gamma,
-            distance=self.distance,
-        )                                                       # [n x n]
+            affinity_type=self.affinity_type,
+        ), nan=0.0)                                                                 # [... x n x n]
         d = features.shape[-1]
         U, L = solve_eig(
             self.A,
             num_eig=d + 1,  # d * (d + 3) // 2 + 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]
+        )                                                                           # [... 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]
 
     def _affinity(self, features: torch.Tensor) -> torch.Tensor:
-        B = affinity_from_features(
-            self.anchor_features,                               # [n x d]
-            features,                                           # [m x d]
+        B = torch.where(self.anchor_mask[..., None], 0.0, affinity_from_features(
+            self.anchor_features,                                                   # [... x n x d]
+            features,                                                               # [... x m x d]
             affinity_focal_gamma=self.affinity_focal_gamma,
-            distance=self.distance,
-        )                                                       # [n x m]
+            affinity_type=self.affinity_type,
+        ))                                                                          # [... x n x m]
         if self.adaptive_scaling:
             diagonal = (
-                einops.rearrange(B, "n m -> m 1 n")             # [m x 1 x n]
-                @ self.Ainv                                     # [n x n]
-                @ einops.rearrange(B, "n m -> m n 1")           # [m x n x 1]
-            ).squeeze(1, 2)                                     # [m]
-            adaptive_scale = diagonal ** -0.5                   # [m]
-            B = B * adaptive_scale
-        return B                                                # [n x m]
+                einops.rearrange(B, "... n m -> ... m 1 n")                         # [... x m x 1 x n]
+                @ self.Ainv                                                         # [... x n x n]
+                @ einops.rearrange(B, "... n m -> ... m n 1")                       # [... x m x n x 1]
+            ).squeeze(-2, -1)                                                       # [... x m]
+            adaptive_scale = diagonal ** -0.5                                       # [... x m]
+            B = B * adaptive_scale[..., None, :]
+        return B                                                                    # [... x n x m]
 
     def update(self, features: torch.Tensor) -> torch.Tensor:
-        B = self._affinity(features)                            # [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]
+        B = self._affinity(features)                                                # [... x n x m]
+        b_r = torch.sum(torch.nan_to_num(B, nan=0.0), dim=-1)                       # [... x n]
+        b_c = torch.sum(B, dim=-2)                                                  # [... x m]
+        self.b_r = self.b_r + b_r                                                   # [... x n]
 
-        row_sum = self.a_r + self.b_r                           # [n]
-        col_sum = b_c + B.mT @ self.Ainv @ self.b_r             # [m]
-        scale = (row_sum[:, None] * col_sum) ** -0.5            # [n x m]
-        return (B * scale).mT                                   # [m x n]
+        row_sum = self.a_r + self.b_r                                               # [... x n]
+        col_sum = b_c + (B.mT @ (self.Ainv @ self.b_r[..., None]))[..., 0]          # [... x m]
+        scale = (row_sum[..., :, None] * col_sum[..., None, :]) ** -0.5             # [... x n x m]
+        return (B * scale).mT                                                       # [... x m x n]
 
     def transform(self, features: torch.Tensor = None) -> torch.Tensor:
-        row_sum = self.a_r + self.b_r                           # [n]
+        row_sum = self.a_r + self.b_r                                               # [... x n]
         if features is None:
-            B = self.A                                          # [n x n]
-            col_sum = row_sum                                   # [n]
+            B = self.A                                                              # [... x n x n]
+            col_sum = row_sum                                                       # [... x n]
         else:
             B = self._affinity(features)
-            b_c = torch.sum(B, dim=-2)                          # [m]
-            col_sum = b_c + B.mT @ self.Ainv @ self.b_r         # [m]
-        scale = (row_sum[:, None] * col_sum) ** -0.5            # [n x m]
-        return (B * scale).mT                                   # [m x n]
+            b_c = torch.sum(B, dim=-2)                                              # [... x m]
+            col_sum = b_c + (B.mT @ (self.Ainv @ self.b_r[..., None]))[..., 0]      # [... x m]
+        scale = (row_sum[..., :, None] * col_sum[..., None, :]) ** -0.5             # [... x n x m]
+        return (B * scale).mT                                                       # [... x m x n]
 
 
 class NCut(OnlineNystromSubsampleFit):
@@ -103,7 +107,7 @@ class NCut(OnlineNystromSubsampleFit):
         self,
         n_components: int = 100,
         affinity_focal_gamma: float = 1.0,
-        distance: DistanceOptions = "cosine",
+        affinity_type: AffinityOptions = "cosine",
         adaptive_scaling: bool = False,
         sample_config: SampleConfig = SampleConfig(),
         eig_solver: EigSolverOptions = "svd_lowrank",
@@ -124,8 +128,8 @@ class NCut(OnlineNystromSubsampleFit):
         OnlineNystromSubsampleFit.__init__(
             self,
             n_components=n_components,
-            kernel=LaplacianKernel(affinity_focal_gamma, distance, adaptive_scaling, eig_solver),
-            distance=distance,
+            kernel=LaplacianKernel(affinity_focal_gamma, affinity_type, adaptive_scaling, eig_solver),
+            distance_type=AFFINITY_TO_DISTANCE[affinity_type],
             sample_config=sample_config,
             eig_solver=eig_solver,
             chunk_size=chunk_size,