Geometry and structural losses · INTEGRATION

Frame-aligned point error (FAPE)

HardRaw PyTorchAbout 30 min2 public examples plus private tests

AlphaFold2AlphaFold-MultimerOpenFold

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

Your task

Measure how similarly predicted and target points sit relative to their own local frames. Express every point in every frame, compare the corresponding local coordinates, and average the scaled distances over valid frame–point pairs. This compares local geometry without first aligning the two structures globally; return one loss per batch item.

Function signature

frame_aligned_point_error(
    pred_R: torch.Tensor,
    pred_t: torch.Tensor,
    target_R: torch.Tensor,
    target_t: torch.Tensor,
    pred_positions: torch.Tensor,
    target_positions: torch.Tensor,
    masks: dict,
    eps_distance: float = 1e-08,
    eps_count: float = 1e-08,
    clamp_distance: None | float = None,
    length_scale: float = 10.0,
) -> torch.Tensor

Worked example

Compare a point in two different frames

Use B=F=P=1, with all three masks True. The predicted frame has rotation Rz below, origin [1,0,0], and world point [1,2,0]. The target frame has identity rotation, origin [0,0,0], and world point [1,0,0].

Rz = [[0,-1,0], [1,0,0], [0,0,1]]
eps_distance=3; eps_count=1
length_scale=2; clamp_distance=None

Local coordinates. The predicted displacement is [0,2,0], so Rzᵀ[0,2,0]=[2,0,0] in its local frame. The target local point is [1,0,0].
Pair distance. Their squared local error is 1, hence the stabilized distance is sqrt(1+3)=2. Divide by length_scale to get 1.
Batch loss. Divide by one valid pair plus eps_count: loss=[1/(1+1)]=[0.5]. The world-space point difference is not the error being averaged.

Parameters

Parameter	Shape and meaning
`pred_R`torch.Tensor	$[B, F, 3, 3]$ Predicted rotations whose columns are local basis axes.
`pred_t`torch.Tensor	$[B, F, 3]$ Predicted frame origins.
`target_R`torch.Tensor	$[B, F, 3, 3]$ Target rotations whose columns are local basis axes.
`target_t`torch.Tensor	$[B, F, 3]$ Target frame origins in global coordinates.
`pred_positions`torch.Tensor	$[B, P, 3]$ Predicted points in global coordinates.
`target_positions`torch.Tensor	$[B, P, 3]$ Target points in global coordinates.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`eps_distance`floatDefault `1e-08`	Strictly positive squared-distance offset inside $\sqrt{\sum e_{\mathrm{local}}^2+\varepsilon_{\mathrm{distance}}}$ .
`eps_count`floatDefault `1e-08`	Strictly positive offset added to the valid frame–point pair count after summation.
`clamp_distance`None \| floatDefault `None`	Optional positive distance cap applied after the distance square root and before length scaling. None means no clamping.
`length_scale`floatDefault `10.0`	Positive coordinate-unit scale. Divide each distance by this value before taking the per-example mean.

Return

torch.Tensor · $[B]$

Return Floating [B] masked, scaled FAPE; no reduction across batch.

Required behavior

Frames and points are separate collections: rotations have shape $(B,F,3,3)$ , origins $(B,F,3)$ , and points $(B,P,3)$ . Rotation columns are local basis axes.
Express each point in each frame using $u_{bfp}=R_{bf}^{\top}(x_{bp}-t_{bf})$ , separately for prediction and target. The local-coordinate arrays have shape $(B,F,P,3)$ .
Compute $d_{bfp}=\sqrt{\sum_c(u^{pred}_{bfpc}-u^{target}_{bfpc})^2+\varepsilon_{distance}}$ . If supplied, clamp this distance to clamp_distance, then divide by length_scale.
A pair is valid only when its frame, point, and pair masks are all True. Use this same combined mask for the error sum and pair count.
Return $\mathrm{loss}_b=\frac{\sum_{f,p}m_{bfp}d_{bfp}/\mathrm{length\_scale}}{\sum_{f,p}m_{bfp}+\varepsilon_{count}}$ , with shape $(B,)$ . Do not average over the batch.
length_scale, both epsilons, and an optional clamp distance are strictly positive. The distance epsilon is inside the square root; the count epsilon is added after counting valid pairs.
Perfectly matching valid geometry has a positive smoothing floor. Do not subtract it. All-invalid examples or empty frame/point axes return exact zero with finite zero gradients.
Gradients are checked through all six floating coordinate/frame tensors, including target tensors and rotation matrix entries. Frames are supplied; do not orthogonalize them or reconstruct them.
Preserve dtype/device and input values; accept arbitrary positive input strides. The full numerical domain and clamp behavior are the same for predicted and target geometry.

Why this matters

Each residue frame supplies its own coordinate system. Comparing local coordinates measures structural agreement without first choosing one global protein alignment; the numbers of frames and points may differ.

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 frame_aligned_point_error(
    pred_R: torch.Tensor,  # (B, F, 3, 3)
    pred_t: torch.Tensor,  # (B, F, 3)
    target_R: torch.Tensor,  # (B, F, 3, 3)
    target_t: torch.Tensor,  # (B, F, 3)
    pred_positions: torch.Tensor,  # (B, P, 3)
    target_positions: torch.Tensor,  # (B, P, 3)
    masks: dict,
    eps_distance: float = 1e-08,
    eps_count: float = 1e-08,
    clamp_distance: None | float = None,
    length_scale: float = 10.0,
) -> torch.Tensor:
    """Return a tensor with shape (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement frame_aligned_point_error")

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.