Molecular graphs, equivariance and docking

Real spherical harmonics from Cartesian polynomial tables

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

Convert each 3D vector into the real angular features defined by supplied polynomial tables. With normalize=True, one output block describes direction independently of vector length; otherwise it also retains the polynomial’s degree-dependent magnitude. Return each degree’s components in their supplied order, with padding zero.

Function signature

real_spherical_harmonics(
    vectors: torch.Tensor,
    valid: torch.Tensor,
    degrees: torch.Tensor,
    coefficients: torch.Tensor,
    powers: torch.Tensor,
    term_valid: torch.Tensor,
    normalization: str = 'component',
    normalize: bool = True,
) -> torch.Tensor

Worked example

Normalize before evaluating a degree-1 block

Take one valid vector [3,4,0], normalize=True, normalization="component", degree 1, and the supplied table $H(v)=\sqrt{3}[v_x,v_y,v_z]$ with no padding. The length is 5, so evaluate at [3/5,4/5,0]. The output block is $[3\sqrt{3}/5,4\sqrt{3}/5,0]\approx[1.039230,1.385641,0]$ . With normalization="norm", that same block is divided by $\sqrt{3}$ and becomes [0.6,0.8,0].

Parameters

Parameter	Shape and meaning
`vectors`torch.Tensor	$[E, 3]$ Cartesian `vectors`, finite at `valid` rows
`valid`torch.Tensor	$[E]$ True marks a usable observation; False marks padding.
`degrees`torch.Tensor	$[Q]$ `int64` $[Q]$ , nonnegative `degrees`; $Q\ge 1$
`coefficients`torch.Tensor	$[Q, D, T]$ Component-normalized fixed polynomial tables
`powers`torch.Tensor	$[Q, T, 3]$ `int64` [ $Q,T$ , $3$ ], nonnegative Cartesian exponents for `valid` terms
`term_valid`torch.Tensor	$[Q, T]$ bool [ $Q,T$ ], table padding mask
`normalization`strDefault `'component'`	"component" uses the supplied polynomial normalization; "norm" divides each degree- $l$ block by $\sqrt{2l+1}$ ; "integral" divides by $\sqrt{4\pi}$ .
`normalize`boolDefault `True`	Normalize input vectors to directions before evaluating harmonics when True.

Return

torch.Tensor · $[E, Q, D]$

Return float [E,Q,D], valid first 2*l+1 components and zero padding.

Required behavior

For each supplied harmonic degree l_q and component $m$ , evaluate the fixed homogeneous harmonic polynomial $H_{\mathrm{qm}}(v)=\sum_{t} \mathrm{coefficients}[q,m,t]\cdot \mathrm{product\_a} v[a]^{\mathrm{powers}}[q,t,a]$ .
The supplied component-normalized tables satisfy $\sum_m H_{lm}(v)^2=(2l+1)\|v\|^{2l}$ ; every monomial has degree l. With normalize=True, evaluate at v/||v|| for nonzero vectors; otherwise evaluate at v directly.
At $v=0$ use $u=0$ , so degree $> 0$ returns zero and $l=0$ is its constant, with no normalized-origin derivative promised. Config normalization='component' default leaves $H$ unchanged; 'norm' divides each degree by $\sqrt{2\cdot l+1}$ ; 'integral' divides by $\sqrt{4\cdot \pi}$ .
Ignore masked rows, invalid terms and components $m\ge 2\cdot l+1$ before arithmetic, including NaN coefficients in padding. Powers on invalid terms are ignored. All outputs are zero for invalid rows. $E=0$ is supported. Tables are supplied constants, but polynomial evaluation and coordinate normalization are the learner's work.

Why this matters

Spherical harmonics encode direction in components that transform together under rotation. The tables fix the real basis; your job is polynomial evaluation and normalization, not generating a different basis.

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 real_spherical_harmonics(
    vectors: torch.Tensor,  # (E, 3)
    valid: torch.Tensor,  # (E,)
    degrees: torch.Tensor,  # (Q,)
    coefficients: torch.Tensor,  # (Q, D, T)
    powers: torch.Tensor,  # (Q, T, 3)
    term_valid: torch.Tensor,  # (Q, T)
    normalization: str = 'component',
    normalize: bool = True,
) -> torch.Tensor:
    """Return a tensor with shape (E, Q, D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement real_spherical_harmonics")

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.