Geometry and structural losses

Masked pairwise distances and hard distance-bin targets

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

AlphaFold2RoseTTAFoldRoseTTAFold2ESMFoldBoltz-2

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Compare every point in one set with every point in another. Return masked distances and discrete bin IDs; the reported-distance smoothing and binning rules are separate.

Function signature

pairwise_distance_encodings(
    x: torch.Tensor,
    y: torch.Tensor,
    mask_x: torch.Tensor,
    mask_y: torch.Tensor,
    boundaries: torch.Tensor,
    squared: bool = False,
    bin_space: str = 'euclidean',
    eps: float = 1e-08,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

One pair, two distance conventions

Take x=[[0,0,0]], y=[[3,4,0]], both masks=[True], boundaries=[3,5,6], squared=False, bin_space="euclidean" and eps=11. The unregularized squared distance is 3²+4²=25, while the reported distance is sqrt(25+11)=6. Only boundaries 3 and 5 satisfy 25≥boundary², so bins=[[2]]. Binning the reported 6 would incorrectly pass boundary 6 as well. distances=[[6]].

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[\ldots, N, 3]$ Cartesian coordinates of the first point set.
`y`torch.Tensor	$[\ldots, M, 3]$ Cartesian coordinates of the second point set; shares the leading batch axes with x.
`mask_x`torch.Tensor	$[\ldots, N]$ Bool $[\ldots,N]$ , `True` means valid.
`mask_y`torch.Tensor	$[\ldots, M]$ Bool $[\ldots,M]$ , `True` means valid.
`boundaries`torch.Tensor	$[K]$ Nonnegative ascending bin `boundaries`; hard metadata, no gradient.
`squared`boolDefault `False`	Return squared rather than smoothed Euclidean distances when True.
`bin_space`strDefault `'euclidean'`	Interpret supplied boundaries in "euclidean" or "squared" distance space, independently of the reported-distance choice.
`eps`floatDefault `1e-08`	Positive value added inside $\sqrt{d^2+\varepsilon}$ for the reported nonsquared distances. Distance-bin labels still use the specified unsmoothed distance domain.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Let $x [\ldots,N,3]$ and $y [\ldots,M,3]$ have identical leading batch axes, with $N$ and $M$ independent. All values, including padding, are finite. True mask_x $[\ldots,N]$ and mask_y $[\ldots,M]$ mean valid atoms; $\mathrm{valid}(i,j)=\mathrm{mask\_x}[i]\land \mathrm{mask\_y}[j]$ .
Compute $\mathrm{d2\_ij}=\sum_{c} (x_{\mathrm{ic}}-y_{\mathrm{jc}})^{2}$ . squared (bool, default False) selects output d2, or $\sqrt{\mathrm{d2}+\varepsilon}$ with $\varepsilon>0$ (default $10^{-8}$ , inside the square root). Invalid distances are exactly zero.
The $K$ boundaries are finite, strictly increasing, nonnegative, and nondifferentiable. Config bin_space is 'euclidean' (default) or 'squared': compare the UNREGULARIZED d2 against $\mathrm{boundaries}^{2}$ or boundaries respectively.
For euclidean bin_space, count boundaries satisfying $d^2\ge b_k^2$ ; for squared bin_space, count $d^2\ge b_k$ . Equality belongs to the upper bin. Invalid pairs use -1; with no boundaries a valid pair has bin 0. Coincident valid points are retained, and empty point axes remain empty.
No coordinates or masks are mutated; no RBF features are requested.

Why this matters

Two singleton point axes create all pair differences; continuous lengths and hard class boundaries then use different treatments of squared distance.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def pairwise_distance_encodings(
    x: torch.Tensor,  # (..., N, 3)
    y: torch.Tensor,  # (..., M, 3)
    mask_x: torch.Tensor,  # (..., N)
    mask_y: torch.Tensor,  # (..., M)
    boundaries: torch.Tensor,  # (K,)
    squared: bool = False,
    bin_space: str = 'euclidean',
    eps: float = 1e-08,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    distances: (..., N, M).
    bins: (..., N, M)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement pairwise_distance_encodings")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.