atomworks 1.0.0__py3-none-any.whl

atomworks/__init__.py

@@ -0,0 +1,48 @@
+"""
+atomworks - Unified package for biological data I/O and machine learning.
+
+This package combines functionality from atomworks.io (I/O operations) and atomworks.ml (ML utilities)
+into a unified interface for biological data processing and machine learning.
+"""
+
+import importlib
+import importlib.metadata
+import logging
+import os
+import warnings
+
+try:
+    __version__ = importlib.metadata.version("atomworks")
+except ImportError:
+    __version__ = "unknown"
+
+# Global logging configuration
+logger = logging.getLogger("atomworks")
+_log_level = os.environ.get("ATOMWORKS_LOG_LEVEL", "WARNING").upper()
+logger.setLevel(_log_level)
+
+# Ensure that deprecation warnings are not repeated
+warnings.filterwarnings("once", category=DeprecationWarning)
+
+# Apply monkey patching to extend AtomArray functionality
+from atomworks.biotite_patch import monkey_patch_biotite  # noqa: E402
+
+monkey_patch_biotite()
+
+
+# Import version information
+# Import subpackages
+from . import io, ml  # noqa: E402
+
+# Re-export key functionality from subpackages for convenience
+# This maintains backward compatibility and provides a clean top-level API
+# Key I/O functionality
+from .io.parser import parse  # noqa: E402
+
+__all__ = [
+    "__version__",
+    "io",
+    "ml",
+    "monkey_patch_atomarray",
+    "parse",
+]

atomworks/biotite_patch.py

