Decomposition Nystrom#

nystrom #

FUNCTION DESCRIPTION
nystrom_extension

Compute Nyström extension for out-of-sample predictions.

Classes#

Functions#

nystrom_extension(X: torch.Tensor, X_train: torch.Tensor, weights: torch.Tensor | None = None, weights_train: torch.Tensor | None = None, *, eigenvectors: torch.Tensor, eigenvalues: torch.Tensor, kernel_fn: Kernel, distance_fn: PairwiseDistance, out_features: int | None = None, scale_by_eigenvalues: bool = True) -> torch.Tensor #

Compute Nyström extension for out-of-sample predictions.

Projects new data points onto the embedding learned from training data using the Nyström approximation method. This enables out-of-sample prediction for kernel-based dimensionality reduction methods.

The function computes the kernel matrix between test (:attr:X) and training samples (:attr:X_train) to obtain :math:K_{\text{test,train}}, then projects onto the training eigenvectors. For kernels with normalization, the normalization effects from training are already captured in the eigenvectors, so no additional normalization is applied to the non-square K_test_train matrix.

Mathematical formulation:

.. math::

\Phi_{\text{test}} = K_{\text{test,train}} \Psi_{\text{train}}

where :math:K_{\text{test,train}} \in \mathbb{R}^{n_{\text{test}} \times n_{\text{train}}} is the kernel matrix between test and training samples, and :math:\Psi_{\text{train}} \in \mathbb{R}^{n_{\text{train}} \times d} are the eigenvectors from training.

For diffusion maps (when scale_by_eigenvalues=True), the embedding is scaled:

.. math::

\Phi_{\text{test}} = K_{\text{test,train}} \Psi_{\text{train}} \Lambda

where :math:\Lambda = \text{diag}(\lambda_1, \ldots, \lambda_d) contains the eigenvalues from training.

PARAMETER DESCRIPTION
X

New data points of shape (n_samples, in_features).

TYPE: Tensor

X_train

Training data points of shape (n_train_samples, in_features).

TYPE: Tensor

weights

New sample weights of shape (n_samples, ). Currently not used because kernel normalization cannot be applied to non-square matrices. Reserved for future extensions.

TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

weights_train

Training sample weights of shape (n_train_samples, ). Currently not used because kernel normalization cannot be applied to non-square matrices. The normalization effects from training are already captured in the eigenvectors. Reserved for future extensions.

TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

eigenvectors

Eigenvectors from training of shape (n_train_samples, out_features).

TYPE: Tensor

eigenvalues

Eigenvalues from training of shape (out_features, ).

TYPE: Tensor

kernel_fn

Kernel function used during training. Must be a Kernel instance.

TYPE: Kernel

distance_fn

Pairwise distance function. Must be a PairwiseDistance instance.

TYPE: PairwiseDistance

out_features

Number of components to return. If None, returns all components.

TYPE: int | None, optional, by default None DEFAULT: None

scale_by_eigenvalues

Whether to scale projection by eigenvalues (for diffusion maps). Set to False for methods like Laplacian eigenmaps.

TYPE: bool, optional, by default True DEFAULT: True

RETURNS DESCRIPTION
Tensor

Embedded coordinates of shape (n_samples, out_features).

Examples:

