nystrom-ncut 0.0.12__py3-none-any.whl → 0.1.0__py3-none-any.whl

nystrom_ncut/__init__.py

@@ -1,5 +1,4 @@
 from .nystrom import (
-    DistanceRealization,
     NCut,
     axis_align,
 )
@@ -15,7 +14,7 @@ from .visualize_utils import (
     rgb_from_tsne_2d,
     rgb_from_umap_3d,
     rgb_from_umap_2d,
-    rgb_from_cosine_tsne_3d,
+    rgb_from_euclidean_tsne_3d,
     rotate_rgb_cube,
     convert_to_lab_color,
     get_mask,

nystrom_ncut/common.py

@@ -24,6 +24,15 @@ def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> t
         return Fn.normalize(x, **normalize_kwargs)
 
 
+def to_euclidean(x: torch.Tensor, disttype: DistanceOptions) -> torch.Tensor:
+    if disttype == "cosine":
+        return lazy_normalize(x, p=2, dim=-1)
+    elif disttype == "rbf":
+        return x
+    else:
+        raise ValueError(f"to_euclidean not implemented for disttype {disttype}.")
+
+
 def quantile_min_max(x, q1=0.01, q2=0.99, n_sample=10000):
     if x.shape[0] > n_sample:
         np.random.seed(0)

nystrom_ncut/nystrom/distance_realization.py

@@ -122,6 +122,7 @@ class DistanceRealization(OnlineNystromSubsampleFit):
             n_components=n_components,
             kernel=GramKernel(distance, eig_solver),
             num_sample=num_sample,
+            distance=distance,
             sample_method=sample_method,
             eig_solver=eig_solver,
             chunk_size=chunk_size,

nystrom_ncut/nystrom/normalized_cut.py

@@ -118,11 +118,11 @@ class NCut(OnlineNystromSubsampleFit):
             n_components=n_components,
             kernel=LaplacianKernel(affinity_focal_gamma, distance, eig_solver),
             num_sample=num_sample,
+            distance=distance,
             sample_method=sample_method,
             eig_solver=eig_solver,
             chunk_size=chunk_size,
         )
-        self.distance: DistanceOptions = distance
 
 
 def axis_align(eigen_vectors: torch.Tensor, max_iter=300):

nystrom_ncut/nystrom/nystrom.py

@@ -4,6 +4,7 @@ from typing import Literal, Tuple
 import torch
 
 from ..common import (
+    DistanceOptions,
     SampleOptions,
     ceildiv,
 )
@@ -145,6 +146,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
         n_components: int,
         kernel: OnlineKernel,
         num_sample: int,
+        distance: DistanceOptions,
         sample_method: SampleOptions,
         eig_solver: EigSolverOptions = "svd_lowrank",
         chunk_size: int = 8192,
@@ -157,6 +159,7 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
             chunk_size=chunk_size,
         )
         self.num_sample: int = num_sample
+        self.distance: DistanceOptions = distance
         self.sample_method: SampleOptions = sample_method
         self.anchor_indices: torch.Tensor = None
 
@@ -176,8 +179,9 @@ class OnlineNystromSubsampleFit(OnlineNystrom):
             self.anchor_indices = precomputed_sampled_indices
         else:
             self.anchor_indices = run_subgraph_sampling(
-                features,
-                self.num_sample,
+                features=features,
+                num_sample=self.num_sample,
+                disttype=self.distance,
                 sample_method=self.sample_method,
             )
         sampled_features = features[self.anchor_indices]

nystrom_ncut/propagation_utils.py

@@ -9,15 +9,17 @@ from .common import (
     SampleOptions,
     ceildiv,
     lazy_normalize,
+    to_euclidean,
 )
 
 
-@torch.no_grad()
+# @torch.no_grad()
 def run_subgraph_sampling(
     features: torch.Tensor,
     num_sample: int,
+    disttype: DistanceOptions,
+    sample_method: SampleOptions,
     max_draw: int = 1000000,
-    sample_method: SampleOptions = "farthest",
 ):
     if num_sample >= features.shape[0]:
         # if too many samples, use all samples and bypass Nystrom-like approximation
