Decomposition Diffusion Map

diffusion_map

CLASS	DESCRIPTION
DiffusionMap	Diffusion map for nonlinear dimensionality reduction.

FUNCTION	DESCRIPTION
diffusion_map	Computes diffusion coordinates based on kernel matrix eigendecomposition

Classes

DiffusionMap(*, kernel_fn: Kernel | str = 'gaussian', kernel_kwargs: dict | None = None, distance_fn: PairwiseDistance | str = 'euclidean', distance_kwargs: dict | None = None, score_fn: Callable | None = None, out_features: int = 0, symmetric_eigendecomposition: bool = True, norm_l2_eigenvectors: bool = True)

Bases: Decomposition

Diffusion map for nonlinear dimensionality reduction.

Computes diffusion coordinates based on kernel matrix eigendecomposition [1, 2]. Provides sklearn-style fit and transform interface. Diffusion maps are particularly effective for manifold learning and nonlinear dimensionality reduction on data with intrinsic low-dimensional structure.

This implementation uses a modular architecture that separates pairwise distance computation from kernel computation, allowing flexible combinations of distance metrics (Euclidean, covariance, Mahalanobis) with various kernel functions (Gaussian, t-distribution, etc.).

PARAMETER	DESCRIPTION
kernel_fn	Kernel function for similarity computation. Can be a Kernel instance, string name of a kernel, or None. Use kernel's normalization parameter to control normalization. TYPE: Kernel | str | None, optional, by default None DEFAULT: 'gaussian'
kernel_kwargs	Additional keyword arguments to pass to the kernel function if kernel_fn is specified as a string. TYPE: dict | None, optional, by default None DEFAULT: None
distance_fn	Pairwise distance function. Can be a PairwiseDistance instance, string name of a distance function, or None. TYPE: PairwiseDistance | str | None, optional, by default None DEFAULT: 'euclidean'
distance_kwargs	Additional keyword arguments to pass to the distance function if distance_fn is specified as a string. TYPE: dict | None, optional, by default None DEFAULT: None
score_fn	Scoring function to evaluate the model. Must be callable if provided. If None, score will raise an error. TYPE: Callable | None, optional, by default None DEFAULT: None
out_features	Number of eigenvectors to compute. If 0, computes all eigenvectors. Must be non-negative. TYPE: int, optional, by default 0 DEFAULT: 0
symmetric_eigendecomposition	Whether to perform symmetric eigendecomposition. TYPE: bool, optional, by default False DEFAULT: True
norm_l2_eigenvectors	Whether to normalize eigenvectors to unit L2 norm. TYPE: bool, optional, by default True DEFAULT: True

ATTRIBUTE	DESCRIPTION
eigenvalues	Eigenvalues from kernel matrix decomposition of shape (out_features, ). TYPE: Tensor
eigenvectors	Eigenvectors from kernel matrix decomposition of shape (n_samples, out_features). TYPE: Tensor
diffusion_coordinates	Diffusion coordinates extracted from the eigenvectors (n_samples, out_features) or (n_samples, out_features - 1). TYPE: Tensor
kernel	Computed kernel matrix of shape (n_samples, n_samples). TYPE: Tensor
explained_variance	Explained variance for each component. TYPE: Tensor
cumulative_explained_variance	Cumulative explained variance for components. TYPE: Tensor

Examples:

Basic usage with Gaussian kernel:

>>> import torch
>>> from spectre.decomposition import DiffusionMap
>>> from spectre.kernel import GaussianKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>> X = torch.randn(100, 5)  # 100 samples, 5 features
>>> kernel_fn = GaussianKernel(bw_method=torch.tensor(0.5))
>>> distance_fn = PairwiseDistanceEuclidean()
>>> dm = DiffusionMap(
...     kernel_fn=kernel_fn,
...     distance_fn=distance_fn,
...     out_features=3,
... )
>>> dm.fit(X)
>>> X_embedded = dm.predict(X)  # Shape: (100, 3)

With kernel normalization:

>>> from spectre.kernel.normalization import KernelNormalizer
>>> normalization = KernelNormalizer(norm_type="markov_symmetric", alpha=0.5)
>>> kernel_fn = GaussianKernel(
...     bw_method=torch.tensor(0.5),
...     normalization=normalization,
... )
>>> dm = DiffusionMap(
...     kernel_fn=kernel_fn,
...     distance_fn=distance_fn,
...     out_features=10,
... )
>>> dm.fit(X)

METHOD	DESCRIPTION
score	Compute score using the provided scoring function.
fit	Fit the diffusion map to the provided data.
predict	Predict diffusion coordinates for new data using Nyström extension.

Source code in spectre/decomposition/diffusion_map.py

