Molecular graphs, equivariance and docking

Real-irrep Clebsch–Gordan paths and channel tensor products

HardRaw PyTorchAbout 30 min1 public example plus private tests

NequIP DiffDockSE(3)-TransformerRFdiffusion

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

Your task

Multiply left and right geometric features through the supplied coupling tables, then mix their channel copies with learned weights. Each output component is the sum of all active paths targeting that output type; a path selects one left type, one right type, and one output type. Return the scaled contribution of every path as well as their sum.

Function signature

clebsch_gordan_tensor_product(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    symmetry: str = 'o3',
    weight_mode: str = 'shared',
    irrep_normalization: str = 'component',
    path_normalization: str = 'fan_in',
) -> dict[str, torch.Tensor]

Worked example

Two vector products produce different representation types

Take one valid pair of polar vectors $a=[1,2,0]$ , $b=[3,1,0]$ , one channel each, and one path to each of an even scalar and an axial vector. Both paths are valid; all learned weights and path scales are 1. Use O(3), shared weights, component normalization, and fan-in normalization.

The supplied scalar CG row computes $(a\cdot b)/\sqrt3=(1\cdot3+2\cdot1)/\sqrt3=5/\sqrt3$ . The vector rows compute $(a\times b)/\sqrt2=[0,0,1\cdot1-2\cdot3]/\sqrt2=[0,0,-5/\sqrt2]$ . Each output has one contributing path and one channel, so fan-in adds no scaling. Both paths satisfy parity $(-1)(-1)=+1$ . features contains scalar $2.886751$ (with zero padding) and vector [0,0,-3.535534].

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`symmetry`strDefault `'o3'`	Choose "o3", "so3". A path is active exactly when `path_valid` is true, $\|l_{a}-l_{b}\|\le l_{c} \le l_{a}+l_{b}$ and, in `symmetry='o3'` default, $\mathrm{parity\_c}=\mathrm{parity\_a}\cdot \mathrm{parity\_b}$ .
`weight_mode`strDefault `'shared'`	"shared" uses one weight tensor for all observations; "sample" supplies a separate weight tensor per observation.
`irrep_normalization`strDefault `'component'`	"component" uses the supplied unit-norm CG rows directly; "norm" multiplies each path by $\sqrt{(2l_{\mathrm{left}}+1)(2l_{\mathrm{right}}+1)/(2l_{\mathrm{out}}+1)}$ .
`path_normalization`strDefault `'fan_in'`	"fan_in" divides each path by $\sqrt{C_{\mathrm{left}}C_{\mathrm{right}}P_{\mathrm{out}}}$ , where $P_{\mathrm{out}}$ counts active paths to this output type; "none" omits that divisor.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

For path $p=(a,b,c)$ , contract $z[n,p,o,m]=\sum_{i,j,u,v} \mathrm{weights}[p,i,j,o]\cdot \mathrm{cg}[p,m,u,v]\cdot \mathrm{left}[n,a,i,u]\cdot \mathrm{right}[n,b,j,v]$ .
In config weight_mode='sample', use weights $[n,p,i,j,o]$ ; default is 'shared'. A path is active exactly when path_valid is true, $|l_{a}-l_{b}|\le l_{c} \le l_{a}+l_{b}$ and, in symmetry='o3' default, $\mathrm{parity\_c}=\mathrm{parity\_a}\cdot \mathrm{parity\_b}$ . symmetry='so3' omits parity selection.
Inactive paths return zero and are excluded from path counts; false path_valid entries may have sentinel indices or NaN tables, ignored before gathers.
For irrep_normalization='component' default, fixed CG rows have unit norm. 'norm' multiplies each path by $\sqrt{(2\cdot l_{a}+1)\cdot (2\cdot l_{b}+1)/(2\cdot l_{c}+1)}$ .
Multiply by supplied path_scale $[p]$ , then, if path_normalization='fan_in' default, divide by $\sqrt{\mathrm{Ca}\cdot \mathrm{Cb}\cdot \mathrm{number\_active\_paths\_to\_c}}$ ; 'none' omits this divisor. The output for type $c$ is the sum of its path outputs.
Tables have axes output,left,right and are fixed intertwining tensors, not arbitrary learned component matrices. Mask invalid rows, input/output component padding and inactive paths, including their nonfinite payloads. $N=0$ and $P=0$ are supported; type lists and channel counts are nonempty.
A packed block $(l,p)$ has $2\cdot l+1$ real components and multiplicity $C$ ; $p$ is $+1$ or $-1$ under inversion. $D$ is a common padded width. Padding is ignored, including NaN padding. Distinct copies of a type are channels, not XYZ coordinates.
The fixture basis is Cartesian $l=1 [x,y,z]$ and $l=2$ orthonormal symmetric-traceless tensors $Q=[(\mathrm{xz}+\mathrm{zx})/\sqrt{2},(\mathrm{xy}+\mathrm{yx})/\sqrt{2},\operatorname{diag}(-1,2,-1)/\sqrt{6},(\mathrm{yz}+\mathrm{zy})/\sqrt{2},\operatorname{diag}(-1,0,1)/\sqrt{2}]$ .
Fixed row-orthonormal CG tensors intertwine these bases; phases are defined by the supplied tables, not inferred from an external checkpoint. Supplied higher-degree tables must obey the same stated normalization and representation law.

Why this matters

Irreducible representations specify how feature components rotate. A Clebsch–Gordan table combines two such component systems into an allowed output type; learned weights mix channels without replacing that geometric rule.

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 clebsch_gordan_tensor_product(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    symmetry: str = 'o3',
    weight_mode: str = 'shared',
    irrep_normalization: str = 'component',
    path_normalization: str = 'fan_in',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    features: (N, O, Co, D).
    path_outputs: (N, P, Co, D).
    active_paths: (P,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement clebsch_gordan_tensor_product")

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.