Molecular graphs, equivariance and docking

RFdiffusion Refinement Graph

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Connect each valid source to the union of its supplied KNN targets and its nearby sequence neighbors. Each fixed candidate slot represents source i → target j and, when accepted, stores receiver-first endpoints [j,i], the source feature, and displacement xyz[j] − xyz[i]. Overlapping neighbor choices produce one edge.

Function signature

rfdiffusion_refinement_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    kmin: int = 32,
) -> dict[str, torch.Tensor]

Worked example

Supplied neighbors add 0→2, 1→2, and 2→0. The repeated sequence index adds local edges 0→1 and 1→0. Receiver counts are [2, 1, 2]; edge_index stores receiver first, while displacement points from source to target.

import torch

activations = {
    'xyz': torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]], dtype=torch.float64, requires_grad=True),
    'features': torch.tensor([[10.0], [20.0], [30.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'node_valid': torch.tensor([True, True, True], dtype=torch.bool),
    'knn_valid': torch.tensor([[True], [True], [True]], dtype=torch.bool),
}
structure = {
    'batch': torch.tensor([0, 0, 0], dtype=torch.int64),
    'sequence_index': torch.tensor([2, 2, 7], dtype=torch.int64),
    'knn_index': torch.tensor([[2], [2], [0]], dtype=torch.int64),
}
kmin = 1

result = rfdiffusion_refinement_graph(activations=activations, masks=masks, structure=structure, kmin=kmin)

Expected return

{
    'edge_index': torch.tensor([[0, 1, 2, 0, 0, 2, 0, 0, 0], [0, 0, 0, 1, 0, 1, 2, 0, 0]], dtype=torch.int64),
    'image_index': torch.full((9,), 0, dtype=torch.int64),
    'edge_valid': torch.tensor([False, True, True, True, False, True, True, False, False], dtype=torch.bool),
    'displacements': torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [-3.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=torch.float64),
    'squared_distances': torch.tensor([0.0, 1.0, 9.0, 1.0, 0.0, 4.0, 9.0, 0.0, 0.0], dtype=torch.float64),
    'neighbor_features': torch.tensor([[0.0], [10.0], [10.0], [20.0], [0.0], [20.0], [30.0], [0.0], [0.0]], dtype=torch.float64),
    'neighbor_count': torch.tensor([2, 1, 2], dtype=torch.int64),
}

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`kmin`intDefault `32`	Minimum sequence-neighborhood width in the RF neighborhood rule.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

xyz is $[N,3]$ and features is $[N,F]$ . node_valid excludes padding; batch identifies independent structures. Sanitize invalid positions/features before calculations. Every retained edge joins valid nodes in the same batch.
Return fixed candidate slots with edge_valid marking retained edges; do not compact or sort them. edge_index is receiver-first. Invalid slots have zero indices, vectors, features, and squared distance. neighbor_count counts retained incoming edges per receiver.
Enumerate source-major, then target, over all N² ordered pairs. RFdiffusion uses source i → target j, so return edge_index=[j,i], displacement=xyz[j]−xyz[i], and features[i]. image_index is zero.
Retain the union of supplied valid KNN choices for each source and off-diagonal pairs with abs(sequence_index[i]−sequence_index[j]) < kmin. kmin defaults to 32 and must lie in [0,999.9). Duplicate KNN choices mark one candidate, not extra edges; a supplied valid KNN self-choice remains eligible.
There is no spatial cutoff, periodic image expansion, or bond prefix in this exercise.

Why this matters

A graph specifies which atoms or residues can exchange messages. The radius mode tests supplied periodic images; RF modes use supplied sequence and neighbor relationships. Their displacement conventions differ, so select the mode before assembling edge geometry.

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_refinement_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    kmin: int = 32,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    edge_index: (2, E).
    image_index: (E,).
    edge_valid: (E,).
    displacements: (E, 3).
    squared_distances: (E,).
    neighbor_features: (E, F).
    neighbor_count: (N,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rfdiffusion_refinement_graph")

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.