def __init__(
    self,
    *,
    kernel_fn: Kernel | str = "gaussian",
    kernel_kwargs: dict | None = None,
    distance_fn: PairwiseDistance | str = "euclidean",
    distance_kwargs: dict | None = None,
    score_fn: Callable | None = None,
    out_features: int = 0,
    symmetric_eigendecomposition: bool = True,
    norm_l2_eigenvectors: bool = True,
) -> None:
    super().__init__(
        kernel_fn=kernel_fn,
        kernel_kwargs=kernel_kwargs,
        distance_fn=distance_fn,
        distance_kwargs=distance_kwargs,
        score_fn=score_fn,
    )
    if not isinstance(out_features, int):
        raise TypeError(
            f"Expected `out_features` to be int, got {type(out_features)}"
        )
    check_in_interval(out_features, "[0, inf)")
    self.out_features = out_features

    if not isinstance(symmetric_eigendecomposition, bool):
        raise ValueError("symmetric_eigendecomposition must be a boolean.")
    self.symmetric_eigendecomposition = symmetric_eigendecomposition

    if not isinstance(norm_l2_eigenvectors, bool):
        raise ValueError("norm_l2_eigenvectors must be a boolean.")
    self.norm_l2_eigenvectors = norm_l2_eigenvectors

    # Internal state.
    self._eigenvalues = None
    self._eigenvectors = None
    self._kernel = None

Attributes

diffusion_coordinates: torch.Tensor property

Return the fitted diffusion coordinates.

explained_variance: torch.Tensor property

Explained variance for each diffusion component.

cumulative_explained_variance: torch.Tensor property

Cumulative explained variance for diffusion components.

eigenvalues: torch.Tensor property

Eigenvalues of the diffusion matrix.

eigenvectors: torch.Tensor property

Eigenvectors of the diffusion matrix.

kernel: torch.Tensor property

Kernel matrix used for diffusion map computation.

Functions

score(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None, **kwargs) -> float | torch.Tensor

Compute score using the provided scoring function.

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, out_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values for scoring. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
float | Tensor	Score computed by the scoring function.

Source code in spectre/decomposition/diffusion_map.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
    **kwargs,
) -> float | torch.Tensor:
    """
    Compute score using the provided scoring function.

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, out_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples, ).

    target : torch.Tensor | None, optional, by default None
        Target values for scoring.

    Returns
    -------
    float | torch.Tensor
        Score computed by the scoring function.
    """
    if self.score_fn is None:
        raise TypeError("To run `score`, set a `score_fn`.")

    self.fit(X, weights=weights)

    return self.score_fn(X, weights=weights, target=target, **kwargs)

