Structure module, templates and recycling

Local 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

Move each movable residue’s frame using an increment expressed in its current local coordinates. Compose the normalized quaternion rotation on the right and rotate the translation increment by the original frame. Return the new frame, its transformed local points, and the normalized increment.

Function signature

af2_local_frame_update(
    rotation: torch.Tensor,
    translation: torch.Tensor,
    local_points: torch.Tensor,
    movable: torch.Tensor,
    single: torch.Tensor,
    w_update: torch.Tensor,
    b_update: torch.Tensor,
    rotation_scale: float = 1.0,
    translation_scale: float = 1.0,
) -> dict[str, torch.Tensor]

Worked example

The original frame rotates the translation

Consider the frame-update stage after projection and scaling, with u=(1,0,0) and v=(1,0,0) for one movable residue. Its old rotation is a positive quarter-turn about z and old translation is zero. Quaternion (1,1,0,0)/sqrt(2) gives a positive quarter-turn about local x. AF2 uses R_new=Rz·Rx; its new origin is Rz·v=(0,1,0). Local point (0,1,0) rotates first to (0,0,1) and stays there under Rz, so the returned point is (0,1,1). The translation uses the original Rz, independently of the new quaternion rotation.

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.
`single`torch.Tensor	$[B, N, C]$ Af2_local: [B,N,C] learned scalar representation.
`w_update`torch.Tensor	$[C, 6]$ Af2_local: [C,6] learned backbone-update head.
`b_update`torch.Tensor	$[6]$ Af2_local: [6].
`rotation_scale`floatDefault `1.0`	Nonnegative multiplier on the local quaternion-vector increment; defaults to1.
`translation_scale`floatDefault `1.0`	Nonnegative multiplier on the local translation increment; defaults to1.

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.
Project single with w_update and b_update; the first three channels are the quaternion-vector increment and the final three the local translation increment.
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 = R @ delta_R and t_new = t + R @ v using the ORIGINAL R. rotation_scale and translation_scale default to 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.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

The learned increment is expressed in each residue frame. Compose the rotation on the right and rotate the translation increment through the original frame before adding it.

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 af2_local_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)
    single: torch.Tensor,  # (B, N, C)
    w_update: torch.Tensor,  # (C, 6)
    b_update: torch.Tensor,  # (6)
    rotation_scale: float = 1.0,
    translation_scale: float = 1.0,
) -> 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 af2_local_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.