torchsx 0.1.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.
- torchsx-0.1.0/LICENSE +21 -0
- torchsx-0.1.0/PKG-INFO +182 -0
- torchsx-0.1.0/README.md +155 -0
- torchsx-0.1.0/pyproject.toml +53 -0
- torchsx-0.1.0/setup.cfg +4 -0
- torchsx-0.1.0/tests/test_amortized_scale.py +854 -0
- torchsx-0.1.0/tests/test_amortized_training.py +426 -0
- torchsx-0.1.0/tests/test_asym_scatter.py +230 -0
- torchsx-0.1.0/tests/test_bg_mean_offset.py +281 -0
- torchsx-0.1.0/tests/test_cli_pipeline.py +658 -0
- torchsx-0.1.0/tests/test_cli_production.py +248 -0
- torchsx-0.1.0/tests/test_core.py +77 -0
- torchsx-0.1.0/tests/test_duckdb_load.py +595 -0
- torchsx-0.1.0/tests/test_ewald_offset.py +72 -0
- torchsx-0.1.0/tests/test_filter.py +121 -0
- torchsx-0.1.0/tests/test_geometry.py +202 -0
- torchsx-0.1.0/tests/test_geometry_reciprocal.py +69 -0
- torchsx-0.1.0/tests/test_grid_tooling.py +95 -0
- torchsx-0.1.0/tests/test_hkl.py +131 -0
- torchsx-0.1.0/tests/test_io.py +119 -0
- torchsx-0.1.0/tests/test_large_batch_diagnostics.py +43 -0
- torchsx-0.1.0/tests/test_levels.py +124 -0
- torchsx-0.1.0/tests/test_loo_pooling.py +436 -0
- torchsx-0.1.0/tests/test_marginal_normalisation.py +125 -0
- torchsx-0.1.0/tests/test_measurement_sources.py +635 -0
- torchsx-0.1.0/tests/test_merge_stats.py +265 -0
- torchsx-0.1.0/tests/test_merging_layers.py +84 -0
- torchsx-0.1.0/tests/test_metadata.py +78 -0
- torchsx-0.1.0/tests/test_native_index.py +102 -0
- torchsx-0.1.0/tests/test_native_parse.py +371 -0
- torchsx-0.1.0/tests/test_observation_filter.py +85 -0
- torchsx-0.1.0/tests/test_pad.py +59 -0
- torchsx-0.1.0/tests/test_peer_cross_attention.py +308 -0
- torchsx-0.1.0/tests/test_peer_features.py +58 -0
- torchsx-0.1.0/tests/test_posterior_nodes.py +232 -0
- torchsx-0.1.0/tests/test_posterior_prior.py +59 -0
- torchsx-0.1.0/tests/test_posterior_smooth_sigma.py +106 -0
- torchsx-0.1.0/tests/test_reciprocal_basis.py +150 -0
- torchsx-0.1.0/tests/test_scatter_mixture.py +488 -0
- torchsx-0.1.0/tests/test_shape_prior.py +128 -0
- torchsx-0.1.0/tests/test_split_gaussian_mean.py +91 -0
- torchsx-0.1.0/tests/test_split_half.py +176 -0
- torchsx-0.1.0/tests/test_split_log_grid.py +212 -0
- torchsx-0.1.0/tests/test_split_mixture.py +232 -0
- torchsx-0.1.0/tests/test_stateless_merge.py +66 -0
- torchsx-0.1.0/tests/test_stats.py +192 -0
- torchsx-0.1.0/tests/test_stream.py +291 -0
- torchsx-0.1.0/tests/test_train_schedule.py +35 -0
- torchsx-0.1.0/tests/test_wilson_surface.py +160 -0
- torchsx-0.1.0/tests/test_writeback.py +60 -0
- torchsx-0.1.0/tests/test_writeback_native.py +122 -0
- torchsx-0.1.0/torchsx/__init__.py +8 -0
- torchsx-0.1.0/torchsx/_categorical.py +58 -0
- torchsx-0.1.0/torchsx/_core.py +602 -0
- torchsx-0.1.0/torchsx/cli/__init__.py +0 -0
- torchsx-0.1.0/torchsx/cli/_pipeline.py +912 -0
- torchsx-0.1.0/torchsx/cli/merge_nn.py +205 -0
- torchsx-0.1.0/torchsx/cli/merge_stats.py +217 -0
- torchsx-0.1.0/torchsx/cli/stats.py +95 -0
- torchsx-0.1.0/torchsx/cli/train.py +278 -0
- torchsx-0.1.0/torchsx/geometry.py +745 -0
- torchsx-0.1.0/torchsx/io/__init__.py +30 -0
- torchsx-0.1.0/torchsx/io/_crystfel.py +105 -0
- torchsx-0.1.0/torchsx/io/_duckdb.py +258 -0
- torchsx-0.1.0/torchsx/io/_metadata.py +187 -0
- torchsx-0.1.0/torchsx/io/_native.py +48 -0
- torchsx-0.1.0/torchsx/io/_parse/__init__.py +8 -0
- torchsx-0.1.0/torchsx/io/_parse/indexing.py +186 -0
- torchsx-0.1.0/torchsx/io/_parse/orchestrator.py +258 -0
- torchsx-0.1.0/torchsx/io/_parse/records.py +520 -0
- torchsx-0.1.0/torchsx/io/_parse/workers.py +35 -0
- torchsx-0.1.0/torchsx/io/_parser.py +40 -0
- torchsx-0.1.0/torchsx/io/_safetensors.py +170 -0
- torchsx-0.1.0/torchsx/io/_writeback.py +463 -0
- torchsx-0.1.0/torchsx/io/sources.py +76 -0
- torchsx-0.1.0/torchsx/io/times.py +438 -0
- torchsx-0.1.0/torchsx/levels.py +306 -0
- torchsx-0.1.0/torchsx/measurement.py +217 -0
- torchsx-0.1.0/torchsx/merge_stats.py +672 -0
- torchsx-0.1.0/torchsx/merging/__init__.py +30 -0
- torchsx-0.1.0/torchsx/merging/config.py +52 -0
- torchsx-0.1.0/torchsx/merging/hkl/__init__.py +75 -0
- torchsx-0.1.0/torchsx/merging/hkl/adjacency.py +71 -0
- torchsx-0.1.0/torchsx/merging/hkl/filter.py +177 -0
- torchsx-0.1.0/torchsx/merging/hkl/hkl.py +118 -0
- torchsx-0.1.0/torchsx/merging/hkl/noise.py +451 -0
- torchsx-0.1.0/torchsx/merging/hkl/shell_fit.py +102 -0
- torchsx-0.1.0/torchsx/merging/hkl/stats.py +294 -0
- torchsx-0.1.0/torchsx/merging/likelihood/__init__.py +10 -0
- torchsx-0.1.0/torchsx/merging/likelihood/base.py +64 -0
- torchsx-0.1.0/torchsx/merging/likelihood/detector_noise.py +104 -0
- torchsx-0.1.0/torchsx/merging/likelihood/poisson.py +173 -0
- torchsx-0.1.0/torchsx/merging/observations.py +699 -0
- torchsx-0.1.0/torchsx/merging/parameters.py +77 -0
- torchsx-0.1.0/torchsx/merging/priors/__init__.py +13 -0
- torchsx-0.1.0/torchsx/merging/priors/reciprocal_basis.py +293 -0
- torchsx-0.1.0/torchsx/merging/quadrature/__init__.py +15 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_fixed_log_grid_triton.py +426 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_fixed_log_grid_tuning.json +44 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_fused_posterior_triton.py +249 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_posterior_nodes_triton.py +378 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_split_log_grid_triton.py +294 -0
- torchsx-0.1.0/torchsx/merging/quadrature/_split_mixture_triton.py +218 -0
- torchsx-0.1.0/torchsx/merging/quadrature/base.py +98 -0
- torchsx-0.1.0/torchsx/merging/quadrature/fixed_log_grid.py +305 -0
- torchsx-0.1.0/torchsx/merging/quadrature/split_log_grid.py +392 -0
- torchsx-0.1.0/torchsx/merging/scale/__init__.py +10 -0
- torchsx-0.1.0/torchsx/merging/scale/amortized.py +452 -0
- torchsx-0.1.0/torchsx/merging/scale/base.py +44 -0
- torchsx-0.1.0/torchsx/merging/scatter/__init__.py +22 -0
- torchsx-0.1.0/torchsx/merging/scatter/base.py +54 -0
- torchsx-0.1.0/torchsx/merging/scatter/gengauss.py +314 -0
- torchsx-0.1.0/torchsx/merging/scatter/mixture.py +269 -0
- torchsx-0.1.0/torchsx/merging/scatter/predicted.py +119 -0
- torchsx-0.1.0/torchsx/merging/scatter_shape.py +213 -0
- torchsx-0.1.0/torchsx/nn/__init__.py +41 -0
- torchsx-0.1.0/torchsx/nn/deepsets.py +579 -0
- torchsx-0.1.0/torchsx/nn/heads.py +370 -0
- torchsx-0.1.0/torchsx/nn/merge_target.py +67 -0
- torchsx-0.1.0/torchsx/nn/metrics.py +397 -0
- torchsx-0.1.0/torchsx/nn/objective.py +380 -0
- torchsx-0.1.0/torchsx/nn/obs_features.py +424 -0
- torchsx-0.1.0/torchsx/nn/pad.py +78 -0
- torchsx-0.1.0/torchsx/nn/pooling.py +395 -0
- torchsx-0.1.0/torchsx/nn/posterior_target.py +662 -0
- torchsx-0.1.0/torchsx/nn/shape_prior.py +98 -0
- torchsx-0.1.0/torchsx/nn/train_amortized.py +1640 -0
- torchsx-0.1.0/torchsx/stats.py +302 -0
- torchsx-0.1.0/torchsx/stream.py +826 -0
- torchsx-0.1.0/torchsx.egg-info/PKG-INFO +182 -0
- torchsx-0.1.0/torchsx.egg-info/SOURCES.txt +133 -0
- torchsx-0.1.0/torchsx.egg-info/dependency_links.txt +1 -0
- torchsx-0.1.0/torchsx.egg-info/entry_points.txt +5 -0
- torchsx-0.1.0/torchsx.egg-info/requires.txt +13 -0
- torchsx-0.1.0/torchsx.egg-info/top_level.txt +1 -0
torchsx-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Hans Peter Seidel
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
torchsx-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,182 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: torchsx
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: GPU-friendly stacked tensors for unmerged serial crystallography data
|
|
5
|
+
Author: Hans Peter Seidel
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Repository, https://github.com/HatPdotS/TorchSX
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
12
|
+
Requires-Python: >=3.10
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
Requires-Dist: torchref<0.8,>=0.7.0
|
|
16
|
+
Requires-Dist: numpy>=2
|
|
17
|
+
Requires-Dist: safetensors>=0.4
|
|
18
|
+
Requires-Dist: matplotlib>=3.7
|
|
19
|
+
Requires-Dist: torchsx-native==0.1.*
|
|
20
|
+
Provides-Extra: duckdb
|
|
21
|
+
Requires-Dist: duckdb>=1.0; extra == "duckdb"
|
|
22
|
+
Provides-Extra: dev
|
|
23
|
+
Requires-Dist: pytest>=6.0; extra == "dev"
|
|
24
|
+
Requires-Dist: black>=21.5; extra == "dev"
|
|
25
|
+
Requires-Dist: isort>=5.9; extra == "dev"
|
|
26
|
+
Dynamic: license-file
|
|
27
|
+
|
|
28
|
+
<p align="center">
|
|
29
|
+
<picture>
|
|
30
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/torchsx-banner-dark.png">
|
|
31
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/torchsx-banner-light.png" alt="TorchSX" width="100%">
|
|
32
|
+
</picture>
|
|
33
|
+
</p>
|
|
34
|
+
|
|
35
|
+
<h1 align="center">A Py<ins>Torch</ins>-based scaling and merging library for serial crystallography (<ins>SX</ins>)</h1>
|
|
36
|
+
|
|
37
|
+
<div align="center">
|
|
38
|
+
|
|
39
|
+
[](https://www.python.org/downloads/)
|
|
40
|
+
[](https://pytorch.org/)
|
|
41
|
+
[](https://github.com/HatPdotS/TorchRef)
|
|
42
|
+
[](https://developer.nvidia.com/cuda-zone)
|
|
43
|
+
[](https://opensource.org/licenses/MIT)
|
|
44
|
+
|
|
45
|
+
</div>
|
|
46
|
+
|
|
47
|
+
In serial crystallography every reflection is seen once per crystal, with an
|
|
48
|
+
unknown partiality, incident intensity and crystal volume. Established programs
|
|
49
|
+
scale each crystal with a few parameters and average. TorchSX instead trains one
|
|
50
|
+
neural network per dataset that predicts, for every observation, a scale factor
|
|
51
|
+
and a scatter distribution from that observation's crystal- and image-level
|
|
52
|
+
context. Merging is a posterior integral rather than an average: each merged
|
|
53
|
+
amplitude and its uncertainty are moments of the posterior of the reflection's
|
|
54
|
+
intensity. The network is trained on the marginal likelihood of the data.
|
|
55
|
+
|
|
56
|
+
TorchSX starts from integrated CrystFEL `.stream` files and writes merged MTZs.
|
|
57
|
+
We are working on integration with dials and deep integration of Probixi.
|
|
58
|
+
|
|
59
|
+
TorchSX is under active development, if you find something behaving unexpectedly please create an issue on GitHub.
|
|
60
|
+
|
|
61
|
+
Refinement and difference maps are done with
|
|
62
|
+
[TorchRef](https://github.com/HatPdotS/TorchRef), or any other program that
|
|
63
|
+
reads amplitudes.
|
|
64
|
+
|
|
65
|
+
## Install
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
pip install torchsx
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
Python ≥ 3.10. This brings TorchRef 0.7, torch, triton, gemmi and
|
|
72
|
+
`torchsx-native`, the compiled `.stream` reader for your platform. TorchSX is intended to be run on a GPU, while performance and memory requirements should be quite low, it relies on triton kernels for core algorithms. We have not investigated CPU runtime or even functionality.
|
|
73
|
+
|
|
74
|
+
From a checkout, build the native package first (needs Rust, see
|
|
75
|
+
[native/README.md](https://github.com/HatPdotS/TorchSX/blob/master/native/README.md)):
|
|
76
|
+
`pip install ./native`, then `pip install -e .`.
|
|
77
|
+
|
|
78
|
+
## Usage
|
|
79
|
+
|
|
80
|
+
```bash
|
|
81
|
+
torchsx.train indexed.stream --spacegroup 'P 63' --dmin 2.0 -o run/
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
This trains the network (400 updates by default) and merges the data with it.
|
|
85
|
+
`run/` then contains:
|
|
86
|
+
|
|
87
|
+
```
|
|
88
|
+
run/
|
|
89
|
+
├── merged.mtz merged amplitudes from all crystals
|
|
90
|
+
├── merged_half1.mtz the same from one half of the crystals
|
|
91
|
+
├── merged_half2.mtz ... and from the other half
|
|
92
|
+
├── merged.metrics.json CC1/2 and R_split between the two halves
|
|
93
|
+
├── best.pt the network at the lowest validation loss; it made the merges
|
|
94
|
+
├── last.pt the network after the last update
|
|
95
|
+
├── prepared.pt the parsed observations, reused by later runs on the same input
|
|
96
|
+
├── run.json the resolved settings
|
|
97
|
+
├── training_summary.json updates, best step and run time
|
|
98
|
+
└── metrics.jsonl training log: one line per update and per validation
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
Each MTZ holds:
|
|
102
|
+
|
|
103
|
+
| Column | MTZ type | Content |
|
|
104
|
+
|---|---|---|
|
|
105
|
+
| `H`, `K`, `L` | H | Miller index; Friedel mates merged |
|
|
106
|
+
| `F-obs` | F | posterior mean of the amplitude \|F\| |
|
|
107
|
+
| `SIGF-obs` | Q | posterior standard deviation of the amplitude |
|
|
108
|
+
| `R-free-flags` | I | refinement free set: 0 = free (5 %), 1 = work |
|
|
109
|
+
|
|
110
|
+
There are no intensity columns. `F-obs` is already the posterior estimate of
|
|
111
|
+
the amplitude, so it must not go through French–Wilson again; refinement
|
|
112
|
+
programs use it as it is. With `--no-merge-friedel` the mates stay apart and the
|
|
113
|
+
files carry `F-obs(+)`, `SIGF-obs(+)`, `F-obs(-)` and `SIGF-obs(-)`.
|
|
114
|
+
`--free-from earlier.mtz` takes the free set from an existing file.
|
|
115
|
+
|
|
116
|
+
The two halves split the crystals, not the training: both are merged with the
|
|
117
|
+
same network. For an independent comparison, train two runs on disjoint sets of
|
|
118
|
+
patterns.
|
|
119
|
+
|
|
120
|
+
### Light and dark
|
|
121
|
+
|
|
122
|
+
```bash
|
|
123
|
+
torchsx.train --split dark=dark.stream --split light=light.stream \
|
|
124
|
+
--spacegroup 'P 63' --dmin 2.0 -o run/
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
One network is trained on both states, and each state is merged on its own:
|
|
128
|
+
`merged_dark.mtz` and `merged_light.mtz`, each with its two halves. A single
|
|
129
|
+
stream with a CrystFEL custom-split file works too (`--split-file`).
|
|
130
|
+
|
|
131
|
+
## Compared with Partialator
|
|
132
|
+
|
|
133
|
+
On the bacteriorhodopsin (bR) data of Weinert et al.<sup>[1](#references)</sup>:
|
|
134
|
+
|
|
135
|
+
<p align="center">
|
|
136
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/readme-br-rfree.png" alt="bR R-free against dark patterns per merge, TorchSX and Partialator" width="65%">
|
|
137
|
+
</p>
|
|
138
|
+
|
|
139
|
+
R<sub>free</sub> of models refined from one starting model against dark merges
|
|
140
|
+
of increasing size. Each size was merged as two independent halves: points show
|
|
141
|
+
their mean, bars their range. R<sub>free</sub> is computed on the same
|
|
142
|
+
reflections for both programs at each size; refinement used TorchRef with one
|
|
143
|
+
protocol for all merges.
|
|
144
|
+
|
|
145
|
+
<p align="center">
|
|
146
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/readme-br-maps.png" alt="bR light-minus-dark difference maps from two independent halves, TorchSX and Partialator" width="100%">
|
|
147
|
+
</p>
|
|
148
|
+
|
|
149
|
+
Light-minus-dark difference maps from two independent halves of the patterns
|
|
150
|
+
(7,800 per state and half), each merged on its own and phased with the dark
|
|
151
|
+
model. Shown: the end of helix E, the EF loop and helix F of the dark model
|
|
152
|
+
(grey) and of the open state 6RPH<sup>[1](#references)</sup> (orange). The
|
|
153
|
+
number under each program is the correlation of its two maps over helix F.
|
|
154
|
+
|
|
155
|
+
## Other commands
|
|
156
|
+
|
|
157
|
+
- `torchsx.merge run/best.pt more.stream -o merged.mtz` merges new or additional
|
|
158
|
+
data with a trained network.
|
|
159
|
+
- `torchsx.merge-stats A/merged.mtz B/merged.mtz -o stats/` compares merges:
|
|
160
|
+
completeness, CC1/2, R_split and, for light/dark pairs, the information in
|
|
161
|
+
ΔF ([docs/merge_stats.md](https://github.com/HatPdotS/TorchSX/blob/master/docs/merge_stats.md)).
|
|
162
|
+
- `torchsx.observation-stats indexed.stream --spacegroup 'P 63' -o figures/`
|
|
163
|
+
plots the raw observations: intensity distributions, skewness, examples
|
|
164
|
+
([docs/stats.md](https://github.com/HatPdotS/TorchSX/blob/master/docs/stats.md)).
|
|
165
|
+
|
|
166
|
+
All options, gain and measurement settings, caching and reproducibility are in
|
|
167
|
+
the [command-line reference](https://github.com/HatPdotS/TorchSX/blob/master/docs/cli.md).
|
|
168
|
+
|
|
169
|
+
In Python, `torchsx.SFXStream.read('indexed.stream', spacegroup='P 63')` reads a
|
|
170
|
+
stream with its geometry and measurement settings.
|
|
171
|
+
|
|
172
|
+
## References
|
|
173
|
+
|
|
174
|
+
1. Weinert, T. et al. Proton uptake mechanism in bacteriorhodopsin captured by
|
|
175
|
+
serial synchrotron crystallography. *Science* **365**, 61–65 (2019).
|
|
176
|
+
[doi:10.1126/science.aaw8634](https://doi.org/10.1126/science.aaw8634)
|
|
177
|
+
2. Seidel, H.-P., Standfuss, J. & Weinert, T. TorchRef: an open-source PyTorch
|
|
178
|
+
framework for crystallographic refinement. *J. Appl. Cryst.*, accepted (2026).
|
|
179
|
+
|
|
180
|
+
## License
|
|
181
|
+
|
|
182
|
+
MIT. See [LICENSE](https://github.com/HatPdotS/TorchSX/blob/master/LICENSE).
|
torchsx-0.1.0/README.md
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
1
|
+
<p align="center">
|
|
2
|
+
<picture>
|
|
3
|
+
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/torchsx-banner-dark.png">
|
|
4
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/torchsx-banner-light.png" alt="TorchSX" width="100%">
|
|
5
|
+
</picture>
|
|
6
|
+
</p>
|
|
7
|
+
|
|
8
|
+
<h1 align="center">A Py<ins>Torch</ins>-based scaling and merging library for serial crystallography (<ins>SX</ins>)</h1>
|
|
9
|
+
|
|
10
|
+
<div align="center">
|
|
11
|
+
|
|
12
|
+
[](https://www.python.org/downloads/)
|
|
13
|
+
[](https://pytorch.org/)
|
|
14
|
+
[](https://github.com/HatPdotS/TorchRef)
|
|
15
|
+
[](https://developer.nvidia.com/cuda-zone)
|
|
16
|
+
[](https://opensource.org/licenses/MIT)
|
|
17
|
+
|
|
18
|
+
</div>
|
|
19
|
+
|
|
20
|
+
In serial crystallography every reflection is seen once per crystal, with an
|
|
21
|
+
unknown partiality, incident intensity and crystal volume. Established programs
|
|
22
|
+
scale each crystal with a few parameters and average. TorchSX instead trains one
|
|
23
|
+
neural network per dataset that predicts, for every observation, a scale factor
|
|
24
|
+
and a scatter distribution from that observation's crystal- and image-level
|
|
25
|
+
context. Merging is a posterior integral rather than an average: each merged
|
|
26
|
+
amplitude and its uncertainty are moments of the posterior of the reflection's
|
|
27
|
+
intensity. The network is trained on the marginal likelihood of the data.
|
|
28
|
+
|
|
29
|
+
TorchSX starts from integrated CrystFEL `.stream` files and writes merged MTZs.
|
|
30
|
+
We are working on integration with dials and deep integration of Probixi.
|
|
31
|
+
|
|
32
|
+
TorchSX is under active development, if you find something behaving unexpectedly please create an issue on GitHub.
|
|
33
|
+
|
|
34
|
+
Refinement and difference maps are done with
|
|
35
|
+
[TorchRef](https://github.com/HatPdotS/TorchRef), or any other program that
|
|
36
|
+
reads amplitudes.
|
|
37
|
+
|
|
38
|
+
## Install
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
pip install torchsx
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
Python ≥ 3.10. This brings TorchRef 0.7, torch, triton, gemmi and
|
|
45
|
+
`torchsx-native`, the compiled `.stream` reader for your platform. TorchSX is intended to be run on a GPU, while performance and memory requirements should be quite low, it relies on triton kernels for core algorithms. We have not investigated CPU runtime or even functionality.
|
|
46
|
+
|
|
47
|
+
From a checkout, build the native package first (needs Rust, see
|
|
48
|
+
[native/README.md](https://github.com/HatPdotS/TorchSX/blob/master/native/README.md)):
|
|
49
|
+
`pip install ./native`, then `pip install -e .`.
|
|
50
|
+
|
|
51
|
+
## Usage
|
|
52
|
+
|
|
53
|
+
```bash
|
|
54
|
+
torchsx.train indexed.stream --spacegroup 'P 63' --dmin 2.0 -o run/
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
This trains the network (400 updates by default) and merges the data with it.
|
|
58
|
+
`run/` then contains:
|
|
59
|
+
|
|
60
|
+
```
|
|
61
|
+
run/
|
|
62
|
+
├── merged.mtz merged amplitudes from all crystals
|
|
63
|
+
├── merged_half1.mtz the same from one half of the crystals
|
|
64
|
+
├── merged_half2.mtz ... and from the other half
|
|
65
|
+
├── merged.metrics.json CC1/2 and R_split between the two halves
|
|
66
|
+
├── best.pt the network at the lowest validation loss; it made the merges
|
|
67
|
+
├── last.pt the network after the last update
|
|
68
|
+
├── prepared.pt the parsed observations, reused by later runs on the same input
|
|
69
|
+
├── run.json the resolved settings
|
|
70
|
+
├── training_summary.json updates, best step and run time
|
|
71
|
+
└── metrics.jsonl training log: one line per update and per validation
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
Each MTZ holds:
|
|
75
|
+
|
|
76
|
+
| Column | MTZ type | Content |
|
|
77
|
+
|---|---|---|
|
|
78
|
+
| `H`, `K`, `L` | H | Miller index; Friedel mates merged |
|
|
79
|
+
| `F-obs` | F | posterior mean of the amplitude \|F\| |
|
|
80
|
+
| `SIGF-obs` | Q | posterior standard deviation of the amplitude |
|
|
81
|
+
| `R-free-flags` | I | refinement free set: 0 = free (5 %), 1 = work |
|
|
82
|
+
|
|
83
|
+
There are no intensity columns. `F-obs` is already the posterior estimate of
|
|
84
|
+
the amplitude, so it must not go through French–Wilson again; refinement
|
|
85
|
+
programs use it as it is. With `--no-merge-friedel` the mates stay apart and the
|
|
86
|
+
files carry `F-obs(+)`, `SIGF-obs(+)`, `F-obs(-)` and `SIGF-obs(-)`.
|
|
87
|
+
`--free-from earlier.mtz` takes the free set from an existing file.
|
|
88
|
+
|
|
89
|
+
The two halves split the crystals, not the training: both are merged with the
|
|
90
|
+
same network. For an independent comparison, train two runs on disjoint sets of
|
|
91
|
+
patterns.
|
|
92
|
+
|
|
93
|
+
### Light and dark
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
torchsx.train --split dark=dark.stream --split light=light.stream \
|
|
97
|
+
--spacegroup 'P 63' --dmin 2.0 -o run/
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
One network is trained on both states, and each state is merged on its own:
|
|
101
|
+
`merged_dark.mtz` and `merged_light.mtz`, each with its two halves. A single
|
|
102
|
+
stream with a CrystFEL custom-split file works too (`--split-file`).
|
|
103
|
+
|
|
104
|
+
## Compared with Partialator
|
|
105
|
+
|
|
106
|
+
On the bacteriorhodopsin (bR) data of Weinert et al.<sup>[1](#references)</sup>:
|
|
107
|
+
|
|
108
|
+
<p align="center">
|
|
109
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/readme-br-rfree.png" alt="bR R-free against dark patterns per merge, TorchSX and Partialator" width="65%">
|
|
110
|
+
</p>
|
|
111
|
+
|
|
112
|
+
R<sub>free</sub> of models refined from one starting model against dark merges
|
|
113
|
+
of increasing size. Each size was merged as two independent halves: points show
|
|
114
|
+
their mean, bars their range. R<sub>free</sub> is computed on the same
|
|
115
|
+
reflections for both programs at each size; refinement used TorchRef with one
|
|
116
|
+
protocol for all merges.
|
|
117
|
+
|
|
118
|
+
<p align="center">
|
|
119
|
+
<img src="https://raw.githubusercontent.com/HatPdotS/TorchSX/master/assets/readme-br-maps.png" alt="bR light-minus-dark difference maps from two independent halves, TorchSX and Partialator" width="100%">
|
|
120
|
+
</p>
|
|
121
|
+
|
|
122
|
+
Light-minus-dark difference maps from two independent halves of the patterns
|
|
123
|
+
(7,800 per state and half), each merged on its own and phased with the dark
|
|
124
|
+
model. Shown: the end of helix E, the EF loop and helix F of the dark model
|
|
125
|
+
(grey) and of the open state 6RPH<sup>[1](#references)</sup> (orange). The
|
|
126
|
+
number under each program is the correlation of its two maps over helix F.
|
|
127
|
+
|
|
128
|
+
## Other commands
|
|
129
|
+
|
|
130
|
+
- `torchsx.merge run/best.pt more.stream -o merged.mtz` merges new or additional
|
|
131
|
+
data with a trained network.
|
|
132
|
+
- `torchsx.merge-stats A/merged.mtz B/merged.mtz -o stats/` compares merges:
|
|
133
|
+
completeness, CC1/2, R_split and, for light/dark pairs, the information in
|
|
134
|
+
ΔF ([docs/merge_stats.md](https://github.com/HatPdotS/TorchSX/blob/master/docs/merge_stats.md)).
|
|
135
|
+
- `torchsx.observation-stats indexed.stream --spacegroup 'P 63' -o figures/`
|
|
136
|
+
plots the raw observations: intensity distributions, skewness, examples
|
|
137
|
+
([docs/stats.md](https://github.com/HatPdotS/TorchSX/blob/master/docs/stats.md)).
|
|
138
|
+
|
|
139
|
+
All options, gain and measurement settings, caching and reproducibility are in
|
|
140
|
+
the [command-line reference](https://github.com/HatPdotS/TorchSX/blob/master/docs/cli.md).
|
|
141
|
+
|
|
142
|
+
In Python, `torchsx.SFXStream.read('indexed.stream', spacegroup='P 63')` reads a
|
|
143
|
+
stream with its geometry and measurement settings.
|
|
144
|
+
|
|
145
|
+
## References
|
|
146
|
+
|
|
147
|
+
1. Weinert, T. et al. Proton uptake mechanism in bacteriorhodopsin captured by
|
|
148
|
+
serial synchrotron crystallography. *Science* **365**, 61–65 (2019).
|
|
149
|
+
[doi:10.1126/science.aaw8634](https://doi.org/10.1126/science.aaw8634)
|
|
150
|
+
2. Seidel, H.-P., Standfuss, J. & Weinert, T. TorchRef: an open-source PyTorch
|
|
151
|
+
framework for crystallographic refinement. *J. Appl. Cryst.*, accepted (2026).
|
|
152
|
+
|
|
153
|
+
## License
|
|
154
|
+
|
|
155
|
+
MIT. See [LICENSE](https://github.com/HatPdotS/TorchSX/blob/master/LICENSE).
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77.0"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "torchsx"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "GPU-friendly stacked tensors for unmerged serial crystallography data"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
authors = [{ name = "Hans Peter Seidel" }]
|
|
13
|
+
requires-python = ">=3.10"
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Programming Language :: Python :: 3 :: Only",
|
|
17
|
+
"Intended Audience :: Science/Research",
|
|
18
|
+
"Topic :: Scientific/Engineering :: Physics",
|
|
19
|
+
]
|
|
20
|
+
dependencies = [
|
|
21
|
+
# Brings torch, triton and gemmi. TorchRef's minor releases change its API,
|
|
22
|
+
# so the bound follows the minor version torchsx is tested against.
|
|
23
|
+
"torchref>=0.7.0,<0.8",
|
|
24
|
+
"numpy>=2", # the Python writer's float32 arithmetic (NEP 50)
|
|
25
|
+
"safetensors>=0.4", # the .sfx working format
|
|
26
|
+
"matplotlib>=3.7", # the figures of torchsx.observation-stats and torchsx.merge-stats
|
|
27
|
+
# The .stream reader/writer (Rust), built from native/ in this repository and
|
|
28
|
+
# released with torchsx. The pin follows native/Cargo.toml's minor version;
|
|
29
|
+
# bump both (and PROTOCOL) when the Rust↔Python interface changes.
|
|
30
|
+
"torchsx-native==0.1.*",
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
[project.urls]
|
|
34
|
+
Repository = "https://github.com/HatPdotS/TorchSX"
|
|
35
|
+
|
|
36
|
+
[project.scripts]
|
|
37
|
+
"torchsx.train" = "torchsx.cli.train:main"
|
|
38
|
+
"torchsx.merge" = "torchsx.cli.merge_nn:main"
|
|
39
|
+
"torchsx.observation-stats" = "torchsx.cli.stats:main"
|
|
40
|
+
"torchsx.merge-stats" = "torchsx.cli.merge_stats:main"
|
|
41
|
+
|
|
42
|
+
[project.optional-dependencies]
|
|
43
|
+
# Read probixi's DuckDB runs (SFXStream.from_duckdb).
|
|
44
|
+
duckdb = ["duckdb>=1.0"]
|
|
45
|
+
|
|
46
|
+
dev = ["pytest>=6.0", "black>=21.5", "isort>=5.9"]
|
|
47
|
+
|
|
48
|
+
[tool.setuptools.packages.find]
|
|
49
|
+
include = ["torchsx*"]
|
|
50
|
+
|
|
51
|
+
[tool.setuptools.package-data]
|
|
52
|
+
# The Triton kernel's per-GPU launch configurations.
|
|
53
|
+
"torchsx.merging.quadrature" = ["_fixed_log_grid_tuning.json"]
|
torchsx-0.1.0/setup.cfg
ADDED