Molecular graphs, equivariance and docking

Origin-safe Bessel channels with a compact polynomial cutoff

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Encode each edge distance as a row of frequency-dependent radial features that smoothly vanish at its cutoff. Each basis[e,k] multiplies the supplied scale, a sinc channel at frequency k, and the polynomial envelope for distance[e] / cutoff[e]. Return both the feature matrix and each edge’s envelope value.

Function signature

bessel_radial_basis(
    distances: torch.Tensor,
    cutoffs: torch.Tensor,
    frequencies: torch.Tensor,
    scales: torch.Tensor,
    valid: torch.Tensor,
    p: int = 6,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

A distance inside the cutoff

For one valid edge, take distance=0.5, cutoff=1, scale=1, one frequency 1, and p=2. Then x=0.5 and the envelope is $1-6x^2+8x^3-3x^4=5/16$ . Since $\operatorname{sinc}(\pi/2)=2/\pi$ , the single basis value is $\frac{5}{8\pi}\approx0.198944$ . At this same cutoff, distance 1 gives exactly zero. The values here are the entries of basis shape (1,1) and envelope shape (1,).

Parameters

Parameter	Shape and meaning
`distances`torch.Tensor	$[E]$ Nonnegative finite at `valid` edges
`cutoffs`torch.Tensor	$[E]$ Positive finite at `valid` edges
`frequencies`torch.Tensor	$[K]$ Finite `frequencies`; $K\ge 1$
`scales`torch.Tensor	$[E]$ Finite caller-supplied normalization at `valid` edges
`valid`torch.Tensor	$[E]$ `True` means observed edge
`p`intDefault `6`	Config $p$ is an integer $\ge 2$ , default $6$ .

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

For valid edge $e$ , $x=r_{e}/R_{e}$ and frequency nu_k, return $b[e,k]=\mathrm{scale}[e]\cdot \nu_{k}\cdot \operatorname{sinc}(\pi\cdot \nu_{k}\cdot x)\cdot c_{p}(x)$ , where $\operatorname{sinc}(z)=\operatorname{sin}(z)/z$ and $\operatorname{sinc}(0)=1$ .
For $0\le x<1$ , $c_{p}(x)=1-((p+1)(p+2)/2)\cdot x^{p}+p\cdot (p+2)\cdot x^{p+1}-(p\cdot (p+1)/2)\cdot x^{p+2}$ ; $c_{p}(x)=0$ for $x\ge 1$ . Config $p$ is an integer $\ge 2$ , default $6$ . NequIP's raw channel is $\nu$ *torch.sinc(nu*x); its global edge-length scale is supplied separately, e.g.
$\mathrm{scale}=2\cdot \pi/r_{\mathrm{max}}^{2}$ , and must not be replaced with $2\cdot \pi/R_{e}^{2}$ for type-dependent cutoffs. Return envelope c_p as well. Frequencies are supplied real learned parameters, including zero; integer initialization is common but not required.
Evaluate the analytic origin limit without NaN gradients, and keep the cutoff stable close to $x=1$ . Mask invalid distances/cutoffs/scales before division, including nonfinite garbage. Invalid rows and all rows at or beyond cutoff are exactly zero. $E=0$ is supported.

Why this matters

A sinc-like radial oscillation carries distance information, while a compact polynomial envelope smoothly limits its spatial range.

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 bessel_radial_basis(
    distances: torch.Tensor,  # (E,)
    cutoffs: torch.Tensor,  # (E,)
    frequencies: torch.Tensor,  # (K,)
    scales: torch.Tensor,  # (E,)
    valid: torch.Tensor,  # (E,)
    p: int = 6,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    basis: (E, K).
    envelope: (E,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement bessel_radial_basis")

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.