Peptide geometry, steric violations and polymer clashes

Your task

Function signature

structural_violation_loss(
    activations: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

Count clashing atoms at the inclusive radius

import torch

activations = {
}
masks = {
    'atom_mask': torch.tensor([[True, True, True, True, True]], dtype=torch.bool),
    'polymer_mask': torch.tensor([[True, True, True, False, True]], dtype=torch.bool),
}
structure = {
    'positions': torch.tensor([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [8.0, 0.0, 0.0]]], dtype=torch.float64),
    'residue_index': torch.tensor([[10, 10, 11, 13, 99]], dtype=torch.int64),
    'chain_id': torch.tensor([[0, 1, 0, 0, 0]], dtype=torch.int64),
}
options = {
    'clash_radius': 1.0,
    'min_clashes': 1,
    'num_chains': 3,
    'min_fraction': 0.75,
    'variant': 'af3_clash',
}

result = structural_violation_loss(activations=activations, masks=masks, structure=structure, options=options)

Expected return

{
    'atom_clash_mask': torch.tensor([[True, True, True, False, False]], dtype=torch.bool),
    'chain_clash_counts': torch.tensor([[2, 1, 0]], dtype=torch.int64),
    'chain_atom_counts': torch.tensor([[3, 1, 0]], dtype=torch.int64),
    'chain_clash_fraction': torch.tensor([[0.6666666666666666, 1.0, 0.0]], dtype=torch.float64),
    'has_clash': torch.tensor([True], dtype=torch.bool),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement four explicitly different variants. B, A, R and C are positive; inputs are CPU \mathrm{float32}/\mathrm{float64} (one common floating dtype), int64 indices and bool masks. positions[B,A,3] are Angstrom coordinates; atom_mask[B,A] is true for present atoms.
• Masked coordinates and ineligible geometric terms are ignored before squaring or normalization. All inputs are finite; eligible geometry/statistics must be within the chosen dtype's arithmetic range. Empty atom axes are outside the contract; all-masked batches, singleton atoms/residues, empty chain slots and noncontiguous tensors are supported.
• variant='peptide' (default). backbone_atoms[B,R,3] indexes positions in N, CA, C order, each index in [0,A). Here residue_index and chain_id have shape [B,R]; is_proline[B,R] is bool. A link r\to r+1 exists only if residue indices differ by \mathrm{exactly} +1 AND chain IDs match.
• For the C_{r}--N_{r+1} length, require those two atoms; for \mathrm{CA\_r}--C_{r}--N_{r+1}, also require CA_r; for C_{r}--N_{r+1}--\mathrm{CA}_{r+1}, also require the next residue's C-alpha atom, \mathrm{CA}_{r+1}. Never borrow a mask from an unrelated angle. Define \operatorname{norm_e}(v)=\sqrt{\sum(v^{2})+\mathrm{geometry\_eps}}.
• The three metrics are \operatorname{norm_e}(N_{\mathrm{next}}-C), \operatorname{dot}(\mathrm{CA}-C,N_{\mathrm{next}}-C)/(\operatorname{norm_e}(\mathrm{CA}-C)\cdot \operatorname{norm_e}(N_{\mathrm{next}}-C)), and \operatorname{dot}(C-N_{\mathrm{next}},\mathrm{CA\_next}-N_{\mathrm{next}})/(\operatorname{norm_e}(C-N_{\mathrm{next}})\cdot \operatorname{norm_e}(\mathrm{CA\_next}-N_{\mathrm{next}})).
• Means/stds for the bond are \mathrm{cn\_mean}[2]/\mathrm{cn\_std}[2], indexed by whether the NEXT residue is proline (0=\mathrm{other},1=\mathrm{proline}). \mathrm{angle\_mean}[2]/\mathrm{angle\_std}[2] contain the two cosine statistics in that order.
• Standard deviations are positive; these are supplied scientific constants, not inferred from the predicted geometry.
• For metric q, mean \mu and std s, \mathrm{error}=\sqrt{(q-\mu)^{2}+\mathrm{error\_eps}}; soft \mathrm{loss}=\operatorname{max}(\mathrm{error}-\mathrm{soft\_tolerance}\cdot s,0), hard \mathrm{flag}=(\mathrm{error}>\mathrm{hard\_tolerance}\cdot s). Apply the term's mask to BOTH its soft loss and hard flag.
• link_losses[B,R -1,3] contains masked losses. mean_losses [B,3] divides each term's sum by its own valid-link \operatorname{count}+\mathrm{count\_eps}, including valid zero-loss links in that count. residue_loss[B,R] splits each link's total equally between its two residues.
• residue_violation[B,R] is true at either endpoint of any flagged term. Defaults \mathrm{geometry\_eps}=\mathrm{error\_eps}=\mathrm{count\_eps}=10^{-6}; \mathrm{soft\_tolerance}=\mathrm{hard\_tolerance}=12.
• variant='within_residue'. Now residue_index and chain_id are [B,A]; together they identify an atom's residue. lower_bound/upper_bound[B,A,A] give directed distance bounds, with 0\le \mathrm{lower}\le \mathrm{upper} on eligible pairs.
• An eligible directed entry (i,j) has both atoms present, i\ne j, and equal residue AND chain IDs. Bounds need not be symmetric (chemical tables normally are). d_{\mathrm{ij}}=\sqrt{\sum((x_{i}-x_{j})^{2})+\mathrm{distance\_eps}}.
• pair_loss[B,A,A] is \operatorname{max}(\mathrm{lower}+\mathrm{tighten}-d,0)+\operatorname{max}(d-\mathrm{upper}+\mathrm{tighten},0) on eligible entries, otherwise zero. Hard violation uses the ORIGINAL bounds: d<\mathrm{lower} OR d>\mathrm{upper}, strictly; tightening never changes the hard threshold.
• variant='between_residue'. radii[B,A] are nonnegative van der Waals radii. pair_exclusion[B,A,A] marks supplied covalent/peptide or disulfide exclusions; either orientation excludes the undirected pair.
• Eligible entries have i<j, both atoms present, different (\mathrm{chain},\mathrm{residue}) identities, and neither exclusion entry set. Use the same \epsilon distance as variant 2. \mathrm{pair\_loss}=\operatorname{max}(r_{i}+r_{j}-\mathrm{overlap\_soft}-d,0); hard violation is d<r_{i}+r_{j}-\mathrm{overlap\_hard}.
• Lower-triangular, diagonal, same-residue and excluded entries are zero. The caller supplies chemical exclusions; do not guess bonds from distance alone.
• Variants 2/3 return pair_loss, \mathrm{mean\_loss}[B]=\sum(\mathrm{pair\_loss})/(number of eligible \mathrm{entries}+\mathrm{count\_eps}), per_atom_loss[B,A]=\mathrm{row\_sum}(\mathrm{pair\_loss})+\mathrm{column\_sum}(\mathrm{pair\_loss}), and atom_violation[B,A]=any incident hard violation.
• Each eligible entry contributes its FULL loss to BOTH endpoints. Within-residue matrices retain both directions: for symmetric bounds an unordered pair contributes twice to each endpoint. Between-residue matrices keep one direction: an unordered pair contributes once to each endpoint.
• The within-residue mean is an added teaching diagnostic, not a separately claimed upstream loss. Defaults \mathrm{distance\_eps}=10^{-10}, \mathrm{count\_eps}=10^{-6}, \mathrm{tighten}=0, \mathrm{overlap\_soft}=\mathrm{overlap\_hard}=1.5. Unused bounds/radii are ignored.
• variant='af3_clash'. residue_index/chain_id are [B,A]; polymer_mask[B,A] selects protein/RNA/DNA atoms. Only atom_mask AND polymer_mask participates. A pair is eligible if its chain IDs differ OR its same-chain residue indices differ by more than 1.
• Test squared Euclidean \mathrm{distance} \le \mathrm{clash\_radius}^{2}, with no \epsilon; the radius boundary is inclusive. atom_clash_mask[B,A] is true if an atom has at least one eligible neighbor, regardless of how many.
• \mathrm{num\_chains}=C (default 1) supplies dense chain slots 0..C-1; participating atoms must have an ID in that range. chain_clash_counts and chain_atom_counts[B,C] are int64; chain_clash_fraction[B,C]=counts/totals, or zero for an empty chain.
• has_clash [B] is true if ANY chain has \operatorname{count}>\mathrm{min\_clashes} OR \mathrm{fraction}>\mathrm{min\_fraction}; both inequalities are strict. Defaults \mathrm{clash\_radius}=1.1 Angstrom, \mathrm{min\_clashes}=100, \mathrm{min\_fraction}=0.5. Non-polymer atoms do not affect numerator or denominator.
• There is no bond/van-der-Waals loss, SASA calculation, or ranking arithmetic in this variant. All epsilons must be finite and strictly positive; other numerical tolerances must be finite and nonnegative. min_fraction must be in [0,1], min_clashes a nonnegative integer, and num_chains a positive integer.
• Unknown variants and these invalid configurations raise ValueError. No geometry/table shape validation is required outside the stated input domain. Variant-specific keys are required only for that variant. All other configuration keys are ignored. Loss variants return exact connected zeros when no term is eligible; bool outputs are false.
• Each variant returns only its listed output subset.

Why this matters

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 structural_violation_loss(
    activations: dict,
    masks: dict,
    structure: dict,
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    link_losses: (B, R-1, 3). Returned only for peptide.
    mean_losses: (B, 3). Returned only for peptide.
    residue_loss: (B, R). Returned only for peptide.
    residue_violation: (B, R). Returned only for peptide.
    pair_loss: (B, A, A). Returned only for within_residue or between_residue.
    mean_loss: (B,). Returned only for within_residue or between_residue.
    per_atom_loss: (B, A). Returned only for within_residue or between_residue.
    atom_violation: (B, A). Returned only for within_residue or between_residue.
    atom_clash_mask: (B, A). Returned only for af3_clash.
    chain_clash_counts: (B, C). Returned only for af3_clash.
    chain_atom_counts: (B, C). Returned only for af3_clash.
    chain_clash_fraction: (B, C). Returned only for af3_clash.
    has_clash: (B,). Returned only for af3_clash."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement structural_violation_loss")

Related exercises

• Builds onPeptide Geometry Violation
• Builds onWithin-Residue Distance Violation
• Builds onBetween-Residue Clash Loss
• Builds onPolymer Clash Summary

Sources and further reading

• www.nature.com
• google-deepmind/alphafold · all_atom.py
• www.nature.com
• google-deepmind/alphafold3 · confidences.py