Geometry and structural losses

Batched rigid transforms

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

AlphaFold2OpenFoldRoseTTAFold2RoseTTAFold All-AtomAlphaFold3

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

Your task

Apply the first supplied rigid transform to points, invert that transform, and compose both transforms. Return the transformed points and inverse/composed rotations and translations using one column-vector convention.

Function signature

rigid_transform_operations(
    R1: torch.Tensor,
    t1: torch.Tensor,
    R2: torch.Tensor,
    t2: torch.Tensor,
    points: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Composition rotates the second translation

Let R1=[[0,-1,0],[1,0,0],[0,0,1]], t1=[1,0,0], R2=identity, t2=[2,0,0], and points=[[1,0,0]]. Applying T1 rotates [1,0,0] to [0,1,0], then translates it to [1,1,0]. The inverse has rotation R1ᵀ and translation [0,1,0]. T1∘T2 has rotation R1 and translation R1[2,0,0]+[1,0,0]=[1,2,0]; this composition applies T2 first.

Parameters

Parameter	Shape and meaning
`R1`torch.Tensor	$[\ldots, 3, 3]$ First proper rotation, mapping local column vectors into the output frame.
`t1`torch.Tensor	$[\ldots, 3]$ Translation of the first transform.
`R2`torch.Tensor	$[\ldots, 3, 3]$ Second proper rotation, applied first in composition.
`t2`torch.Tensor	$[\ldots, 3]$ Translation of the second transform.
`points`torch.Tensor	$[\ldots, P, 3]$ Point rows to transform; batch axes match the transforms.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Use COLUMN-vector geometry: $T(x)=Rx+t$ , although coordinates are stored in $[\ldots,P,3]$ rows. Inputs R1,R2 $[\ldots,3,3]$ are proper orthonormal rotations; translations t1,t2 $[\ldots,3]$ and points $[\ldots,P,3]$ have the same leading axes.
Return all three operations: $\mathrm{applied}=\mathrm{T1}(\mathrm{points})$ ; $\mathrm{inverse\_rotation}=R_1^\top$ and $\mathrm{inverse\_translation}=-R_1^\top t_1$ ; $\mathrm{composed\_rotation}=R_1R_2$ and $\mathrm{composed\_translation}=R_1t_2+t_1$ . Composition applies T2 FIRST and then T1.
Do not change the matrices, project onto $\mathrm{SO}(3)$ , invert them numerically, or add a residual to coordinates. There are no masks, epsilons, or options. $P=0$ is valid and gives empty applied coordinates while inverse/composition are still computed.

Why this matters

Composition order matters because rotations do not commute, and a translation must be expressed in the frame where it is added.

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 rigid_transform_operations(
    R1: torch.Tensor,  # (..., 3, 3)
    t1: torch.Tensor,  # (..., 3)
    R2: torch.Tensor,  # (..., 3, 3)
    t2: torch.Tensor,  # (..., 3)
    points: torch.Tensor,  # (..., P, 3)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    applied: (..., P, 3).
    inverse_rotation: (..., 3, 3).
    inverse_translation: (..., 3).
    composed_rotation: (..., 3, 3).
    composed_translation: (..., 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rigid_transform_operations")

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.

Sources and further reading

aqlaboratory/openfold · rigid_utils.py

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