Rigid Permutation Alignment

Your task

Function signature

rigid_permutation_alignment(
    pred: torch.Tensor,
    target: torch.Tensor,
    weights: torch.Tensor,
    mask: torch.Tensor,
    permutations: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Worked example

import torch

pred = torch.tensor([[[3.0, 2.0, 3.0], [1.0, 2.0, 3.0], [1.0, 5.0, 3.0], [1.0, 2.0, 7.0]]], dtype=torch.float64, requires_grad=True)
target = 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]]], dtype=torch.float64, requires_grad=True)
weights = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([[True, True, True, True]], dtype=torch.bool)
permutations = torch.tensor([[[0, 1, 2, 3], [1, 0, 2, 3]]], dtype=torch.int64)

result = rigid_permutation_alignment(pred=pred, target=target, weights=weights, mask=mask, permutations=permutations)

Expected return

{
    'candidate_cost': torch.tensor([[1.9999999999999996, 1.5284180038657104e-30]], dtype=torch.float64),
    'cost': torch.tensor([1.5284180038657104e-30], dtype=torch.float64),
    'matched_target': torch.tensor([[[3.0, 2.0, 3.0000000000000004], [0.9999999999999998, 1.9999999999999998, 3.0], [1.0000000000000007, 5.000000000000001, 2.999999999999999], [0.9999999999999996, 1.9999999999999993, 7.000000000000002]]], dtype=torch.float64),
    'rotation': torch.tensor([[[1.0, 2.7755575615628914e-16, -5.551115123125783e-17], [5.551115123125783e-17, 1.0000000000000004, -1.1102230246251565e-16], [1.3877787807814457e-16, -3.608224830031759e-16, 1.0000000000000004]]], dtype=torch.float64),
    'translation': torch.tensor([[0.9999999999999998, 1.9999999999999998, 3.0]], dtype=torch.float64),
    'permutation_index': torch.tensor([1], dtype=torch.int64),
    'permutation': torch.tensor([[1, 0, 2, 3]], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• pred/target are [B,N,3]; permutations is [B,K,N], with K positive. Each candidate maps target coordinates as target[permutation]. Supplied permutations preserve chemical admissibility and mask validity; do not search for additional permutations.
• Mask excluded coordinates before computing candidates. weights is a nonnegative per-atom weight and contributes only where mask=True.
• For each candidate, fit target to pred with weighted centroids and Kabsch SVD, correcting reflections so det(rotation)=+1. Return the weighted mean squared coordinate error, dividing by the exact positive weight sum.
• A zero-weight structure returns zero candidate cost and matched coordinates, identity rotation, and zero translation.
• Pick the first candidate with the minimum cost. Return all candidate costs, selected cost/index/permutation, and selected matched coordinates/frame. Rotation uses column-vector form: matched = target @ rotation.T + translation.
• Selection is discrete, but keep gradients through the selected candidate’s floating computation. Do not mutate coordinates or modify supplied permutations.
• 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.
• Derivative tests use separated candidate minima and nondegenerate weighted covariance singular values. Discrete selection and degenerate SVD branches do not define general smooth derivatives.

Why this matters

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_permutation_alignment(
    pred: torch.Tensor,  # (B, N, 3)
    target: torch.Tensor,  # (B, N, 3)
    weights: torch.Tensor,  # (B, N)
    mask: torch.Tensor,  # (B, N)
    permutations: torch.Tensor,  # (B, K, N)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    candidate_cost: (B, K).
    cost: (B,).
    matched_target: (B, N, 3).
    rotation: (B, 3, 3).
    translation: (B, 3).
    permutation_index: (B,).
    permutation: (B, N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rigid_permutation_alignment")

Related exercises

• Builds onWeighted proper rigid alignment
• Combined exerciseSymmetry assignment: rigid fits and ambiguous atom renaming

Sources and further reading

• www.nature.com
• www.nature.com
• www.biorxiv.org
• aqlaboratory/openfold · multi_chain_permutation.py
• aqlaboratory/openfold · multi_chain_permutation.py
• aqlaboratory/openfold · loss.py