SO(3) heat-kernel scores and supplied-variate angle sampling

Your task

Function signature

so3_heat_kernel_score(
    omega: torch.Tensor,
    sigma: torch.Tensor,
    uniform: torch.Tensor,
    axis: torch.Tensor,
    valid: torch.Tensor,
    series_terms: int = 16,
    grid_points: int = 65,
    distribution: str = 'heat_kernel',
) -> dict[str, torch.Tensor]

Worked example

Group score at a nonzero rotation

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs omega and axis are same-dtype floating [B,K,3], uniform is floating [B,K], valid is bool [B,K], and sigma is floating [B], with B\ge 1 and K\ge 0. True means a real query/sample. All floating tensors share float32 or float64 CPU dtype.
• Valid omega is a principal rotation vector in radians: \sum(\omega^{2})\le \pi^{2}+10^{-6}, where the fixed tolerance only permits endpoint rounding and no vector is clamped. Every sigma is finite in [0.5,2], including empty or fully masked rows.
• Valid uniform lies in [0,1], and valid axis is finite with Euclidean norm in [10^{-6},10^{6}]. Axis may be a supplied unit direction or nonzero Gaussian vector; always normalize it, never resample it. Missing omega/axis/uniform may contain NaN/Inf and must be masked before any operation.
• Compute internally in float64, return all outputs in the input floating dtype, and never mutate inputs.
• Config series_terms defaults to 16 and is integer L in [16,64]; grid_points defaults to 65 and is integer G in [\operatorname{max}(65,2\cdot L+1),513]; distribution defaults to 'heat_kernel' and is either 'heat_kernel' or 'haar'. Bool values are not valid integer configurations.
• For \theta=\|\omega\| and \mathrm{ell}=0,...,L-1, define a_\ell=(2\ell+1)e^{-\ell(\ell+1)\sigma^2} and \chi_\ell(\theta)=1+2\sum_{m=1}^{\ell}\cos(m\theta). H_L(\theta;\sigma)=\sum_\ell a_\ell\chi_\ell(\theta).
• This is ORIGINAL DiffDock v1.0's no-one-half exponent; RFdiffusion's \operatorname{exp}(-\mathrm{ell}\cdot (\mathrm{ell}+1)\cdot t/2) matches under t=2\cdot \sigma^{2}.
• Here t is HEAT time, not a discrete or normalized schedule time: the original RFdiffusion table passes t=\sigma_{\mathrm{RF}}^{2}, so \sigma_{\mathrm{RF}}=\sqrt{2}\cdot \sigma.
• Its density helper returns H rather than H/(8\cdot \pi^{2}); multiply group_density by 8\cdot \pi^{2} to match that normalized-Haar convention, while the log score is unchanged.
• In distribution='haar', set H=1 and all derivatives to zero exactly, independent of sigma: a finite heat-kernel sigma is not silently treated as Haar.
• Return \mathrm{group\_density}=H/(8\cdot \pi^{2}), with the original DiffDock volume convention, and \mathrm{angle\_density}=H\cdot (1-\operatorname{cos}(\theta))/\pi, with respect to dtheta on [0,\pi]. Use 2\cdot \operatorname{sin}(\theta/2)^{2} for the latter factor near identity.
• Return \texttt{radial\_score}=\frac{H^{\prime}(\theta)}{H} and \texttt{vector\_score}=\frac{H^{\prime}(\theta)}{\theta H}\omega. These differentiate log GROUP density, never the angle marginal or its Haar Jacobian.
• At identity set \mathrm{radial\_score}=0 and \mathrm{vector\_score}=0, using \lim_{\theta\to 0}\frac{H^{\prime}(\theta)}{\theta}=H^{\prime\prime}(0) in the vector-score coefficient so its Jacobian remains correct. Stable cosine characters and a stable \operatorname{sin}(x)/x limit avoid singular sin-ratio evaluation.
• Define output \mathrm{angle\_grid}[j]=j\cdot \pi/(G-1), j=0,...,G-1. Evaluate p_{j}=\mathrm{angle\_density}(\mathrm{angle\_grid}[j]) and r_{j}=\mathrm{radial\_score}(\mathrm{angle\_grid}[j]); let d_{j}=\mathrm{grid}[j+1]-\mathrm{grid}[j] and m_j=\frac{p_j+p_{j+1}}2\,d_j.
• Return \mathrm{angle\_cdf}[:,0]=0 and \texttt{angle\_cdf}[:,j]=\frac{\sum_{i<j}m_i}{\sum_i m_i}, with the last element normalized to one.
• Return \texttt{score\_norm\_sq}=\frac{\sum_j(p_jr_j^2+p_{j+1}r_{j+1}^2)d_j/2}{\sum_jm_j}; \mathrm{score\_rms}=\sqrt{\mathrm{score\_norm\_sq}}; and \mathrm{diffdock\_score\_norm}=\sqrt{\mathrm{score\_norm\_sq}/\pi}. All three are exact zero in the Haar branch with zero sigma gradients.
• For each supplied u, choose i=\operatorname{clamp}(\sum_{j} [u\ge \mathrm{CDF\_j}]-1,0,G-2), interpolate \texttt{sample\_angle}=\texttt{grid}_i+\frac{u-\mathrm{CDF}_i}{\mathrm{CDF}_{i+1}-\mathrm{CDF}_i}(\texttt{grid}_{i+1}-\texttt{grid}_i), and return \texttt{sample\_rotvec}=\texttt{sample\_angle}\frac{\texttt{axis}}{\|\texttt{axis}\|}.
• Interior CDF knots use the interval on their right; u=0 and u=1 give angles 0 and \pi. This inverse CDF is LINEAR BETWEEN GRID KNOTS: it is not inversion of a continuous linearly varying PDF. Per-query outputs (the four density/score tensors plus the two sample tensors) are zero for invalid entries.
• Grid, CDF and normalizers are not masked by query count and exist for empty K. No RNG, cached tables, e3nn, matrix logarithm or rotation-matrix conversion is used. Invalid supported configuration values, observed ranges or listed shapes raise ValueError.
• The finite series, trapezoid moments and piecewise-linear inverse CDF define the exercise; no accuracy or positivity guarantee is made outside this moderate-noise domain, and even Haar sampling here is a finite-grid approximation rather than an exact Haar rotation generator.

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 so3_heat_kernel_score(
    omega: torch.Tensor,  # (B, K, 3)
    sigma: torch.Tensor,  # (B,)
    uniform: torch.Tensor,  # (B, K)
    axis: torch.Tensor,  # (B, K, 3)
    valid: torch.Tensor,  # (B, K)
    series_terms: int = 16,
    grid_points: int = 65,
    distribution: str = 'heat_kernel',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    group_density: (B, K).
    angle_density: (B, K).
    radial_score: (B, K).
    vector_score: (B, K, 3).
    angle_grid: (G,).
    angle_cdf: (B, G).
    score_norm_sq: (B,).
    score_rms: (B,).
    diffdock_score_norm: (B,).
    sample_angle: (B, K).
    sample_rotvec: (B, K, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement so3_heat_kernel_score")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · so3.py
• RosettaCommons/RFdiffusion · igso3.py
• www.nature.com