Molecular graphs, equivariance and docking

Periodic Radius Graph

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

For every receiver, sender, and supplied periodic image, decide whether the sender image lies strictly inside the receiver’s type-dependent radius. Keep a fixed slot for every candidate: each accepted slot stores the sender feature and sender-image-minus-receiver displacement, and edge_valid distinguishes real edges from zeroed rejected slots.

Function signature

periodic_radius_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    bond_prefix: bool = False,
) -> dict[str, torch.Tensor]

Worked example

The nearest periodic image can cross the cell boundary

Two valid nodes in one graph lie at x=[0,1.5], with types [0,1] and features [10,20]. Use a cubic cell of width 2, supplied x-image shifts [0,-1,+1], all images valid, cutoffs [[1,1],[0.75,1]], and no bond prefix. Candidate order is receiver, sender, then image.

Receiver 0 sees sender 1's negative image at $1.5-2=-0.5$ . Its slot is $(0\cdot2+1)\cdot3+1=4$ : displacement [-0.5,0,0], squared distance 0.25 < 1, sender feature 20. Receiver 1 sees sender 0's positive image at slot 8: displacement [0.5,0,0], squared distance $0.25<0.75^2$ , feature 10. All other slots fail the self-edge or cutoff rule. The 12-slot mask is true only at 4 and 8, with incoming counts [1,1].

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.
`bond_prefix`boolDefault `False`	Prepend supplied bond candidates when True.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

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 candidates receiver-major, then sender, then supplied image index. Translation is image_shift @ cell using lattice vectors as cell ROWS. Displacement is sender + translation − receiver.
Use the receiver/sender type-specific cutoff for that batch and retain squared_distance < cutoff² strictly. Remove only a same-node edge at zero image shift; periodic self-images can remain. An image_valid=False slot is always excluded.
bond_prefix defaults to False. When enabled, prepend supplied valid directed bond edges in their given order, using image_index=−1 and unshifted sender−receiver displacement. Keep same-batch distinct valid endpoints without a distance test; duplicates with radius edges remain separate.
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

A periodic radius graph tests each supplied lattice image as a separate candidate. The fixed output layout makes image choice and masking explicit; optional bond candidates are prepended without deduplication.

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 periodic_radius_graph(
    activations: dict,
    masks: dict,
    structure: dict,
    bond_prefix: bool = False,
) -> 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 periodic_radius_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.