torch-calculate-electrostatic-potential 0.6.0__tar.gz

torch_calculate_electrostatic_potential-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_calculate_electrostatic_potential-0.6.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020-2026, 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_calculate_electrostatic_potential-0.6.0/PKG-INFO

@@ -0,0 +1,223 @@
+Metadata-Version: 2.5
+Name: torch-calculate-electrostatic-potential
+Version: 0.6.0
+Summary: Differentiable Cryo-EM electrostatic potential calculator with PyTorch.
+Project-URL: homepage, https://github.com/teamtomo/teamtomo
+Project-URL: repository, https://github.com/teamtomo/teamtomo
+Author-email: Volodymyr Masalitin <volodymyr.masalitin@ista.ac.at>, Matthew Giammar <mdgiammar@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: einops>=0.6.0
+Requires-Dist: gemmi
+Requires-Dist: numpy>=1.21.0
+Requires-Dist: setuptools
+Requires-Dist: torch-structure-manipulation
+Requires-Dist: torch>=2.0.0
+Requires-Dist: tqdm>=4.60.0
+Description-Content-Type: text/markdown
+
+# torch-calculate-electrostatic-potential
+
+Differentiable 2D projected and 3D electrostatic potentials from Peng 1996
+electron-scattering factors.
+
+The high-level API consumes
+`torch_structure_manipulation.AtomicStructure`. The tensor-only
+`calculate_scattering_potential_2d` and `calculate_scattering_potential_3d`
+kernels remain public and support arbitrary leading batch dimensions.
+
+Coordinates and spacing are in Angstroms. Axis order is ZYX in 3D and YX in 2D.
+
+## Units and normalization
+
+`peng1996_element_params.json` contains Peng et al. (1996) **elastic electron
+scattering factors**, not X-ray form factors:
+
+```text
+f_e(s) = sum_i a_i exp(-b_i s^2),  s = sin(theta) / wavelength
+```
+
+The amplitudes `a_i` and `f_e` are in Angstroms and `b_i` is in Angstroms
+squared. The X-ray-to-electron Mott-Bethe conversion
+`f_e(s) = 0.023934 (Z - f_X(s)) / s^2` is therefore already incorporated in the
+tabulated coefficients and must not be applied again.
+
+An electron scattering factor is not itself a real-space potential in volts.
+The package converts it to the Fourier transform of the electrostatic potential
+using
+
+```text
+V_tilde(g) = C f_e(g / 2),  g = 2s
+C = 2 pi hbar^2 / (m_e e) = 47.877647... V Angstrom^2
+```
+
+The inverse transform returned by `calculate_scattering_potential_3d` and
+`potential_from_structure_3d` is therefore in **volts**. The 2D functions
+analytically integrate the 3D potential over the omitted spatial axis and
+return a projected potential in **volt-Angstroms**.
+
+The bonded coefficients come from
+[Shtyrov et al. (2026)](https://pmc.ncbi.nlm.nih.gov/articles/PMC13167779/)
+and use the equivalent convention `f_e(g) = sum_i a_i exp(-b_i g^2 / 4)`.
+Protein and RNA currently share the same coefficient table because RNA-specific
+factors have not yet been measured.
+
+## Installation
+
+```sh
+# From PyPI (after first release)
+pip install torch-calculate-electrostatic-potential
+```
+
+```sh
+# Development install from the monorepo
+pip install -e packages/primitives/torch-calculate-electrostatic-potential
+```
+
+With [uv](https://github.com/astral-sh/uv): `uv pip install torch-calculate-electrostatic-potential`.
+
+## Usage
+
+```python
+from torch_calculate_electrostatic_potential import (
+    GridConfig,
+    potential_from_structure_2d,
+    potential_from_structure_3d,
+)
+from torch_structure_manipulation import AtomicStructure
+
+structure = AtomicStructure.from_dataframe(atoms, device="cuda")
+
+grid_3d = GridConfig.from_grid_shape_and_voxel_size(
+    grid_shape=(128, 128, 128),
+    voxel_size=(1.0, 1.0, 1.0),
+    center_zyx=(0.0, 0.0, 0.0),
+    sublattice_radius=5.0,
+)
+volume = potential_from_structure_3d(
+    structure,
+    grid_3d,
+    scattering_factors="peng_bonded",
+    bonded_fallback="elemental",
+)
+
+grid_2d = GridConfig.from_grid_shape_and_voxel_size(
+    grid_shape=(128, 128),
+    voxel_size=(1.0, 1.0),
+    center_yx=(0.0, 0.0),
+)
+projected = potential_from_structure_2d(structure, grid_2d)
+```
+
+`scattering_factors="peng_elemental"` is the default and ignores bonding
+metadata. `"peng_bonded"` must be selected explicitly and requires
+`bonded_environments` and per-atom `molecule_types`. Unsupported `other`
+molecules and absent keys either emit one warning and use elemental values
+(`bonded_fallback="elemental"`) or raise (`bonded_fallback="error"`).
+
+The molecule type is the scattering-factor provider key, not merely descriptive
+metadata. Custom providers can therefore supply different tables for protein,
+RNA, or any additional molecule type:
+
+## Batched structures and bonded factors
+
+`AtomicStructure` may carry broadcast-compatible batch dimensions on positions
+and other numerical fields. The tensor kernels and elemental Peng lookup support
+that directly.
+
+Bonded factors are different:
+
+- `bonded_environments` and `molecule_types` are **flat tuples** (one string per
+  atom), shared across the whole batch.
+- `resolve_scattering_parameters(..., scattering_factors="peng_bonded")` requires
+  **one-dimensional** `atomic_numbers` with shape `(n_atoms,)`.
+
+Practical guidance:
+
+| Use case | Elemental | Bonded |
+|----------|-----------|--------|
+| Single structure | yes | yes |
+| Multiple poses, same chemistry (`positions` batched, `atomic_numbers` `(n,)`) | yes | yes |
+| Batched `atomic_numbers` with shape `(batch, n_atoms)` | yes | no — raises |
+| Different chemistry per batch member | N/A | no — not representable |
+
+For different structures, call `potential_from_structure_3d` once per
+`AtomicStructure` (or loop over batch indices).
+
+```python
+from torch_calculate_electrostatic_potential import BondedScatteringFactorTable
+
+custom_factors = {
+    "protein": BondedScatteringFactorTable(
+        parameters_a=protein_parameters_a,
+        parameters_b=protein_parameters_b,
+    ),
+    "rna": BondedScatteringFactorTable(
+        parameters_a=rna_parameters_a,
+        parameters_b=rna_parameters_b,
+    ),
+}
+volume = potential_from_structure_3d(
+    structure,
+    grid_3d,
+    scattering_factors=custom_factors,
+    bonded_fallback="error",
+)
+```
+
+Each parameter mapping is keyed by the structure's `bonded_environments`
+strings. The low-level tensor API remains available for callers that have
+already resolved arbitrary per-atom `a` and `b` tensors.
+
+The lower-level route exposes parameter tensors directly:
+
+```python
+from torch_calculate_electrostatic_potential import (
+    calculate_scattering_potential_3d,
+    get_peng_scattering_parameters,
+)
+
+atom_params_a, atom_params_b = get_peng_scattering_parameters(atomic_numbers)
+potential_volume = calculate_scattering_potential_3d(
+    atom_pos_zyx,
+    atom_bfactors,
+    atom_params_a,
+    atom_params_b,
+    grid_3d,
+    atom_occupancies=occupancies,
+)
+```
+
+Positions, B-factors, occupancies, and explicit parameter tensors remain
+differentiable. `sublattice_radius` controls the finite local stencil; increase
+it for broad Gaussians.
+
+## Testing
+
+Install the package together with test dependencies:
+
+```sh
+pip install "torch-calculate-electrostatic-potential[test]" @ git+https://github.com/teamtomo/torch-calculate-electrostatic-potential.git
+pytest
+```
+
+With coverage: `pytest --cov=torch_calculate_electrostatic_potential --cov-report=html`.
+
+## Requirements
+
+- Python >= 3.11
+- PyTorch >= 2.0
+- torch-structure-manipulation, numpy, einops, tqdm
+
+## License
+
+BSD 3-Clause License

torch_calculate_electrostatic_potential-0.6.0/README.md

@@ -0,0 +1,196 @@
+# torch-calculate-electrostatic-potential
+
+Differentiable 2D projected and 3D electrostatic potentials from Peng 1996
+electron-scattering factors.
+
+The high-level API consumes
+`torch_structure_manipulation.AtomicStructure`. The tensor-only
+`calculate_scattering_potential_2d` and `calculate_scattering_potential_3d`
+kernels remain public and support arbitrary leading batch dimensions.
+
+Coordinates and spacing are in Angstroms. Axis order is ZYX in 3D and YX in 2D.
+
+## Units and normalization
+
+`peng1996_element_params.json` contains Peng et al. (1996) **elastic electron
+scattering factors**, not X-ray form factors:
+
+```text
+f_e(s) = sum_i a_i exp(-b_i s^2),  s = sin(theta) / wavelength
+```
+
+The amplitudes `a_i` and `f_e` are in Angstroms and `b_i` is in Angstroms
+squared. The X-ray-to-electron Mott-Bethe conversion
+`f_e(s) = 0.023934 (Z - f_X(s)) / s^2` is therefore already incorporated in the
+tabulated coefficients and must not be applied again.
+
+An electron scattering factor is not itself a real-space potential in volts.
+The package converts it to the Fourier transform of the electrostatic potential
+using
+
+```text
+V_tilde(g) = C f_e(g / 2),  g = 2s
+C = 2 pi hbar^2 / (m_e e) = 47.877647... V Angstrom^2
+```
+
+The inverse transform returned by `calculate_scattering_potential_3d` and
+`potential_from_structure_3d` is therefore in **volts**. The 2D functions
+analytically integrate the 3D potential over the omitted spatial axis and
+return a projected potential in **volt-Angstroms**.
+
+The bonded coefficients come from
+[Shtyrov et al. (2026)](https://pmc.ncbi.nlm.nih.gov/articles/PMC13167779/)
+and use the equivalent convention `f_e(g) = sum_i a_i exp(-b_i g^2 / 4)`.
+Protein and RNA currently share the same coefficient table because RNA-specific
+factors have not yet been measured.
+
+## Installation
+
+```sh
+# From PyPI (after first release)
+pip install torch-calculate-electrostatic-potential
+```
+
+```sh
+# Development install from the monorepo
+pip install -e packages/primitives/torch-calculate-electrostatic-potential
+```
+
+With [uv](https://github.com/astral-sh/uv): `uv pip install torch-calculate-electrostatic-potential`.
+
+## Usage
+
+```python
+from torch_calculate_electrostatic_potential import (
+    GridConfig,
+    potential_from_structure_2d,
+    potential_from_structure_3d,
+)
+from torch_structure_manipulation import AtomicStructure
+
+structure = AtomicStructure.from_dataframe(atoms, device="cuda")
+
+grid_3d = GridConfig.from_grid_shape_and_voxel_size(
+    grid_shape=(128, 128, 128),
+    voxel_size=(1.0, 1.0, 1.0),
+    center_zyx=(0.0, 0.0, 0.0),
+    sublattice_radius=5.0,
+)
+volume = potential_from_structure_3d(
+    structure,
+    grid_3d,
+    scattering_factors="peng_bonded",
+    bonded_fallback="elemental",
+)
+
+grid_2d = GridConfig.from_grid_shape_and_voxel_size(
+    grid_shape=(128, 128),
+    voxel_size=(1.0, 1.0),
+    center_yx=(0.0, 0.0),
+)
+projected = potential_from_structure_2d(structure, grid_2d)
+```
+
+`scattering_factors="peng_elemental"` is the default and ignores bonding
+metadata. `"peng_bonded"` must be selected explicitly and requires
+`bonded_environments` and per-atom `molecule_types`. Unsupported `other`
+molecules and absent keys either emit one warning and use elemental values
+(`bonded_fallback="elemental"`) or raise (`bonded_fallback="error"`).
+
+The molecule type is the scattering-factor provider key, not merely descriptive
+metadata. Custom providers can therefore supply different tables for protein,
+RNA, or any additional molecule type:
+
+## Batched structures and bonded factors
+
+`AtomicStructure` may carry broadcast-compatible batch dimensions on positions
+and other numerical fields. The tensor kernels and elemental Peng lookup support
+that directly.
+
+Bonded factors are different:
+
+- `bonded_environments` and `molecule_types` are **flat tuples** (one string per
+  atom), shared across the whole batch.
+- `resolve_scattering_parameters(..., scattering_factors="peng_bonded")` requires
+  **one-dimensional** `atomic_numbers` with shape `(n_atoms,)`.
+
+Practical guidance:
+
+| Use case | Elemental | Bonded |
+|----------|-----------|--------|
+| Single structure | yes | yes |
+| Multiple poses, same chemistry (`positions` batched, `atomic_numbers` `(n,)`) | yes | yes |
+| Batched `atomic_numbers` with shape `(batch, n_atoms)` | yes | no — raises |
+| Different chemistry per batch member | N/A | no — not representable |
+
+For different structures, call `potential_from_structure_3d` once per
+`AtomicStructure` (or loop over batch indices).
+
+```python
+from torch_calculate_electrostatic_potential import BondedScatteringFactorTable
+
+custom_factors = {
+    "protein": BondedScatteringFactorTable(
+        parameters_a=protein_parameters_a,
+        parameters_b=protein_parameters_b,
+    ),
+    "rna": BondedScatteringFactorTable(
+        parameters_a=rna_parameters_a,
+        parameters_b=rna_parameters_b,
+    ),
+}
+volume = potential_from_structure_3d(
+    structure,
+    grid_3d,
+    scattering_factors=custom_factors,
+    bonded_fallback="error",
+)
+```
+
+Each parameter mapping is keyed by the structure's `bonded_environments`
+strings. The low-level tensor API remains available for callers that have
+already resolved arbitrary per-atom `a` and `b` tensors.
+
+The lower-level route exposes parameter tensors directly:
+
+```python
+from torch_calculate_electrostatic_potential import (
+    calculate_scattering_potential_3d,
+    get_peng_scattering_parameters,
+)
+
+atom_params_a, atom_params_b = get_peng_scattering_parameters(atomic_numbers)
+potential_volume = calculate_scattering_potential_3d(
+    atom_pos_zyx,
+    atom_bfactors,
+    atom_params_a,
+    atom_params_b,
+    grid_3d,
+    atom_occupancies=occupancies,
+)
+```
+
+Positions, B-factors, occupancies, and explicit parameter tensors remain
+differentiable. `sublattice_radius` controls the finite local stencil; increase
+it for broad Gaussians.
+
+## Testing
+
+Install the package together with test dependencies:
+
+```sh
+pip install "torch-calculate-electrostatic-potential[test]" @ git+https://github.com/teamtomo/torch-calculate-electrostatic-potential.git
+pytest
+```
+
+With coverage: `pytest --cov=torch_calculate_electrostatic_potential --cov-report=html`.
+
+## Requirements
+
+- Python >= 3.11
+- PyTorch >= 2.0
+- torch-structure-manipulation, numpy, einops, tqdm
+
+## License
+
+BSD 3-Clause License