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.
Files changed (135) hide show
  1. torchsx-0.1.0/LICENSE +21 -0
  2. torchsx-0.1.0/PKG-INFO +182 -0
  3. torchsx-0.1.0/README.md +155 -0
  4. torchsx-0.1.0/pyproject.toml +53 -0
  5. torchsx-0.1.0/setup.cfg +4 -0
  6. torchsx-0.1.0/tests/test_amortized_scale.py +854 -0
  7. torchsx-0.1.0/tests/test_amortized_training.py +426 -0
  8. torchsx-0.1.0/tests/test_asym_scatter.py +230 -0
  9. torchsx-0.1.0/tests/test_bg_mean_offset.py +281 -0
  10. torchsx-0.1.0/tests/test_cli_pipeline.py +658 -0
  11. torchsx-0.1.0/tests/test_cli_production.py +248 -0
  12. torchsx-0.1.0/tests/test_core.py +77 -0
  13. torchsx-0.1.0/tests/test_duckdb_load.py +595 -0
  14. torchsx-0.1.0/tests/test_ewald_offset.py +72 -0
  15. torchsx-0.1.0/tests/test_filter.py +121 -0
  16. torchsx-0.1.0/tests/test_geometry.py +202 -0
  17. torchsx-0.1.0/tests/test_geometry_reciprocal.py +69 -0
  18. torchsx-0.1.0/tests/test_grid_tooling.py +95 -0
  19. torchsx-0.1.0/tests/test_hkl.py +131 -0
  20. torchsx-0.1.0/tests/test_io.py +119 -0
  21. torchsx-0.1.0/tests/test_large_batch_diagnostics.py +43 -0
  22. torchsx-0.1.0/tests/test_levels.py +124 -0
  23. torchsx-0.1.0/tests/test_loo_pooling.py +436 -0
  24. torchsx-0.1.0/tests/test_marginal_normalisation.py +125 -0
  25. torchsx-0.1.0/tests/test_measurement_sources.py +635 -0
  26. torchsx-0.1.0/tests/test_merge_stats.py +265 -0
  27. torchsx-0.1.0/tests/test_merging_layers.py +84 -0
  28. torchsx-0.1.0/tests/test_metadata.py +78 -0
  29. torchsx-0.1.0/tests/test_native_index.py +102 -0
  30. torchsx-0.1.0/tests/test_native_parse.py +371 -0
  31. torchsx-0.1.0/tests/test_observation_filter.py +85 -0
  32. torchsx-0.1.0/tests/test_pad.py +59 -0
  33. torchsx-0.1.0/tests/test_peer_cross_attention.py +308 -0
  34. torchsx-0.1.0/tests/test_peer_features.py +58 -0
  35. torchsx-0.1.0/tests/test_posterior_nodes.py +232 -0
  36. torchsx-0.1.0/tests/test_posterior_prior.py +59 -0
  37. torchsx-0.1.0/tests/test_posterior_smooth_sigma.py +106 -0
  38. torchsx-0.1.0/tests/test_reciprocal_basis.py +150 -0
  39. torchsx-0.1.0/tests/test_scatter_mixture.py +488 -0
  40. torchsx-0.1.0/tests/test_shape_prior.py +128 -0
  41. torchsx-0.1.0/tests/test_split_gaussian_mean.py +91 -0
  42. torchsx-0.1.0/tests/test_split_half.py +176 -0
  43. torchsx-0.1.0/tests/test_split_log_grid.py +212 -0
  44. torchsx-0.1.0/tests/test_split_mixture.py +232 -0
  45. torchsx-0.1.0/tests/test_stateless_merge.py +66 -0
  46. torchsx-0.1.0/tests/test_stats.py +192 -0
  47. torchsx-0.1.0/tests/test_stream.py +291 -0
  48. torchsx-0.1.0/tests/test_train_schedule.py +35 -0
  49. torchsx-0.1.0/tests/test_wilson_surface.py +160 -0
  50. torchsx-0.1.0/tests/test_writeback.py +60 -0
  51. torchsx-0.1.0/tests/test_writeback_native.py +122 -0
  52. torchsx-0.1.0/torchsx/__init__.py +8 -0
  53. torchsx-0.1.0/torchsx/_categorical.py +58 -0
  54. torchsx-0.1.0/torchsx/_core.py +602 -0
  55. torchsx-0.1.0/torchsx/cli/__init__.py +0 -0
  56. torchsx-0.1.0/torchsx/cli/_pipeline.py +912 -0
  57. torchsx-0.1.0/torchsx/cli/merge_nn.py +205 -0
  58. torchsx-0.1.0/torchsx/cli/merge_stats.py +217 -0
  59. torchsx-0.1.0/torchsx/cli/stats.py +95 -0
  60. torchsx-0.1.0/torchsx/cli/train.py +278 -0
  61. torchsx-0.1.0/torchsx/geometry.py +745 -0
  62. torchsx-0.1.0/torchsx/io/__init__.py +30 -0
  63. torchsx-0.1.0/torchsx/io/_crystfel.py +105 -0
  64. torchsx-0.1.0/torchsx/io/_duckdb.py +258 -0
  65. torchsx-0.1.0/torchsx/io/_metadata.py +187 -0
  66. torchsx-0.1.0/torchsx/io/_native.py +48 -0
  67. torchsx-0.1.0/torchsx/io/_parse/__init__.py +8 -0
  68. torchsx-0.1.0/torchsx/io/_parse/indexing.py +186 -0
  69. torchsx-0.1.0/torchsx/io/_parse/orchestrator.py +258 -0
  70. torchsx-0.1.0/torchsx/io/_parse/records.py +520 -0
  71. torchsx-0.1.0/torchsx/io/_parse/workers.py +35 -0
  72. torchsx-0.1.0/torchsx/io/_parser.py +40 -0
  73. torchsx-0.1.0/torchsx/io/_safetensors.py +170 -0
  74. torchsx-0.1.0/torchsx/io/_writeback.py +463 -0
  75. torchsx-0.1.0/torchsx/io/sources.py +76 -0
  76. torchsx-0.1.0/torchsx/io/times.py +438 -0
  77. torchsx-0.1.0/torchsx/levels.py +306 -0
  78. torchsx-0.1.0/torchsx/measurement.py +217 -0
  79. torchsx-0.1.0/torchsx/merge_stats.py +672 -0
  80. torchsx-0.1.0/torchsx/merging/__init__.py +30 -0
  81. torchsx-0.1.0/torchsx/merging/config.py +52 -0
  82. torchsx-0.1.0/torchsx/merging/hkl/__init__.py +75 -0
  83. torchsx-0.1.0/torchsx/merging/hkl/adjacency.py +71 -0
  84. torchsx-0.1.0/torchsx/merging/hkl/filter.py +177 -0
  85. torchsx-0.1.0/torchsx/merging/hkl/hkl.py +118 -0
  86. torchsx-0.1.0/torchsx/merging/hkl/noise.py +451 -0
  87. torchsx-0.1.0/torchsx/merging/hkl/shell_fit.py +102 -0
  88. torchsx-0.1.0/torchsx/merging/hkl/stats.py +294 -0
  89. torchsx-0.1.0/torchsx/merging/likelihood/__init__.py +10 -0
  90. torchsx-0.1.0/torchsx/merging/likelihood/base.py +64 -0
  91. torchsx-0.1.0/torchsx/merging/likelihood/detector_noise.py +104 -0
  92. torchsx-0.1.0/torchsx/merging/likelihood/poisson.py +173 -0
  93. torchsx-0.1.0/torchsx/merging/observations.py +699 -0
  94. torchsx-0.1.0/torchsx/merging/parameters.py +77 -0
  95. torchsx-0.1.0/torchsx/merging/priors/__init__.py +13 -0
  96. torchsx-0.1.0/torchsx/merging/priors/reciprocal_basis.py +293 -0
  97. torchsx-0.1.0/torchsx/merging/quadrature/__init__.py +15 -0
  98. torchsx-0.1.0/torchsx/merging/quadrature/_fixed_log_grid_triton.py +426 -0
  99. torchsx-0.1.0/torchsx/merging/quadrature/_fixed_log_grid_tuning.json +44 -0
  100. torchsx-0.1.0/torchsx/merging/quadrature/_fused_posterior_triton.py +249 -0
  101. torchsx-0.1.0/torchsx/merging/quadrature/_posterior_nodes_triton.py +378 -0
  102. torchsx-0.1.0/torchsx/merging/quadrature/_split_log_grid_triton.py +294 -0
  103. torchsx-0.1.0/torchsx/merging/quadrature/_split_mixture_triton.py +218 -0
  104. torchsx-0.1.0/torchsx/merging/quadrature/base.py +98 -0
  105. torchsx-0.1.0/torchsx/merging/quadrature/fixed_log_grid.py +305 -0
  106. torchsx-0.1.0/torchsx/merging/quadrature/split_log_grid.py +392 -0
  107. torchsx-0.1.0/torchsx/merging/scale/__init__.py +10 -0
  108. torchsx-0.1.0/torchsx/merging/scale/amortized.py +452 -0
  109. torchsx-0.1.0/torchsx/merging/scale/base.py +44 -0
  110. torchsx-0.1.0/torchsx/merging/scatter/__init__.py +22 -0
  111. torchsx-0.1.0/torchsx/merging/scatter/base.py +54 -0
  112. torchsx-0.1.0/torchsx/merging/scatter/gengauss.py +314 -0
  113. torchsx-0.1.0/torchsx/merging/scatter/mixture.py +269 -0
  114. torchsx-0.1.0/torchsx/merging/scatter/predicted.py +119 -0
  115. torchsx-0.1.0/torchsx/merging/scatter_shape.py +213 -0
  116. torchsx-0.1.0/torchsx/nn/__init__.py +41 -0
  117. torchsx-0.1.0/torchsx/nn/deepsets.py +579 -0
  118. torchsx-0.1.0/torchsx/nn/heads.py +370 -0
  119. torchsx-0.1.0/torchsx/nn/merge_target.py +67 -0
  120. torchsx-0.1.0/torchsx/nn/metrics.py +397 -0
  121. torchsx-0.1.0/torchsx/nn/objective.py +380 -0
  122. torchsx-0.1.0/torchsx/nn/obs_features.py +424 -0
  123. torchsx-0.1.0/torchsx/nn/pad.py +78 -0
  124. torchsx-0.1.0/torchsx/nn/pooling.py +395 -0
  125. torchsx-0.1.0/torchsx/nn/posterior_target.py +662 -0
  126. torchsx-0.1.0/torchsx/nn/shape_prior.py +98 -0
  127. torchsx-0.1.0/torchsx/nn/train_amortized.py +1640 -0
  128. torchsx-0.1.0/torchsx/stats.py +302 -0
  129. torchsx-0.1.0/torchsx/stream.py +826 -0
  130. torchsx-0.1.0/torchsx.egg-info/PKG-INFO +182 -0
  131. torchsx-0.1.0/torchsx.egg-info/SOURCES.txt +133 -0
  132. torchsx-0.1.0/torchsx.egg-info/dependency_links.txt +1 -0
  133. torchsx-0.1.0/torchsx.egg-info/entry_points.txt +5 -0
  134. torchsx-0.1.0/torchsx.egg-info/requires.txt +13 -0
  135. 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
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
40
+ [![PyTorch](https://img.shields.io/badge/PyTorch-2.4+-ee4c2c.svg)](https://pytorch.org/)
41
+ [![TorchRef](https://img.shields.io/badge/TorchRef-0.7-f97316.svg)](https://github.com/HatPdotS/TorchRef)
42
+ [![CUDA](https://img.shields.io/badge/CUDA-supported-76b900.svg)](https://developer.nvidia.com/cuda-zone)
43
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](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).
@@ -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
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
13
+ [![PyTorch](https://img.shields.io/badge/PyTorch-2.4+-ee4c2c.svg)](https://pytorch.org/)
14
+ [![TorchRef](https://img.shields.io/badge/TorchRef-0.7-f97316.svg)](https://github.com/HatPdotS/TorchRef)
15
+ [![CUDA](https://img.shields.io/badge/CUDA-supported-76b900.svg)](https://developer.nvidia.com/cuda-zone)
16
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](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"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+