Molecular graphs, equivariance and docking

Distance basis expansion with smooth finite support

EasyRaw PyTorchAbout 10 min1 public example plus private tests

SchNetPaiNNTorchMD-Net

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

Your task

Encode each valid distance with Gaussian radial basis functions multiplied by a cosine cutoff. Return both the basis features and their cutoff envelope.

Function signature

radial_basis_cutoff(
    distances: torch.Tensor,
    centers: torch.Tensor,
    widths: torch.Tensor,
    valid: torch.Tensor,
    cutoff: float,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Inside and beyond the cutoff

With cutoff 1, the envelope is 1 at distance 0 and 1/2 at distance 1/2. At distance 1/2, the channel centered there is exactly 1/2; the channel centered at 0 is exp(-1/8)/2. At or beyond the cutoff, every channel is zero.

import torch

distances = torch.tensor([0.0, 0.5, 1.0, 1.5], dtype=torch.float64, requires_grad=True)
centers = torch.tensor([0.0, 0.5], dtype=torch.float64, requires_grad=True)
widths = torch.tensor([1.0, 2.0], dtype=torch.float64, requires_grad=True)
valid = torch.tensor([True, True, True, True], dtype=torch.bool)
cutoff = 1.0

result = radial_basis_cutoff(distances=distances, centers=centers, widths=widths, valid=valid, cutoff=cutoff)

Expected return

(torch.tensor([[1.0, 0.9692332344763441], [0.4412484512922977, 0.5], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float64), torch.tensor([1.0, 0.5, 0.0, 0.0], dtype=torch.float64))

Parameters

Parameter	Shape and meaning
`distances`torch.Tensor	$[E]$ Nonnegative for `valid` edges.
`centers`torch.Tensor	$[K]$ Supplied `centers`, graded as learnable parameters.
`widths`torch.Tensor	$[K]$ Strictly positive supplied `widths`.
`valid`torch.Tensor	$[E]$ `True` means real edge.
`cutoff`float	Required positive radius R. Use 0.5(1+cos(pir/R)) only for valid r<R; equality and larger distances return zero.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Inputs distances $[E]$ , centers $[K]$ , positive widths $[K]$ , valid $[E]$ use $E\ge 0$ and $K\ge 1$ . Boolean True means a real edge; real distances are finite and nonnegative, centers finite, widths strictly positive. cutoff is a required positive scalar $R$ .
For valid $r<R$ , return envelope $c(r)=0.5\cdot (1+\operatorname{cos}(\pi\cdot r/R))$ and $\mathrm{features}[e,k]=\operatorname{exp}(-0.5\cdot ((r-\mathrm{centers}[k])/\mathrm{widths}[k])^{2})\cdot c(r)$ . At $r=R$ , beyond $R$ , or for invalid edges, both outputs are exactly zero.
The cutoff predicate is strictly $<$ ; at equality define the first derivative as zero, matching the continuous one-sided limits. Widths are used directly: do not square them twice, transform them, clamp them or insert $\varepsilon$ . Mask invalid/outside-cutoff distances to a safe value before arithmetic; padded $r$ may be NaN/Inf.
Return features $[E,K]$ and envelope $[E]$ , preserving dtype/device and accepting strided inputs; an empty edge axis returns empty outputs. There is no sine/Bessel basis or learned cutoff in this launch variant.

Why this matters

Gaussian channels describe distance ranges, while a cosine envelope turns every channel off smoothly at a shared cutoff radius.

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 radial_basis_cutoff(
    distances: torch.Tensor,  # (E,)
    centers: torch.Tensor,  # (K,)
    widths: torch.Tensor,  # (K,)
    valid: torch.Tensor,  # (E,)
    cutoff: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    features: (E, K).
    envelope: (E,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement radial_basis_cutoff")

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.