Geometry and structural losses

Between-Residue Clash Loss

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Penalize between-residue atom pairs that overlap their supplied van der Waals radii beyond the allowed tolerance. Return upper-triangle pair losses, their normalized mean, and incident per-atom loss and clash flags.

Function signature

between_residue_clash(
    activations: dict,
    masks: dict,
    structure: dict,
    count_eps: float = 1e-06,
    distance_eps: float = 1e-10,
    overlap_soft: float = 1.5,
    overlap_hard: float = 1.5,
) -> dict[str, torch.Tensor]

Worked example

Only upper-triangle pairs (0,1) and (0,2) are eligible; excluding (2,1) also excludes (1,2). The illustrative distance_eps=3 makes the first distance 2 and its loss 1. The other loss is zero. With count_eps=1, mean_loss=1/(2+1)=1/3; incident atom losses are [1,1,0].

import torch

activations = {
    'radii': torch.tensor([[2.0, 2.0, 2.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'atom_mask': torch.tensor([[True, True, True]], dtype=torch.bool),
    'pair_exclusion': torch.tensor([[[False, False, False], [False, False, False], [False, True, False]]], dtype=torch.bool),
}
structure = {
    'positions': torch.tensor([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [5.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True),
    'residue_index': torch.tensor([[0, 1, 2]], dtype=torch.int64),
    'chain_id': torch.tensor([[0, 0, 0]], dtype=torch.int64),
}
count_eps = 1.0
distance_eps = 3.0
overlap_soft = 1.0
overlap_hard = 0.5

result = between_residue_clash(activations=activations, masks=masks, structure=structure, count_eps=count_eps, distance_eps=distance_eps, overlap_soft=overlap_soft, overlap_hard=overlap_hard)

Expected return

{
    'pair_loss': torch.tensor([[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]], dtype=torch.float64),
    'mean_loss': torch.tensor([0.3333333333333333], dtype=torch.float64),
    'per_atom_loss': torch.tensor([[1.0, 1.0, 0.0]], dtype=torch.float64),
    'atom_violation': torch.tensor([[True, True, False]], dtype=torch.bool),
}

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`count_eps`floatDefault `1e-06`	Positive value added to the number of eligible upper-triangle pairs in the mean_loss denominator.
`distance_eps`floatDefault `1e-10`	Positive value added to squared atom separation before sqrt.
`overlap_soft`floatDefault `1.5`	Allowed overlap subtracted before the soft between-residue penalty.
`overlap_hard`floatDefault `1.5`	Allowed overlap used for the hard between-residue flag.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Atoms belong to the same residue only when BOTH residue_index and chain_id match. Exclude absent atoms and sanitize their coordinates before distance calculations.
Include only different-residue pairs with flattened atom index i<j. Exclude pair_exclusion[i,j] OR pair_exclusion[j,i]; count each unordered pair once.
Use stabilized distance $d=\sqrt{\|r_i-r_j\|^2+distance\_eps}$ , distance_eps default $10^{-10}$ .
Let $r=radii_i+radii_j$ . Soft loss is $\max(r-overlap\_soft-d,0)$ ; a hard clash requires $d<r-overlap\_hard$ . Both overlap tolerances default to 1.5.
Zero excluded pair losses. mean_loss divides their sum by the eligible pair count plus count_eps (default 1e-6). per_atom_loss adds row and column sums; atom_violation is any flagged pair incident to the atom.
Tolerances are finite nonnegative numbers; distance_eps/count_eps must be positive. Even when no pair is eligible, return correctly shaped exact-zero losses and False flags.

Why this matters

Bad peptide geometry, within-residue distortions and overlaps are different problems. Soft variants provide penalties for optimization; the AF3 clash variant reports a hard structural flag. Select one variant before deciding what to count.

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 between_residue_clash(
    activations: dict,
    masks: dict,
    structure: dict,
    count_eps: float = 1e-06,
    distance_eps: float = 1e-10,
    overlap_soft: float = 1.5,
    overlap_hard: float = 1.5,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pair_loss: (B, A, A).
    mean_loss: (B,).
    per_atom_loss: (B, A).
    atom_violation: (B, A)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement between_residue_clash")

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.

Related exercises

Sources and further reading

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