torch-structure-manipulation 0.6.0__tar.gz

torch_structure_manipulation-0.6.0/.gitignore

@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+.DS_Store
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Files downloaded for unit tests
+**/tests/tmp/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# ruff
+.ruff_cache/
+
+# IDEs
+.idea/
+.vscode/
+
+# Mojo extension compile cache (experimental torch-fourier-slice kernels)
+__mojocache__/
+*.mojopkg
+# experimental demo outputs
+packages/primitives/torch-fourier-slice/examples/*.npz
+packages/primitives/torch-fourier-slice/examples/*.png
+packages/primitives/torch-fourier-slice/examples/*.gif
+# Personal local notes (not for commit)
+notes/*.local.md
+
+lightning_logs/

torch_structure_manipulation-0.6.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020, TeamTomo
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

torch_structure_manipulation-0.6.0/PKG-INFO

@@ -0,0 +1,64 @@
+Metadata-Version: 2.5
+Name: torch-structure-manipulation
+Version: 0.6.0
+Summary: Typed molecular structure data, bonding annotations, and tensor transforms
+Project-URL: homepage, https://github.com/teamtomo/teamtomo
+Project-URL: repository, https://github.com/teamtomo/teamtomo
+Author: TeamTomo Developers
+Author-email: Davide Torre <davidetorre99@gmail.com>
+License: BSD-3-Clause
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: gemmi
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: roma
+Requires-Dist: torch
+Description-Content-Type: text/markdown
+
+# torch-structure-manipulation
+
+Typed, in-memory utilities for atomic structures:
+
+- `AtomicStructure.from_dataframe` converts mmdf-compatible DataFrames into
+  device-aware tensors, with positions stored in `(z, y, x)` array order.
+- `AtomicStructure.from_annotated_dataframe` annotates bonding metadata and
+  then constructs the structure in one step.
+- `annotate_bonding_environments` adds template-derived bonding and
+  `protein`/`rna`/`other` molecule annotations without reading files.
+- `classify_structure_composition` and `get_scattering_provider_keys` expose
+  aggregate composition labels and per-atom Peng provider keys respectively.
+- Centering, rotation, translation, atom selection, and coordinate conversion
+  helpers are re-exported from the package root (also available under
+  ``structure_transforms``).
+
+```python
+import mmdf
+
+from torch_structure_manipulation import (
+    AtomicStructure,
+    annotate_bonding_environments,
+    center_structure,
+)
+
+# File I/O belongs to the caller; mmdf is not a runtime dependency.
+atoms = mmdf.read("structure.cif")
+structure = AtomicStructure.from_annotated_dataframe(atoms, include_hydrogens=False)
+
+# Or annotate first when you still need the DataFrame:
+annotated = annotate_bonding_environments(atoms)
+centered = center_structure(annotated, center_point=(0.0, 0.0, 0.0), zyx=False)
+structure = AtomicStructure.from_dataframe(centered)
+```
+
+## License
+
+This project is licensed under the BSD 3-Clause License; see `LICENSE`.

torch_structure_manipulation-0.6.0/README.md

@@ -0,0 +1,38 @@
+# torch-structure-manipulation
+
+Typed, in-memory utilities for atomic structures:
+
+- `AtomicStructure.from_dataframe` converts mmdf-compatible DataFrames into
+  device-aware tensors, with positions stored in `(z, y, x)` array order.
+- `AtomicStructure.from_annotated_dataframe` annotates bonding metadata and
+  then constructs the structure in one step.
+- `annotate_bonding_environments` adds template-derived bonding and
+  `protein`/`rna`/`other` molecule annotations without reading files.
+- `classify_structure_composition` and `get_scattering_provider_keys` expose
+  aggregate composition labels and per-atom Peng provider keys respectively.
+- Centering, rotation, translation, atom selection, and coordinate conversion
+  helpers are re-exported from the package root (also available under
+  ``structure_transforms``).
+
+```python
+import mmdf
+
+from torch_structure_manipulation import (
+    AtomicStructure,
+    annotate_bonding_environments,
+    center_structure,
+)
+
+# File I/O belongs to the caller; mmdf is not a runtime dependency.
+atoms = mmdf.read("structure.cif")
+structure = AtomicStructure.from_annotated_dataframe(atoms, include_hydrogens=False)
+
+# Or annotate first when you still need the DataFrame:
+annotated = annotate_bonding_environments(atoms)
+centered = center_structure(annotated, center_point=(0.0, 0.0, 0.0), zyx=False)
+structure = AtomicStructure.from_dataframe(centered)
+```
+
+## License
+
+This project is licensed under the BSD 3-Clause License; see `LICENSE`.

torch_structure_manipulation-0.6.0/pyproject.toml

@@ -0,0 +1,95 @@
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+[tool.hatch.version]
+source = "vcs"
+tag-pattern = "^torch-structure-manipulation@v(?P<version>.+)$"
+fallback-version = "0.1.0"
+
+[tool.hatch.version.raw-options]
+search_parent_directories = true
+tag_regex = "^torch-structure-manipulation@v(?P<version>\\d+\\.\\d+\\.\\d+.*)$"
+git_describe_command = "git describe --dirty --tags --long --match 'torch-structure-manipulation@v[0-9]*.[0-9]*.[0-9]*'"
+
+[tool.hatch.build.targets.wheel]
+only-include = ["src"]
+sources = ["src"]
+
+[project]
+name = "torch-structure-manipulation"
+dynamic = ["version"]
+description = "Typed molecular structure data, bonding annotations, and tensor transforms"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "BSD-3-Clause" }
+authors = [
+    { name = "Davide Torre", email = "davidetorre99@gmail.com" },
+    { name = "TeamTomo Developers" },
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Typing :: Typed",
+]
+dependencies = ["gemmi", "numpy", "pandas", "roma", "torch"]
+
+[project.urls]
+homepage = "https://github.com/teamtomo/teamtomo"
+repository = "https://github.com/teamtomo/teamtomo"
+
+[dependency-groups]
+test = ["pytest", "pytest-cov"]
+dev = [
+    { include-group = "test" },
+    "ipython",
+    "mypy",
+    "pdbpp",
+    "rich",
+    "ruff",
+]
+
+[tool.ruff]
+line-length = 88
+target-version = "py311"
+src = ["src"]
+
+[tool.ruff.lint]
+pydocstyle = { convention = "numpy" }
+select = ["E", "W", "F", "D", "D417", "I", "UP", "C4", "B", "A001", "RUF", "TC", "TID"]
+ignore = ["D401"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*.py" = ["D", "S"]
+
+[tool.mypy]
+files = "src/torch_structure_manipulation"
+strict = true
+disallow_any_generics = false
+disallow_subclassing_any = false
+show_error_codes = true
+pretty = true
+
+[[tool.mypy.overrides]]
+module = ["pandas", "pandas.*"]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = ["roma", "roma.*"]
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+testpaths = ["tests"]
+filterwarnings = ["error"]
+
+[tool.coverage.run]
+source = ["torch_structure_manipulation"]
+
+[tool.check-manifest]
+ignore = [".ruff_cache/**/*", "tests/**/*"]

torch_structure_manipulation-0.6.0/src/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-0.6.0/src/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),
+        )