Within-Residue Distance Violation

Your task

Function signature

within_residue_violation(
    activations: dict,
    masks: dict,
    structure: dict,
    count_eps: float = 1e-06,
    distance_eps: float = 1e-10,
    tighten: float = 0.0,
) -> dict[str, torch.Tensor]

Worked example

Worked example

import torch

activations = {
    'lower_bound': torch.full((1, 3, 3), 3.0, dtype=torch.float64, requires_grad=True),
    'upper_bound': torch.full((1, 3, 3), 4.0, dtype=torch.float64, requires_grad=True),
}
masks = {
    'atom_mask': torch.tensor([[True, True, True]], 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, 0, 1]], dtype=torch.int64),
    'chain_id': torch.tensor([[0, 0, 0]], dtype=torch.int64),
}
count_eps = 1.0
distance_eps = 3.0
tighten = 0.0

result = within_residue_violation(activations=activations, masks=masks, structure=structure, count_eps=count_eps, distance_eps=distance_eps, tighten=tighten)

Expected return

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

Parameters

Return

dict[str, torch.Tensor]

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 all directed same-residue pairs except the diagonal. Both(i,j) and (j,i) are present; do not halve the returned loss.
• Use stabilized distance d=\sqrt{\|r_i-r_j\|^2+distance\_eps}, distance_eps default 10^{-10}.
• Soft loss is \max(lower+tighten-d,0)+\max(d-upper+tighten,0), with tighten default 0. Hard flags compare d strictly to the ORIGINAL lower/upper bounds, without tighten.
• 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.
• Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
• Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

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 within_residue_violation(
    activations: dict,
    masks: dict,
    structure: dict,
    count_eps: float = 1e-06,
    distance_eps: float = 1e-10,
    tighten: float = 0.0,
) -> 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 within_residue_violation")

Related exercises

• Combined exercisePeptide geometry, steric violations and polymer clashes

Sources and further reading

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