Pairwise Distance Euclidean

euclidean

CLASS	DESCRIPTION
PairwiseDistanceEuclidean	PyTorch module for computing Euclidean pairwise distances.

FUNCTION	DESCRIPTION
pairwise_distance_euclidean	Compute Euclidean pairwise distances between tensors.

Classes

PairwiseDistanceEuclidean(reduce: Literal['mean', 'sum', 'none', None] = None, compute_mode: Literal['donot_use_mm_for_euclid_dist', 'use_mm_for_euclid_dist', 'use_mm_for_euclid_dist_if_necessary'] = 'donot_use_mm_for_euclid_dist', normalize: bool = False)

Bases: PairwiseDistance

PyTorch module for computing Euclidean pairwise distances.

This module wraps the pairwise_distance_euclidean function as a PyTorch module for integration into neural network architectures.

PARAMETER	DESCRIPTION
reduce	Reduction method for the computed distances. "mean": Average distance across specified dimensions "sum": Sum distances across specified dimensions "none" or None: Return full distance matrix without reduction TYPE: Literal["mean", "sum", "none", None], optional, by default None DEFAULT: None
compute_mode	Computation strategy for Euclidean distances. "donot_use_mm_for_euclid_dist": Use element-wise operations "use_mm_for_euclid_dist": Force matrix multiplication approach "use_mm_for_euclid_dist_if_necessary": Adaptive selection based on data TYPE: str, optional, by default "donot_use_mm_for_euclid_dist" DEFAULT: 'donot_use_mm_for_euclid_dist'
normalize	Whether to normalize distances to unit scale. TYPE: bool, optional, by default False DEFAULT: False

METHOD	DESCRIPTION
forward_step	Compute Euclidean pairwise distances.

Source code in spectre/pairwise_distance/euclidean.py

def __init__(
    self,
    reduce: Literal["mean", "sum", "none", None] = None,
    compute_mode: Literal[
        "donot_use_mm_for_euclid_dist",
        "use_mm_for_euclid_dist",
        "use_mm_for_euclid_dist_if_necessary",
    ] = "donot_use_mm_for_euclid_dist",
    normalize: bool = False,
) -> None:
    super().__init__(reduce=reduce, compute_mode=compute_mode, normalize=normalize)

Functions

forward_step(X: torch.Tensor, **kwargs) -> torch.Tensor

Compute Euclidean pairwise distances.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, n_features). TYPE: Tensor
**kwargs	Additional keyword arguments. Y (torch.Tensor | None): second tensor for cross-distances TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Pairwise distance matrix.

Source code in spectre/pairwise_distance/euclidean.py

def forward_step(self, X: torch.Tensor, **kwargs) -> torch.Tensor:
    """
    Compute Euclidean pairwise distances.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor with shape (n_samples, n_features).

    **kwargs : dict
        Additional keyword arguments.

        - Y (torch.Tensor | None): second tensor for cross-distances

    Returns
    -------
    torch.Tensor
        Pairwise distance matrix.
    """
    return pairwise_distance_euclidean(
        X,
        Y=kwargs.get("Y", None),
        reduce=self.reduce,
        compute_mode=self.compute_mode,
        normalize=self.normalize,
    )

Functions

pairwise_distance_euclidean(X: torch.Tensor, Y: torch.Tensor | None = None, reduce: Literal['mean', 'sum', 'none', None] = None, compute_mode: Literal['donot_use_mm_for_euclid_dist', 'use_mm_for_euclid_dist', 'use_mm_for_euclid_dist_if_necessary'] = 'donot_use_mm_for_euclid_dist', normalize: bool = False) -> torch.Tensor

Compute Euclidean pairwise distances between tensors.

