Geometry and structural losses

Right-handed local frame from three points

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

AlphaFold2OpenFoldRoseTTAFold2RoseTTAFold All-Atom

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

Your task

Construct a right-handed coordinate frame from an origin and two supporting points. Return its rotation and translation, together with a mask identifying nondegenerate frames.

Function signature

local_frame_from_three_points(
    p_neg_x: torch.Tensor,
    origin: torch.Tensor,
    p_xy: torch.Tensor,
    mask: torch.Tensor,
    eps: float = 1e-06,
) -> dict[str, torch.Tensor]

Worked example

Build a translated quarter-turn frame

Take origin=[[1,2,0]], p_neg_x=[[1,1,0]], p_xy=[[-2,4,0]], mask=[True] and eps=1e-6. The positive local x axis is [0,1,0]. Removing the projection [0,2,0] from the third-point displacement [-3,2,0] leaves [-3,0,0], so local y is [-1,0,0] and their cross product is [0,0,1]. Store these axes as columns:

R=\begin{bmatrix}0&-1&0\\1&0&0\\0&0&1\end{bmatrix}.

Return this rotation with shape (1,3,3), t=[[1,2,0]] and frame_valid=[True]. Transposing this matrix would reverse the turn.

Parameters

Parameter	Shape and meaning
`p_neg_x`torch.Tensor	$[\ldots, 3]$ A point toward negative local $x$ .
`origin`torch.Tensor	$[\ldots, 3]$ Global location of the local frame origin.
`p_xy`torch.Tensor	$[\ldots, 3]$ Point in the positive- $y$ half of the local xy plane after projection.
`mask`torch.Tensor	$[\ldots]$ `True` means observed triplet.
`eps`floatDefault `1e-06`	Positive length threshold. Clamp each squared normalization denominator below by eps²; a frame is valid only when both defining squared lengths are strictly greater than eps².

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Inputs p_neg_x,origin,p_xy $[\ldots,3]$ have matching leading axes and finite values; mask $[\ldots]$ is True for observed triplets.
With $\varepsilon>0$ (default $10^{-6}$ ), set $a=\mathrm{origin}-p_{\mathrm{neg\_x}}$ , $\mathrm{e0}=a/\sqrt{\operatorname{max}(\sum(a^{2}),\varepsilon^{2})}$ ; $v=p_{\mathrm{xy}}-\mathrm{origin}$ ; $b=v-(v\cdot e_0)e_0$ ; $\mathrm{e1}=b/\sqrt{\operatorname{max}(\sum(b^{2}),\varepsilon^{2})}$ ; $e_2=e_0\times e_1$ .
A frame is valid exactly when mask is True AND $\sum(a^{2})>\varepsilon^{2}$ AND $\sum(b^{2})>\varepsilon^{2}$ . Return $R$ with $e_0,e_1,e_2$ as columns and $t=\mathrm{origin}$ on valid frames.
Invalid frames return $R=\mathrm{identity}$ , $t=\mathrm{zero}$ , frame_valid=False, including coincident/collinear or sub- $\epsilon$ directions. Clamp denominators before selection so invalid branches have finite gradients. The $\epsilon$ is a minimum norm, not an additive bias: valid frames are orthonormal.
Empty leading axes return empty output tensors. Equivariance is required only on valid frames; the fixed invalid fallback has no orientation.

Why this matters

One point defines the negative- $x$ direction, the origin anchors the frame, and a third point selects the $xy$ plane; a cross product fixes handedness.

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 local_frame_from_three_points(
    p_neg_x: torch.Tensor,  # (..., 3)
    origin: torch.Tensor,  # (..., 3)
    p_xy: torch.Tensor,  # (..., 3)
    mask: torch.Tensor,  # (...)
    eps: float = 1e-06,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    R: (..., 3, 3).
    t: (..., 3).
    frame_valid: (...)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement local_frame_from_three_points")

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.