Structure module, templates and recycling

Global Quaternion Frame Update

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

Update each movable residue frame using a quaternion rotation and translation expressed in global coordinates. Left-compose the rotation and add the translation directly to the old origin, then transform the supplied local points through the resulting frame.

Function signature

rfdiffusion_global_frame_update(
    rotation: torch.Tensor,
    translation: torch.Tensor,
    local_points: torch.Tensor,
    movable: torch.Tensor,
    offsets: torch.Tensor,
    rotation_scale: float = 0.01,
    translation_scale: float = 0.1,
) -> 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 local-to-world transforms.
`translation`torch.Tensor	$[B, N, 3]$ World origins.
`local_points`torch.Tensor	$[B, N, P, 3]$ Local residue atoms/points.
`movable`torch.Tensor	$[B, N]$ True=update allowed.
`offsets`torch.Tensor	$[B, N, 2, 3]$ Supplied equivariant network update; slot 0 is translation and slot 1 the rotation-vector increment.
`rotation_scale`floatDefault `0.01`	Nonnegative multiplier on the quaternion-vector increment; defaults to 0.01.
`translation_scale`floatDefault `0.1`	Nonnegative multiplier on the world-coordinate translation increment; defaults to 0.1.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Frames use column vectors: a local point maps to rotation @ point + translation. movable is boolean; False keeps the original frame.
offsets[...,0,:] is the global translation increment v before multiplication by translation_scale. offsets[...,1,:] is the quaternion-vector increment before multiplication by rotation_scale.
Apply rotation_scale to the vector increment and normalize quaternion $q=(1,u)/\sqrt{1+\|u\|^2}$ , using scalar-first order. Convert it to its proper 3×3 rotation matrix; do not use an axis-angle exponential.
Compose R_new = delta_R @ R and t_new = t + v in global coordinates. rotation_scale defaults to 0.01 and translation_scale to 0.1.
Both scales must be nonnegative. Zero both increments where movable is False, so delta_quaternion is the identity quaternion there. Transform every supplied local point with the final frame.
For normalized quaternion $(w,x,y,z)$ , use $R=\begin{bmatrix}1-2(y^2+z^2)&2(xy-wz)&2(xz+wy)\\2(xy+wz)&1-2(x^2+z^2)&2(yz-wx)\\2(xz-wy)&2(yz+wx)&1-2(x^2+y^2)\end{bmatrix}$ .

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 rfdiffusion_global_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)
    offsets: torch.Tensor,  # (B, N, 2, 3)
    rotation_scale: float = 0.01,
    translation_scale: float = 0.1,
) -> 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 rfdiffusion_global_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.