Ambiguous Atom Renaming

Your task

Function signature

ambiguous_atom_renaming(
    pred: torch.Tensor,
    target: torch.Tensor,
    weights: torch.Tensor,
    mask: torch.Tensor,
    permutations: torch.Tensor,
    ambiguous: torch.Tensor,
    rename_group: torch.Tensor,
    distance_eps: float = 1e-10,
    error_eps: float = 1e-10,
) -> dict[str, torch.Tensor]

Worked example

Choose a swap from anchor distances

import torch

pred = torch.tensor([[[3.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
target = torch.tensor([[[1.0, 0.0, 0.0], [3.0, 0.0, 0.0], [0.0, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
weights = torch.tensor([[1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([[True, True, True]], dtype=torch.bool)
permutations = torch.tensor([[[0, 1, 2], [1, 0, 2]]], dtype=torch.int64)
ambiguous = torch.tensor([[True, True, False]], dtype=torch.bool)
rename_group = torch.tensor([[True, True, False]], dtype=torch.bool)
distance_eps = 1e-10
error_eps = 1e-10

result = ambiguous_atom_renaming(pred=pred, target=target, weights=weights, mask=mask, permutations=permutations, ambiguous=ambiguous, rename_group=rename_group, distance_eps=distance_eps, error_eps=error_eps)

Expected return

{
    'candidate_cost': torch.tensor([[3.9999999999833333, 2e-05]], dtype=torch.float64),
    'cost': torch.tensor([2e-05], dtype=torch.float64),
    'matched_target': torch.tensor([[[3.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]], dtype=torch.float64),
    'rotation': torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]], dtype=torch.float64),
    'translation': torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64),
    'permutation_index': torch.tensor([1], dtype=torch.int64),
    'permutation': torch.tensor([[1, 0, 2]], 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.
• Compute stabilized pair distances \sqrt{\|r_i-r_j\|^2+distance\_eps}, then error \sqrt{(d^{pred}_{ij}-d^{target}_{ij})^2+error\_eps}. Both epsilons default to 10^{-10} and are positive.
• Sum errors with weights_i × weights_j only where ambiguous[i] AND rename_group[i] AND NOT ambiguous[j]. This is a SUM, without dividing by weight or pair count. Do not rigidly align in this objective; rotation is identity, translation zero, and matched_target is the chosen permuted target.
• 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.

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 ambiguous_atom_renaming(
    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,  # (B, N)
    rename_group: torch.Tensor,  # (B, N)
    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 ambiguous_atom_renaming")

Related exercises

• Builds onMasked pairwise distances and hard distance-bin targets
• 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