Pairwise Distance Covariance

covariance

CLASS	DESCRIPTION
PairwiseDistanceCov	PyTorch module for computing covariance-based pairwise distances.

FUNCTION	DESCRIPTION
pairwise_distance_cov	Compute covariance-based pairwise distances between data samples.

Classes

PairwiseDistanceCov(cov_k_neighbor: int, cov_mode: Literal['batchwise', 'sequential'] = 'batchwise', 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 covariance-based pairwise distances.

This module computes pairwise distances using local covariance matrices estimated from nearest neighbors, implementing a heterogeneous Gaussian distance metric.

PARAMETER	DESCRIPTION
cov_k_neighbor	Number of nearest neighbors to use for covariance estimation. TYPE: int
cov_mode	Computation mode for covariance calculation. TYPE: Literal["batchwise", "sequential"], optional, by default "batchwise" DEFAULT: 'batchwise'
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 covariance-based pairwise distances.

Source code in spectre/pairwise_distance/covariance.py

def __init__(
    self,
    cov_k_neighbor: int,
    cov_mode: Literal["batchwise", "sequential"] = "batchwise",
    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)

    if not isinstance(cov_k_neighbor, int):
        raise TypeError(
            f"Metric `cov` requires `cov_k_neighbor` to be int, got {type(cov_k_neighbor)}."
        )
    self.cov_k_neighbor = cov_k_neighbor

    avail_cov_mode = ["batchwise", "sequential"]
    if cov_mode not in avail_cov_mode:
        raise ValueError(
            f"Expected `cov_mode` to be {avail_cov_mode}, got {cov_mode}."
        )
    self.cov_mode = cov_mode

Functions

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

Compute covariance-based pairwise distances.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, n_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Pairwise distance matrix.

Source code in spectre/pairwise_distance/covariance.py

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

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

    Returns
    -------
    torch.Tensor
        Pairwise distance matrix.
    """
    return pairwise_distance_cov(
        X,
        cov_k_neighbor=self.cov_k_neighbor,
        cov_mode=self.cov_mode,
        reduce=self.reduce,
        compute_mode=self.compute_mode,
        normalize=self.normalize,
    )

Functions

pairwise_distance_cov(X: torch.Tensor, cov_k_neighbor: int = 64, cov_mode: Literal['batchwise', 'sequential'] = 'batchwise', 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, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute covariance-based pairwise distances between data samples.

This function implements a heterogeneous Gaussian distance metric using local covariance matrices estimated from nearest neighbors. The distance between two samples is computed using their local covariance structure.

The heterogeneous Gaussian distance is computed as: \(d(x_k, x_l) = \sqrt{(x_k - x_l)^\top (C_{k}^{-1} + C_{l}^{-1}) (x_k - x_l)}\), where \(C_k\) and \(C_l\) are the local covariance matrices at samples \(x_k\) and \(x_l\).

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, in_features). TYPE: Tensor
cov_k_neighbor	Number of nearest neighbors to use for covariance estimation. TYPE: int, optional, by default 64 DEFAULT: 64
cov_mode	Computation mode for covariance calculation. "batchwise": Vectorized computation for all samples "sequential": Loop-based computation sample by sample TYPE: Literal["batchwise", "sequential"], optional, By default "batchwise" DEFAULT: 'batchwise'
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
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. If reduce is None: (n_samples, n_samples) If reduce is "mean" or "sum": (n_samples, )

Examples:

>>> import torch
>>> X = torch.randn(100, 5)
>>> distances = pairwise_distance_cov(X, cov_k_neighbor=20)
>>> distances.shape
torch.Size([100, 100])

Source code in spectre/pairwise_distance/covariance.py

def pairwise_distance_cov(
    X: torch.Tensor,
    cov_k_neighbor: int = 64,
    cov_mode: Literal["batchwise", "sequential"] = "batchwise",
    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,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute covariance-based pairwise distances between data samples.

    This function implements a heterogeneous Gaussian distance metric using
    local covariance matrices estimated from nearest neighbors. The distance
    between two samples is computed using their local covariance structure.

    The heterogeneous Gaussian distance is computed as:
    $d(x_k, x_l) = \\sqrt{(x_k - x_l)^\\top (C_{k}^{-1} + C_{l}^{-1}) (x_k - x_l)}$,
    where $C_k$ and $C_l$ are the local covariance matrices at samples $x_k$ and $x_l$.

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

    cov_k_neighbor : int, optional, by default 64
        Number of nearest neighbors to use for covariance estimation.

    cov_mode : Literal["batchwise", "sequential"], optional, By default "batchwise"
        Computation mode for covariance calculation.

        - "batchwise": Vectorized computation for all samples
        - "sequential": Loop-based computation sample by sample

    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

    compute_mode : str, optional, by default "donot_use_mm_for_euclid_dist"
        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

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

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

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

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

    Examples
    --------
    >>> import torch
    >>> X = torch.randn(100, 5)
    >>> distances = pairwise_distance_cov(X, cov_k_neighbor=20)
    >>> distances.shape
    torch.Size([100, 100])
    """
    check_2d(X)

    # Find `cov_n_neigbor` nearest neighbors.
    _, indices = pairwise_distance_euclidean(X, compute_mode=compute_mode).topk(
        cov_k_neighbor + 1, dim=-1, largest=False
    )

    if cov_mode == "batchwise":
        inv_cov = _pinv_cov_batchwise(X[indices], rowvar=False, bias=False)
    elif cov_mode == "sequential":
        inv_cov = _pinv_cov_sequential(X, indices)

    # Heterogeneous Gaussian distance.
    x_diff = X.unsqueeze(dim=1) - X.unsqueeze(dim=0)
    cov_inv_sum = inv_cov.unsqueeze(dim=1) + inv_cov.unsqueeze(dim=0)
    X = torch.sqrt(torch.einsum("nmi,nmik,nmk->nm", x_diff, cov_inv_sum, x_diff) + eps)

    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}."
        )