@@ -0,0 +1,421 @@
+"""Collection of monkey patches for biotite."""
+
+from typing import Callable
+
+
+import biotite
+from biotite.structure import AtomArray, AtomArrayStack, Atom
+import numpy as np
+import biotite.structure as struc
+
+__all__ = [
+    "monkey_patch_biotite",
+]
+
+_HAS_BEEN_PATCHED = False
+
+
+def apply_if_version_lt(version: str, min_version: str) -> Callable:
+    """
+    Decorator to apply a function only if the given version is less than the given minimal version.
+
+    Args:
+        version: Version to check.
+        min_version: Minimal semantic version (e.g. "0.38.0"). If the given version is lower, the
+            decorated function is called; otherwise, it is a no-op.
+
+    Example:
+        @apply_if_version_lt(biotite.__version__, "0.38.0")
+        def patch_bug():
+            # Patch code here
+            ...
+
+    Returns:
+        Decorator that conditionally applies the function.
+    """
+    from functools import wraps
+
+    def version_tuple(version: str) -> tuple[int, ...]:
+        # Only consider numeric parts, ignore pre/post-release tags
+        return tuple(int(part) for part in version.split(".") if part.isdigit())
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            current = version_tuple(version)
+            minimum = version_tuple(min_version)
+            if current < minimum:
+                return func(*args, **kwargs)
+            return None
+
+        return wrapper
+
+    return decorator
+
+
+def _add_query_mask_idxs_methods() -> None:
+    """Add `query`, `mask`, and `idxs` methods to `AtomArray` and `AtomArrayStack`."""
+    from atomworks.io.utils.query import query, mask, idxs
+
+    def query_method(self: AtomArray | AtomArrayStack, expr: str) -> AtomArray | AtomArrayStack:
+        """
+        Query the AtomArray using pandas-like syntax.
+
+        Examples
+        --------
+        >>> # Using function calls
+        >>> array.query("~has_nan_coord() & has_bonds()")
+
+        >>> # Combining with regular attributes
+        >>> array.query("has_bonds() & (chain_id == 'A') & (atom_name == 'CA')")
+        """
+        return query(self, expr)  # type: ignore
+
+    def mask_method(self: AtomArray | AtomArrayStack, expr: str) -> np.ndarray:
+        """
+        Query the AtomArray using pandas-like syntax and return a boolean mask.
+        """
+        return mask(self, expr)  # type: ignore
+
+    def idxs_method(self: AtomArray | AtomArrayStack, expr: str) -> np.ndarray:
+        """
+        Query the AtomArray using pandas-like syntax and return the indices of the matching atoms.
+        """
+        return idxs(self, expr)  # type: ignore
+
+    struc.AtomArray.query = query_method
+    struc.AtomArrayStack.query = query_method
+    struc.AtomArray.mask = mask_method
+    struc.AtomArrayStack.mask = mask_method
+    struc.AtomArray.idxs = idxs_method
+    struc.AtomArrayStack.idxs = idxs_method
+
+
+def _enable_lean_atom_array_repr() -> None:
+    """Improve the AtomArray representation to be leaner (only shows at most 20 atoms), for debugging."""
+    if not getattr(struc.AtomArray, "_repr_lean", False):
+        original_repr = struc.AtomArray.__repr__
+
+        def lean_atom_array_repr(self: struc.AtomArray) -> str:
+            """Lean AtomArray representation that only shows at most 20 atoms (first 10 and last 10)."""
+            atoms = ""
+            n_atoms = self.array_length()
+            for i in range(0, n_atoms):
+                if len(atoms) == 0:
+                    atoms = "\n\t" + self.get_atom(i).__repr__()
+                elif i >= 10 and i < (n_atoms - 10):
+                    if i == 10:
+                        atoms += "\n\t... (" + str(n_atoms - 21) + " not shown) ..."
+                    continue
+                else:
+                    atoms = atoms + ",\n\t" + self.get_atom(i).__repr__()
+            return f"AtomArray([{atoms}\n])"
+
+        setattr(struc.AtomArray, "__repr__", lean_atom_array_repr)
+        setattr(struc.AtomArray, "_repr_original", original_repr)
+        setattr(struc.AtomArray, "_repr_lean", True)
+
+
+def _enable_segment_slices_in_atom_arrays() -> None:
+    """Enable `SegmentSlice` in `AtomArray` slicing."""
+    from atomworks.io.utils.selection import SegmentSlice
+
+    if not getattr(struc.AtomArray, "_getitem_new", False):
+        original_getitem = struc.AtomArray.__getitem__
+
+        def getitem_with_segment_slices(self, item):
+            if isinstance(item, SegmentSlice):
+                item = item(self)
+            return original_getitem(self, item)
+
+        setattr(struc.AtomArray, "__getitem__", getitem_with_segment_slices)
+        setattr(struc.AtomArray, "_getitem_original", original_getitem)
+        setattr(struc.AtomArray, "_getitem_new", True)
+
+
+def _update_get_residue_starts() -> None:
+    """Improve the `get_residue_starts` function to disambiguate symmetry copies."""
+    from atomworks.io.utils.selection import get_residue_starts  # noqa: E402
+
+    struc.get_residue_starts = get_residue_starts
+
+
+def _update_array() -> None:
+    """Improve the `array` function to not truncate the datatype of annotations."""
+
+    def array(atoms: list[Atom]) -> AtomArray:
+        """Patch of Biotite's `array` function to not truncate the datatype of annotations.
+
+        Args:
+            atoms: The atoms to be combined in an array. All atoms must share the same
+                annotation categories.
+
+        Returns:
+            The listed atoms as array.
+
+        Raises:
+            ValueError: If atoms do not share the same annotation categories.
+
+        Examples:
+            Creating an atom array from atoms:
+
+            >>> atom1 = Atom([1, 2, 3], chain_id="A")
+            >>> atom2 = Atom([2, 3, 4], chain_id="A")
+            >>> atom3 = Atom([3, 4, 5], chain_id="B")
+            >>> atom_array = array([atom1, atom2, atom3])
+            >>> print(atom_array)
+                A       0                       1.000    2.000    3.000
+                A       0                       2.000    3.000    4.000
+                B       0                       3.000    4.000    5.000
+        """
+        # Check if all atoms have the same annotation names
+        # Equality check requires sorting
+        names = sorted(atoms[0]._annot.keys())
+        for i, atom in enumerate(atoms):
+            if sorted(atom._annot.keys()) != names:
+                raise ValueError(
+                    f"The atom at index {i} does not share the same " f"annotation categories as the atom at index 0"
+                )
+        array = AtomArray(len(atoms))
+
+        for name in names:
+            if hasattr(atoms[0]._annot[name], "dtype"):
+                # (Preserve dtype if possible)
+                dtype = atoms[0]._annot[name].dtype
+            else:
+                dtype = type(atoms[0]._annot[name])
+            annotation_values = [atom._annot[name] for atom in atoms]
+            annotation_values = np.array(annotation_values, dtype=dtype)  # maintain dtype
+            array.set_annotation(name, annotation_values)
+        array._coord = np.stack([atom.coord for atom in atoms])
+        return array
+
+    struc.array = array
+
+
+def _update_pdbx_set_structure() -> None:
+    """Improve the `set_structure` function to handle altloc atoms."""
+
+    # fmt: off
+    # ruff: noqa
+    import biotite.structure.io.pdbx as pdbx
+    from biotite.structure.io.pdbx.convert import (
+        MaskValue,
+        _check_non_empty,
+        _determine_entity_id,
+        _get_or_create_block,
+        _repeat,
+        _set_inter_residue_bonds,
+        _set_intra_residue_bonds,
+        unitcell_from_vectors,
+    )
+
+    def set_structure(
+        pdbx_file,
+        array,
+        data_block=None,
+        include_bonds=False,
+        extra_fields=[],
+    ):
+        """
+        Set the ``atom_site`` category with atom information from an
+        :class:`AtomArray` or :class:`AtomArrayStack`.
+
+        This will save the coordinates, the mandatory annotation categories
+        and the optional annotation categories
+        ``atom_id``, ``b_factor``, ``occupancy`` and ``charge``.
+        If the atom array (stack) contains the annotation ``'atom_id'``,
+        these values will be used for atom numbering instead of continuous
+        numbering.
+        Furthermore, inter-residue bonds will be written into the
+        ``struct_conn`` category.
+
+        Parameters
+        ----------
+        pdbx_file : CIFFile or CIFBlock or BinaryCIFFile or BinaryCIFBlock
+            The file object.
+        array : AtomArray or AtomArrayStack
+            The structure to be written. If a stack is given, each array in
+            the stack will be in a separate model.
+        data_block : str, optional
+            The name of the data block.
+            Default is the first (and most times only) data block of the
+            file.
+            If the data block object is passed directly to `pdbx_file`,
+            this parameter is ignored.
+            If the file is empty, a new data block will be created.
+        include_bonds : bool, optional
+            If set to true and `array` has associated ``bonds`` , the
+            intra-residue bonds will be written into the ``chem_comp_bond``
+            category.
+            Inter-residue bonds will be written into the ``struct_conn``
+            independent of this parameter.
+        extra_fields : list of str, optional
+            List of additional fields from the ``atom_site`` category
+            that should be written into the file.
+            Default is an empty list.
+
+        Notes
+        -----
+        In some cases, the written inter-residue bonds cannot be read again
+        due to ambiguity to which atoms the bond refers.
+        This is the case, when two equal residues in the same chain have
+        the same (or a masked) `res_id`.
+
+        Examples
+        --------
+
+        >>> import os.path
+        >>> file = CIFFile()
+        >>> set_structure(file, atom_array)
+        >>> file.write(os.path.join(path_to_directory, "structure.cif"))
+
+        """
+        _check_non_empty(array)
+
+        block = _get_or_create_block(pdbx_file, data_block)
+        Category = block.subcomponent_class()
+        Column = Category.subcomponent_class()
+
+        # Fill PDBx columns from information
+        # in structures' attribute arrays as good as possible
+        atom_site = Category()
+        atom_site["group_PDB"] = np.where(array.hetero, "HETATM", "ATOM")
+        atom_site["type_symbol"] = np.copy(array.element)
+        atom_site["label_atom_id"] = np.copy(array.atom_name)
+        if "altloc_id" in array.get_annotation_categories():
+            atom_site["label_alt_id"] = np.copy(array.altloc_id)
+        else:
+            atom_site["label_alt_id"] = Column(
+                # AtomArrays do not store altloc atoms
+                np.full(array.array_length(), "."),
+                np.full(array.array_length(), MaskValue.INAPPLICABLE),
+            )
+        atom_site["label_comp_id"] = np.copy(array.res_name)
+        atom_site["label_asym_id"] = np.copy(array.chain_id)
+        if "chain_entity" in array.get_annotation_categories():
+            atom_site["label_entity_id"] = np.copy(array.chain_entity)
+        else:
+            atom_site["label_entity_id"] = _determine_entity_id(array.chain_id)
+        atom_site["label_seq_id"] = np.copy(array.res_id)
+        atom_site["pdbx_PDB_ins_code"] = Column(
+            np.copy(array.ins_code),
+            np.where(array.ins_code == "", MaskValue.INAPPLICABLE, MaskValue.PRESENT),
+        )
+        atom_site["auth_seq_id"] = atom_site["label_seq_id"]
+        atom_site["auth_comp_id"] = atom_site["label_comp_id"]
+        atom_site["auth_asym_id"] = atom_site["label_asym_id"]
+        atom_site["auth_atom_id"] = atom_site["label_atom_id"]
+
+        annot_categories = array.get_annotation_categories()
+        if "atom_id" in annot_categories:
+            atom_site["id"] = np.copy(array.atom_id)
+        if "b_factor" in annot_categories:
+            atom_site["B_iso_or_equiv"] = np.copy(array.b_factor)
+        if "occupancy" in annot_categories:
+            atom_site["occupancy"] = np.copy(array.occupancy)
+        if "charge" in annot_categories:
+            atom_site["pdbx_formal_charge"] = Column(
+                np.array([f"{int(c):+d}" if c != 0 else "?" for c in array.charge]),
+                np.where(array.charge == 0, MaskValue.MISSING, MaskValue.PRESENT),
+            )
+
+        # Handle all remaining custom fields
+        if len(extra_fields) > 0:
+            # ... check to avoid clashes with standard annotations
+            _standard_annotations = [
+                "hetero",
+                "element",
+                "atom_name",
+                "res_name",
+                "chain_id",
+                "res_id",
+                "ins_code",
+                "atom_id",
+                "b_factor",
+                "occupancy",
+                "charge",
+            ]
+            _reserved_annotation_names = list(atom_site.keys()) + _standard_annotations
+
+            for annot in extra_fields:
+                if annot in _reserved_annotation_names:
+                    raise ValueError(
+                        f"Annotation name '{annot}' is reserved and cannot be written to as extra field. "
+                        "Please choose another name."
+                    )
+                atom_site[annot] = np.copy(array.get_annotation(annot))
+
+        if array.bonds is not None:
+            struct_conn = _set_inter_residue_bonds(array, atom_site)
+            if struct_conn is not None:
+                block["struct_conn"] = struct_conn
+            if include_bonds:
+                chem_comp_bond = _set_intra_residue_bonds(array, atom_site)
+                if chem_comp_bond is not None:
+                    block["chem_comp_bond"] = chem_comp_bond
+
+        # In case of a single model handle each coordinate
+        # simply like a flattened array
+        if isinstance(array, AtomArray) or (isinstance(array, AtomArrayStack) and array.stack_depth() == 1):
+            # 'ravel' flattens coord without copy
+            # in case of stack with stack_depth = 1
+            atom_site["Cartn_x"] = np.copy(np.ravel(array.coord[..., 0]))
+            atom_site["Cartn_y"] = np.copy(np.ravel(array.coord[..., 1]))
+            atom_site["Cartn_z"] = np.copy(np.ravel(array.coord[..., 2]))
+            atom_site["pdbx_PDB_model_num"] = np.ones(array.array_length(), dtype=np.int32)
+        # In case of multiple models repeat annotations
+        # and use model-specific coordinates
+        else:
+            atom_site = _repeat(atom_site, array.stack_depth())
+            coord = np.reshape(array.coord, (array.stack_depth() * array.array_length(), 3))
+            atom_site["Cartn_x"] = np.copy(coord[:, 0])
+            atom_site["Cartn_y"] = np.copy(coord[:, 1])
+            atom_site["Cartn_z"] = np.copy(coord[:, 2])
+            atom_site["pdbx_PDB_model_num"] = np.repeat(
+                np.arange(1, array.stack_depth() + 1, dtype=np.int32),
+                repeats=array.array_length(),
+            )
+        if "atom_id" not in annot_categories:
+            # Count from 1
+            atom_site["id"] = np.arange(1, len(atom_site["group_PDB"]) + 1)
+        block["atom_site"] = atom_site
+
+        # Write box into file
+        if array.box is not None:
+            # PDBx files can only store one box for all models
+            # -> Use first box
+            if array.box.ndim == 3:
+                box = array.box[0]
+            else:
+                box = array.box
+            len_a, len_b, len_c, alpha, beta, gamma = unitcell_from_vectors(box)
+            cell = Category()
+            cell["length_a"] = len_a
+            cell["length_b"] = len_b
+            cell["length_c"] = len_c
+            cell["angle_alpha"] = np.rad2deg(alpha)
+            cell["angle_beta"] = np.rad2deg(beta)
+            cell["angle_gamma"] = np.rad2deg(gamma)
+            block["cell"] = cell
+
+    pdbx.set_structure = set_structure
+    # fmt: on
+
+
+def monkey_patch_biotite() -> None:
+    """Monkey-patch biotite to add query, mask, and idxs methods to AtomArray and AtomArrayStack."""
+    global _HAS_BEEN_PATCHED
+
+    if _HAS_BEEN_PATCHED:
+        # ... ensure that the monkey patching is only applied once
+        return
+
+    _add_query_mask_idxs_methods()
+    _enable_lean_atom_array_repr()
+    _enable_segment_slices_in_atom_arrays()
+    _update_get_residue_starts()
+    _update_array()
+    _update_pdbx_set_structure()
+
+    _HAS_BEEN_PATCHED = True

atomworks/cli/__init__.py

@@ -0,0 +1,17 @@
+"""AtomWorks command-line interface."""
+
+from __future__ import annotations
+
+import typer
+
+app = typer.Typer(help="AtomWorks command-line interface")
+
+# Import commands to register them with the root app and expose sub-apps
+from . import ccd as _ccd  # noqa: E402
+from . import pdb as _pdb  # noqa: E402, T100
+from . import setup as _setup  # noqa: E402
+
+# Expose namespaced groups: `atomworks ccd ...`, ...
+app.add_typer(_ccd.app, name="ccd")
+app.add_typer(_pdb.app, name="pdb")
+app.add_typer(_setup.app, name="setup")

atomworks/cli/__main__.py

@@ -0,0 +1,13 @@
+"""Entry point for the AtomWorks command-line interface."""
+
+from __future__ import annotations
+
+from . import app
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()