Pairwise Distance Mahalanobis

mahalanobis

CLASS	DESCRIPTION
PairwiseDistanceMahalanobis	PyTorch module for computing Mahalanobis pairwise distances.

FUNCTION	DESCRIPTION
pairwise_distance_mahalanobis	Compute Mahalanobis pairwise distances between data samples.

Classes

PairwiseDistanceMahalanobis(reduce: Literal['mean', 'sum', 'none', None] = None, normalize: bool = False)

Bases: PairwiseDistance

PyTorch module for computing Mahalanobis pairwise distances.

This module computes pairwise Mahalanobis distances using provided inverse covariance matrices for each data sample. Supports computing distances within a single dataset (X vs X) or between two datasets (X vs Y).

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
normalize	Whether to normalize distances to unit scale. TYPE: bool, optional, by default False DEFAULT: False

Examples:

Computing distances within a single dataset:

>>> import torch
>>> from spectre.pairwise_distance import PairwiseDistanceMahalanobis
>>> X = torch.randn(50, 3)
>>> metric_inv = torch.eye(3).unsqueeze(0).repeat(50, 1, 1)
>>> mahal_dist = PairwiseDistanceMahalanobis()
>>> distances = mahal_dist(X, metric_inv=metric_inv)
>>> distances.shape
... torch.Size([50, 50])

Computing cross-distances between two datasets:

>>> Y = torch.randn(30, 3)
>>> metric_inv_y = torch.eye(3).unsqueeze(0).repeat(30, 1, 1)
>>> distances = mahal_dist(X, metric_inv=metric_inv, Y=Y, metric_inv_y=metric_inv_y)
>>> distances.shape
... torch.Size([50, 30])

METHOD	DESCRIPTION
forward_step	Compute Mahalanobis pairwise distances.

Source code in spectre/pairwise_distance/mahalanobis.py

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

Functions

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

Compute Mahalanobis pairwise distances.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, in_features). TYPE: Tensor
**kwargs	Additional keyword arguments. metric_inv : torch.Tensor (inverse covariance matrices for X) Y : torch.Tensor | None (optional second dataset) metric_inv_y : torch.Tensor | None (inverse covariance matrices for Y) TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Pairwise distance matrix.

