Symmetry assignment: rigid fits and ambiguous atom renaming

Your task

Function signature

symmetry_permutation_alignment(
    pred: torch.Tensor,
    target: torch.Tensor,
    weights: torch.Tensor,
    mask: torch.Tensor,
    permutations: torch.Tensor,
    ambiguous: torch.Tensor | None = None,
    rename_group: torch.Tensor | None = None,
    objective: str = 'rigid_mse',
    distance_eps: float = 1e-10,
    error_eps: float = 1e-10,
) -> dict[str, torch.Tensor]

Worked example

A labeling swap changes anchor distances

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• For B\ge 1 batches, N\ge 0 points and K\ge 1 supplied admissible permutations, select the best permuted TARGET against the PREDICTION. pred,target are[B,N,3], weights[B,N], mask[B,N], permutations[B,K,N]. permutations [b,k,i] gives the target index matched to prediction slot i.
• Every candidate is a bijection of0..N-1, already checked for chemical/chain admissibility, and preserves mask: \mathrm{mask}[b,\mathrm{permutations}[b,k,i]]=\mathrm{mask}[b,i]. Thus candidates cannot win by discarding observations.
• Weights are finite positive at observed slots, arbitrary finite at masked slots; effective w=\operatorname{where}(\mathrm{mask},\mathrm{weights},0). Masked coordinates may be nonfinite and must be sanitized before arithmetic.
• objective='rigid_mse' (default): For each b,k solve \min_{R\in\mathrm{SO}(3),\,t}\frac{\sum_i w_i\|R\,\texttt{target}_{p_k(i)}+t-\texttt{pred}_i\|^2}{\sum_i w_i}. Rotation uses COLUMN vectors.
• Center the two clouds with the SAME prediction-slot weights and form H=\sum_{i} w_{i}(\mathrm{target\_p}(i)-\mu_{x})(\mathrm{pred\_i}-\mu_{y})^{\top}. The narrow torch.linalg.svd call is allowed, but implement centering, covariance, determinant/reflection correction, translation and loss yourself. Reflection is not an allowed rotation.
• matched_target is the selected rotated/translated target at observed slots and0 at masked slots; return that selected R,t. Cost has units of squared angstroms.
• objective='ambiguous_distance': additionally receive ambiguous[B,N] (all chemically ambiguous atoms, including those outside the queried group) and rename_group[B,N] (the residue/group whose naming is being chosen). Candidates may change only indices that are BOTH ambiguous and in rename_group; every other index stays fixed.
• Define \mathrm{dp\_ij}=\sqrt{\mathrm{distance\_eps}+\Vert \mathrm{pred\_i}-\mathrm{pred\_j}\Vert ^{2}} and \mathrm{dt\_kij}=\sqrt{\mathrm{distance\_eps}+\Vert \mathrm{target\_p\_k}(i)-\mathrm{target\_p\_k}(j)\Vert ^{2}}.
• Candidate cost is SUM_ij w_{i}\cdot w_{j}\cdot A_{\mathrm{ij}}\cdot \sqrt{\mathrm{error\_eps}+(\mathrm{dp\_ij}-\mathrm{dt\_kij})^{2}}, where A_{\mathrm{ij}}=\mathrm{ambiguous\_i} AND rename_group_i AND not ambiguous_j. This is a SUM, not a mean, squared error or rigid-aligned RMSD.
• Unit weights reproduce the masked ambiguous-to-unambiguous comparison in AF2 \mathrm{Algorithm26}/\mathrm{OpenFold}; extra ambiguous atoms elsewhere are not fixed anchors. Both epsilons default 10^{-10}, matching that common source convention.
• Their locations are independent: distance regularization versus absolute-error regularization. Perfect matching still has the explicit \sqrt{\mathrm{error\_eps}} floor for each selected pair. No eligible pairs gives0. This mode does not use SVD or require a nondegenerate cloud.
• matched_target is the selected renamed target without geometric fitting; returned R=\mathrm{identity},t=0. Cost has weighted angstrom-error units, not squared-distance units. Return candidate_cost[B,K] and choose the FIRST minimum candidate per batch.
• Return selected cost, matched_target[B,N,3], rotation [B,3,3], translation [B,3], permutation[B,N], permutation_index [B]. Selection is exact only within the SUPPLIED candidate set.
• It does not solve chemical graph isomorphism, unconstrained factorial chain assignment, Hungarian matching or upstream greedy-anchor candidate discovery. Do not choose independent permutations for different downstream losses.
• For rigid_mse the supported numerical domain is either no observations (including N=0), or positive observed weights with every candidate's weighted centered covariance full rank and having distinct positive singular values. Nonempty rank-deficient/repeated-singular-value fits are outside that mode's contract.
• Smooth selection fixtures have a unique winner. With no observations either mode returns candidate costs0, index0, identity R,zero t and zero/empty matched coordinates. Duplicate candidates choose the first; no derivative is claimed at a general changing-assignment boundary.
• Empty B/K, unknown objective, or nonpositive distance_eps/error_eps raise ValueError; other shape/domain violations are outside the contract. The continuous costs and selected transform remain differentiable on smooth branches, including zero VJPs for absent observations.
• Chemical tables, candidate enumeration and experimental coordinate acquisition are supplied resources.

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 symmetry_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)
    ambiguous: torch.Tensor | None = None,  # (B, N)
    rename_group: torch.Tensor | None = None,  # (B, N)
    objective: str = 'rigid_mse',
    distance_eps: float = 1e-10,
    error_eps: float = 1e-10,
) -> 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 symmetry_permutation_alignment")

Related exercises

• Builds onRigid Permutation Alignment
• Builds onAmbiguous 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