RFdiffusion rotational reverse step

Your task

Function signature

rfdiffusion_rotation_step(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    steps: int = 50,
    series_terms: int = 32,
) -> dict[str, torch.Tensor]

Worked example

A score correction and noise both change orientation

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Compute one rotational reverse-SDE step for original RFdiffusion from supplied current and predicted-clean residue frames. This operation updates orientations and rotates backbone atoms about their current Cα positions; it does not compute the translation posterior or the outer runner's terminal return.
• The floating inputs noisy_rotation, clean_rotation and coordinates have shape [B,N,3,3], with B,N\ge1. The first two are proper rotation matrices R_t,R_0; coordinate rows are N, Cα and C. noise has shape [B,N,3] and contains supplied Gaussian variates.
• Boolean valid and fixed have shape [B,N]; fixed=True protects a valid residue. The active set is valid AND not fixed. Matrices and coordinates are independently supplied; the caller constructs consistent frames from its backbone. All real tensors share one float32/float64 CPU dtype, and no input may be mutated.
• Config steps is an integer T\in[2,1000], default 50. Integer step[B] lies in [1,T]. The per-example floating sigma_min, min_b, max_b and noise_scale all have shape [B]. Require finite \sigma_{\min}>0, 0<b_{\min}\le b_{\max} and noise scale a\ge0.
• These are always checked, even when every residue is fixed or invalid. Use the source linear schedule, with u=\mathrm{step}/T: \sigma=\sigma_{\min}+b_{\min}u+\tfrac12(b_{\max}-b_{\min})u^2,\qquad g=\sqrt{2\sigma\,[b_{\min}+(b_{\max}-b_{\min})u]}. Require current \sigma\in[0.75,2]. This is the RF width convention.
• Its heat-kernel exponent contains 1/2; it corresponds to width \sigma/\sqrt2 in the separate DiffDock-convention kernel exercise.
• For each active residue form Q=R_tR_0^\top and \omega=\operatorname{Log}(Q), requiring relative angle \theta=\|\omega\|\le2.8 radians. Use the principal logarithm away from the branch cut.
• For explicit matrix gradients define v=\operatorname{vee}(Q-Q^\top)/2, c=(\operatorname{tr}Q-1)/2, q=v^\top v. When q<10^{-8} and c>0, use \omega=(1+q/6+3q^2/40+5q^3/112)v; otherwise use \omega=\operatorname{atan2}(\sqrt q,c)v/\sqrt q. This fixes the local algebraic extension for matrix-input derivatives.
• Check used matrices for orthogonality and determinant one to absolute tolerance 10^{-3}; do not project or detach them. Config series_terms is an integer L\in[16,64], default 32.
• Evaluate the finite normalized-Haar kernel H(\theta,\sigma)=\sum_{\ell=0}^{L-1}(2\ell+1)e^{-\ell(\ell+1)\sigma^2/2}\left[1+2\sum_{k=1}^{\ell}\cos(k\theta)\right]. Let s=(H'(\theta)/(\theta H(\theta)))\omega. At \theta=0, use the analytic factor H''(0)/H(0), so the score is exactly zero with its nonzero linear derivative.
• Avoid differentiating an unguarded square root or dividing by zero in an unused branch. Float64 working arithmetic is permitted for both input dtypes; the finite series and its derivatives, rather than source cache interpolation, define the answer.
• With \Delta t=1/T and supplied noise \varepsilon, compute \delta=g^2\Delta t\,s+g\sqrt{\Delta t}\,a\varepsilon,\qquad P=\operatorname{Exp}(\delta),\qquad R_{\mathrm{next}}=PR_t. The perturbation is LEFT-composed.
• For each N/Cα/C point x, return P(x-x_{\mathrm{CA}})+x_{\mathrm{CA}}; the Cα pivot is unchanged. Use the Rodrigues exponential with its analytic zero limit. Return rotation, coordinates, relative_vector, score, tangent, sigma and diffusion_coefficient with the declared shapes and original dtype.
• Valid fixed residues return their input noisy rotation and coordinates exactly, and zero relative vector, score and tangent. Invalid residues return zeros for all residue-shaped outputs. Only active clean rotations and noise are read; fixed/invalid clean/noise values and all invalid residue values may be NaN/Inf.
• Sanitize them before arithmetic, with zero gradients. Used noisy rotations, coordinates and active noise must be finite. A zero noise scale does not disable active-noise validity or its scale derivative. Reject invalid listed shapes/dtypes/types, schedule ranges, used rotations, relative-angle bounds and nonfinite used values with ValueError.
• Outputs and intermediate arithmetic must remain finite and representable. No random draw, source lookup, network execution or implicit terminal override occurs here.
• For v=(x,y,z), \widehat v=\begin{bmatrix}0&-z&y\\z&0&-x\\-y&x&0\end{bmatrix}. The inverse vee reads entries (2,1), (0,2), (1,0).
• Use Exp(v)=I+(\sin\theta/\theta)\widehat v+[(1-\cos\theta)/\theta^2]\widehat v^2, with \theta=\|v\|. The coefficient limits at zero are 1 and 1/2; use safe analytic series near zero.

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 rfdiffusion_rotation_step(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    steps: int = 50,
    series_terms: int = 32,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    rotation: (B, N, 3, 3).
    coordinates: (B, N, 3, 3).
    relative_vector: (B, N, 3).
    score: (B, N, 3).
    tangent: (B, N, 3).
    sigma: (B,).
    diffusion_coefficient: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rfdiffusion_rotation_step")

Sources and further reading

• www.nature.com
• RosettaCommons/RFdiffusion · diffusion.py
• RosettaCommons/RFdiffusion · diffusion.py
• RosettaCommons/RFdiffusion · utils.py