Wrapped torsion corruption, density and score normalization

Your task

Function signature

wrapped_torsion_diffusion(
    angle0: torch.Tensor,
    noise: torch.Tensor,
    uniform: torch.Tensor,
    valid: torch.Tensor,
    sigma: torch.Tensor,
    image_radius: int = 4,
    quadrature_points: int = 128,
) -> dict[str, torch.Tensor]

Worked example

Cross the angular boundary

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs angle0, noise and uniform are same-dtype 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 torsion. All floating inputs use float32 or float64 on the same CPU device.
• Valid angle0 values are finite with absolute \mathrm{value} \le 100\cdot \pi; valid noise values are finite with absolute \mathrm{value} \le 32; valid uniform lies in [0,1). Masked payloads may contain NaN/Inf and must be replaced before arithmetic.
• Every sigma is finite in [0.35,3], with range comparisons interpreted in its input dtype; sigma remains meaningful even for an empty or fully masked molecule. Work internally in float64 and cast returned floating tensors to the input dtype. Let W(x)=x-2\pi\left\lfloor\frac{x+\pi}{2\pi}\right\rfloor, mapping to [-\pi,\pi).
• Return \mathrm{displacement}=W(\sigma\cdot \mathrm{noise}), \mathrm{noisy\_angle}=W(\mathrm{angle0}+\sigma\cdot \mathrm{noise}), and \mathrm{initial\_angle}=2\cdot \pi\cdot \mathrm{uniform}-\pi. The latter is a supplied-uniform torus prior, distinct from Gaussian corruption.
• Config image_radius has default 4 and is an integer J in [3,8]; quadrature_points has default 128 and is an integer Q in [64,1024]. Bool values are not accepted as integer configuration values.
• For d=\mathrm{displacement} define z_{k}=d+2\cdot \pi\cdot k for k=-J,...,J, Z(s)=\operatorname{erf}\left(\frac{(2J+1)\pi}{\sqrt2\,s}\right), and p_J(d;s)=\frac{\sum_k e^{-z_k^2/(2s^2)}}{\sqrt{2\pi}\,s Z(s)}.
• Return \mathrm{density}=p_{J}, \mathrm{log\_density}=\operatorname{log}(p_{J}), and \operatorname{score}=\frac{d\log p_J}{d d}=-\frac{\sum_k z_k e^{-z_k^2/(2s^2)}}{s^2\sum_k e^{-z_k^2/(2s^2)}}. Evaluate log_density stably, without a density floor or \epsilon.
• Z exactly normalizes the truncated image sum over the principal interval; this finite-image likelihood approximates, but is not defined as, the infinite wrapped-normal likelihood of the supplied corruption.
• On midpoint grid g_{j}=-\pi+(j+1/2)\cdot 2\cdot \pi/Q, j=0,...,Q-1, return \texttt{score\_norm\_sq}[B]=\frac{\sum_j p_J(g_j;s)\operatorname{score}(g_j;s)^2}{\sum_j p_J(g_j;s)}. This is a quadrature-defined expected SQUARED score, not its square root.
• All six per-torsion outputs are zero where valid=False, including log_density as a bookkeeping convention; score_norm_sq is not masked by torsion count. Empty K gives [B,0] per-torsion tensors and a defined [B] normalizer. Inputs are not mutated and no random numbers are drawn.
• Invalid supported configuration values, sigma ranges, observed finite/range checks or listed shapes raise ValueError. No narrower-noise or larger-image-range accuracy is promised.

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 wrapped_torsion_diffusion(
    angle0: torch.Tensor,  # (B, K)
    noise: torch.Tensor,  # (B, K)
    uniform: torch.Tensor,  # (B, K)
    valid: torch.Tensor,  # (B, K)
    sigma: torch.Tensor,  # (B,)
    image_radius: int = 4,
    quadrature_points: int = 128,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    displacement: (B, K).
    noisy_angle: (B, K).
    initial_angle: (B, K).
    density: (B, K).
    log_density: (B, K).
    score: (B, K).
    score_norm_sq: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement wrapped_torsion_diffusion")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · torus.py