fastplms/models/esmfold2/esmfold2_output.py

"""Convert ESMFold2 coordinate tensors into molecular-complex records."""

from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass, field
from itertools import groupby
from typing import Any

import numpy as np
import torch

from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER
from .esmfold2_molecular_complex import MolecularComplex, MolecularComplexMetadata


def get_element_symbol(atomic_number: int) -> str:
    """Map a training-time atomic number to a chemical symbol."""

    return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_number, "X")


def _decode_atom_name(encoded_name: Any) -> str:
    values = encoded_name.tolist() if hasattr(encoded_name, "tolist") else encoded_name
    return "".join(chr(int(value) + 32) for value in values if int(value)).strip()


@dataclass
class _ComplexRecords:
    sequence: list[str] = field(default_factory=list)
    chain_ids: list[int] = field(default_factory=list)
    token_to_atoms: list[list[int]] = field(default_factory=list)
    confidence: list[float] = field(default_factory=list)
    positions: list[list[float]] = field(default_factory=list)
    elements: list[str] = field(default_factory=list)
    atom_names: list[str] = field(default_factory=list)
    atom_hetero: list[bool] = field(default_factory=list)
    chain_lookup: dict[int, str] = field(default_factory=dict)
    entity_lookup: dict[int, str] = field(default_factory=dict)

    def add_token(
        self,
        *,
        residue_name: str,
        asym_id: int,
        plddt: float,
        atoms: Iterable[tuple[list[float], str, str]],
        hetero: bool,
    ) -> None:
        atom_start = len(self.positions)
        for position, element, atom_name in atoms:
            self.positions.append(position)
            self.elements.append(element)
            self.atom_names.append(atom_name)
            self.atom_hetero.append(hetero)
        self.sequence.append(residue_name)
        self.chain_ids.append(asym_id)
        self.confidence.append(plddt)
        self.token_to_atoms.append([atom_start, len(self.positions)])

    def build(self, complex_id: str) -> MolecularComplex:
        return MolecularComplex(
            id=complex_id,
            sequence=self.sequence,
            atom_positions=np.asarray(self.positions, dtype=np.float32).reshape(-1, 3),
            atom_elements=np.asarray(self.elements, dtype=object),
            token_to_atoms=np.asarray(self.token_to_atoms, dtype=np.int32).reshape(-1, 2),
            chain_id=np.asarray(self.chain_ids, dtype=np.int64),
            plddt=np.asarray(self.confidence, dtype=np.float32),
            atom_names=np.asarray(self.atom_names, dtype=object),
            atom_hetero=np.asarray(self.atom_hetero, dtype=bool),
            metadata=MolecularComplexMetadata(
                entity_lookup=self.entity_lookup,
                chain_lookup=self.chain_lookup,
                assembly_composition=None,
            ),
        )


def build_molecular_complex_from_features(
    coords: torch.Tensor,
    plddt: torch.Tensor,
    atom_mask: torch.Tensor,
    ref_element: torch.Tensor,
    ref_atom_name_chars: torch.Tensor,
    chain_infos: list[Any],
    complex_id: str,
) -> MolecularComplex:
    """Decode model features into one complex without intermediate structure files.

    Protein, DNA, and RNA tokens are grouped by residue index. Ligand atom
    tokens are collapsed into one non-polymer residue per chain.
    """

    M = atom_mask.bool().cpu().numpy()
    X = coords.float().cpu().numpy()
    atom_names = ref_atom_name_chars.cpu().numpy()
    elements = ref_element.cpu().numpy()
    confidence = plddt.float().cpu().numpy()
    records = _ComplexRecords()

    def decode_atoms(tokens: Iterable[Any]):
        for token in tokens:
            for atom_index in range(token.atom_start, token.atom_start + token.atom_count):
                if M[atom_index]:
                    yield (
                        X[atom_index].tolist(),
                        get_element_symbol(int(elements[atom_index])),
                        _decode_atom_name(atom_names[atom_index]),
                    )

    for chain in chain_infos:
        is_nonpolymer = chain.mol_type == MOL_TYPE_NONPOLYMER
        records.chain_lookup[chain.asym_id] = chain.chain_id
        records.entity_lookup[chain.entity_id] = "non-polymer" if is_nonpolymer else "polymer"

        if is_nonpolymer:
            mean_confidence = (
                float(np.mean([confidence[token.token_index] for token in chain.tokens]))
                if chain.tokens
                else 0.0
            )
            records.add_token(
                residue_name=chain.tokens[0].residue_name if chain.tokens else "LIG",
                asym_id=chain.asym_id,
                plddt=mean_confidence,
                atoms=decode_atoms(chain.tokens),
                hetero=True,
            )
            continue

        residue_groups = groupby(chain.tokens, key=lambda token: token.residue_index)
        for _residue_index, group in residue_groups:
            residue_tokens = list(group)
            records.add_token(
                residue_name=residue_tokens[0].residue_name,
                asym_id=chain.asym_id,
                plddt=float(np.mean([confidence[token.token_index] for token in residue_tokens])),
                atoms=decode_atoms(residue_tokens),
                hetero=False,
            )

    return records.build(complex_id)


def build_molecular_complex(
    structure: Any,
    coords: torch.Tensor,
    plddt: torch.Tensor,
    complex_id: str,
) -> MolecularComplex:
    """Decode coordinates using the atom and residue arrays of a prepared structure."""

    records = _ComplexRecords()
    coordinate_index = 0
    confidence_index = 0

    for chain in structure.chains:
        asym_id = int(chain["asym_id"])
        mol_type = int(chain["mol_type"])
        is_nonpolymer = mol_type == MOL_TYPE_NONPOLYMER
        records.chain_lookup[asym_id] = str(chain["name"])
        records.entity_lookup[int(chain["entity_id"])] = (
            "non-polymer" if is_nonpolymer else "polymer"
        )

        residue_start = int(chain["res_idx"])
        residue_stop = residue_start + int(chain["res_num"])
        for residue in structure.residues[residue_start:residue_stop]:
            atom_start = int(residue["atom_idx"])
            atom_stop = atom_start + int(residue["atom_num"])
            decoded_atoms: list[tuple[list[float], str, str]] = []
            for atom in structure.atoms[atom_start:atom_stop]:
                if not atom["is_present"]:
                    continue
                decoded_atoms.append(
                    (
                        coords[coordinate_index].tolist(),
                        get_element_symbol(int(atom["element"].item())),
                        _decode_atom_name(atom["name"]),
                    )
                )
                coordinate_index += 1

            records.add_token(
                residue_name=str(residue["name"]),
                asym_id=asym_id,
                plddt=float(plddt[confidence_index].item()),
                atoms=decoded_atoms,
                hetero=is_nonpolymer,
            )
            confidence_index += 1

    return records.build(complex_id)


__all__ = [
    "build_molecular_complex",
    "build_molecular_complex_from_features",
    "get_element_symbol",
]