This function uses torch.cdist to efficiently compute pairwise Euclidean distances with various computation modes and optional normalization.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, n_features). TYPE: Tensor
Y	Second tensor with same feature dimension. If None, computes distances within X. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
reduce	Reduction method to apply to distance matrix. TYPE: Literal["mean", "sum", "none", None], optional, by default None DEFAULT: None
compute_mode	Computation mode for torch.cdist. Option "donot_use_mm_for_euclid_dist" avoids matrix multiplication to prevent potential non-zero diagonal elements. TYPE: str, optional, by default "donot_use_mm_for_euclid_dist" DEFAULT: 'donot_use_mm_for_euclid_dist'
normalize	Whether to L2-normalize the distance matrix along dimension 1. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
Tensor	Pairwise distance matrix. If reduce is None: (n_samples_X, n_samples_Y) If reduce is "mean" or "sum": (n_samples_X, )

References

• https://pytorch.org/docs/stable/generated/torch.cdist.html
• https://github.com/pytorch/pytorch/issues/57690

Examples:

>>> import torch
>>> X = torch.randn(50, 10)
>>> distances = pairwise_distance_euclidean(X)
>>> distances.shape
torch.Size([50, 50])

>>> # With reduction
>>> mean_distances = pairwise_distance_euclidean(X, reduce="mean")
>>> mean_distances.shape
torch.Size([50])

Source code in spectre/pairwise_distance/euclidean.py

def pairwise_distance_euclidean(
    X: torch.Tensor,
    Y: torch.Tensor | None = None,
    reduce: Literal["mean", "sum", "none", None] = None,
    compute_mode: Literal[
        "donot_use_mm_for_euclid_dist",
        "use_mm_for_euclid_dist",
        "use_mm_for_euclid_dist_if_necessary",
    ] = "donot_use_mm_for_euclid_dist",
    normalize: bool = False,
) -> torch.Tensor:
    """
    Compute Euclidean pairwise distances between tensors.

    This function uses `torch.cdist` to efficiently compute pairwise Euclidean
    distances with various computation modes and optional normalization.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor with shape (n_samples, n_features).

    Y : torch.Tensor | None, optional, by default None
        Second tensor with same feature dimension. If None, computes distances
        within X.

    reduce : Literal["mean", "sum", "none", None], optional, by default None
        Reduction method to apply to distance matrix.

    compute_mode : str, optional, by default "donot_use_mm_for_euclid_dist"
        Computation mode for `torch.cdist`. Option "donot_use_mm_for_euclid_dist"
        avoids matrix multiplication to prevent potential non-zero diagonal
        elements.

    normalize : bool, optional, by default False
        Whether to L2-normalize the distance matrix along dimension 1.

    Returns
    -------
    torch.Tensor
        Pairwise distance matrix.

        - If reduce is None: (n_samples_X, n_samples_Y)
        - If reduce is "mean" or "sum": (n_samples_X, )

    References
    ----------
    - <https://pytorch.org/docs/stable/generated/torch.cdist.html>
    - <https://github.com/pytorch/pytorch/issues/57690>

    Examples
    --------
    >>> import torch
    >>> X = torch.randn(50, 10)
    >>> distances = pairwise_distance_euclidean(X)
    >>> distances.shape
    torch.Size([50, 50])

    >>> # With reduction
    >>> mean_distances = pairwise_distance_euclidean(X, reduce="mean")
    >>> mean_distances.shape
    torch.Size([50])
    """
    check_2d(X)

    if Y is not None:
        check_2d(Y)
        check_same_shape(X, Y, axis=1)

    # https://docs.pytorch.org/docs/stable/generated/torch.cdist.html:
    # `compute_mode=‘donot_use_mm_for_euclid_dist’` - will never use matrix
    # multiplication approach to calculate euclidean distance.
    # Apparently, when mm is used, it may sometimes result in non-zero diagonal
    # elements: https://github.com/pytorch/pytorch/issues/57690.
    X = torch.cdist(X, Y if Y is not None else X, compute_mode=compute_mode)

    if normalize:
        X = torch.nn.functional.normalize(X, dim=1, p=2)

    if reduce == "mean":
        return X.mean(dim=-1)
    elif reduce == "sum":
        return X.sum(dim=-1)
    elif reduce in ["none", None]:
        return X
    else:
        raise ValueError(
            f"Expected `reduce` to be 'mean', 'sum', 'none', or None, got {reduce}."
        )