Geometry and structural losses

Weighted proper rigid alignment

HardRaw PyTorchAbout 30 min2 public examples plus private tests

OpenFold3Boltz-2ProtenixESMFold2

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

Your task

Fit a proper rotation and translation that aligns corresponding moving points to reference points under supplied weights. Return the fitted rotation, translation, and every transformed moving point.

Function signature

weighted_kabsch_alignment(
    moving: torch.Tensor,
    reference: torch.Tensor,
    weights: torch.Tensor,
    mask: torch.Tensor,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

Recover a known rigid transform

A positive $90^\circ$ z rotation maps [2, 0, 0] to [0, 2, 0]; adding [4, -2, 1] gives its reference point [4, 0, 1]. The other points follow the same transform. Recover this rotation and translation up to floating-point SVD roundoff, and apply them to every moving point.

import torch

moving = torch.tensor([[[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0], [1.0, 2.0, 3.0]]], dtype=torch.float64, requires_grad=True)
reference = torch.tensor([[[4.0, -2.0, 1.0], [4.0, 0.0, 1.0], [1.0, -2.0, 1.0], [4.0, -2.0, 5.0], [2.0, -1.0, 4.0]]], dtype=torch.float64, requires_grad=True)
weights = torch.tensor([[1.0, 2.0, 1.0, 3.0, 1.0]], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([[True, True, True, True, True]], dtype=torch.bool)
eps = 1e-08

result = weighted_kabsch_alignment(moving=moving, reference=reference, weights=weights, mask=mask, eps=eps)

Expected return

{
    'aligned': torch.tensor([[[4.0, -2.0, 0.9999999999999996], [4.0, 4.440892098500626e-16, 0.9999999999999998], [1.0, -2.000000000000001, 0.9999999999999996], [3.999999999999999, -2.0, 5.0], [1.9999999999999996, -1.0000000000000004, 4.0]]], dtype=torch.float64),
    'R': torch.tensor([[[5.551115123125783e-17, -1.0, -1.942890293094024e-16], [1.0000000000000002, -2.220446049250313e-16, -5.551115123125783e-17], [8.326672684688674e-17, 0.0, 1.0000000000000002]]], dtype=torch.float64),
    't': torch.tensor([[4.0, -2.0, 0.9999999999999996]], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`moving`torch.Tensor	$[\ldots, N, 3]$ Floating $[\ldots,N,3]$ cloud to transform.
`reference`torch.Tensor	$[\ldots, N, 3]$ Desired corresponding cloud.
`weights`torch.Tensor	$[\ldots, N]$ Floating nonnegative $[\ldots,N]$ fitting `weights`, differentiable on active smooth cases.
`mask`torch.Tensor	$[\ldots, N]$ $\mathrm{True}=\mathrm{include}$ weight in fit.
`eps`floatDefault `1e-08`	A fit is active only when its total valid weight is strictly greater than eps. This is a support threshold, not an additive distance or covariance stabilizer.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

moving,reference $[\ldots,N,3]$ , $\mathrm{weights}[\ldots,N]\ge 0$ and bool mask $[\ldots,N]$ have matching leading axes and finite values. True mask means include in FIT; use $w=\operatorname{where}(\mathrm{mask},\mathrm{weights},0)$ .
If $\sum(w)\le \varepsilon$ (positive, default $10^{-8}$ ), return $R=\mathrm{identity}$ , $t=\mathrm{zero}$ and $\mathrm{aligned}=\mathrm{moving}$ , including $N=0$ .
Otherwise use exact weighted centroids $c_x=\frac{\sum_i w_i\,\texttt{moving}_i}{\sum_iw_i}$ , $c_y=\frac{\sum_i w_i\,\texttt{reference}_i}{\sum_iw_i}$ , with no $\epsilon$ bias. Form $H=\sum_iw_i(\texttt{moving}_i-c_x)(\texttt{reference}_i-c_y)^\top$ .
Use the allowed SVD kernel $H=U\operatorname{diag}(s)V^T$ . Set $D=\operatorname{diag}(1,1,\operatorname{sign}\det(VU^T))$ , choosing +1 at zero. Then $R=VDU^T$ and $t=c_y-Rc_x$ .
Return $\mathrm{aligned\_i}=R \mathrm{moving\_i}+t$ for EVERY row, including masked rows (the mask controls fitting, not output suppression). Column-vector convention applies; stored rows therefore use $\mathrm{moving} \mathbin{@} R^{\top}$ . $R$ must have $\mathrm{determinant} +1$ .
Use a raw $3$ x3 determinant formula, not another linalg API. Positive-fit graded fixtures have full rank and separated singular values. Rank-deficient fits can have nonunique rotations: only forward $\mathrm{fit}/\operatorname{SO}(3)$ properties are meaningful there, and no gradient or exact-choice contract is imposed.
Use a nonsingular, distinct-spectrum placeholder before SVD on empty-fit branches so inactive SVD backward cannot introduce NaNs.

Why this matters

Centering removes translation. A small covariance matrix then determines the rotation that best aligns the two clouds. A determinant correction prevents a mirror reflection from being mistaken for a rotation.

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 weighted_kabsch_alignment(
    moving: torch.Tensor,  # (..., N, 3)
    reference: torch.Tensor,  # (..., N, 3)
    weights: torch.Tensor,  # (..., N)
    mask: torch.Tensor,  # (..., N)
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    aligned: (..., N, 3).
    R: (..., 3, 3).
    t: (..., 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement weighted_kabsch_alignment")

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

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