Structure module, templates and recycling

RFdiffusion Template Pair Features

ExpertRaw PyTorchAbout 45 min2 public examples plus private tests

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

Your task

Reconstruct a virtual Cβ from each template backbone and describe each ordered residue pair by its distance, orientation angles and residue identities. Encode close-pair orientations as sine/cosine channels, then project the exact feature layout into a template embedding.

Function signature

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

Worked example

One translated backbone pair

The second backbone is the first translated by [0, 0, 2]. The first virtual C-beta is approximately [-0.5407, -0.5680, -0.5827]; the second adds the same translation. Their distance is $2$ , so edges [1, 3] give histogram [0, 1, 0]. The omega orientation has sine $0$ and cosine $1$ . Only pair (0, 1) is enabled; selecting its middle histogram channel gives embedding $1$ , while other pair embeddings are zero.

import torch

activations = {
    'distance_edges': torch.tensor([1.0, 3.0], dtype=torch.float64),
}
parameters = {
    'w_features': torch.tensor([[0.0], [1.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'b_features': torch.tensor([0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'atom_mask': torch.full((1, 1, 2, 4), True, dtype=torch.bool),
    'pair_mask': torch.tensor([[[False, True], [False, False]]], dtype=torch.bool),
}
structure = {
    'atom_positions': torch.tensor([[[[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 1.0, 2.0], [0.0, 0.0, 2.0], [1.0, 0.0, 2.0], [0.0, 0.0, 2.0]]]]], dtype=torch.float64, requires_grad=True),
    'residue_type': torch.tensor([[[0, 0]]], dtype=torch.int64),
}
eps = 1e-08
type_count = 1

result = rfdiffusion_template_pair_features(activations=activations, parameters=parameters, masks=masks, structure=structure, eps=eps, type_count=type_count)

Expected return

{
    'pseudo_beta': torch.tensor([[[[-0.54067466, -0.56802827, -0.58273431], [-0.54067466, -0.56802827, 1.41726569]]]], dtype=torch.float64),
    'pair_mask': torch.tensor([[[[False, True], [False, False]]]], dtype=torch.bool),
    'features': torch.tensor([[[[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, -0.8473820602168685, 0.8026564473784078, 1.0, -0.5309836570202663, 0.5964416379511691, 1.0, 1.0, 1.0]], [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]]]], dtype=torch.float64),
    'embedding': torch.tensor([[[[[0.0], [1.0]], [[0.0], [0.0]]]]], dtype=torch.float64),
}

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 geometric squared-norm floor for frame/angle construction and an additive squared-distance stabilizer for local unit directions; use the distinct placements in the selected feature rule.
`type_count`intDefault `22`	Size of the residue-type vocabulary.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

atom_positions is $[B,T,N,4,3]$ , with atom slots N, CA, C, CB. atom_mask marks available atoms. residue_type uses type_count categories (default 22); one-hot feature order follows these supplied integer IDs.
Construct virtual Cβ as $CA-0.58273431[(CA-N)\times(C-CA)]+0.56802827(CA-N)-0.54067466(C-CA)$ , requiring N/CA/C.
A candidate pair requires both complete N/CA/C backbones, both valid pseudo-beta atoms, and pair_mask. Squared-distance bin index counts distance_edges² strictly below squared pseudo-beta distance.
A close pair additionally has distance² < last_edge² and i≠j. Nonclose pairs use the final distance bin. Build omega=dihedral(CA_i,CB_i,CB_j,CA_j), theta=dihedral(N_i,CA_i,CB_i,CB_j), and phi=angle(CA_i−CB_i,CB_j−CB_i).
Unit vectors divide by sqrt(max(norm²,eps)), with a valid norm² strictly greater than eps (default 1e-8). A dihedral removes the central-axis component of its endpoint vectors; sin is dot(cross(axis,first_unit),last_unit), cos is their dot product. For phi, cos is the unit-vector dot and sin is sqrt(max(squared_cross_norm,eps)).
Orientation features are zero unless the pair is close and every needed unit vector is valid. Concatenate: distance one-hot, sin(omega), sin(theta), sin(phi), cos(omega), cos(theta), cos(phi), pair-valid indicator, right residue one-hot, left residue one-hot.
Zero invalid feature rows, then project with w_features and b_features. Return pseudo_beta, the final pair_mask, raw concatenated features, and learned embedding; zero invalid embeddings after every bias.
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

Templates contribute distances, local directions and residue identities. Missing coordinates determine whether a pair is usable; the selected AF2, AF3 or RFdiffusion mode determines which geometric features are included.

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 rfdiffusion_template_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
    type_count: int = 22,
) -> 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,K+2V+8).
    embedding: (B, T, N, N, Ct)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rfdiffusion_template_pair_features")

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.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.