Shared rigid coordinate augmentation and Haar rotations

Your task

Function signature

rigid_coordinate_augmentation(
    coordinates: torch.Tensor,
    companion: torch.Tensor,
    rotation_seed: torch.Tensor,
    translation_noise: torch.Tensor,
    translation_scale: torch.Tensor,
    valid_mask: torch.Tensor,
    companion_mask: torch.Tensor,
    eps: float = 1e-08,
    rotation: str = 'two_vectors',
) -> dict[str, torch.Tensor]

Worked example

Rotate two streams about one center

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs coordinates,companion [B,N,3] and independent bool valid_mask,companion_mask [B,N] describe two coordinate streams. Let center be the arithmetic mean of valid primary coordinates, or 0 when none are valid. The companion uses THIS SAME center, never its own.
• \mathrm{translation}=\mathrm{translation\_scale}\cdot \mathrm{translation\_noise} with supplied translation_noise [B,3] and scalar \mathrm{translation\_scale}\ge 0.
• Construct a RIGHT-ACTING row-vector rotation Q; outputs are (\mathrm{coordinates}-\mathrm{center})\mathbin{@}Q+\mathrm{translation} and (\mathrm{companion}-\mathrm{center})\mathbin{@}Q+\mathrm{translation}, then independently zero invalid rows.
• rotation='two_vectors' (default) accepts rotation_seed [B,2,3]: e_0=a/\lVert a\rVert, b_{\perp}=b-e_0(e_0\cdot b), e_1=b_{\perp}/\lVert b_{\perp}\rVert, e_2=e_0\times e_1, Q=\mathrm{stack}(\mathrm{e0},\mathrm{e1},\mathrm{e2}) as ROWS. This matches the AF3 source's right action.
• The alternate 'quaternion' accepts rotation_seed [B,4] in (w,x,y,z) order, normalizes it and constructs the standard active quaternion matrix Q, STILL right-multiplying row coordinates as in the inspected ESMFold2 release. Do not silently transpose either source convention.
• Independent isotropic Gaussian seeds generate Haar rotations; arbitrary deterministic seeds are valid inputs but carry no random-distribution claim. No internal random draws. eps is a positive norm threshold, default 10^{-8}.
• If \Vert a\Vert \le \varepsilon or \Vert b_{\mathrm{perp}}\Vert \le \varepsilon, or quaternion \mathrm{norm}\le \varepsilon, set Q=I and rotation_valid=False; otherwise True. This total degenerate fallback is a declared teaching adaptation, not the source's clamped normalization. Avoid unsafe inactive divisions.
• All seeds, translation inputs and observed coordinates are finite; missing coordinates can contain NaN and are ignored. B\ge 1,N\ge 0, including empty/all-masked primary streams. If only the companion has valid rows it rotates about center 0. Return center,Q,translation and both transformed streams.
• The only keys are rotation and \varepsilon; unsupported modes or nonpositive \varepsilon raise ValueError.

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_coordinate_augmentation(
    coordinates: torch.Tensor,  # (B, N, 3)
    companion: torch.Tensor,  # (B, N, 3)
    rotation_seed: torch.Tensor,  # (B, 2, 3) or (B, 4)
    translation_noise: torch.Tensor,  # (B, 3)
    translation_scale: torch.Tensor,  # Scalar tensor, shape ().
    valid_mask: torch.Tensor,  # (B, N)
    companion_mask: torch.Tensor,  # (B, N)
    eps: float = 1e-08,
    rotation: str = 'two_vectors',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    coordinates_augmented: (B, N, 3).
    companion_augmented: (B, N, 3).
    centroid: (B, 3).
    rotation: (B, 3, 3).
    translation: (B, 3).
    rotation_valid: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rigid_coordinate_augmentation")

Sources and further reading

• www.nature.com
• google-deepmind/alphafold3 · diffusion_head.py
• biohub.ai
• huggingface/transformers · generation_esmfold2.py
• arxiv.org
• gcorso/DiffDock · sampling.py