Ligand translation, rotation, bond torsion and confidence heads

Your task

Function signature

geometric_score_readouts(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 0.0001,
    parity: str = 'so3',
    scale_by_sigma: bool = True,
) -> dict[str, torch.Tensor]

Worked example

Follow the raw vectors through calibration, then compute a bond score

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs describe N\ge 1 ligand atoms partitioned into B\ge 1 graphs by batch [N], plus M\ge 0 bonds and F\ge 0 atom-to-bond edges (F=0 if M=0). node_valid [N], bond_valid [M], torsion_edge_valid [F] use True for enabled items.
• features[N,T,C,5] contains real irreps with node_degrees [T] in {0,1,2}, node_parities [T] in {−1,+1}; only the first 2\cdot l+1 components are real. Type 0 must be degree-0 even scalars. xyz [N,3] is Cartesian row-coordinate storage.
• Invalid nodes/padded components may be NaN/Inf and are ignored before arithmetic. All observed values and parameters are finite. Valid node batch IDs lie in [0,B); invalid ones can be garbage. All graph tensors, including empty-graph entries, must be finite; all radial_widths, translation_sigma and score scales are positive.
• Use \mathrm{eps}>0 (default 10^{-4}) in BOTH geometric and score-vector lengths.
• For a displacement d define r_{\mathrm{eps}}=\sqrt{\sum(d\cdot d)+\varepsilon^{2}}, u=d/r_{\mathrm{eps}}. This is an explicitly regularized NON-UNIT direction, not e3nn's exact unit-normalized harmonic convention. The supplied fixed quadrupole_basis [5,3,3] contains orthonormal symmetric-traceless matrices Q_a.
• Define padded \mathrm{H0}=[1,0,0,0,0], \mathrm{H1}=[\sqrt{3}\cdot u_{x},\sqrt{3}\cdot u_{y},\sqrt{3}\cdot u_{z},0,0], H_2[a]=\sqrt{\frac{15}{2}}\,u^\top Q_a u. The graded convention is Q in order \frac{xz+zx}{\sqrt2}, \frac{xy+yx}{\sqrt2}, \frac{-xx+2yy-zz}{\sqrt6}, \frac{yz+zy}{\sqrt2}, \frac{zz-xx}{\sqrt2}.
• Fixed CG tensors must use this same basis and zero padding. A zero displacement gives \mathrm{H1}=\mathrm{H2}=0, \mathrm{H0}=1, r_{\mathrm{eps}}=\varepsilon. Geometry remains differentiable there. Let center [b] be the arithmetic mean of valid xyz in graph b (zero when empty).
• For each valid atom i evaluate H on \mathrm{xyz\_i}-\mathrm{center}[\mathrm{batch\_i}], and bare Gaussian RBF \rho_{k}=\operatorname{exp}(-0.5\cdot ((r_{\mathrm{eps}}-\mathrm{radial\_centers}[k])/\mathrm{radial\_widths}[k])^{2}), with no cutoff.
• Concatenate [\rho,graph_time [\mathrm{batch\_i}],features [i,0,:,0]] in that exact order.
• Apply z=\operatorname{SiLU}(\mathrm{attr}\mathbin{@}\mathrm{center\_weight1}+\mathrm{center\_bias1}), where \operatorname{SiLU}(x)=x/(1+\operatorname{exp}(-x)); then coefficients A[i,p,c,o]=\sum_{u} z[i,u]\cdot \mathrm{center\_weight2}[u,p,c,o]+\mathrm{center\_bias2}[p,c,o], o=0,1.
• For global_paths [p]=(node_type t,harmonic degree l,output parity slot s), contract \texttt{global\_path\_scale}[p]\sum_{c,a,b} A[i,p,c,o]\,\texttt{features}[i,t,c,a]\,H_\ell[b]\,\texttt{global\_cg}[p,m,a,b]. Only m<3 is used. Slot 0 is a polar (1 o) vector and slot 1 an axial (1 e) vector.
• Every path obeys angular triangle and parity selection. Average these atom contributions per graph, obtaining polar [B,2,3] and axial [B,2,3]. Learned radial/channel weights, not geometric basis values or final scores, are computed in this step. parity defaults to 'so3'.
• In that source-motivated mixed-parity mode, \mathrm{translation\_raw}=\mathrm{polar}[:,0]+\mathrm{axial}[:,0] and \mathrm{rotation\_raw}=\mathrm{polar}[:,1]+\mathrm{axial}[:,1]. These are equivariant under proper rotations, not given a polar/axial reflection law.
• In 'o3' mode use \mathrm{translation\_raw}=\mathrm{polar}[:,0], \mathrm{rotation\_raw}=\mathrm{axial}[:,1].
• For each head h=(\mathrm{translation},\mathrm{rotation}), set \mathrm{length\_h}=\sqrt{\operatorname{norm}(\mathrm{raw\_h})^{2}+\varepsilon^{2}}, z_{h}=\operatorname{SiLU}(\operatorname{concat}([\mathrm{length\_h}],\mathrm{graph\_time}) \mathbin{@} \mathrm{magnitude\_weight1}[h]+\mathrm{magnitude\_bias1}[h]), \mathrm{amplitude\_h}=z_{h} \mathbin{@} \mathrm{magnitude\_weight2}[h]+\mathrm{magnitude\_bias2}[h], and \mathrm{vector\_h}=\mathrm{raw\_h}/\mathrm{length\_h}\cdot \mathrm{amplitude\_h}.
• A zero raw vector yields zero, without NaN. scale_by_sigma defaults True: divide translation by translation_sigma [B], multiply rotation by rotation_score_scale [B]. False skips both and the torsion scale below. Empty graphs have all vector outputs and confidence zero. bond_index [2,M] stores endpoint atoms.
• A bond is active iff bond_valid, both endpoints are valid, and they have the same graph. Invalid bonds/edges may contain arbitrary indices; active endpoints/atoms must be in range. torsion_edge_index [2,F] is (\mathrm{bond\_receiver},\mathrm{atom\_sender}).
• An edge is active iff enabled, its bond and atom are active, and its atom belongs to the bond's graph. Cross-graph edges are ignored, not used in the denominator. Repeated active edges are separate neighbors. Neighborhood/radius and chemical rotatability decisions are supplied discrete inputs.
• Let bond \mathrm{midpoint}=(\mathrm{xyz\_a}+\mathrm{xyz\_b})/2 and A_{\mathrm{bond}}=\mathrm{H2}(\mathrm{xyz\_b}-\mathrm{xyz\_a}).
• For each active atom-to-bond edge compute edge harmonics H_l(xyz_atom-midpoint) and a filter for each \mathrm{filter\_paths}[q]=(l,L): \texttt{filter}[q,m]=\texttt{filter\_path\_scale}[q]\sum_{a,b}\texttt{filter\_cg}[q,m,a,b]H_\ell[a]A_{\mathrm{bond}}[b]. L in {0,1,2}; |l-2|\le L \le l+2.
• This is a genuine tensor product with the bond's \mathrm{rank}-2 irrep, not a dot product with the directed bond vector. Consequently reversing bond endpoints leaves this readout unchanged. Higher-L intermediate filters unable to couple to a scalar from l\le 2 node features need not be represented.
• Concatenate [edge_rho,features [\mathrm{atom},0,:,0],\mathrm{features}[\mathrm{first\_endpoint},0,:,0]+\mathrm{features}[\mathrm{second\_endpoint},0,:,0]].
• Apply the same specified SiLU first-layer formula with \mathrm{torsion\_weight1}/\mathrm{bias1}, then \mathrm{torsion\_weight2}/\mathrm{bias2} to produce coefficients D [f,p,c,o], o=0,...,O-1.
• torsion_paths [p]=(node_type t,filter q,output scalar slot s), where slot 0 means odd pseudoscalar 0\mathrm{o} and slot 1 even scalar 0\mathrm{e}. Require node degree equals filter L and \mathrm{node\_parity}\cdot (-1)^{l} equals the output parity. Compute
• \texttt{edge\_scalar}[f,s,o]=\sum_{\substack{p:\mathrm{slot}(p)=s\\c,a,b}}\texttt{torsion\_path\_scale}[p]D[f,p,c,o]\,\texttt{features}[\texttt{atom},t,c,a]\,\texttt{torsion\_cg}[p,a,b]\,\texttt{filter}[q,b] . Average over active neighbors of each bond; empty neighborhoods give zeros.
• Concatenate [\mathrm{odd\_channels},\mathrm{even\_channels}], each O wide; in o3 mode replace even_channels by zero.
• The final \mathrm{torsion}= \sum(\operatorname{tanh}(\operatorname{concat} \mathbin{@} \mathrm{torsion\_head\_weight1})\cdot \mathrm{torsion\_head\_weight2}), with no biases so strict-parity torsions remain pseudoscalars. If scaling, multiply by torsion_score_scale [M]. Inactive bonds and empty neighborhoods return zero.
• Supplied rotation/torsion scales are deterministic normalization values from other exercises, not supplied head outputs.
• For confidence, confidence_types [J] selects only degree-0 node types.
• Extract features [:,\mathrm{confidence\_types},:,0], flatten type then channel, take the valid-atom mean separately for each graph, and apply \operatorname{SiLU}(\mathrm{pool}\mathbin{@}\mathrm{confidence\_weight1}+\mathrm{confidence\_bias1})\mathbin{@}\mathrm{confidence\_weight2}+\mathrm{confidence\_bias2}.
• This is an unbounded confidence logit, no sigmoid.
• In o3 mode zero selected odd-scalar channels before the MLP.
• In so3 mode odd and even scalar channels may mix, matching proper-rotation rather than reflection invariance. Confidence does not update atom features. Return center [B,3], translation_raw [B,3], rotation_raw [B,3], translation [B,3], rotation [B,3], torsion [M], confidence [B]. No inputs are mutated.
• supports only \varepsilon, parity, scale_by_sigma. Nonpositive \varepsilon/widths, unsupported parity strings/degrees, wrong \mathrm{type}-0 irrep, invalid triangle/parity paths, and nonscalar confidence selections raise ValueError; other malformed shapes or invalid fixed intertwiner tables are outside scope.
• Fixed basis/path metadata are not learned. Defaults do not reproduce upstream checkpoints: the MLPs here are shallow SiLU without batch normalization/dropout, and \varepsilon regularizes geometry and norm calibration.
• translation_sigma is required. Supply positive values even when scale_by_sigma=false.
• rotation_score_scale is required. Supply positive values even when scale_by_sigma=false.
• torsion_score_scale is required. Supply positive values even when scale_by_sigma=false.

Why this matters

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 geometric_score_readouts(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 0.0001,
    parity: str = 'so3',
    scale_by_sigma: bool = True,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    center: (B, 3).
    translation_raw: (B, 3).
    rotation_raw: (B, 3).
    translation: (B, 3).
    rotation: (B, 3).
    torsion: (M,).
    confidence: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement geometric_score_readouts")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · score_model.py
• gcorso/DiffDock · all_atom_score_model.py
• gcorso/DiffDock · cg_model.py
• gcorso/DiffDock · tensor_layers.py