RFdiffusion Full Graph

Your task

Function signature

rfdiffusion_full_graph(
    activations: dict,
    masks: dict,
    structure: dict,
) -> dict[str, torch.Tensor]

Worked example

Worked example

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),
}
structure = {
    'batch': torch.tensor([0, 0, 0], dtype=torch.int64),
    'sequence_index': torch.tensor([2, 2, 7], dtype=torch.int64),
}

result = rfdiffusion_full_graph(activations=activations, masks=masks, structure=structure)

Expected return

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

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• xyz is [N,3] and features is [N,C]. 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.
• Keep pairs with unequal sequence_index values. Distinct nodes with identical sequence indices are excluded.
• There is no spatial cutoff, periodic image expansion, or bond prefix in this exercise.
• 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 rfdiffusion_full_graph(
    activations: dict,
    masks: dict,
    structure: dict,
) -> 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_full_graph")

Related exercises

• Combined exerciseTyped radius, bond-prefix and sequence-neighbor graph assembly

Sources and further reading

• www.nature.com
• mir-group/nequip · _edge.py
• arxiv.org
• gcorso/DiffDock · score_model.py
• gcorso/DiffDock · all_atom_score_model.py
• gcorso/DiffDock · cg_model.py
• www.nature.com
• RosettaCommons/RFdiffusion · util_module.py