@@ -28,6 +30,7 @@ def run_subgraph_sampling(
     else:
         # sample subgraph
         if sample_method == "farthest":  # default
+            features = to_euclidean(features, disttype)
             if num_sample > max_draw:
                 logging.warning(
                     f"num_sample is larger than max_draw, apply farthest point sampling on random sampled {max_draw} samples"
@@ -144,12 +147,12 @@ def extrapolate_knn(
     anchor_features: torch.Tensor,          # [n x d]
     anchor_output: torch.Tensor,            # [n x d']
     extrapolation_features: torch.Tensor,   # [m x d]
+    distance: DistanceOptions,
     knn: int = 10,                          # k
-    distance: DistanceOptions = "cosine",
     affinity_focal_gamma: float = 1.0,
     chunk_size: int = 8192,
     device: str = None,
-    move_output_to_cpu: bool = False
+    move_output_to_cpu: bool = False,
 ) -> torch.Tensor:                          # [m x d']
     """A generic function to propagate new nodes using KNN.
 
@@ -168,7 +171,7 @@ def extrapolate_knn(
         >>> old_eigenvectors = torch.randn(3000, 20)
         >>> old_features = torch.randn(3000, 100)
         >>> new_features = torch.randn(200, 100)
-        >>> new_eigenvectors = extrapolate_knn(old_features,old_eigenvectors,new_features,knn=3)
+        >>> new_eigenvectors = extrapolate_knn(old_features, old_eigenvectors, new_features, knn=3)
         >>> # new_eigenvectors.shape = (200, 20)
 
     """
@@ -197,21 +200,24 @@ def extrapolate_knn(
             _V = _V.cpu()
         V_list.append(_V)
 
-    anchor_output = torch.cat(V_list, dim=0)
-    return anchor_output
+    extrapolation_output = torch.cat(V_list, dim=0)
+    return extrapolation_output
 
 
 # wrapper functions for adding new nodes to existing graph
 def extrapolate_knn_with_subsampling(
-    full_features: torch.Tensor,
-    full_output: torch.Tensor,
-    extrapolation_features: torch.Tensor,
-    knn: int,
-    num_sample: int,
+    full_features: torch.Tensor,            # [n x d]
+    full_output: torch.Tensor,              # [n x d']
+    extrapolation_features: torch.Tensor,   # [m x d]
+    num_sample: int,                        # n'
     sample_method: SampleOptions,
-    chunk_size: int,
-    device: str
-):
+    distance: DistanceOptions,
+    knn: int = 10,                          # k
+    affinity_focal_gamma: float = 1.0,
+    chunk_size: int = 8192,
+    device: str = None,
+    move_output_to_cpu: bool = False,
+) -> torch.Tensor:                          # [m x d']
     """Propagate eigenvectors to new nodes using KNN. Note: this is equivalent to the class API `NCUT.tranform(new_features)`, expect for the sampling is re-done in this function.
     Args:
         full_output (torch.Tensor): eigenvectors from existing nodes, shape (num_sample, num_eig)
@@ -237,8 +243,9 @@ def extrapolate_knn_with_subsampling(
 
     # sample subgraph
     anchor_indices = run_subgraph_sampling(
-        full_features,
-        num_sample,
+        features=full_features,
+        num_sample=num_sample,
+        disttype=distance,
         sample_method=sample_method,
     )
 
@@ -247,12 +254,15 @@ def extrapolate_knn_with_subsampling(
     extrapolation_features = extrapolation_features.to(device)
 
     # propagate eigenvectors from subgraph to new nodes
-    new_eigenvectors = extrapolate_knn(
+    extrapolation_output = extrapolate_knn(
         anchor_features,
         anchor_output,
         extrapolation_features,
+        distance,
         knn=knn,
+        affinity_focal_gamma=affinity_focal_gamma,
         chunk_size=chunk_size,
-        device=device
+        device=device,
+        move_output_to_cpu=move_output_to_cpu,
     )
-    return new_eigenvectors
+    return extrapolation_output

nystrom_ncut/visualize_utils.py

@@ -8,6 +8,7 @@ from sklearn.base import BaseEstimator
 
 from .common import (
     lazy_normalize,
+    to_euclidean,
     quantile_min_max,
     quantile_normalize,
 )
@@ -20,76 +21,72 @@ from .propagation_utils import (
 )
 
 
-def _identity(X: torch.Tensor) -> torch.Tensor:
-    return X
-
-
 def _rgb_with_dimensionality_reduction(
     features: torch.Tensor,
     num_sample: int,
-    metric: Literal["cosine", "euclidean"],
+    disttype: Literal["cosine", "euclidean"],
     rgb_func: Callable[[torch.Tensor, float], torch.Tensor],
     q: float,
     knn: int,
     reduction: Callable[..., BaseEstimator],
     reduction_dim: int,
     reduction_kwargs: Dict[str, Any],
-    transform_func: Callable[[torch.Tensor], torch.Tensor],
     seed: int,
     device: str,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> torch.Tensor:
 
     if True:
         _subgraph_indices = run_subgraph_sampling(
-            features,
+            features=features,
             num_sample=10000,
+            disttype=disttype,
             sample_method="farthest",
         )
         features = extrapolate_knn(
-            features[_subgraph_indices],
-            features[_subgraph_indices],
-            features,
-            distance="cosine",
+            anchor_features=features[_subgraph_indices],
+            anchor_output=features[_subgraph_indices],
+            extrapolation_features=features,
+            distance=disttype,
         )
 
     subgraph_indices = run_subgraph_sampling(
-        features,
-        num_sample,
+        features=features,
+        num_sample=num_sample,
+        disttype=disttype,
         sample_method="farthest",
     )
 
     _inp = features[subgraph_indices].numpy(force=True)
     _subgraph_embed = reduction(
         n_components=reduction_dim,
-        metric=metric,
+        metric=disttype,
         random_state=seed,
         **reduction_kwargs
     ).fit_transform(_inp)
 
     _subgraph_embed = torch.tensor(_subgraph_embed, dtype=torch.float32)
-    X_nd = transform_func(extrapolate_knn(
+    rgb = rgb_func(extrapolate_knn(
         features[subgraph_indices],
         _subgraph_embed,
         features,
+        disttype,
         knn=knn,
-        distance=metric,
         device=device,
         move_output_to_cpu=True
-    ))
-    rgb = rgb_func(X_nd, q)
-    return X_nd, rgb
+    ), q)
+    return rgb
 
 
 def rgb_from_tsne_2d(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     perplexity: int = 150,
-    metric: Literal["cosine", "euclidean"] = "cosine",
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
     device: str = None,
-):
+) -> torch.Tensor:
     """
     Returns:
         (torch.Tensor): Embedding in 2D, shape (n_samples, 2)
@@ -108,32 +105,32 @@ def rgb_from_tsne_2d(
         )
         perplexity = num_sample // 2
 
-    x2d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric=metric,
+        disttype=disttype,
         rgb_func=rgb_from_2d_colormap,
         q=q,
         knn=knn,
         reduction=TSNE, reduction_dim=2, reduction_kwargs={
             "perplexity": perplexity,
-        }, transform_func=_identity,
+        },
         seed=seed,
         device=device,
     )
