torch-calculate-electrostatic-potential 0.6.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- torch_calculate_electrostatic_potential-0.6.0/.gitignore +115 -0
- torch_calculate_electrostatic_potential-0.6.0/LICENSE +29 -0
- torch_calculate_electrostatic_potential-0.6.0/PKG-INFO +223 -0
- torch_calculate_electrostatic_potential-0.6.0/README.md +196 -0
- torch_calculate_electrostatic_potential-0.6.0/pyproject.toml +159 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/__init__.py +37 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/atom_stack.py +193 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/grid.py +436 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/potential.py +326 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/py.typed +0 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/structure.py +148 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/utils/__init__.py +1 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/utils/elastic_scattering_bonding_protein.json +198 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/utils/elastic_scattering_bonding_rna.json +198 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/utils/peng1996_element_params.json +1392 -0
- torch_calculate_electrostatic_potential-0.6.0/src/torch_calculate_electrostatic_potential/utils/peng_model.py +255 -0
- torch_calculate_electrostatic_potential-0.6.0/tests/conftest.py +26 -0
- torch_calculate_electrostatic_potential-0.6.0/tests/test_atom_stack.py +171 -0
- torch_calculate_electrostatic_potential-0.6.0/tests/test_grid.py +238 -0
- torch_calculate_electrostatic_potential-0.6.0/tests/test_potential.py +343 -0
- torch_calculate_electrostatic_potential-0.6.0/tests/test_structure.py +275 -0
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# C extensions
|
|
7
|
+
*.so
|
|
8
|
+
|
|
9
|
+
# Distribution / packaging
|
|
10
|
+
.Python
|
|
11
|
+
env/
|
|
12
|
+
build/
|
|
13
|
+
develop-eggs/
|
|
14
|
+
dist/
|
|
15
|
+
downloads/
|
|
16
|
+
eggs/
|
|
17
|
+
.eggs/
|
|
18
|
+
lib/
|
|
19
|
+
lib64/
|
|
20
|
+
parts/
|
|
21
|
+
sdist/
|
|
22
|
+
var/
|
|
23
|
+
wheels/
|
|
24
|
+
*.egg-info/
|
|
25
|
+
.installed.cfg
|
|
26
|
+
*.egg
|
|
27
|
+
|
|
28
|
+
.DS_Store
|
|
29
|
+
|
|
30
|
+
# PyInstaller
|
|
31
|
+
*.manifest
|
|
32
|
+
*.spec
|
|
33
|
+
|
|
34
|
+
# Installer logs
|
|
35
|
+
pip-log.txt
|
|
36
|
+
pip-delete-this-directory.txt
|
|
37
|
+
|
|
38
|
+
# Unit test / coverage reports
|
|
39
|
+
htmlcov/
|
|
40
|
+
.tox/
|
|
41
|
+
.coverage
|
|
42
|
+
.coverage.*
|
|
43
|
+
.cache
|
|
44
|
+
nosetests.xml
|
|
45
|
+
coverage.xml
|
|
46
|
+
*.cover
|
|
47
|
+
.hypothesis/
|
|
48
|
+
.pytest_cache/
|
|
49
|
+
|
|
50
|
+
# Files downloaded for unit tests
|
|
51
|
+
**/tests/tmp/
|
|
52
|
+
|
|
53
|
+
# Translations
|
|
54
|
+
*.mo
|
|
55
|
+
*.pot
|
|
56
|
+
|
|
57
|
+
# Django stuff:
|
|
58
|
+
*.log
|
|
59
|
+
local_settings.py
|
|
60
|
+
|
|
61
|
+
# Flask stuff:
|
|
62
|
+
instance/
|
|
63
|
+
.webassets-cache
|
|
64
|
+
|
|
65
|
+
# Scrapy stuff:
|
|
66
|
+
.scrapy
|
|
67
|
+
|
|
68
|
+
# Sphinx documentation
|
|
69
|
+
docs/_build/
|
|
70
|
+
|
|
71
|
+
# PyBuilder
|
|
72
|
+
target/
|
|
73
|
+
|
|
74
|
+
# Jupyter Notebook
|
|
75
|
+
.ipynb_checkpoints
|
|
76
|
+
|
|
77
|
+
# dotenv
|
|
78
|
+
.env
|
|
79
|
+
|
|
80
|
+
# virtualenv
|
|
81
|
+
.venv
|
|
82
|
+
venv/
|
|
83
|
+
ENV/
|
|
84
|
+
|
|
85
|
+
# Spyder project settings
|
|
86
|
+
.spyderproject
|
|
87
|
+
.spyproject
|
|
88
|
+
|
|
89
|
+
# Rope project settings
|
|
90
|
+
.ropeproject
|
|
91
|
+
|
|
92
|
+
# mkdocs documentation
|
|
93
|
+
/site
|
|
94
|
+
|
|
95
|
+
# mypy
|
|
96
|
+
.mypy_cache/
|
|
97
|
+
|
|
98
|
+
# ruff
|
|
99
|
+
.ruff_cache/
|
|
100
|
+
|
|
101
|
+
# IDEs
|
|
102
|
+
.idea/
|
|
103
|
+
.vscode/
|
|
104
|
+
|
|
105
|
+
# Mojo extension compile cache (experimental torch-fourier-slice kernels)
|
|
106
|
+
__mojocache__/
|
|
107
|
+
*.mojopkg
|
|
108
|
+
# experimental demo outputs
|
|
109
|
+
packages/primitives/torch-fourier-slice/examples/*.npz
|
|
110
|
+
packages/primitives/torch-fourier-slice/examples/*.png
|
|
111
|
+
packages/primitives/torch-fourier-slice/examples/*.gif
|
|
112
|
+
# Personal local notes (not for commit)
|
|
113
|
+
notes/*.local.md
|
|
114
|
+
|
|
115
|
+
lightning_logs/
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
BSD 3-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2020-2026, TeamTomo
|
|
4
|
+
All rights reserved.
|
|
5
|
+
|
|
6
|
+
Redistribution and use in source and binary forms, with or without
|
|
7
|
+
modification, are permitted provided that the following conditions are met:
|
|
8
|
+
|
|
9
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
|
10
|
+
list of conditions and the following disclaimer.
|
|
11
|
+
|
|
12
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
13
|
+
this list of conditions and the following disclaimer in the documentation
|
|
14
|
+
and/or other materials provided with the distribution.
|
|
15
|
+
|
|
16
|
+
3. Neither the name of the copyright holder nor the names of its
|
|
17
|
+
contributors may be used to endorse or promote products derived from
|
|
18
|
+
this software without specific prior written permission.
|
|
19
|
+
|
|
20
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
21
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
22
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
23
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
24
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
25
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
26
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
27
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
28
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
29
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
@@ -0,0 +1,223 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: torch-calculate-electrostatic-potential
|
|
3
|
+
Version: 0.6.0
|
|
4
|
+
Summary: Differentiable Cryo-EM electrostatic potential calculator with PyTorch.
|
|
5
|
+
Project-URL: homepage, https://github.com/teamtomo/teamtomo
|
|
6
|
+
Project-URL: repository, https://github.com/teamtomo/teamtomo
|
|
7
|
+
Author-email: Volodymyr Masalitin <volodymyr.masalitin@ista.ac.at>, Matthew Giammar <mdgiammar@gmail.com>
|
|
8
|
+
License: BSD-3-Clause
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: License :: OSI Approved :: BSD License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
17
|
+
Classifier: Typing :: Typed
|
|
18
|
+
Requires-Python: >=3.11
|
|
19
|
+
Requires-Dist: einops>=0.6.0
|
|
20
|
+
Requires-Dist: gemmi
|
|
21
|
+
Requires-Dist: numpy>=1.21.0
|
|
22
|
+
Requires-Dist: setuptools
|
|
23
|
+
Requires-Dist: torch-structure-manipulation
|
|
24
|
+
Requires-Dist: torch>=2.0.0
|
|
25
|
+
Requires-Dist: tqdm>=4.60.0
|
|
26
|
+
Description-Content-Type: text/markdown
|
|
27
|
+
|
|
28
|
+
# torch-calculate-electrostatic-potential
|
|
29
|
+
|
|
30
|
+
Differentiable 2D projected and 3D electrostatic potentials from Peng 1996
|
|
31
|
+
electron-scattering factors.
|
|
32
|
+
|
|
33
|
+
The high-level API consumes
|
|
34
|
+
`torch_structure_manipulation.AtomicStructure`. The tensor-only
|
|
35
|
+
`calculate_scattering_potential_2d` and `calculate_scattering_potential_3d`
|
|
36
|
+
kernels remain public and support arbitrary leading batch dimensions.
|
|
37
|
+
|
|
38
|
+
Coordinates and spacing are in Angstroms. Axis order is ZYX in 3D and YX in 2D.
|
|
39
|
+
|
|
40
|
+
## Units and normalization
|
|
41
|
+
|
|
42
|
+
`peng1996_element_params.json` contains Peng et al. (1996) **elastic electron
|
|
43
|
+
scattering factors**, not X-ray form factors:
|
|
44
|
+
|
|
45
|
+
```text
|
|
46
|
+
f_e(s) = sum_i a_i exp(-b_i s^2), s = sin(theta) / wavelength
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
The amplitudes `a_i` and `f_e` are in Angstroms and `b_i` is in Angstroms
|
|
50
|
+
squared. The X-ray-to-electron Mott-Bethe conversion
|
|
51
|
+
`f_e(s) = 0.023934 (Z - f_X(s)) / s^2` is therefore already incorporated in the
|
|
52
|
+
tabulated coefficients and must not be applied again.
|
|
53
|
+
|
|
54
|
+
An electron scattering factor is not itself a real-space potential in volts.
|
|
55
|
+
The package converts it to the Fourier transform of the electrostatic potential
|
|
56
|
+
using
|
|
57
|
+
|
|
58
|
+
```text
|
|
59
|
+
V_tilde(g) = C f_e(g / 2), g = 2s
|
|
60
|
+
C = 2 pi hbar^2 / (m_e e) = 47.877647... V Angstrom^2
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
The inverse transform returned by `calculate_scattering_potential_3d` and
|
|
64
|
+
`potential_from_structure_3d` is therefore in **volts**. The 2D functions
|
|
65
|
+
analytically integrate the 3D potential over the omitted spatial axis and
|
|
66
|
+
return a projected potential in **volt-Angstroms**.
|
|
67
|
+
|
|
68
|
+
The bonded coefficients come from
|
|
69
|
+
[Shtyrov et al. (2026)](https://pmc.ncbi.nlm.nih.gov/articles/PMC13167779/)
|
|
70
|
+
and use the equivalent convention `f_e(g) = sum_i a_i exp(-b_i g^2 / 4)`.
|
|
71
|
+
Protein and RNA currently share the same coefficient table because RNA-specific
|
|
72
|
+
factors have not yet been measured.
|
|
73
|
+
|
|
74
|
+
## Installation
|
|
75
|
+
|
|
76
|
+
```sh
|
|
77
|
+
# From PyPI (after first release)
|
|
78
|
+
pip install torch-calculate-electrostatic-potential
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
```sh
|
|
82
|
+
# Development install from the monorepo
|
|
83
|
+
pip install -e packages/primitives/torch-calculate-electrostatic-potential
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
With [uv](https://github.com/astral-sh/uv): `uv pip install torch-calculate-electrostatic-potential`.
|
|
87
|
+
|
|
88
|
+
## Usage
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
from torch_calculate_electrostatic_potential import (
|
|
92
|
+
GridConfig,
|
|
93
|
+
potential_from_structure_2d,
|
|
94
|
+
potential_from_structure_3d,
|
|
95
|
+
)
|
|
96
|
+
from torch_structure_manipulation import AtomicStructure
|
|
97
|
+
|
|
98
|
+
structure = AtomicStructure.from_dataframe(atoms, device="cuda")
|
|
99
|
+
|
|
100
|
+
grid_3d = GridConfig.from_grid_shape_and_voxel_size(
|
|
101
|
+
grid_shape=(128, 128, 128),
|
|
102
|
+
voxel_size=(1.0, 1.0, 1.0),
|
|
103
|
+
center_zyx=(0.0, 0.0, 0.0),
|
|
104
|
+
sublattice_radius=5.0,
|
|
105
|
+
)
|
|
106
|
+
volume = potential_from_structure_3d(
|
|
107
|
+
structure,
|
|
108
|
+
grid_3d,
|
|
109
|
+
scattering_factors="peng_bonded",
|
|
110
|
+
bonded_fallback="elemental",
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
grid_2d = GridConfig.from_grid_shape_and_voxel_size(
|
|
114
|
+
grid_shape=(128, 128),
|
|
115
|
+
voxel_size=(1.0, 1.0),
|
|
116
|
+
center_yx=(0.0, 0.0),
|
|
117
|
+
)
|
|
118
|
+
projected = potential_from_structure_2d(structure, grid_2d)
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
`scattering_factors="peng_elemental"` is the default and ignores bonding
|
|
122
|
+
metadata. `"peng_bonded"` must be selected explicitly and requires
|
|
123
|
+
`bonded_environments` and per-atom `molecule_types`. Unsupported `other`
|
|
124
|
+
molecules and absent keys either emit one warning and use elemental values
|
|
125
|
+
(`bonded_fallback="elemental"`) or raise (`bonded_fallback="error"`).
|
|
126
|
+
|
|
127
|
+
The molecule type is the scattering-factor provider key, not merely descriptive
|
|
128
|
+
metadata. Custom providers can therefore supply different tables for protein,
|
|
129
|
+
RNA, or any additional molecule type:
|
|
130
|
+
|
|
131
|
+
## Batched structures and bonded factors
|
|
132
|
+
|
|
133
|
+
`AtomicStructure` may carry broadcast-compatible batch dimensions on positions
|
|
134
|
+
and other numerical fields. The tensor kernels and elemental Peng lookup support
|
|
135
|
+
that directly.
|
|
136
|
+
|
|
137
|
+
Bonded factors are different:
|
|
138
|
+
|
|
139
|
+
- `bonded_environments` and `molecule_types` are **flat tuples** (one string per
|
|
140
|
+
atom), shared across the whole batch.
|
|
141
|
+
- `resolve_scattering_parameters(..., scattering_factors="peng_bonded")` requires
|
|
142
|
+
**one-dimensional** `atomic_numbers` with shape `(n_atoms,)`.
|
|
143
|
+
|
|
144
|
+
Practical guidance:
|
|
145
|
+
|
|
146
|
+
| Use case | Elemental | Bonded |
|
|
147
|
+
|----------|-----------|--------|
|
|
148
|
+
| Single structure | yes | yes |
|
|
149
|
+
| Multiple poses, same chemistry (`positions` batched, `atomic_numbers` `(n,)`) | yes | yes |
|
|
150
|
+
| Batched `atomic_numbers` with shape `(batch, n_atoms)` | yes | no — raises |
|
|
151
|
+
| Different chemistry per batch member | N/A | no — not representable |
|
|
152
|
+
|
|
153
|
+
For different structures, call `potential_from_structure_3d` once per
|
|
154
|
+
`AtomicStructure` (or loop over batch indices).
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
from torch_calculate_electrostatic_potential import BondedScatteringFactorTable
|
|
158
|
+
|
|
159
|
+
custom_factors = {
|
|
160
|
+
"protein": BondedScatteringFactorTable(
|
|
161
|
+
parameters_a=protein_parameters_a,
|
|
162
|
+
parameters_b=protein_parameters_b,
|
|
163
|
+
),
|
|
164
|
+
"rna": BondedScatteringFactorTable(
|
|
165
|
+
parameters_a=rna_parameters_a,
|
|
166
|
+
parameters_b=rna_parameters_b,
|
|
167
|
+
),
|
|
168
|
+
}
|
|
169
|
+
volume = potential_from_structure_3d(
|
|
170
|
+
structure,
|
|
171
|
+
grid_3d,
|
|
172
|
+
scattering_factors=custom_factors,
|
|
173
|
+
bonded_fallback="error",
|
|
174
|
+
)
|
|
175
|
+
```
|
|
176
|
+
|
|
177
|
+
Each parameter mapping is keyed by the structure's `bonded_environments`
|
|
178
|
+
strings. The low-level tensor API remains available for callers that have
|
|
179
|
+
already resolved arbitrary per-atom `a` and `b` tensors.
|
|
180
|
+
|
|
181
|
+
The lower-level route exposes parameter tensors directly:
|
|
182
|
+
|
|
183
|
+
```python
|
|
184
|
+
from torch_calculate_electrostatic_potential import (
|
|
185
|
+
calculate_scattering_potential_3d,
|
|
186
|
+
get_peng_scattering_parameters,
|
|
187
|
+
)
|
|
188
|
+
|
|
189
|
+
atom_params_a, atom_params_b = get_peng_scattering_parameters(atomic_numbers)
|
|
190
|
+
potential_volume = calculate_scattering_potential_3d(
|
|
191
|
+
atom_pos_zyx,
|
|
192
|
+
atom_bfactors,
|
|
193
|
+
atom_params_a,
|
|
194
|
+
atom_params_b,
|
|
195
|
+
grid_3d,
|
|
196
|
+
atom_occupancies=occupancies,
|
|
197
|
+
)
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
Positions, B-factors, occupancies, and explicit parameter tensors remain
|
|
201
|
+
differentiable. `sublattice_radius` controls the finite local stencil; increase
|
|
202
|
+
it for broad Gaussians.
|
|
203
|
+
|
|
204
|
+
## Testing
|
|
205
|
+
|
|
206
|
+
Install the package together with test dependencies:
|
|
207
|
+
|
|
208
|
+
```sh
|
|
209
|
+
pip install "torch-calculate-electrostatic-potential[test]" @ git+https://github.com/teamtomo/torch-calculate-electrostatic-potential.git
|
|
210
|
+
pytest
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
With coverage: `pytest --cov=torch_calculate_electrostatic_potential --cov-report=html`.
|
|
214
|
+
|
|
215
|
+
## Requirements
|
|
216
|
+
|
|
217
|
+
- Python >= 3.11
|
|
218
|
+
- PyTorch >= 2.0
|
|
219
|
+
- torch-structure-manipulation, numpy, einops, tqdm
|
|
220
|
+
|
|
221
|
+
## License
|
|
222
|
+
|
|
223
|
+
BSD 3-Clause License
|
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
# torch-calculate-electrostatic-potential
|
|
2
|
+
|
|
3
|
+
Differentiable 2D projected and 3D electrostatic potentials from Peng 1996
|
|
4
|
+
electron-scattering factors.
|
|
5
|
+
|
|
6
|
+
The high-level API consumes
|
|
7
|
+
`torch_structure_manipulation.AtomicStructure`. The tensor-only
|
|
8
|
+
`calculate_scattering_potential_2d` and `calculate_scattering_potential_3d`
|
|
9
|
+
kernels remain public and support arbitrary leading batch dimensions.
|
|
10
|
+
|
|
11
|
+
Coordinates and spacing are in Angstroms. Axis order is ZYX in 3D and YX in 2D.
|
|
12
|
+
|
|
13
|
+
## Units and normalization
|
|
14
|
+
|
|
15
|
+
`peng1996_element_params.json` contains Peng et al. (1996) **elastic electron
|
|
16
|
+
scattering factors**, not X-ray form factors:
|
|
17
|
+
|
|
18
|
+
```text
|
|
19
|
+
f_e(s) = sum_i a_i exp(-b_i s^2), s = sin(theta) / wavelength
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
The amplitudes `a_i` and `f_e` are in Angstroms and `b_i` is in Angstroms
|
|
23
|
+
squared. The X-ray-to-electron Mott-Bethe conversion
|
|
24
|
+
`f_e(s) = 0.023934 (Z - f_X(s)) / s^2` is therefore already incorporated in the
|
|
25
|
+
tabulated coefficients and must not be applied again.
|
|
26
|
+
|
|
27
|
+
An electron scattering factor is not itself a real-space potential in volts.
|
|
28
|
+
The package converts it to the Fourier transform of the electrostatic potential
|
|
29
|
+
using
|
|
30
|
+
|
|
31
|
+
```text
|
|
32
|
+
V_tilde(g) = C f_e(g / 2), g = 2s
|
|
33
|
+
C = 2 pi hbar^2 / (m_e e) = 47.877647... V Angstrom^2
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
The inverse transform returned by `calculate_scattering_potential_3d` and
|
|
37
|
+
`potential_from_structure_3d` is therefore in **volts**. The 2D functions
|
|
38
|
+
analytically integrate the 3D potential over the omitted spatial axis and
|
|
39
|
+
return a projected potential in **volt-Angstroms**.
|
|
40
|
+
|
|
41
|
+
The bonded coefficients come from
|
|
42
|
+
[Shtyrov et al. (2026)](https://pmc.ncbi.nlm.nih.gov/articles/PMC13167779/)
|
|
43
|
+
and use the equivalent convention `f_e(g) = sum_i a_i exp(-b_i g^2 / 4)`.
|
|
44
|
+
Protein and RNA currently share the same coefficient table because RNA-specific
|
|
45
|
+
factors have not yet been measured.
|
|
46
|
+
|
|
47
|
+
## Installation
|
|
48
|
+
|
|
49
|
+
```sh
|
|
50
|
+
# From PyPI (after first release)
|
|
51
|
+
pip install torch-calculate-electrostatic-potential
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
```sh
|
|
55
|
+
# Development install from the monorepo
|
|
56
|
+
pip install -e packages/primitives/torch-calculate-electrostatic-potential
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
With [uv](https://github.com/astral-sh/uv): `uv pip install torch-calculate-electrostatic-potential`.
|
|
60
|
+
|
|
61
|
+
## Usage
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
from torch_calculate_electrostatic_potential import (
|
|
65
|
+
GridConfig,
|
|
66
|
+
potential_from_structure_2d,
|
|
67
|
+
potential_from_structure_3d,
|
|
68
|
+
)
|
|
69
|
+
from torch_structure_manipulation import AtomicStructure
|
|
70
|
+
|
|
71
|
+
structure = AtomicStructure.from_dataframe(atoms, device="cuda")
|
|
72
|
+
|
|
73
|
+
grid_3d = GridConfig.from_grid_shape_and_voxel_size(
|
|
74
|
+
grid_shape=(128, 128, 128),
|
|
75
|
+
voxel_size=(1.0, 1.0, 1.0),
|
|
76
|
+
center_zyx=(0.0, 0.0, 0.0),
|
|
77
|
+
sublattice_radius=5.0,
|
|
78
|
+
)
|
|
79
|
+
volume = potential_from_structure_3d(
|
|
80
|
+
structure,
|
|
81
|
+
grid_3d,
|
|
82
|
+
scattering_factors="peng_bonded",
|
|
83
|
+
bonded_fallback="elemental",
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
grid_2d = GridConfig.from_grid_shape_and_voxel_size(
|
|
87
|
+
grid_shape=(128, 128),
|
|
88
|
+
voxel_size=(1.0, 1.0),
|
|
89
|
+
center_yx=(0.0, 0.0),
|
|
90
|
+
)
|
|
91
|
+
projected = potential_from_structure_2d(structure, grid_2d)
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
`scattering_factors="peng_elemental"` is the default and ignores bonding
|
|
95
|
+
metadata. `"peng_bonded"` must be selected explicitly and requires
|
|
96
|
+
`bonded_environments` and per-atom `molecule_types`. Unsupported `other`
|
|
97
|
+
molecules and absent keys either emit one warning and use elemental values
|
|
98
|
+
(`bonded_fallback="elemental"`) or raise (`bonded_fallback="error"`).
|
|
99
|
+
|
|
100
|
+
The molecule type is the scattering-factor provider key, not merely descriptive
|
|
101
|
+
metadata. Custom providers can therefore supply different tables for protein,
|
|
102
|
+
RNA, or any additional molecule type:
|
|
103
|
+
|
|
104
|
+
## Batched structures and bonded factors
|
|
105
|
+
|
|
106
|
+
`AtomicStructure` may carry broadcast-compatible batch dimensions on positions
|
|
107
|
+
and other numerical fields. The tensor kernels and elemental Peng lookup support
|
|
108
|
+
that directly.
|
|
109
|
+
|
|
110
|
+
Bonded factors are different:
|
|
111
|
+
|
|
112
|
+
- `bonded_environments` and `molecule_types` are **flat tuples** (one string per
|
|
113
|
+
atom), shared across the whole batch.
|
|
114
|
+
- `resolve_scattering_parameters(..., scattering_factors="peng_bonded")` requires
|
|
115
|
+
**one-dimensional** `atomic_numbers` with shape `(n_atoms,)`.
|
|
116
|
+
|
|
117
|
+
Practical guidance:
|
|
118
|
+
|
|
119
|
+
| Use case | Elemental | Bonded |
|
|
120
|
+
|----------|-----------|--------|
|
|
121
|
+
| Single structure | yes | yes |
|
|
122
|
+
| Multiple poses, same chemistry (`positions` batched, `atomic_numbers` `(n,)`) | yes | yes |
|
|
123
|
+
| Batched `atomic_numbers` with shape `(batch, n_atoms)` | yes | no — raises |
|
|
124
|
+
| Different chemistry per batch member | N/A | no — not representable |
|
|
125
|
+
|
|
126
|
+
For different structures, call `potential_from_structure_3d` once per
|
|
127
|
+
`AtomicStructure` (or loop over batch indices).
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
from torch_calculate_electrostatic_potential import BondedScatteringFactorTable
|
|
131
|
+
|
|
132
|
+
custom_factors = {
|
|
133
|
+
"protein": BondedScatteringFactorTable(
|
|
134
|
+
parameters_a=protein_parameters_a,
|
|
135
|
+
parameters_b=protein_parameters_b,
|
|
136
|
+
),
|
|
137
|
+
"rna": BondedScatteringFactorTable(
|
|
138
|
+
parameters_a=rna_parameters_a,
|
|
139
|
+
parameters_b=rna_parameters_b,
|
|
140
|
+
),
|
|
141
|
+
}
|
|
142
|
+
volume = potential_from_structure_3d(
|
|
143
|
+
structure,
|
|
144
|
+
grid_3d,
|
|
145
|
+
scattering_factors=custom_factors,
|
|
146
|
+
bonded_fallback="error",
|
|
147
|
+
)
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
Each parameter mapping is keyed by the structure's `bonded_environments`
|
|
151
|
+
strings. The low-level tensor API remains available for callers that have
|
|
152
|
+
already resolved arbitrary per-atom `a` and `b` tensors.
|
|
153
|
+
|
|
154
|
+
The lower-level route exposes parameter tensors directly:
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
from torch_calculate_electrostatic_potential import (
|
|
158
|
+
calculate_scattering_potential_3d,
|
|
159
|
+
get_peng_scattering_parameters,
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
atom_params_a, atom_params_b = get_peng_scattering_parameters(atomic_numbers)
|
|
163
|
+
potential_volume = calculate_scattering_potential_3d(
|
|
164
|
+
atom_pos_zyx,
|
|
165
|
+
atom_bfactors,
|
|
166
|
+
atom_params_a,
|
|
167
|
+
atom_params_b,
|
|
168
|
+
grid_3d,
|
|
169
|
+
atom_occupancies=occupancies,
|
|
170
|
+
)
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
Positions, B-factors, occupancies, and explicit parameter tensors remain
|
|
174
|
+
differentiable. `sublattice_radius` controls the finite local stencil; increase
|
|
175
|
+
it for broad Gaussians.
|
|
176
|
+
|
|
177
|
+
## Testing
|
|
178
|
+
|
|
179
|
+
Install the package together with test dependencies:
|
|
180
|
+
|
|
181
|
+
```sh
|
|
182
|
+
pip install "torch-calculate-electrostatic-potential[test]" @ git+https://github.com/teamtomo/torch-calculate-electrostatic-potential.git
|
|
183
|
+
pytest
|
|
184
|
+
```
|
|
185
|
+
|
|
186
|
+
With coverage: `pytest --cov=torch_calculate_electrostatic_potential --cov-report=html`.
|
|
187
|
+
|
|
188
|
+
## Requirements
|
|
189
|
+
|
|
190
|
+
- Python >= 3.11
|
|
191
|
+
- PyTorch >= 2.0
|
|
192
|
+
- torch-structure-manipulation, numpy, einops, tqdm
|
|
193
|
+
|
|
194
|
+
## License
|
|
195
|
+
|
|
196
|
+
BSD 3-Clause License
|