>>> import torch
>>> from spectre.decomposition.nystrom import nystrom_extension
>>> from spectre.kernel import GaussianKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>> X_train = torch.randn(100, 5)
>>> X = torch.randn(20, 5)
>>> # Assume eigenvectors and eigenvalues from training
>>> kernel_fn = GaussianKernel(bw_method=0.5)
>>> distance_fn = PairwiseDistanceEuclidean()
>>> coords = nystrom_extension(
...     X,
...     X_train,
...     eigenvectors,
...     eigenvalues,
...     kernel_fn,
...     distance_fn,
...     out_features=3,
... )
>>> coords.shape
torch.Size([20, 3])
Source code in spectre/decomposition/nystrom.py
def nystrom_extension(
    X: torch.Tensor,
    X_train: torch.Tensor,
    weights: torch.Tensor | None = None,
    weights_train: torch.Tensor | None = None,
    *,
    eigenvectors: torch.Tensor,
    eigenvalues: torch.Tensor,
    kernel_fn: Kernel,
    distance_fn: PairwiseDistance,
    out_features: int | None = None,
    scale_by_eigenvalues: bool = True,
) -> torch.Tensor:
    """
    Compute Nyström extension for out-of-sample predictions.

    Projects new data points onto the embedding learned from training data using
    the Nyström approximation method. This enables out-of-sample prediction for
    kernel-based dimensionality reduction methods.

    The function computes the kernel matrix between test (:attr:`X`) and training
    samples (:attr:`X_train`) to obtain :math:`K_{\\text{test,train}}`, then projects
    onto the training eigenvectors. For kernels with normalization, the normalization
    effects from training are already captured in the eigenvectors, so no additional
    normalization is applied to the non-square K_test_train matrix.

    Mathematical formulation:

    .. math::

        \\Phi_{\\text{test}} = K_{\\text{test,train}} \\Psi_{\\text{train}}

    where :math:`K_{\\text{test,train}} \\in \\mathbb{R}^{n_{\\text{test}}
    \\times n_{\\text{train}}}` is the kernel matrix between test and training
    samples, and :math:`\\Psi_{\\text{train}} \\in \\mathbb{R}^{n_{\\text{train}}
    \\times d}` are the eigenvectors from training.

    For diffusion maps (when ``scale_by_eigenvalues=True``), the embedding is scaled:

    .. math::

        \\Phi_{\\text{test}} = K_{\\text{test,train}} \\Psi_{\\text{train}} \\Lambda

    where :math:`\\Lambda = \\text{diag}(\\lambda_1, \\ldots, \\lambda_d)` contains
    the eigenvalues from training.


    Parameters
    ----------
    X : torch.Tensor
        New data points of shape (n_samples, in_features).

    X_train : torch.Tensor
        Training data points of shape (n_train_samples, in_features).

    weights : torch.Tensor | None, optional, by default None
        New sample weights of shape (n_samples, ). Currently not used because
        kernel normalization cannot be applied to non-square matrices.
        Reserved for future extensions.

    weights_train : torch.Tensor | None, optional, by default None
        Training sample weights of shape (n_train_samples, ). Currently not used
        because kernel normalization cannot be applied to non-square matrices.
        The normalization effects from training are already captured in the
        eigenvectors. Reserved for future extensions.

    eigenvectors : torch.Tensor
        Eigenvectors from training of shape (n_train_samples, out_features).

    eigenvalues : torch.Tensor
        Eigenvalues from training of shape (out_features, ).

    kernel_fn : Kernel
        Kernel function used during training. Must be a Kernel instance.

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

    out_features : int | None, optional, by default None
        Number of components to return. If None, returns all components.

    scale_by_eigenvalues : bool, optional, by default True
        Whether to scale projection by eigenvalues (for diffusion maps).
        Set to False for methods like Laplacian eigenmaps.


    Returns
    -------
    torch.Tensor
        Embedded coordinates of shape (n_samples, out_features).


    Examples
    --------
    >>> import torch
    >>> from spectre.decomposition.nystrom import nystrom_extension
    >>> from spectre.kernel import GaussianKernel
    >>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
    >>> X_train = torch.randn(100, 5)
    >>> X = torch.randn(20, 5)
    >>> # Assume eigenvectors and eigenvalues from training
    >>> kernel_fn = GaussianKernel(bw_method=0.5)
    >>> distance_fn = PairwiseDistanceEuclidean()
    >>> coords = nystrom_extension(
    ...     X,
    ...     X_train,
    ...     eigenvectors,
    ...     eigenvalues,
    ...     kernel_fn,
    ...     distance_fn,
    ...     out_features=3,
    ... )
    >>> coords.shape
    torch.Size([20, 3])
    """
    # TODO: Check if this works correctly.
    if out_features is None:
        out_features = len(eigenvalues)

    # Compute distance matrix between test and training samples (n_test x n_train).
    pdist = distance_fn(X, Y=X_train)

    # Compute unnormalized kernel matrix.
    # We pass weights=None because `forward_step` does not use weights
    # (weights are only used by normalization, which expects square matrices).
    kernel = kernel_fn.forward_step(pdist, weights=None)

    # For Nyström extension with normalized kernels, we cannot directly apply
    # the normalization because it expects square matrices. Instead, normalization
    # effects are already captured in the eigenvectors from training.
    # The kernel matrix K_test_train is used as-is for projection.

    # Project onto eigenvectors.
    projection = torch.matmul(kernel, eigenvectors[:, :out_features])

    # Scale by eigenvalues if requested (diffusion maps yes, Laplacian eigenmaps no).
    if scale_by_eigenvalues:
        projection = projection * eigenvalues[:out_features].unsqueeze(0)

    return projection