Template geometry and pair features

Your task

Function signature

template_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'af2',
    eps: float = 1e-08,
    type_count: int = 22,
    norm_eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]

Worked example

An off-diagonal direction in a rotated frame

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Construct template pair features and their embedding. Positive B,T,N,A,\mathrm{Ct}; T\ge 1, so no-template input is represented by all-False masks. atom_positions[B,T,N,4,3] is ordered N,\mathrm{CA},C,\mathrm{CB}, with True atom_mask meaning present. All supplied floats are finite, including masked entries.
• residue_type[B,T,N] is in [0,A); A=\mathrm{type\_count} (default 22). pair_mask[B,N,N] allows a directed pair, e.g. an intrachain restriction. distance_edges [K] has K\ge 1 positive strictly increasing length cutoffs.
• An exhaustive distance bin is k=\sum(\mathrm{distance\_squared} > \mathrm{edge\_squared}): equality stays in the lower bin. This deliberate complete-bucket convention is not AF2's strict finite-interval distogram. \varepsilon>0 (default 10^{-8}) is a squared-norm threshold.
• mode='af2' (default) or 'af3': pseudo_beta is CA for \mathrm{residue\_type}=\mathrm{glycine\_index} (default 7), otherwise CB. Missing pseudo-\beta coordinates return zero.
• Build x=\operatorname{normalize}(C-\mathrm{CA}), y=\operatorname{normalize}((N-\mathrm{CA})-\operatorname{dot}(N-\mathrm{CA},x)x), z=\operatorname{cross}(x,y), with squared norms clamped below by \varepsilon before division.
• A backbone frame is valid only if N/\mathrm{CA}/C exist and both tested squared norms exceed \varepsilon.
• For pair i,j, u=\frac{R_i^\top(\mathrm{CA}_j-\mathrm{CA}_i)}{\sqrt{\|R_i^\top(\mathrm{CA}_j-\mathrm{CA}_i)\|^2+\varepsilon}}, where R columns are x,y,z. Its diagonal is zero. pair_valid requires both frames, both pseudo-betas and pair_mask.
• Concatenate in this exact order: [distance one-hot K+1, pair_valid1, residue-j one-hot A, residue-i one-hot A, u3, pair_valid1], giving F=K+2A+6. Zero every feature of invalid pairs. \mathrm{embedding}=\mathrm{features}\mathbin{@}w_{\mathrm{features}}+b_{\mathrm{features}}, also zeroed after bias for invalid pairs.
• mode='af3' additionally layer-normalizes query[B,N,N,\mathrm{Cq}] over C_q with population variance and \sqrt{\mathrm{var}+\mathrm{norm\_eps}}, \mathrm{norm\_eps}>0 default 10^{-5}, then applies query_scale/query_bias and bias-free w_query. Add this projected query to every template embedding before final masking.
• This is the protein-template, query-conditioned feature construction; independent input projections in AF3 can be represented by blocks of the concatenated projection.
• mode='rfdiffusion': ignore CB and reconstruct \beta=\mathrm{CA}-0.58273431\cdot \operatorname{cross}(CA-N,C-CA)+0.56802827\cdot (\mathrm{CA}-N)-0.54067466\cdot (C-\mathrm{CA}). Require N/\mathrm{CA}/C at both ends and pair_mask.
• Nonself pairs with \beta distance STRICTLY below the final distance edge have orientations \omega=\operatorname{dih}(\mathrm{CA\_i},\beta_{i},\beta_{j},\mathrm{CA\_j}), \theta=\operatorname{dih}(N_{i},\mathrm{CA\_i},\beta_{i},\beta_{j}), \phi=\operatorname{angle}(\mathrm{CA\_i},\beta_{i},\beta_{j}).
• For \operatorname{dih}(a,b,c,d), let \mathrm{axis}=\operatorname{normalize}(c-b), v=(a-b)-\bigl((a-b)\cdot\mathrm{axis}\bigr)\mathrm{axis}, w=(d-c)-\bigl((d-c)\cdot\mathrm{axis}\bigr)\mathrm{axis}; \operatorname{angle}=\operatorname{atan2}(\operatorname{dot}(\operatorname{cross}(\mathrm{axis},v),w),\operatorname{dot}(v,w)).
• If axis, v or w has squared \operatorname{norm}\le \varepsilon, all six orientation channels are zero. Otherwise concatenate [\sin\omega,\sin\theta,\sin\phi,\cos\omega,\cos\theta,\cos\phi]. Diagonal/long-range orientations are zero and distance one-hot is the final (overflow) bin.
• Concatenate [distance one-hot K+1, orientations6, pair_valid1, type_j A,type_i A], F=K+2A+8; mask and project as above.
• The \phi sine uses \sqrt{\operatorname{max}(\Vert \operatorname{cross}(\operatorname{unit}(\mathrm{CA\_i}-\beta_{i}),\operatorname{unit}(\beta_{j}-\beta_{i}))\Vert ^{2},\varepsilon)}; smooth fixtures stay away from this floor. Unlike upstream sentinel/NaN conventions, absent geometry has exact-zero features.
• Both alphafold modes require proper nondegenerate frame geometry when marked valid; all modes exclude exact bin/cutoff boundaries from differentiability claims. No template pair stack is hidden inside this function.

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 template_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'af2',
    eps: float = 1e-08,
    type_count: int = 22,
    norm_eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pseudo_beta: (B, T, N, 3).
    pair_mask: (B, T, N, N).
    features: (B, T, N, N, F).
    embedding: (B, T, N, N, Ct)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement template_pair_features")

Related exercises

• Builds onAlphaFold 2 Template Pair Features
• Builds onAlphaFold 3 Template Pair Features
• Builds onRFdiffusion Template Pair Features

Sources and further reading

• www.nature.com
• www.nature.com
• www.nature.com
• google-deepmind/alphafold · modules.py
• google-deepmind/alphafold3 · template_modules.py
• RosettaCommons/RFdiffusion · kinematics.py