Pairwise Distance Periodic

periodic

CLASS	DESCRIPTION
PairwiseDistancePeriodic	PyTorch module for computing periodic pairwise distances.

FUNCTION	DESCRIPTION
pairwise_distance_periodic	Compute periodic pairwise distances between tensors.

Classes

PairwiseDistancePeriodic(reduce: Literal['mean', 'sum', 'none', None] = None, period: float = 2 * pi, normalize: bool = False)

Bases: PairwiseDistance

PyTorch module for computing periodic pairwise distances.

Computes distances for periodic variables (e.g., dihedral angles) where differences are wrapped to [-pi, pi]. This ensures that angles near 0 and 2 pi are treated as close in distance space.

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
period	Period of the periodic variables. Default assumes angles in radians. Use 360.0 for degrees. TYPE: float, optional, by default 2 pi DEFAULT: 2 * pi
normalize	Whether to normalize distances to unit scale. TYPE: bool, optional, by default False DEFAULT: False

METHOD	DESCRIPTION
forward_step	Compute periodic pairwise distances.

Source code in spectre/pairwise_distance/periodic.py

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

    if not isinstance(period, float):
        raise TypeError(f"`period` must be a float, got {type(period)}.")
    check_in_interval(period, "(0, inf)")
    self.period = period

Functions

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

Compute periodic pairwise distances.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, in_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/periodic.py

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

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

    **kwargs : dict
        Additional keyword arguments.

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

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

Functions

pairwise_distance_periodic(X: torch.Tensor, Y: torch.Tensor | None = None, reduce: Literal['mean', 'sum', 'none', None] = None, period: float = 2 * torch.pi, normalize: bool = False) -> torch.Tensor

Compute periodic pairwise distances between tensors.

Computes distances for periodic variables (e.g., dihedral angles) where differences are wrapped to [-pi, pi]. This ensures that angles near 0 and 2 pi are treated as close in distance space.

PARAMETER	DESCRIPTION
X	Input tensor with shape (n_samples, in_features) containing periodic variables. 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
period	Period of the periodic variables. Use 2π for radians, 360.0 for degrees. TYPE: float, optional, by default 2 pi DEFAULT: 2 * pi
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.

Source code in spectre/pairwise_distance/periodic.py

def pairwise_distance_periodic(
    X: torch.Tensor,
    Y: torch.Tensor | None = None,
    reduce: Literal["mean", "sum", "none", None] = None,
    period: float = 2 * torch.pi,
    normalize: bool = False,
) -> torch.Tensor:
    """
    Compute periodic pairwise distances between tensors.

    Computes distances for periodic variables (e.g., dihedral angles) where
    differences are wrapped to [-pi, pi]. This ensures that angles near 0 and 2 pi
    are treated as close in distance space.

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

    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.

    period : float, optional, by default 2 pi
        Period of the periodic variables. Use 2π for radians, 360.0 for degrees.

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

    Returns
    -------
    torch.Tensor
        Pairwise distance matrix.
    """
    check_2d(X)

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

    diff = X.unsqueeze(1) - Y.unsqueeze(0)
    diff = 2 * pi * diff / period

    # Wrap to [-pi, pi] using atan2
    wrapped_diff = torch.atan2(torch.sin(diff), torch.cos(diff))

    # Convert back to original period
    wrapped_diff = wrapped_diff * period / (2 * pi)

    distances = torch.sqrt((wrapped_diff**2).sum(dim=-1))

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

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