Molecular graphs, equivariance and docking

Masked batched kNN graphs and neighbor gather

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

Select the nearest valid neighbors of each node within its batch item. Return exactly the requested number of slots with indices, masks, gathered features, and directed geometry.

Function signature

batched_neighbor_graph(
    xyz: torch.Tensor,
    features: torch.Tensor,
    valid: torch.Tensor,
    k: int,
) -> dict[str, torch.Tensor]

Worked example

Distance order, ties and padding in one row

Use one batch with xyz=[[0,0,0],[3,0,0],[1,0,0],[-1,0,0]], features=[[10],[20],[30],[40]], all nodes valid and k=4. For query 0, other-node squared distances are [9,1,1]. Sort by distance first, then smaller destination index: indices=[2,3,1,0], neighbor_valid=[True,True,True,False], squared_distances=[1,1,9,0], and neighbor_features=[[30],[40],[20],[0]]. Matching displacements are [[1,0,0],[-1,0,0],[3,0,0],[0,0,0]]. Node 1 appears after nodes 2 and 3 despite its smaller ID; the final zero index is padding, not a self edge.

Parameters

Parameter	Shape and meaning
`xyz`torch.Tensor	$[B, N, 3]$ Last axis Cartesian `xyz`.
`features`torch.Tensor	$[B, N, C]$ Neighbor `features` to gather.
`valid`torch.Tensor	$[B, N]$ `True` means real node.
`k`int	Required nonnegative number of neighbor slots per query. Keep this width even when k exceeds the number of available valid neighbors.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Given xyz $[B,N,3]$ , features $[B,N,C]$ and valid $[B,N]$ , with $B,N,C\ge 1$ , build directed neighbors from query $i$ to destination $j$ within the same batch element. Boolean True means a real node. Self edges $i=j$ are excluded, including when other distinct nodes coincide.
k is a required integer $K\ge 0$ ; return exactly $K$ slots per query, even when $K>N$ or fewer valid destinations exist. Select the smallest squared Euclidean distances, resolving equal distances by smaller destination index; torch.argsort(...,stable=True) is an allowed numerical exception.
Put valid neighbors first and pad every remaining slot with index $0$ , neighbor_valid=False, zero gathered features, zero squared distance, and zero displacement. A masked query has only padding.
Displacement is $\mathrm{xyz}[b,j]-\mathrm{xyz}[b,i]$ , not its negative; squared_distances is its squared xyz norm, with no sqrt or $\varepsilon$ . Gather features from the neighbor $j$ , never query $i$ or another batch.
Valid-node coordinates/features are finite; padded values may be arbitrary, including NaN/Inf, so mask before arithmetic/gather use.
Return indices $[B,N,K]$ int64, neighbor_valid $[B,N,K]$ bool, squared_distances $[B,N,K]$ , neighbor_features $[B,N,K,C]$ , and displacements $[B,N,K,3]$ , preserving floating dtype/device and strided inputs. The graph is not symmetrized. There is no radius, chain, bond, or hybrid-union variant in this launch contract.

Why this matters

Neighbor selection is discrete, but gathered features and coordinate differences must still preserve their valid gradient paths.

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 batched_neighbor_graph(
    xyz: torch.Tensor,  # (B, N, 3)
    features: torch.Tensor,  # (B, N, C)
    valid: torch.Tensor,  # (B, N)
    k: int,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B, N, K).
    neighbor_valid: (B, N, K).
    squared_distances: (B, N, K).
    neighbor_features: (B, N, K, C).
    displacements: (B, N, K, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement batched_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.

Sources and further reading

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