Source code in spectre/pairwise_distance/mahalanobis.py

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

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

    **kwargs : dict
        Additional keyword arguments.

        - metric_inv : torch.Tensor (inverse covariance matrices for X)
        - Y : torch.Tensor | None (optional second dataset)
        - metric_inv_y : torch.Tensor | None (inverse covariance matrices for Y)

    Returns
    -------
    torch.Tensor
        Pairwise distance matrix.
    """
    metric_inv = kwargs.get("metric_inv", None)
    if metric_inv is None:
        raise ValueError("metric_inv must be provided as keyword argument.")

    Y = kwargs.get("Y", None)
    metric_inv_y = kwargs.get("metric_inv_y", None)

    # Validate: if Y provided, metric_inv_y must also be provided (and vice versa).
    if (Y is None) != (metric_inv_y is None):
        raise ValueError(
            "`Y` and `metric_inv_y` must both be provided or both be None. Got "
            f"Y={'provided' if Y is not None else 'None'}, "
            f"metric_inv_y={'provided' if metric_inv_y is not None else 'None'}."
        )

    return pairwise_distance_mahalanobis(
        X,
        metric_inv=metric_inv,
        Y=Y,
        metric_inv_y=metric_inv_y,
        reduce=self.reduce,
        normalize=self.normalize,
    )

Functions

pairwise_distance_mahalanobis(X: torch.Tensor, metric_inv: torch.Tensor, Y: torch.Tensor | None = None, metric_inv_y: torch.Tensor | None = None, normalize: bool = False, reduce: Literal['mean', 'sum', 'none', None] = None, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute Mahalanobis pairwise distances between data samples.

The Mahalanobis distance between two samples \(x_k\) and \(y_l\) is computed as: \(d(x_k, y_l) = \sqrt{(x_k - y_l)^\top M_{kl}^{-1} (x_k - y_l)}\), where \(M_{kl}^{-1}\) is the symmetrized inverse covariance or diffusion matrix: \(M_{kl}^{-1} = (M_{X_k}^{-1} + M_{Y_l}^{-1}) / 2\).

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples_x, in_features). TYPE: Tensor
metric_inv	Inverse covariance matrices for each sample in X with shape (n_samples_x, in_features, in_features). TYPE: Tensor
Y	Second tensor with shape (n_samples_y, in_features). If None, computes distances within X. If provided, metric_inv_y must also be provided. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
metric_inv_y	Inverse covariance matrices for each sample in Y with shape (n_samples_y, in_features, in_features). Must be provided if Y is provided. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
normalize	Whether to normalize distances to unit scale. TYPE: bool, optional, by default False DEFAULT: False
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
eps	Small value added for numerical stability. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Pairwise distance matrix. Shape depends on Y and reduce parameters: If Y is None and reduce is None: (n_samples_x, n_samples_x) If Y is provided and reduce is None: (n_samples_x, n_samples_y) If reduce is "mean" or "sum": (n_samples_x, )

Examples:

Computing distances within a single dataset:

>>> import torch
>>> X = torch.randn(50, 3)
>>> metric_inv = torch.eye(3).unsqueeze(0).repeat(50, 1, 1)  # Identity matrices
>>> distances = pairwise_distance_mahalanobis(X, metric_inv)
>>> distances.shape
... torch.Size([50, 50])

Computing cross-distances between two datasets:

>>> X = torch.randn(50, 3)
>>> Y = torch.randn(30, 3)
>>> metric_inv_X = torch.eye(3).unsqueeze(0).repeat(50, 1, 1)
>>> metric_inv_y = torch.eye(3).unsqueeze(0).repeat(30, 1, 1)
>>> distances = pairwise_distance_mahalanobis(
...     X, metric_inv_X, Y=Y, metric_inv_y=metric_inv_y
... )
>>> distances.shape
... torch.Size([50, 30])

Source code in spectre/pairwise_distance/mahalanobis.py

def pairwise_distance_mahalanobis(
    X: torch.Tensor,
    metric_inv: torch.Tensor,
    Y: torch.Tensor | None = None,
    metric_inv_y: torch.Tensor | None = None,
    normalize: bool = False,
    reduce: Literal["mean", "sum", "none", None] = None,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute Mahalanobis pairwise distances between data samples.

    The Mahalanobis distance between two samples $x_k$ and $y_l$ is
    computed as:
    $d(x_k, y_l) = \\sqrt{(x_k - y_l)^\\top  M_{kl}^{-1} (x_k - y_l)}$,
    where $M_{kl}^{-1}$ is the symmetrized inverse covariance or diffusion matrix:
    $M_{kl}^{-1} = (M_{X_k}^{-1} + M_{Y_l}^{-1}) / 2$.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor with shape (n_samples_x, in_features).

    metric_inv : torch.Tensor
        Inverse covariance matrices for each sample in X with shape
        (n_samples_x, in_features, in_features).

    Y : torch.Tensor | None, optional, by default None
        Second tensor with shape (n_samples_y, in_features). If None, computes
        distances within X. If provided, metric_inv_y must also be provided.

    metric_inv_y : torch.Tensor | None, optional, by default None
        Inverse covariance matrices for each sample in Y with shape
        (n_samples_y, in_features, in_features). Must be provided if Y is provided.

    normalize : bool, optional, by default False
        Whether to normalize distances to unit scale.

    reduce : Literal["mean", "sum", "none", None], optional, by default None
        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

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value added for numerical stability.

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

        Shape depends on Y and reduce parameters:

        - If Y is None and reduce is None: (n_samples_x, n_samples_x)
        - If Y is provided and reduce is None: (n_samples_x, n_samples_y)
        - If reduce is "mean" or "sum": (n_samples_x, )

    Examples
    --------
    Computing distances within a single dataset:

    >>> import torch
    >>> X = torch.randn(50, 3)
    >>> metric_inv = torch.eye(3).unsqueeze(0).repeat(50, 1, 1)  # Identity matrices
    >>> distances = pairwise_distance_mahalanobis(X, metric_inv)
    >>> distances.shape
    ... torch.Size([50, 50])

    Computing cross-distances between two datasets:

    >>> X = torch.randn(50, 3)
    >>> Y = torch.randn(30, 3)
    >>> metric_inv_X = torch.eye(3).unsqueeze(0).repeat(50, 1, 1)
    >>> metric_inv_y = torch.eye(3).unsqueeze(0).repeat(30, 1, 1)
    >>> distances = pairwise_distance_mahalanobis(
    ...     X, metric_inv_X, Y=Y, metric_inv_y=metric_inv_y
    ... )
    >>> distances.shape
    ... torch.Size([50, 30])
    """
    check_2d(X)
    check_same_len(X, metric_inv)

    if Y is not None:
        check_2d(Y)
        check_same_len(X, Y)

    if metric_inv_y is not None:
        check_same_len(Y, metric_inv_y)

    Y = Y if Y is not None else X
    metric_inv_y = metric_inv_y if metric_inv_y is not None else metric_inv

    # (n_X, n_Y, in_features).
    diffs = X[:, None, :] - Y[None, :, :]

    # Symmetrize inverse covariance matrices: (n_X, n_Y, in_features, in_features).
    metric_inv_sym = (metric_inv[:, None, :, :] + metric_inv_y[None, :, :, :]) / 2

    # Compute Mahalanobis distance.
    D = torch.sqrt(torch.einsum("ija,ijab,ijb->ij", diffs, metric_inv_sym, diffs) + eps)

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

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