Structure module, templates and recycling · INTEGRATION

Normalized-quaternion backbone updates

HardRaw PyTorchAbout 30 min2 public examples plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Update residue frames and place their local points in world coordinates. The selected mode determines whether the increment is expressed in the residue’s local frame or in the global frame; this changes both rotation composition and translation. Preserve each immovable frame and return the normalized quaternion increment.

Function signature

quaternion_frame_update(
    rotation: torch.Tensor,
    translation: torch.Tensor,
    local_points: torch.Tensor,
    movable: torch.Tensor,
    options: dict,
    single: torch.Tensor | None = None,
    w_update: torch.Tensor | None = None,
    b_update: torch.Tensor | None = None,
    offsets: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]

Worked example

Global and local rotation order differ

At the frame-update stage, take one movable residue whose old rotation is a positive quarter-turn about z and whose old translation is zero. Suppose the supplied or projected increments, after the declared scales, are u=(1,0,0) and v=(1,0,0). Quaternion (1,1,0,0)/sqrt(2) gives Rx, a positive quarter-turn about x. RFdiffusion global mode uses R_new=Rx·Rz and new translation (1,0,0). Local point (0,1,0) becomes (-1,0,0) under Rz, stays there under Rx, and ends at (0,0,0) after translation. For the same scaled increments, AF2 local mode instead gives (0,1,1) because it uses Rz·Rx and rotates v with the old Rz.

Parameters

Parameter	Shape and meaning
`rotation`torch.Tensor	$[B, N, 3, 3]$ Proper rotation mapping local vectors into world coordinates.
`translation`torch.Tensor	$[B, N, 3]$ World origin of each frame.
`local_points`torch.Tensor	$[B, N, P, 3]$ Points expressed in each frame’s local coordinates.
`movable`torch.Tensor	$[B, N]$ $\mathrm{True}=\mathrm{update}$ allowed
`single`torch.Tensor \| NoneDefault `None`	$[B, N, C]$ Learned scalar representation used only by the AF2 update projection.
`w_update`torch.Tensor \| NoneDefault `None`	$[C, 6]$ `af2_local`: $[C,6]$ learned backbone-update head
`b_update`torch.Tensor \| NoneDefault `None`	$[6]$ Supplied update affine offset; add it only at its named projection or normalization stage.
`offsets`torch.Tensor \| NoneDefault `None`	$[B, N, 2, 3]$ Supplied global translation offset in slot 0 and rotation-vector offset in slot 1.
`options`dict	Scalar settings. Omitted keys use the defaults below.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Update each supplied proper rotation[ $B,N$ , $3$ , $3$ ], translation[ $B,N$ , $3$ ] and transform local_points[ $B,N,P$ , $3$ ] with column-vector convention. $B,N,P$ and feature width $C$ are positive; all floats finite. movable[ $B,N$ ] True permits an increment; protected residues still return their existing transformed points.
mode='af2_local' (default): compute $\mathrm{update}=\mathrm{single}$ [ $B,N,C$ ]@ $w_{\mathrm{update}}[C,6]+b_{\mathrm{update}}[6]$ . First3 channels are rotation vector $u$ , final3 translation vector $v$ . Default $\mathrm{rotation\_scale}=\mathrm{translation\_scale}=1$ .
mode='rfdiffusion_global': take offsets[ $B,N$ , $2$ , $3$ ] from the separately learned SE( $3$ ) network; slot0 is translation $v$ , slot1 rotation $u$ . Default $\mathrm{translation\_scale}=0.1$ and $\mathrm{rotation\_scale}=0.01$ , following original RFdiffusion's / $10$ and / $100$ . scales are finite nonnegative numbers.
In either mode multiply $u/v$ by these scales and zero both where movable=False. Form a WXYZ quaternion $q=(1,u)/\sqrt{1+\sum(u^{2})}$ ; no $\epsilon$ is necessary. Convert $q$ to the proper active rotation delta_R.
af2_local applies $\mathrm{rotation\_new}=\mathrm{rotation}\mathbin{@}\delta_{R}$ and $\mathrm{translation\_new}=\mathrm{translation}+\mathrm{rotation}\mathbin{@}v$ , where translation uses the OLD rotation.
rfdiffusion_global applies $\mathrm{rotation\_new}=\delta_{R}\mathbin{@}\mathrm{rotation}$ and $\mathrm{translation\_new}=\mathrm{translation}+v$ ; do not rotate the old translation.
Return $\mathrm{points}=\mathrm{rotation\_new}\mathbin{@}\mathrm{local\_points}+\mathrm{translation\_new}$ and delta_quaternion[ $B,N$ , $4$ ] as well as the new transform. Zero increment gives the identity update; a protected residue has $\delta_{\mathrm{quaternion}}=(1,0,0,0)$ .
No detach, stochastic integrator, axis-angle exponential, or post-update quaternion sign-flip convention is included. Original RFdiffusion's motif_mask=True means protected; the supplied movable mask intentionally has the opposite polarity.

Why this matters

The same rotation increment can act in the residue’s local coordinates or in the global frame. The selected mode determines both the multiplication order and how the translation increment is interpreted.

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 quaternion_frame_update(
    rotation: torch.Tensor,  # (B, N, 3, 3)
    translation: torch.Tensor,  # (B, N, 3)
    local_points: torch.Tensor,  # (B, N, P, 3)
    movable: torch.Tensor,  # (B, N)
    options: dict,
    single: torch.Tensor | None = None,  # (B, N, C)
    w_update: torch.Tensor | None = None,  # (C, 6)
    b_update: torch.Tensor | None = None,  # (6)
    offsets: torch.Tensor | None = None,  # (B, N, 2, 3)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    rotation: (B, N, 3, 3).
    translation: (B, N, 3).
    points: (B, N, P, 3).
    delta_quaternion: (B, N, 4)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement quaternion_frame_update")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.