-    return x2d, rgb
+    return rgb
 
 
 def rgb_from_tsne_3d(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     perplexity: int = 150,
-    metric: Literal["cosine", "euclidean"] = "cosine",
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
     device: str = None,
-):
+) -> torch.Tensor:
     """
     Returns:
         (torch.Tensor): Embedding in 3D, shape (n_samples, 3)
@@ -152,31 +149,32 @@ def rgb_from_tsne_3d(
         )
         perplexity = num_sample // 2
 
-    x3d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric=metric,
+        disttype=disttype,
         rgb_func=rgb_from_3d_rgb_cube,
         q=q,
         knn=knn,
         reduction=TSNE, reduction_dim=3, reduction_kwargs={
             "perplexity": perplexity,
-        }, transform_func=_identity,
+        },
         seed=seed,
         device=device,
     )
-    return x3d, rgb
+    return rgb
 
 
-def rgb_from_cosine_tsne_3d(
+def rgb_from_euclidean_tsne_3d(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     perplexity: int = 150,
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
     device: str = None
-):
+) -> torch.Tensor:
     """
     Returns:
         (torch.Tensor): Embedding in 3D, shape (n_samples, 3)
@@ -195,40 +193,36 @@ def rgb_from_cosine_tsne_3d(
         )
         perplexity = num_sample // 2
 
