Polymer Clash Summary

Your task

Function signature

polymer_clash_summary(
    masks: dict,
    structure: dict,
    clash_radius: float = 1.1,
    min_clashes: int = 100,
    num_chains: int = 1,
    min_fraction: float = 0.5,
) -> dict[str, torch.Tensor]

Worked example

Count each clashing atom once

import torch

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),
}
clash_radius = 1.0
min_clashes = 1
num_chains = 3
min_fraction = 0.75

result = polymer_clash_summary(masks=masks, structure=structure, clash_radius=clash_radius, min_clashes=min_clashes, num_chains=num_chains, min_fraction=min_fraction)

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

• An eligible atom has atom_mask and polymer_mask both True. Compare atoms in different chains, or in the same chain with absolute residue_index separation >1; immediate neighbors and same-residue pairs are excluded.
• A clash uses the inclusive boundary squared_distance ≤ clash_radius² (default radius 1.1). No distance epsilon is added. Count an atom once if it has any eligible clashing partner, regardless of partner count.
• num_chains is a positive integer (default 1); output chain slots correspond to IDs 0…num_chains−1. chain_atom_counts counts eligible polymer atoms; fractions divide clash counts by total.clamp_min(1), so an empty chain has fraction 0.
• has_clash is True when ANY chain has count > min_clashes OR fraction > min_fraction. These comparisons are strict; defaults are100 and 0.5. min_clashes is a nonnegative integer; min_fraction lies in[0,1].
• Return exact integer counts and boolean flags alongside the floating fractions; this thresholding task does not require differentiating the discrete decisions.
• 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 polymer_clash_summary(
    masks: dict,
    structure: dict,
    clash_radius: float = 1.1,
    min_clashes: int = 100,
    num_chains: int = 1,
    min_fraction: float = 0.5,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    atom_clash_mask: (B, A).
    chain_clash_counts: (B, C).
    chain_atom_counts: (B, C).
    chain_clash_fraction: (B, C).
    has_clash: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement polymer_clash_summary")

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