Molecular graphs, equivariance and docking · INTEGRATION

Typed radius, bond-prefix and sequence-neighbor graph assembly

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

Build a directed graph while retaining a slot for every candidate edge. An accepted slot stores its receiver and sender, displacement, and the sender’s feature; rejected slots are zeroed and marked by edge_valid. Radius mode tests supplied periodic sender images against a strict distance cutoff. RFdiffusion modes use sequence/KNN rules and the separately specified displacement sign.

Function signature

radius_neighbor_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    neighborhood: str = 'radius',
    bond_prefix: bool = False,
    kmin: int = 32,
) -> dict[str, torch.Tensor]

Worked example

One periodic candidate

In radius mode, let receiver 0 be at [0,0,0], sender 1 at [1,0,0], and sender feature be [7,2]. They are valid and in the same batch. With cell rows [[3,0,0],[0,3,0],[0,0,3]], supplied image shift [-1,0,0], and directed cutoff 2.5, the sender image is [-2,0,0]. This candidate has displacement [-2,0,0], squared distance 4 < 6.25, endpoints [0,1], and neighbor feature [7,2]. It contributes one incoming neighbor to receiver 0; it stays in its original candidate slot.

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.
`neighborhood`strDefault `'radius'`	Select periodic radius candidates, RF refinement kNN/sequence union, or RF full nonzero-sequence-separation pairs.
`bond_prefix`boolDefault `False`	Prepend supplied bond candidates when True.
`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

Config neighborhood='radius' (default): enumerate $E=N\cdot N\cdot S$ candidates in lexicographic order $(\mathrm{receiver} i, \mathrm{sender} j, \mathrm{image} s)$ , last index fastest.
For receiver batch $b$ , shift $u=\mathrm{image\_shifts}[b,s]$ and row-vector lattice cell $[b]$ , compute $r=\mathrm{xyz}[j]+u\mathbin{@}\mathrm{cell}[b]-\mathrm{xyz}[i]$ .
Keep exactly observed $i$ and $j$ in the same batch, an observed image, not ( $i=j$ and $u=0$ ), and $\|r\|^2<\texttt{cutoffs}[b,\texttt{types}[i],\texttt{types}[j]]^2$ . Nonzero self images are allowed. A duplicate supplied image remains a duplicate edge. Directed type cutoffs need not be symmetric.
Nonnegative cutoffs are allowed: cutoff zero disables that directed radius relation. Return fixed-width masked edge arrays, not a data-dependent compact list: every rejected candidate has indices/image_index $0$ and zero geometry/features. neighbor_count counts accepted candidates per receiver. Config bond_prefix=False by default.
If True, prepend $L$ supplied directed bond_edge_index columns to radius candidates, in their supplied order. Accept a bond independently of radius only if bond_edge_valid, both nodes are valid in the same batch and endpoints differ.
Its displacement is $\mathrm{xyz}[\mathrm{sender}]-\mathrm{xyz}[\mathrm{receiver}]$ without an image shift, and accepted $\mathrm{image\_index}=-1$ . False bonds have all-zero outputs. $E=L+N\cdot N\cdot S$ in this mode.
Do not deduplicate: duplicate supplied bonds and separate bond/radius copies remain separate messages, as in DiffDock's concatenation. $L=0$ is required if $N=0$ . Invalid nodes may carry nonfinite xyz/features and sentinel batch/type indices; invalid images may carry arbitrary shifts, all ignored before arithmetic/gather.
Cell/cutoffs are finite globally. $N=0$ or $S=0$ gives $E=0$ . Selection is exact on the supplied finite list; no nearest-image approximation, automatic image generation, completeness claim or sorting is implied.
Config neighborhood='rfd_refinement' selects a different nonperiodic graph: each source row $i$ supplies knn_index $[i,:]$ with knn_valid. Form a Boolean adjacency from those choices, then OR it with $i\ne j$ and $\left|\mathrm{sequence\_index}[j]-\mathrm{sequence\_index}[i]\right|<\mathrm{kmin}$ .
Config $\mathrm{kmin}=32$ by default, $0\le \mathrm{kmin}<999.9$ . The inequality is STRICT, matching the released code rather than $\mathrm{its} \le \mathrm{comment}$ . Both nodes must be valid and share batch. Duplicate kNN entries and overlaps with sequence neighbors produce ONE pair.
An explicitly supplied self kNN choice is retained (the upstream $\mathrm{top\_k}\ge L$ branch can include self); the sequence term alone never introduces self. False knn_valid entries and all choices from invalid source rows may have sentinel indices and are ignored before scatter.
Config neighborhood='rfd_full' ignores kNN and selects same-batch valid pairs with nonzero sequence-index difference, matching make_full_graph.
RF modes enumerate $N\cdot N (\mathrm{source} i,\mathrm{target} j)$ candidates and return $\mathrm{edge\_index}=[j,i]$ in the common $\mathrm{RECEIVER},\mathrm{SENDER}$ order, $\mathrm{image\_index}=0$ , sender features features $[i]$ , and incoming receiver counts.
Their displacement is $\mathrm{xyz}[j]-\mathrm{xyz}[i]$ , RECEIVER minus SENDER, matching RFdiffusion DGL rel_pos; this sign differs from the radius mode. Pair attributes from upstream pair $[i,j]$ therefore gather as pair $[\mathrm{sender},\mathrm{receiver}]$ .
RF modes ignore cell/types/radius/image tensors and reject bond_prefix=True. kNN selection is supplied, not recomputed here: the existing batched_neighbor_graph covers unperturbed distance kNN.
Upstream's diagonal $999.9$ penalty and $\mathrm{distance}+\varepsilon\cdot \mathrm{sequence\_separation}$ ranking must be reflected in supplied choices for exact upstream selection; that $\epsilon$ ranking is not silently attributed to the old distance-only exercise.
No coordinate detach is allowed here, unlike upstream's detached basis geometry.

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 radius_neighbor_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    neighborhood: str = 'radius',
    bond_prefix: bool = False,
    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 radius_neighbor_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.