torch-structure-manipulation 0.6.0__py3-none-any.whl

torch_structure_manipulation/__init__.py

@@ -0,0 +1,54 @@
+"""Atomic structure data, bonding annotations, and structure transforms."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .atomic_structure import AtomicStructure
+from .bonding import (
+    annotate_bonding_environments,
+    classify_structure_composition,
+    get_scattering_provider_keys,
+)
+from .structure_transforms import (
+    apply_rotation,
+    apply_rotation_to_coords,
+    apply_translation,
+    apply_translation_to_coords,
+    ball_query_atoms,
+    calculate_center_from_tensors,
+    center_structure,
+    center_structure_from_coords,
+    df_to_atomxyz,
+    df_to_atomzyx,
+    find_atoms_in_ball,
+    get_nucleic_acid_residues,
+    get_protein_residues,
+    remove_sidechains,
+    separate_protein_rna,
+)
+
+try:
+    __version__ = version("torch-structure-manipulation")
+except PackageNotFoundError:
+    __version__ = "uninstalled"
+
+__all__ = [
+    "AtomicStructure",
+    "annotate_bonding_environments",
+    "apply_rotation",
+    "apply_rotation_to_coords",
+    "apply_translation",
+    "apply_translation_to_coords",
+    "ball_query_atoms",
+    "calculate_center_from_tensors",
+    "center_structure",
+    "center_structure_from_coords",
+    "classify_structure_composition",
+    "df_to_atomxyz",
+    "df_to_atomzyx",
+    "find_atoms_in_ball",
+    "get_nucleic_acid_residues",
+    "get_protein_residues",
+    "get_scattering_provider_keys",
+    "remove_sidechains",
+    "separate_protein_rna",
+]

torch_structure_manipulation/atomic_structure.py

@@ -0,0 +1,228 @@
+"""Tensor-backed atomic structure data."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
+
+import gemmi
+import torch
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+@dataclass(frozen=True, slots=True)
+class AtomicStructure:
+    """Lightweight numerical representation of an atomic structure.
+
+    Positions use array-friendly ``(z, y, x)`` order. Numerical fields may have
+    arbitrary broadcast-compatible batch dimensions before their atom dimension.
+    Text metadata is shared across batches and held in immutable tuples.
+
+    Bonding metadata (``bonded_environments``, ``molecule_types``) is not batched:
+    one tuple is shared by every batch member. That matches ensemble-of-poses use
+    cases (same chemistry, different coordinates) but not batched structures with
+    different chemistry. See :meth:`from_annotated_dataframe` and the bonded-factor
+    notes on
+    :func:`torch_calculate_electrostatic_potential.potential_from_structure_3d`.
+    """
+
+    positions_zyx: torch.Tensor
+    atomic_numbers: torch.Tensor
+    elements: tuple[str, ...]
+    atom_names: tuple[str, ...]
+    b_factors: torch.Tensor
+    occupancies: torch.Tensor
+    bonded_environments: tuple[str, ...] | None = None
+    molecule_types: tuple[str, ...] | None = None
+
+    def __post_init__(self) -> None:
+        """Validate atom dimensions and batch broadcasting."""
+        if self.positions_zyx.ndim < 2 or self.positions_zyx.shape[-1] != 3:
+            raise ValueError("positions_zyx must have shape (..., n_atoms, 3)")
+        n_atoms = self.positions_zyx.shape[-2]
+        if self.atomic_numbers.ndim < 1 or self.atomic_numbers.shape[-1] != n_atoms:
+            raise ValueError(
+                "atomic_numbers must have shape (..., n_atoms) with the same "
+                "number of atoms as positions_zyx"
+            )
+        if len(self.elements) != n_atoms or len(self.atom_names) != n_atoms:
+            raise ValueError("text metadata must have one value per atom")
+        numerical_fields = {
+            "b_factors": self.b_factors,
+            "occupancies": self.occupancies,
+        }
+        for name, value in numerical_fields.items():
+            if value.ndim > 0 and value.shape[-1] != n_atoms:
+                raise ValueError(f"{name} must be scalar or have shape (..., n_atoms)")
+        batch_shapes = [
+            self.positions_zyx.shape[:-2],
+            self.atomic_numbers.shape[:-1],
+            *[
+                value.shape[:-1] if value.ndim > 0 else ()
+                for value in numerical_fields.values()
+            ],
+        ]
+        try:
+            torch.broadcast_shapes(*batch_shapes)  # type: ignore[no-untyped-call]
+        except RuntimeError as error:
+            raise ValueError(
+                "numerical AtomicStructure fields have incompatible batch shapes"
+            ) from error
+        if (
+            self.bonded_environments is not None
+            and len(self.bonded_environments) != n_atoms
+        ):
+            raise ValueError("bonded_environments must have one value per atom")
+        if self.molecule_types is not None and len(self.molecule_types) != n_atoms:
+            raise ValueError("molecule_types must have one value per atom")
+
+    @classmethod
+    def from_dataframe(
+        cls,
+        df: pd.DataFrame,
+        *,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float32,
+    ) -> AtomicStructure:
+        """Construct from an mmdf-compatible DataFrame.
+
+        Required columns are ``x``, ``y``, ``z``, and ``element``. Atom names
+        come from ``atom`` when present. ``b_isotropic`` and ``occupancy``
+        default to zero and one, respectively. Optional ``bonded_environments``
+        and ``molecule_type`` columns are preserved when present.
+        """
+        required = {"x", "y", "z", "element"}
+        missing = sorted(required.difference(df.columns))
+        if missing:
+            raise ValueError(f"missing required structure columns: {missing}")
+
+        elements = tuple(str(value).strip().upper() for value in df["element"])
+        if "atomic_number" in df:
+            atomic_number_values = [int(value) for value in df["atomic_number"]]
+        else:
+            atomic_number_values = [
+                gemmi.Element(element).atomic_number for element in elements
+            ]
+        unknown = sorted(
+            element
+            for element, atomic_number in zip(
+                elements, atomic_number_values, strict=True
+            )
+            if atomic_number == 0
+        )
+        if unknown:
+            raise ValueError(f"unknown element symbols: {unknown}")
+
+        positions = torch.as_tensor(
+            df.loc[:, ["z", "y", "x"]].to_numpy(copy=True),
+            dtype=dtype,
+            device=device,
+        )
+        atomic_numbers = torch.tensor(
+            atomic_number_values,
+            dtype=torch.int64,
+            device=device,
+        )
+        b_values = df["b_isotropic"] if "b_isotropic" in df else [0.0] * len(df)
+        occupancy_values = df["occupancy"] if "occupancy" in df else [1.0] * len(df)
+        atom_names = (
+            tuple(str(value).strip() for value in df["atom"])
+            if "atom" in df
+            else ("",) * len(df)
+        )
+        bonded = (
+            tuple(str(value) for value in df["bonded_environments"])
+            if "bonded_environments" in df
+            else None
+        )
+        molecule_types = (
+            tuple(str(value) for value in df["molecule_type"])
+            if "molecule_type" in df
+            else None
+        )
+        return cls(
+            positions_zyx=positions,
+            atomic_numbers=atomic_numbers,
+            elements=elements,
+            atom_names=atom_names,
+            b_factors=torch.as_tensor(b_values, dtype=dtype, device=device),
+            occupancies=torch.as_tensor(occupancy_values, dtype=dtype, device=device),
+            bonded_environments=bonded,
+            molecule_types=molecule_types,
+        )
+
+    @classmethod
+    def from_annotated_dataframe(
+        cls,
+        df: pd.DataFrame,
+        *,
+        include_hydrogens: bool = True,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float32,
+    ) -> AtomicStructure:
+        """Annotate bonding metadata, then construct from the result.
+
+        This is the usual entry point for Peng bonded scattering factors: it
+        calls :func:`~torch_structure_manipulation.annotate_bonding_environments`
+        to add ``bonded_environments`` and ``molecule_type`` columns, then
+        delegates to :meth:`from_dataframe`.
+
+        The input must include ``chain``, ``residue_id``, ``residue``, ``atom``,
+        and ``element`` in addition to the coordinate columns required by
+        :meth:`from_dataframe`.
+        """
+        from .bonding import annotate_bonding_environments
+
+        annotated = annotate_bonding_environments(
+            df, include_hydrogens=include_hydrogens
+        )
+        return cls.from_dataframe(annotated, device=device, dtype=dtype)
+
+    @property
+    def num_atoms(self) -> int:
+        """Number of atoms in each structure."""
+        return self.positions_zyx.shape[-2]
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        """Broadcasted batch shape of all numerical fields."""
+        batch_shapes = [
+            self.positions_zyx.shape[:-2],
+            self.atomic_numbers.shape[:-1],
+            self.b_factors.shape[:-1] if self.b_factors.ndim > 0 else (),
+            self.occupancies.shape[:-1] if self.occupancies.ndim > 0 else (),
+        ]
+        return torch.Size(
+            torch.broadcast_shapes(*batch_shapes)  # type: ignore[no-untyped-call]
+        )
+
+    @property
+    def device(self) -> torch.device:
+        """Device containing the atomic positions."""
+        return self.positions_zyx.device
+
+    def with_positions(self, positions_zyx: torch.Tensor) -> AtomicStructure:
+        """Return a copy with replacement, broadcast-compatible positions."""
+        if positions_zyx.ndim < 2 or positions_zyx.shape[-2:] != (self.num_atoms, 3):
+            raise ValueError(
+                "replacement positions must have shape (..., n_atoms, 3) with the "
+                "same number of atoms"
+            )
+        return replace(self, positions_zyx=positions_zyx)
+
+    def to(
+        self,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> AtomicStructure:
+        """Return a copy with numerical tensors moved to a device and dtype."""
+        floating_dtype = self.positions_zyx.dtype if dtype is None else dtype
+        return replace(
+            self,
+            positions_zyx=self.positions_zyx.to(device=device, dtype=floating_dtype),
+            atomic_numbers=self.atomic_numbers.to(device=device),
+            b_factors=self.b_factors.to(device=device, dtype=floating_dtype),
+            occupancies=self.occupancies.to(device=device, dtype=floating_dtype),
+        )

torch_structure_manipulation/bonding.py

@@ -0,0 +1,260 @@
+"""Annotate atomic bonding environments from packaged residue templates."""
+
+from __future__ import annotations
+
+import json
+from importlib.resources import files
+from itertools import pairwise
+from typing import Any
+
+import pandas as pd
+
+from .structure_transforms import (
+    get_nucleic_acid_residues,
+    get_protein_residues,
+)
+
+_REQUIRED_COLUMNS = {"atom", "chain", "element", "residue", "residue_id"}
+NUCLEIC_ACID_RESIDUES = get_nucleic_acid_residues()
+PROTEIN_RESIDUES = get_protein_residues()
+
+
+def _load_templates() -> tuple[
+    dict[str, dict[str, list[str]]], dict[str, dict[str, list[str]]]
+]:
+    resource = files(__package__).joinpath("bonding_data.json")
+    data: dict[str, Any] = json.loads(resource.read_text(encoding="utf-8"))
+    return data["protein"], data["rna"]
+
+
+_PROTEIN_BONDING, _RNA_BONDING = _load_templates()
+
+
+def annotate_bonding_environments(
+    df: pd.DataFrame, include_hydrogens: bool = True
+) -> pd.DataFrame:
+    """Return a copy with ``bonded_environments`` and per-atom ``molecule_type``.
+
+    The input follows mmdf naming conventions but is not tied to mmdf itself.
+    Bonding templates are keyed by residue and atom names; adjacent numeric
+    residue IDs are used to add peptide and phosphodiester bonds.
+    """
+    missing = sorted(_REQUIRED_COLUMNS.difference(df.columns))
+    if missing:
+        raise ValueError(f"missing required structure columns: {missing}")
+
+    result = df.copy()
+    if result.empty:
+        result["bonded_environments"] = pd.Series(dtype="object", index=result.index)
+        result["molecule_type"] = pd.Series(dtype="object", index=result.index)
+        return result
+
+    residues = [str(value).strip().upper() for value in result["residue"]]
+    atoms = [str(value).strip().upper() for value in result["atom"]]
+    elements = [str(value).strip().upper() for value in result["element"]]
+    residue_ids = [str(value) for value in result["residue_id"]]
+    chains = [str(value) for value in result["chain"]]
+
+    residue_lookup: dict[tuple[str, str], dict[str, str]] = {}
+    residue_names: dict[tuple[str, str], str] = {}
+    for residue, atom, element, residue_id, chain in zip(
+        residues, atoms, elements, residue_ids, chains, strict=True
+    ):
+        key = (chain, residue_id)
+        residue_lookup.setdefault(key, {})[atom] = element
+        residue_names.setdefault(key, residue)
+
+    next_residue, previous_residue = _build_residue_order(residue_names)
+    molecule_types = [_molecule_type(residue) for residue in residues]
+    environments = [
+        _environment_for_atom(
+            residue=residue,
+            atom=atom,
+            element=element,
+            key=(chain, residue_id),
+            residue_lookup=residue_lookup,
+            next_residue=next_residue,
+            previous_residue=previous_residue,
+            include_hydrogens=include_hydrogens,
+        )
+        for residue, atom, element, residue_id, chain in zip(
+            residues, atoms, elements, residue_ids, chains, strict=True
+        )
+    ]
+    result["bonded_environments"] = environments
+    result["molecule_type"] = molecule_types
+    return result
+
+
+def classify_structure_composition(df: pd.DataFrame) -> str:
+    """Summarize residue composition for the whole table.
+
+    Returns an aggregate label such as ``"protein"``, ``"rna"``,
+    ``"rna+protein"``, or ``"other"``. This is descriptive metadata only and
+    is not a valid per-atom scattering-provider key.
+    """
+    residues = {str(value).strip().upper() for value in df["residue"]}
+    has_protein = bool(residues & PROTEIN_RESIDUES)
+    has_rna = bool(residues & NUCLEIC_ACID_RESIDUES)
+    if has_protein and has_rna:
+        return "rna+protein"
+    if has_protein:
+        return "protein"
+    if has_rna:
+        return "rna"
+    return "other"
+
+
+def get_scattering_provider_keys(df: pd.DataFrame) -> list[str]:
+    """Return the Peng scattering-provider key for each atom.
+
+    Each value is ``"protein"``, ``"rna"``, or ``"other"`` and matches the
+    ``molecule_type`` column written by :func:`annotate_bonding_environments`.
+    """
+    return [_molecule_type(str(residue).strip().upper()) for residue in df["residue"]]
+
+
+def _environment_for_atom(
+    *,
+    residue: str,
+    atom: str,
+    element: str,
+    key: tuple[str, str],
+    residue_lookup: dict[tuple[str, str], dict[str, str]],
+    next_residue: dict[tuple[str, str], tuple[str, str]],
+    previous_residue: dict[tuple[str, str], tuple[str, str]],
+    include_hydrogens: bool,
+) -> str:
+    template = _template_for(residue)
+    bonded_names = list(template.get(atom, ()))
+    if residue in _RNA_BONDING and atom == "O3'" and key in next_residue:
+        bonded_names = [name for name in bonded_names if name != "HO3'"]
+
+    bonded_elements: list[str] = []
+    residue_atoms = residue_lookup.get(key, {})
+    for bonded_name in bonded_names:
+        bonded_element = _find_element(bonded_name, residue_atoms)
+        if bonded_element:
+            if include_hydrogens or bonded_element != "H":
+                bonded_elements.append(bonded_element)
+        elif include_hydrogens and _is_hydrogen_name(bonded_name):
+            bonded_elements.append("H")
+
+    if residue in _PROTEIN_BONDING and atom == "C" and key in next_residue:
+        _append_atom_element(
+            bonded_elements, "N", residue_lookup[next_residue[key]], include_hydrogens
+        )
+    elif residue in _PROTEIN_BONDING and atom == "N" and key in previous_residue:
+        _append_atom_element(
+            bonded_elements,
+            "C",
+            residue_lookup[previous_residue[key]],
+            include_hydrogens,
+        )
+    elif residue in _RNA_BONDING and atom == "O3'" and key in next_residue:
+        _append_atom_element(
+            bonded_elements, "P", residue_lookup[next_residue[key]], include_hydrogens
+        )
+    elif residue in _RNA_BONDING and atom == "P" and key in previous_residue:
+        _append_atom_element(
+            bonded_elements,
+            "O3'",
+            residue_lookup[previous_residue[key]],
+            include_hydrogens,
+        )
+
+    bonded_key = "".join(sorted(bonded_elements))
+    category = _oxygen_carbon_category(
+        residue, atom, element, bonded_key, key, residue_lookup, next_residue
+    )
+    suffix = f", {category}" if category is not None else ""
+    return f"{element}({bonded_key}{suffix})"
+
+
+def _oxygen_carbon_category(
+    residue: str,
+    atom: str,
+    element: str,
+    bonded_key: str,
+    key: tuple[str, str],
+    residue_lookup: dict[tuple[str, str], dict[str, str]],
+    next_residue: dict[tuple[str, str], tuple[str, str]],
+) -> str | None:
+    if element != "O" or bonded_key != "C":
+        return None
+    if atom == "OXT" or (residue == "ASP" and atom in {"OD1", "OD2"}):
+        return "carboxyl"
+    if residue == "GLU" and atom in {"OE1", "OE2"}:
+        return "carboxyl"
+    if (
+        residue in _PROTEIN_BONDING
+        and atom == "O"
+        and key in next_residue
+        and _find_element("N", residue_lookup[next_residue[key]])
+    ):
+        return "amide"
+    return None
+
+
+def _build_residue_order(
+    residue_names: dict[tuple[str, str], str],
+) -> tuple[
+    dict[tuple[str, str], tuple[str, str]],
+    dict[tuple[str, str], tuple[str, str]],
+]:
+    by_chain: dict[str, list[tuple[float, tuple[str, str]]]] = {}
+    for key in residue_names:
+        try:
+            numeric_id = float(key[1])
+        except ValueError:
+            continue
+        by_chain.setdefault(key[0], []).append((numeric_id, key))
+
+    next_residue: dict[tuple[str, str], tuple[str, str]] = {}
+    previous_residue: dict[tuple[str, str], tuple[str, str]] = {}
+    for residues in by_chain.values():
+        residues.sort(key=lambda item: item[0])
+        for (left_id, left_key), (right_id, right_key) in pairwise(residues):
+            if abs(right_id - left_id - 1.0) < 1e-6:
+                next_residue[left_key] = right_key
+                previous_residue[right_key] = left_key
+    return next_residue, previous_residue
+
+
+def _template_for(residue: str) -> dict[str, list[str]]:
+    if residue in _PROTEIN_BONDING:
+        return _PROTEIN_BONDING[residue]
+    return _RNA_BONDING.get(residue, {})
+
+
+def _molecule_type(residue: str) -> str:
+    if residue in PROTEIN_RESIDUES:
+        return "protein"
+    if residue in NUCLEIC_ACID_RESIDUES:
+        return "rna"
+    return "other"
+
+
+def _append_atom_element(
+    elements: list[str],
+    atom_name: str,
+    residue_atoms: dict[str, str],
+    include_hydrogens: bool,
+) -> None:
+    element = _find_element(atom_name, residue_atoms)
+    if element and (include_hydrogens or element != "H"):
+        elements.append(element)
+
+
+def _find_element(atom_name: str, residue_atoms: dict[str, str]) -> str:
+    if atom_name in residue_atoms:
+        return residue_atoms[atom_name]
+    normalized = atom_name.replace("'", "").replace("*", "")
+    for candidate, element in residue_atoms.items():
+        if candidate.replace("'", "").replace("*", "") == normalized:
+            return element
+    return ""
+
+
+def _is_hydrogen_name(atom_name: str) -> bool:
+    return atom_name.upper().startswith("H")