src/datasets/pdb_dataset.py

"""Prepare RNA coordinates from a Biotite ``AtomArray`` for RiboSphere.

The output atom-channel order matches ``get_backbone_coords()``:

    P, C5', C4', C3', C2', C1', O5', O4', O3', O2', N9/N1

The final channel is N9 for purines (A/G) and N1 for pyrimidines (C/U).
Coordinates read from a PDB file remain in Angstrom.
"""

from __future__ import annotations

from typing import Any, Iterator

import numpy as np
import torch


COMMON_BACKBONE_ATOMS = (
    "P",
    "C5'",
    "C4'",
    "C3'",
    "C2'",
    "C1'",
    "O5'",
    "O4'",
    "O3'",
    "O2'",
)

BASE_ANCHOR_ATOM = {
    "A": "N9",
    "G": "N9",
    "C": "N1",
    "U": "N1",
}

RIBOSPHERE_A11_CHANNELS = COMMON_BACKBONE_ATOMS + ("N9/N1",)

ATOM_SELECTION = {
    "a1": (COMMON_BACKBONE_ATOMS.index("C4'"),),
    "a10": tuple(range(len(COMMON_BACKBONE_ATOMS))),
    "a11": tuple(range(len(COMMON_BACKBONE_ATOMS) + 1)),
}

STANDARD_RNA_RESIDUES = frozenset(BASE_ANCHOR_ATOM)


def _residue_slices(atom_array: Any) -> Iterator[slice]:
    """Yield contiguous slices for individual residues.

    Residues are distinguished by chain ID, residue ID, and insertion code.
    The input must be a single-model ``AtomArray`` whose atoms are ordered
    contiguously by residue.
    """
    if np.asarray(atom_array.coord).ndim != 2:
        raise ValueError(
            "Expected a single-model AtomArray; call "
            "PDBFile.get_structure(model=1)"
        )

    start = 0
    num_atoms = len(atom_array)
    for index in range(1, num_atoms):
        residue_changed = (
            atom_array.chain_id[index] != atom_array.chain_id[start]
            or atom_array.res_id[index] != atom_array.res_id[start]
            or atom_array.ins_code[index] != atom_array.ins_code[start]
        )
        if residue_changed:
            yield slice(start, index)
            start = index

    if num_atoms > 0:
        yield slice(start, num_atoms)


def prepare_rna_coordinates(
    atom_array: Any,
    *,
    atoms: str = "a11",
    max_residues: int | None = 512,
) -> torch.Tensor:
    """Extract complete standard RNA residues for RiboSphere.

    Parameters
    ----------
    atom_array
        A single-model Biotite ``AtomArray``. When reading a PDB file, use
        ``PDBFile.get_structure(model=1, altloc="first")``.
    atoms
        RiboSphere representation: ``"a1"`` (C4'), ``"a10"`` (ten backbone
        atoms), or ``"a11"`` (ten backbone atoms plus N9/N1).
    max_residues
        Maximum number of retained complete residues. ``None`` keeps all
        residues.

    Returns
    -------
    torch.Tensor
        Centered float32 coordinates with shape ``[1, L, A, 3]``. Coordinates
        remain in Angstrom. Residues with non-standard names or missing
        required A11 atoms are discarded.
    """
    if atoms not in ATOM_SELECTION:
        supported = ", ".join(sorted(ATOM_SELECTION))
        raise ValueError(
            f"Unsupported atom representation {atoms!r}; expected one of "
            f"{supported}"
        )
    if max_residues is not None and max_residues <= 0:
        raise ValueError("max_residues must be positive or None")

    complete_residues: list[np.ndarray] = []

    for residue_slice in _residue_slices(atom_array):
        residue = atom_array[residue_slice]
        residue_name = str(residue.res_name[0]).strip().upper()
        if residue_name not in STANDARD_RNA_RESIDUES:
            continue

        # PDBFile.get_structure(..., altloc="first") should already resolve
        # alternate locations. setdefault() makes the first remaining atom with
        # a given name deterministic if duplicates are nevertheless present.
        coordinates_by_name: dict[str, np.ndarray] = {}
        for atom_name, coordinate, is_hetero in zip(
            residue.atom_name,
            residue.coord,
            residue.hetero,
        ):
            if bool(is_hetero):
                continue
            normalized_name = str(atom_name).strip()
            coordinates_by_name.setdefault(
                normalized_name,
                np.asarray(coordinate),
            )

        base_anchor = BASE_ANCHOR_ATOM[residue_name]
        required_atoms = COMMON_BACKBONE_ATOMS + (base_anchor,)
        if not all(
            atom_name in coordinates_by_name
            for atom_name in required_atoms
        ):
            continue

        complete_residues.append(
            np.stack(
                [
                    coordinates_by_name[atom_name]
                    for atom_name in required_atoms
                ],
                axis=0,
            )
        )

    if not complete_residues:
        raise ValueError("No complete standard RNA residues were found")

    coordinates = torch.from_numpy(
        np.stack(complete_residues, axis=0).astype(
            np.float32,
            copy=False,
        )
    )

    if max_residues is not None:
        coordinates = coordinates[:max_residues]

    # Select the representation before centering so the returned A1/A10/A11
    # coordinates are each independently centered.
    atom_indices = torch.tensor(
        ATOM_SELECTION[atoms],
        dtype=torch.long,
        device=coordinates.device,
    )
    coordinates = coordinates.index_select(1, atom_indices)
    coordinates = coordinates - coordinates.mean(
        dim=(0, 1),
        keepdim=True,
    )

    return coordinates.unsqueeze(0)


def kabsch_rmsd(predicted: object, target: object) -> float:
    """Return translation- and rotation-aligned RMSD.

    Both inputs must contain the same atoms in the same order and end in an
    xyz dimension. The result uses the same length unit as the inputs.
    """
    predicted_tensor = torch.as_tensor(
        predicted,
        dtype=torch.float64,
    ).reshape(-1, 3)
    target_tensor = torch.as_tensor(
        target,
        dtype=torch.float64,
        device=predicted_tensor.device,
    ).reshape(-1, 3)

    if predicted_tensor.shape != target_tensor.shape:
        raise ValueError(
            "Coordinate shapes differ: "
            f"{predicted_tensor.shape} vs {target_tensor.shape}"
        )
    if predicted_tensor.shape[0] == 0:
        raise ValueError("Cannot compute RMSD for empty coordinates")

    predicted_centered = predicted_tensor - predicted_tensor.mean(
        dim=0,
        keepdim=True,
    )
    target_centered = target_tensor - target_tensor.mean(
        dim=0,
        keepdim=True,
    )

    covariance = predicted_centered.T @ target_centered
    left_vectors, _, right_vectors_transposed = torch.linalg.svd(covariance)

    correction = torch.eye(
        3,
        dtype=predicted_centered.dtype,
        device=predicted_centered.device,
    )
    if torch.det(left_vectors @ right_vectors_transposed) < 0:
        correction[-1, -1] = -1

    rotation = (
        left_vectors
        @ correction
        @ right_vectors_transposed
    )
    aligned = predicted_centered @ rotation

    return torch.sqrt(
        ((aligned - target_centered) ** 2).sum(dim=-1).mean()
    ).item()