Structure module, templates and recycling

Learned torsions and hierarchical atom reconstruction

ExpertRaw PyTorchAbout 45 min1 public example plus private tests

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

Your task

Predict torsion angle pairs from supplied residue features, normalize them, and compose a hierarchy of rigid-group frames. Use those frames to reconstruct the requested atom coordinates.

Function signature

torsion_frame_reconstruction(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

Place an atom through a torsion frame

For the reconstruction stage, suppose the angle head produced one valid raw torsion pair [3,0] with default eps=1e-8. Normalization gives [sin,cos]=[1,0], a positive quarter-turn about local x. Use two groups with parent=[-1,0]: group 0 has identity default rotation/zero translation; group 1 has identity default rotation and local translation (1,0,0). The backbone rotates by a positive quarter-turn about global z and translates by (10,0,0). A present atom assigned to group 1 at local position (0,1,0) first rotates to (0,0,1), then shifts to (1,0,1) in the backbone frame. The backbone maps this to world coordinate (10,1,1). This calculation starts after the learned head; the full task also computes that raw torsion pair.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`eps`floatDefault `1e-08`	Positive squared-length threshold for sine/cosine angle pairs. Normalize by $\sqrt{\max(\lVert v\rVert^2,\varepsilon)}$ ; use the identity-angle fallback when $\lVert v\rVert^2\leq\varepsilon$ .

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Implement a ReLU residual angle head and topologically ordered torsion-frame reconstruction. All dimensions $B,N,C,H,K,L,A,P$ are positive. There are $K$ torsions and $G=K+1$ rigid groups, and A residue types. Every floating input is finite and shares dtype/device.
residue_type[ $B,N$ ] is in $[0,A)$ , residue_mask[ $B,N$ ] $\mathrm{True}=\mathrm{usable}$ . $\operatorname{ReLU}(t)=\operatorname{max}(t,0)$ . First replace masked single/initial_single[ $B,N,C$ ] by zero.
$h=\operatorname{ReLU}(\mathrm{single})\mathbin{@}w_{\mathrm{single}}+b_{\mathrm{single}}+\operatorname{ReLU}(\mathrm{initial\_single})\mathbin{@}w_{\mathrm{initial}}+b_{\mathrm{initial}}$ .
For $\mathrm{ell}=0$ .. $L-1$ , set $h=h+(\operatorname{ReLU}(\operatorname{ReLU}(h)\mathbin{@}w_{\mathrm{res1}}[\mathrm{ell}]+b_{\mathrm{res1}}[\mathrm{ell}])\mathbin{@}w_{\mathrm{res2}}[\mathrm{ell}]+b_{\mathrm{res2}}[\mathrm{ell}])$ .
$\mathrm{raw\_angles}=\operatorname{reshape}$ ( $\operatorname{ReLU}(h)\mathbin{@}w_{\mathrm{angles}}+b_{\mathrm{angles}}$ ,[ $B,N,K$ , $2$ ]), ordered $[\operatorname{sin},\operatorname{cos}]$ . Mask raw_angles to zero on invalid residues.
With $\varepsilon>0$ default $10^{-8}$ , normalize each pair by $\sqrt{\max(\sum r_{\mathrm{raw}}^2,\varepsilon)}$ ; if $\sum r_{\mathrm{raw}}^2\le\varepsilon$ return $[0,1]$ exactly instead.
This explicit degenerate-angle identity fallback differs from an upstream zero-angle vector and has zero normalized-angle derivative inside the fallback region. Return both raw and normalized angles. Prepend [ $\operatorname{sin}=0$ , $\operatorname{cos}=1$ ] for rigid group0.
Each group's $x$ -axis rotation is $\mathrm{Rx}=[[1,0,0],[0,\operatorname{cos},-\operatorname{sin}],[0,\operatorname{sin},\operatorname{cos}]]$ .
Gather default_rotation[ $A,G$ , $3$ , $3$ ] and default_translation[ $A,G$ , $3$ ] by residue type, then local group transform is ( $\mathrm{default\_rotation}\mathbin{@}\mathrm{Rx}$ ,default_translation): torsion acts FIRST in that local group.
parent $[G]$ is int64: $-1$ means the supplied backbone transform, otherwise $0\le \mathrm{parent}[g]<g$ . Group0 must have $\mathrm{parent}=-1$ , identity default rotation and zero default translation; its extra $x$ rotation is identity.
Canonical AF2 has $K=7$ , $G=8$ and $\mathrm{parent}=[-1,0,0,0,0,4,5,6]$ , representing independent $\mathrm{backbone}/\mathrm{chi1}$ frames and the $\mathrm{chi2}\to \mathrm{chi3}\to \mathrm{chi4}$ chain. More compact topological examples use the same algorithm.
Using column-vector actions $R\cdot x+t$ , compose $\mathrm{world\_R}[g]=\mathrm{parent\_R}\mathbin{@}\mathrm{local\_R}[g]$ , $\mathrm{world\_t}[g]=\mathrm{parent\_t}+\mathrm{parent\_R}\mathbin{@}\mathrm{local\_t}[g]$ .
For each atom a gather atom_group[ $A,P$ ] in $[0,G)$ , atom_positions[ $A,P$ , $3$ ] in that group's local frame, and atom_exists[ $A,P$ ]. Its world coordinate is $\mathrm{world\_R}[\mathrm{group}]\mathbin{@}\mathrm{local\_atom}+\mathrm{world\_t}[\mathrm{group}]$ .
Atom validity is table existence AND residue_mask; invalid atoms are exact zero. Return all group frames, replacing invalid-residue frames by identity rotations and zero translations. Table transforms and backbone rotations are proper rigid transforms on forward inputs.
Ideal-geometry tables, residue/atom assignment and arbitrary absent-atom padding are supplied; no residue constants, nn module, pretrained head, relaxation or atom-layout library may be imported. No global rotation acts on angle features.

Why this matters

A protein residue can be represented as small rigid groups connected by rotatable bonds. The angle head predicts sine/cosine pairs; a parent tree turns those local rotations into global atom positions.

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 torsion_frame_reconstruction(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    raw_angles: (B, N, K, 2).
    angles: (B, N, K, 2).
    frame_rotation: (B, N, K+1, 3, 3).
    frame_translation: (B, N, K+1, 3).
    atom_positions: (B, N, P, 3).
    atom_mask: (B, N, P)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement torsion_frame_reconstruction")

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.