fit(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> DiffusionMap

Fit the diffusion map to the provided data.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, out_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Not used in diffusion maps (unsupervised method). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
self	Fitted estimator. TYPE: DiffusionMap

Source code in spectre/decomposition/diffusion_map.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "DiffusionMap":
    """
    Fit the diffusion map to the provided data.

    Parameters
    ----------
    X : torch.Tensor
        Training data of shape (n_samples, out_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples, ).

    target : torch.Tensor | None, optional, by default None
        Not used in diffusion maps (unsupervised method).

    Returns
    -------
    self : DiffusionMap
        Fitted estimator.
    """
    check_2d(X)

    self.in_features = X.shape[1]

    if self.out_features == 0:
        self.out_features = min(X.shape[0], X.shape[1])

    # Store training data for out-of-sample extension
    self._X_fit = X.clone()
    self._weights_fit = weights.clone() if weights is not None else None

    result = _diffusion_map_impl(
        X=X,
        weights=weights,
        kernel_fn=self.kernel_fn,
        distance_fn=self.distance_fn,
        out_features=0,  # Compute all, slice later
        symmetric_eigendecomposition=self.symmetric_eigendecomposition,
        norm_l2_eigenvectors=self.norm_l2_eigenvectors,
    )

    self.is_fitted = True

    self._eigenvalues = result.eigenvalues
    self._eigenvectors = result.eigenvectors
    self._kernel = result.kernel

    return self

predict(X: torch.Tensor) -> torch.Tensor

Predict diffusion coordinates for new data using Nyström extension.

For training data, returns the fitted diffusion coordinates. For new data, uses Nyström extension to approximate the embedding.

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, n_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Diffusion coordinates of shape (n_samples, out_features) or (n_samples, out_features - 1) depending on out_features setting.

Notes

The Nyström extension approximates the embedding of new points by computing kernel similarities to training points and projecting onto the learned eigenvectors. This is an approximation that works well when the new data lies on or near the manifold learned from training data.

Examples:

>>> import torch
>>> from spectre.decomposition import DiffusionMap
>>> from spectre.kernel import GaussianKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>> X_train = torch.randn(100, 5)
>>> X_test = torch.randn(20, 5)
>>> kernel_fn = GaussianKernel(bw_method=torch.tensor(0.5))
>>> dm = DiffusionMap(
...     kernel_fn=kernel_fn,
...     distance_fn=PairwiseDistanceEuclidean(),
...     out_features=3,
... )
>>> dm.fit(X_train)
>>> X_test_embedded = dm.predict(X_test)  # Shape: (20, 3)

Source code in spectre/decomposition/diffusion_map.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict diffusion coordinates for new data using Nyström extension.

    For training data, returns the fitted diffusion coordinates. For new data,
    uses Nyström extension to approximate the embedding.

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, n_features).

    Returns
    -------
    torch.Tensor
        Diffusion coordinates of shape (n_samples, out_features) or
        (n_samples, out_features - 1) depending on out_features setting.

    Notes
    -----
    The Nyström extension approximates the embedding of new points by computing
    kernel similarities to training points and projecting onto the learned
    eigenvectors. This is an approximation that works well when the new data
    lies on or near the manifold learned from training data.

    Examples
    --------
    >>> import torch
    >>> from spectre.decomposition import DiffusionMap
    >>> from spectre.kernel import GaussianKernel
    >>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
    >>> X_train = torch.randn(100, 5)
    >>> X_test = torch.randn(20, 5)
    >>> kernel_fn = GaussianKernel(bw_method=torch.tensor(0.5))
    >>> dm = DiffusionMap(
    ...     kernel_fn=kernel_fn,
    ...     distance_fn=PairwiseDistanceEuclidean(),
    ...     out_features=3,
    ... )
    >>> dm.fit(X_train)
    >>> X_test_embedded = dm.predict(X_test)  # Shape: (20, 3)
    """
    self.check_fitted()

    if X.shape == self._X_fit.shape and torch.allclose(X, self._X_fit):
        return self.diffusion_coordinates

    return nystrom_extension(
        X=X,
        X_train=self._X_fit,
        eigenvectors=self._eigenvectors,
        eigenvalues=self._eigenvalues,
        kernel_fn=self.kernel_fn,
        distance_fn=self.distance_fn,
        out_features=self.out_features,
        scale_by_eigenvalues=True,
    )

Functions

diffusion_map(X: torch.Tensor, weights: torch.Tensor | None = None, *, kernel_fn: Kernel, distance_fn: PairwiseDistance, out_features: int = 0, symmetric_eigendecomposition: bool = True, norm_l2_eigenvectors: bool = True) -> DecompositionResult

Computes diffusion coordinates based on kernel matrix eigendecomposition [1, 2]. Provides sklearn-style fit and transform interface. Diffusion maps are particularly effective for manifold learning and nonlinear dimensionality reduction on data with intrinsic low-dimensional structure.

This implementation uses a modular architecture that separates pairwise distance computation from kernel computation, allowing flexible combinations of distance metrics (Euclidean, covariance, Mahalanobis) with various kernel functions (Gaussian, t-distribution, etc.).

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, out_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). If None, weights are assigned ones. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
kernel_fn	Kernel function for similarity computation. Must be a Kernel instance. Use kernel's normalization parameter to control normalization. TYPE: Kernel
distance_fn	Pairwise distance function. Must be a PairwiseDistance instance. TYPE: PairwiseDistance
out_features	Number of eigenvectors to compute. If 0, computes all eigenvectors. Must be non-negative. TYPE: int, optional, by default 0 DEFAULT: 0
symmetric_eigendecomposition	Whether to compute symmetric eigendecomposition. TYPE: bool, optional, by default True DEFAULT: True
norm_l2_eigenvectors	Whether to normalize eigenvectors to unit L2 norm. TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
DecompositionResult	A tuple containing: eigenvalues : Eigenvalues of shape (out_features,) eigenvectors : Eigenvectors of shape (n_samples, out_features) kernel : Normalized kernel matrix of shape (n_samples, n_samples)

Examples:

Basic usage with Gaussian kernel:

>>> import torch
>>> from spectre.decomposition.diffusion_map import diffusion_map
>>> from spectre.kernel import GaussianKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>> X = torch.randn(50, 10)
>>> kernel_fn = GaussianKernel(bw_method=torch.tensor(1.0))
>>> pairwise_dist_fn = PairwiseDistanceEuclidean()
>>> eigenvals, eigenvecs, kernel_matrix = diffusion_map(
...     X=X,
...     kernel_fn=kernel_fn,
...     distance_fn=pairwise_dist_fn,
...     out_features=5,
... )

With kernel normalization:

>>> from spectre.kernel.normalization import KernelNormalizer
>>> normalization = KernelNormalizer(norm_type="markov_symmetric", alpha=0.5)
>>> kernel_fn = GaussianKernel(
...     bw_method=torch.tensor(1.0),
...     normalization=normalization,
... )
>>> eigenvals, eigenvecs, K = diffusion_map(
...     X=X,
...     kernel_fn=kernel_fn,
...     distance_fn=pairwise_dist_fn,
...     out_features=5,
... )

Source code in spectre/decomposition/diffusion_map.py

def diffusion_map(
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    *,
    kernel_fn: Kernel,
    distance_fn: PairwiseDistance,
    out_features: int = 0,
    symmetric_eigendecomposition: bool = True,
    norm_l2_eigenvectors: bool = True,
) -> DecompositionResult:
    """
    Computes diffusion coordinates based on kernel matrix eigendecomposition
    [@coifman2008diffusion;@coifman2005geometric].
    Provides sklearn-style fit and transform interface. Diffusion maps are
    particularly effective for manifold learning and nonlinear dimensionality
    reduction on data with intrinsic low-dimensional structure.

    This implementation uses a modular architecture that separates pairwise
    distance computation from kernel computation, allowing flexible combinations
    of distance metrics (Euclidean, covariance, Mahalanobis) with various
    kernel functions (Gaussian, t-distribution, etc.).

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, out_features).

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,). If None, weights are assigned ones.

    kernel_fn : Kernel
        Kernel function for similarity computation. Must be a Kernel instance.
        Use kernel's `normalization` parameter to control normalization.

    distance_fn : PairwiseDistance
        Pairwise distance function. Must be a PairwiseDistance instance.

    out_features : int, optional, by default 0
        Number of eigenvectors to compute. If 0, computes all eigenvectors.
        Must be non-negative.

    symmetric_eigendecomposition : bool, optional, by default True
        Whether to compute symmetric eigendecomposition.

    norm_l2_eigenvectors : bool, optional, by default True
        Whether to normalize eigenvectors to unit L2 norm.

    Returns
    -------
    DecompositionResult
        A tuple containing:

        - eigenvalues : Eigenvalues of shape (out_features,)
        - eigenvectors : Eigenvectors of shape (n_samples, out_features)
        - kernel : Normalized kernel matrix of shape (n_samples, n_samples)

    Examples
    --------
    Basic usage with Gaussian kernel:

    >>> import torch
    >>> from spectre.decomposition.diffusion_map import diffusion_map
    >>> from spectre.kernel import GaussianKernel
    >>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
    >>> X = torch.randn(50, 10)
    >>> kernel_fn = GaussianKernel(bw_method=torch.tensor(1.0))
    >>> pairwise_dist_fn = PairwiseDistanceEuclidean()
    >>> eigenvals, eigenvecs, kernel_matrix = diffusion_map(
    ...     X=X,
    ...     kernel_fn=kernel_fn,
    ...     distance_fn=pairwise_dist_fn,
    ...     out_features=5,
    ... )

    With kernel normalization:

    >>> from spectre.kernel.normalization import KernelNormalizer
    >>> normalization = KernelNormalizer(norm_type="markov_symmetric", alpha=0.5)
    >>> kernel_fn = GaussianKernel(
    ...     bw_method=torch.tensor(1.0),
    ...     normalization=normalization,
    ... )
    >>> eigenvals, eigenvecs, K = diffusion_map(
    ...     X=X,
    ...     kernel_fn=kernel_fn,
    ...     distance_fn=pairwise_dist_fn,
    ...     out_features=5,
    ... )
    """
    check_2d(X)

    if not isinstance(kernel_fn, Kernel):
        raise TypeError(
            f"kernel_fn must be Kernel instance, got {type(kernel_fn).__name__}"
        )

    if not isinstance(distance_fn, PairwiseDistance):
        raise TypeError(
            f"distance_fn must be PairwiseDistance instance, "
            f"got {type(distance_fn).__name__}"
        )

    if not isinstance(out_features, int):
        raise TypeError(f"Expected `out_features` to be int, got {type(out_features)}")
    check_in_interval(out_features, "[0, inf)")

    if out_features == 0:
        out_features = min(X.shape[0], X.shape[1])

    if not isinstance(symmetric_eigendecomposition, bool):
        raise ValueError("symmetric_eigendecomposition must be a boolean.")

    return _diffusion_map_impl(
        X=X,
        weights=weights,
        kernel_fn=kernel_fn,
        distance_fn=distance_fn,
        out_features=out_features,
        symmetric_eigendecomposition=symmetric_eigendecomposition,
        norm_l2_eigenvectors=norm_l2_eigenvectors,
    )