-    def cosine_to_rbf(X: torch.Tensor) -> torch.Tensor:
-        dr = DistanceRealization(n_components=3, num_sample=20000, distance="cosine", eig_solver="svd_lowrank")
-        return dr.fit_transform(X)
-
-    def rgb_from_cosine(X_3d: torch.Tensor, q: float) -> torch.Tensor:
-        return rgb_from_3d_rgb_cube(cosine_to_rbf(X_3d), q=q)
+    def rgb_func(X_3d: torch.Tensor, q: float) -> torch.Tensor:
+        return rgb_from_3d_rgb_cube(to_euclidean(X_3d, disttype), q=q)
 
-    x3d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric="cosine",
-        rgb_func=rgb_from_cosine,
+        disttype="cosine",
+        rgb_func=rgb_func,
         q=q,
         knn=knn,
         reduction=TSNE, reduction_dim=3, reduction_kwargs={
             "perplexity": perplexity,
-        }, transform_func=_identity,
+        },
         seed=seed,
         device=device,
     )
-    return x3d, rgb
+    return rgb
 
 
 def rgb_from_umap_2d(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     n_neighbors: int = 150,
     min_dist: float = 0.1,
-    metric: Literal["cosine", "euclidean"] = "cosine",
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
     device: str = None,
-):
+) -> torch.Tensor:
     """
     Returns:
         (torch.Tensor): Embedding in 2D, shape (n_samples, 2)
@@ -239,34 +233,34 @@ def rgb_from_umap_2d(
     except ImportError:
         raise ImportError("umap import failed, please install `pip install umap-learn`")
 
-    x2d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric=metric,
+        disttype=disttype,
         rgb_func=rgb_from_2d_colormap,
         q=q,
         knn=knn,
         reduction=UMAP, reduction_dim=2, reduction_kwargs={
             "n_neighbors": n_neighbors,
             "min_dist": min_dist,
-        }, transform_func=_identity,
+        },
         seed=seed,
         device=device,
     )
-    return x2d, rgb
+    return rgb
 
 
 def rgb_from_umap_sphere(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     n_neighbors: int = 150,
     min_dist: float = 0.1,
-    metric: Literal["cosine", "euclidean"] = "cosine",
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
     device: str = None,
-):
+) -> torch.Tensor:
     """
     Returns:
         (torch.Tensor): Embedding in 2D, shape (n_samples, 2)
@@ -277,37 +271,37 @@ def rgb_from_umap_sphere(
     except ImportError:
         raise ImportError("umap import failed, please install `pip install umap-learn`")
 
-    def transform_func(X: torch.Tensor) -> torch.Tensor:
-        return torch.stack((
+    def rgb_func(X: torch.Tensor, q: float) -> torch.Tensor:
+        return rgb_from_3d_rgb_cube(torch.stack((
             torch.sin(X[:, 0]) * torch.cos(X[:, 1]),
             torch.sin(X[:, 0]) * torch.sin(X[:, 1]),
             torch.cos(X[:, 0]),
-        ), dim=1)
+        ), dim=1), q=q)
 
-    x3d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric=metric,
-        rgb_func=rgb_from_3d_rgb_cube,
+        disttype=disttype,
+        rgb_func=rgb_func,
         q=q,
         knn=knn,
         reduction=UMAP, reduction_dim=2, reduction_kwargs={
             "n_neighbors": n_neighbors,
             "min_dist": min_dist,
             "output_metric": "haversine",
-        }, transform_func=transform_func,
+        },
         seed=seed,
         device=device,
     )
