Three-space docking initialization and reverse diffusion step

Your task

Function signature

pose_reverse_diffusion_step(
    activations: dict,
    masks: dict,
    structure: dict,
    mode: str = 'sde',
    no_final_step_noise: bool = False,
    suppress_noise: bool = False,
) -> dict[str, torch.Tensor]

Worked example

A decreasing-time translation step

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Space order is translation,rotation,torsion. t,dt [B,3] satisfy 0\le \mathrm{dt}\le t\le 1; time decreases to t_{\mathrm{next}}=t-\mathrm{dt}. sigma_min,sigma_max [3] satisfy 0<\operatorname{min}<\operatorname{max}.
• For space j, \sigma_{j}=\mathrm{min\_j}^{1-t_{j}}\cdot \mathrm{max\_j}^{t_{j}}.
• Use ORIGINAL DiffDock v1.0 coefficients g^{2}=\sigma^{2}\cdot \operatorname{log}(\operatorname{max}/\operatorname{min})\cdot (2,4,2); the rotation factor 4 is intentional and corresponds to its heat kernel exp[-l(l+1)\cdot \sigma^{2}], not the maintained code's changed convention.
• translation_score,rotation_score [B,3] and torsion_score [B,K] are already physically normalized scores, not raw neural-head outputs.
• In mode='sde' (default), \delta_{j}=g_{j}^{2}\cdot \mathrm{dt\_j}\cdot \mathrm{score\_j}+g_{j}\cdot \sqrt{\mathrm{dt\_j}}\cdot \mathrm{noise\_j}. 'ode' uses \delta_{j}=0.5\cdot g_{j}^{2}\cdot \mathrm{dt\_j}\cdot \mathrm{score\_j} and ignores all noise; suppress_noise=True in SDE instead drops only noise, not the factor 1 drift.
• no_final_step_noise=False by default; True suppresses stochastic terms for supplied last_step[B]=True. All noise tensors have the corresponding score shape. At \mathrm{dt\_j}=0 ignore that space's score/noise payloads and return zero increment, without evaluating an unsafe sqrt derivative.
• Torsion increments are real radians, not wrapped here; directed_bond_torsion_rotation applies them separately. Let c be the mean of valid current positions [B,N,3]. \operatorname{Exp}(\delta_{\mathrm{rotation}}) is a GLOBAL active rotation, so:
• \mathrm{rigid\_positions}=(\mathrm{positions}-c)\mathbin{@}\operatorname{Exp}(\delta_{\mathrm{rotation}})^{\top}+c+\delta_{\mathrm{translation}}. It is not the local/body right-composition convention of lie_euler.
• Return the increments and this intermediate rigid pose; a separate torsion action and equal-weight Kabsch complete DiffDock's conformer update.
• For initialization, initial_conformer [B,N,3] is already torsion-randomized by the separate angular/rotation exercises.
• \mathrm{initial\_positions}=(\mathrm{initial\_conformer}-\mathrm{its\_valid\_mean})\mathbin{@}\mathrm{initial\_rotation}^{\top}+\sigma_{\mathrm{max}}[0]\cdot \mathrm{initial\_translation\_noise}+\mathrm{reference\_center}.
• initial_rotation [B,3,3] is a supplied proper rotation (Haar for the original prior), initial_translation_noise and reference_center are [B,3]; original centered-receptor coordinates use \mathrm{reference\_center}=0.
• Masks: valid_mask [B,N] observed atoms; diffuse_mask [B] whole movable ligand; torsion_mask [B,K] active torsions. Fixed ligands retain their current and initial coordinates, with zero increments and identity rotation_update. An example without valid atoms is inactive; invalid atoms are zero.
• Inactive scores/noises and initialization transforms may be nonfinite and are ignored. All active payloads and schedule parameters must be finite. B\ge 1,N\ge 0,K\ge 0; no chemistry parsing or denoiser is inside this function.
• Rodrigues uses a degree-six Taylor branch at squared \mathrm{angle}<10^{-6}, with no re-orthogonalization or angular clipping. keys are mode,suppress_noise,no_final_step_noise; unsupported mode, invalid scales or invalid time bounds raise ValueError.

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 pose_reverse_diffusion_step(
    activations: dict,
    masks: dict,
    structure: dict,
    mode: str = 'sde',
    no_final_step_noise: bool = False,
    suppress_noise: bool = False,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    sigma: (B, 3).
    g: (B, 3).
    t_next: (B, 3).
    translation_update: (B, 3).
    rotation_vector_update: (B, 3).
    torsion_update: (B, K).
    rotation_update: (B, 3, 3).
    rigid_positions: (B, N, 3).
    initial_positions: (B, N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement pose_reverse_diffusion_step")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · sampling.py
• gcorso/DiffDock · diffusion_utils.py
• gcorso/DiffDock · so3.py