-    return x3d, rgb
+    return rgb
 
 
 def rgb_from_umap_3d(
     features: torch.Tensor,
     num_sample: int = 1000,
+    disttype: Literal["cosine", "euclidean"] = "cosine",
     n_neighbors: int = 150,
     min_dist: float = 0.1,
-    metric: Literal["cosine", "euclidean"] = "cosine",
     q: float = 0.95,
     knn: int = 10,
     seed: int = 0,
@@ -323,21 +317,21 @@ def rgb_from_umap_3d(
     except ImportError:
         raise ImportError("umap import failed, please install `pip install umap-learn`")
 
-    x3d, rgb = _rgb_with_dimensionality_reduction(
+    rgb = _rgb_with_dimensionality_reduction(
         features=features,
         num_sample=num_sample,
-        metric=metric,
+        disttype=disttype,
         rgb_func=rgb_from_3d_rgb_cube,
         q=q,
         knn=knn,
         reduction=UMAP, reduction_dim=3, reduction_kwargs={
             "n_neighbors": n_neighbors,
             "min_dist": min_dist,
-        }, transform_func=_identity,
+        },
         seed=seed,
         device=device,
     )
-    return x3d, rgb
+    return rgb
 
 
 def flatten_sphere(X_3d):

{nystrom_ncut-0.0.12.dist-info → nystrom_ncut-0.1.0.dist-info}/METADATA

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

@@ -0,0 +1,14 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nystrom_ncut/__init__.py,sha256=ffExLdGTaPsUweHcYc61Ose6a5A5Tfo9hm48zjEl6ho,441
+nystrom_ncut/common.py,sha256=l5kjF6neukdtbMrt5bE69pN0JM9r-93JYltLlYJuBik,2227
+nystrom_ncut/propagation_utils.py,sha256=79M61iJfp_RWj_xLOn51PHiextWcEWTQ7NWl2T51-3Y,10907
+nystrom_ncut/visualize_utils.py,sha256=uGfBBkETQ8uRJ-UXPWpiXT8KEYMLmEHXPKJMLvN0c34,16543
+nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
+nystrom_ncut/nystrom/distance_realization.py,sha256=FGH7VjbtRrSROH0d8OPuCUxLQy5j7Z8BuE4hrSGGZG4,6031
+nystrom_ncut/nystrom/normalized_cut.py,sha256=s9ZS3-tQbWnxAlPc01v9l7fqBhl28lvOalaCO2y-Gd8,7175
+nystrom_ncut/nystrom/nystrom.py,sha256=OV5o9UL9fkrz9HdsD6rXh7MTsenPKrtCNRIczMuDS_4,12779
+nystrom_ncut-0.1.0.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
+nystrom_ncut-0.1.0.dist-info/METADATA,sha256=wAhGONU0ZM1VWoLqEwtfAs3_GUAt4CItkHq3ISuFyVE,6058
+nystrom_ncut-0.1.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nystrom_ncut-0.1.0.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
+nystrom_ncut-0.1.0.dist-info/RECORD,,

nystrom_ncut-0.0.12.dist-info/RECORD

@@ -1,14 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nystrom_ncut/__init__.py,sha256=JKfF6atok5T9V692RhlhgeRO5a2cN-bfAVa9irmTLfs,463
-nystrom_ncut/common.py,sha256=RMPQvg9R2s7V-q7zAStN9YCZt7gpc5Ut-KSKtvELBQ4,1934
-nystrom_ncut/propagation_utils.py,sha256=WeWKxRBm01ITILMgjsit5_fCe9oW1kJOPmAjjcmliMo,10340
-nystrom_ncut/visualize_utils.py,sha256=JkDyWML6k7k6S2Z7xnpbUvMWssEXcXqXu7gBy8wnids,16809
-nystrom_ncut/nystrom/__init__.py,sha256=4EpxD3Cmc8Fif4vo8DG-6FpTfCnNanD5zCZxK3WrMwQ,121
-nystrom_ncut/nystrom/distance_realization.py,sha256=MWSdfPfUnr7BdiKFkogjQvcGagvj7OzLQklnVp0fPx8,6000
-nystrom_ncut/nystrom/normalized_cut.py,sha256=_U3zrbe6V-5TQ4uWmqckxs2JTIhygQlnRDTFBI1ghD4,7194
-nystrom_ncut/nystrom/nystrom.py,sha256=nL-zxbEE_ygJEZEPmeNrUpVeffvxdrsTcbxFanFuXQY,12613
-nystrom_ncut-0.0.12.dist-info/LICENSE,sha256=2bm9uFabQZ3Ykb_SaSU_uUbAj2-htc6WJQmS_65qD00,1073
-nystrom_ncut-0.0.12.dist-info/METADATA,sha256=pM-WT6Ly-IKYJ3DV2d-oOyc--K4VOyArB0sT5gHfHL4,6059
-nystrom_ncut-0.0.12.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nystrom_ncut-0.0.12.dist-info/top_level.txt,sha256=gM8IWWHYysIRTCvCTcdS4RShOyl9pxpylgSwPUZR2XM,22
-nystrom_ncut-0.0.12.dist-info/RECORD,,