torchsx-native 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.
@@ -0,0 +1 @@
1
+ target/
@@ -0,0 +1,287 @@
1
+ # This file is automatically @generated by Cargo.
2
+ # It is not intended for manual editing.
3
+ version = 4
4
+
5
+ [[package]]
6
+ name = "autocfg"
7
+ version = "1.5.1"
8
+ source = "registry+https://github.com/rust-lang/crates.io-index"
9
+ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
10
+
11
+ [[package]]
12
+ name = "crossbeam-deque"
13
+ version = "0.8.8"
14
+ source = "registry+https://github.com/rust-lang/crates.io-index"
15
+ checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a"
16
+ dependencies = [
17
+ "crossbeam-epoch",
18
+ "crossbeam-utils",
19
+ ]
20
+
21
+ [[package]]
22
+ name = "crossbeam-epoch"
23
+ version = "0.9.21"
24
+ source = "registry+https://github.com/rust-lang/crates.io-index"
25
+ checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d"
26
+ dependencies = [
27
+ "crossbeam-utils",
28
+ ]
29
+
30
+ [[package]]
31
+ name = "crossbeam-utils"
32
+ version = "0.8.23"
33
+ source = "registry+https://github.com/rust-lang/crates.io-index"
34
+ checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
35
+
36
+ [[package]]
37
+ name = "either"
38
+ version = "1.18.0"
39
+ source = "registry+https://github.com/rust-lang/crates.io-index"
40
+ checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34"
41
+
42
+ [[package]]
43
+ name = "heck"
44
+ version = "0.5.0"
45
+ source = "registry+https://github.com/rust-lang/crates.io-index"
46
+ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
47
+
48
+ [[package]]
49
+ name = "libc"
50
+ version = "0.2.189"
51
+ source = "registry+https://github.com/rust-lang/crates.io-index"
52
+ checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
53
+
54
+ [[package]]
55
+ name = "matrixmultiply"
56
+ version = "0.3.11"
57
+ source = "registry+https://github.com/rust-lang/crates.io-index"
58
+ checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
59
+ dependencies = [
60
+ "autocfg",
61
+ "rawpointer",
62
+ ]
63
+
64
+ [[package]]
65
+ name = "memchr"
66
+ version = "2.8.3"
67
+ source = "registry+https://github.com/rust-lang/crates.io-index"
68
+ checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
69
+
70
+ [[package]]
71
+ name = "ndarray"
72
+ version = "0.17.2"
73
+ source = "registry+https://github.com/rust-lang/crates.io-index"
74
+ checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
75
+ dependencies = [
76
+ "matrixmultiply",
77
+ "num-complex",
78
+ "num-integer",
79
+ "num-traits",
80
+ "portable-atomic",
81
+ "portable-atomic-util",
82
+ "rawpointer",
83
+ ]
84
+
85
+ [[package]]
86
+ name = "num-complex"
87
+ version = "0.4.6"
88
+ source = "registry+https://github.com/rust-lang/crates.io-index"
89
+ checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
90
+ dependencies = [
91
+ "num-traits",
92
+ ]
93
+
94
+ [[package]]
95
+ name = "num-integer"
96
+ version = "0.1.47"
97
+ source = "registry+https://github.com/rust-lang/crates.io-index"
98
+ checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b"
99
+ dependencies = [
100
+ "num-traits",
101
+ ]
102
+
103
+ [[package]]
104
+ name = "num-traits"
105
+ version = "0.2.19"
106
+ source = "registry+https://github.com/rust-lang/crates.io-index"
107
+ checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
108
+ dependencies = [
109
+ "autocfg",
110
+ ]
111
+
112
+ [[package]]
113
+ name = "numpy"
114
+ version = "0.29.0"
115
+ source = "registry+https://github.com/rust-lang/crates.io-index"
116
+ checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4"
117
+ dependencies = [
118
+ "libc",
119
+ "ndarray",
120
+ "num-complex",
121
+ "num-integer",
122
+ "num-traits",
123
+ "pyo3",
124
+ "pyo3-build-config",
125
+ "rustc-hash",
126
+ ]
127
+
128
+ [[package]]
129
+ name = "once_cell"
130
+ version = "1.21.4"
131
+ source = "registry+https://github.com/rust-lang/crates.io-index"
132
+ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
133
+
134
+ [[package]]
135
+ name = "portable-atomic"
136
+ version = "1.15.0"
137
+ source = "registry+https://github.com/rust-lang/crates.io-index"
138
+ checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
139
+
140
+ [[package]]
141
+ name = "portable-atomic-util"
142
+ version = "0.2.8"
143
+ source = "registry+https://github.com/rust-lang/crates.io-index"
144
+ checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715"
145
+ dependencies = [
146
+ "portable-atomic",
147
+ ]
148
+
149
+ [[package]]
150
+ name = "proc-macro2"
151
+ version = "1.0.107"
152
+ source = "registry+https://github.com/rust-lang/crates.io-index"
153
+ checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
154
+ dependencies = [
155
+ "unicode-ident",
156
+ ]
157
+
158
+ [[package]]
159
+ name = "pyo3"
160
+ version = "0.29.2"
161
+ source = "registry+https://github.com/rust-lang/crates.io-index"
162
+ checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
163
+ dependencies = [
164
+ "libc",
165
+ "once_cell",
166
+ "portable-atomic",
167
+ "pyo3-build-config",
168
+ "pyo3-ffi",
169
+ "pyo3-macros",
170
+ ]
171
+
172
+ [[package]]
173
+ name = "pyo3-build-config"
174
+ version = "0.29.2"
175
+ source = "registry+https://github.com/rust-lang/crates.io-index"
176
+ checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
177
+ dependencies = [
178
+ "target-lexicon",
179
+ ]
180
+
181
+ [[package]]
182
+ name = "pyo3-ffi"
183
+ version = "0.29.2"
184
+ source = "registry+https://github.com/rust-lang/crates.io-index"
185
+ checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
186
+ dependencies = [
187
+ "libc",
188
+ "pyo3-build-config",
189
+ ]
190
+
191
+ [[package]]
192
+ name = "pyo3-macros"
193
+ version = "0.29.2"
194
+ source = "registry+https://github.com/rust-lang/crates.io-index"
195
+ checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
196
+ dependencies = [
197
+ "proc-macro2",
198
+ "pyo3-macros-backend",
199
+ "quote",
200
+ "syn",
201
+ ]
202
+
203
+ [[package]]
204
+ name = "pyo3-macros-backend"
205
+ version = "0.29.2"
206
+ source = "registry+https://github.com/rust-lang/crates.io-index"
207
+ checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
208
+ dependencies = [
209
+ "heck",
210
+ "proc-macro2",
211
+ "quote",
212
+ "syn",
213
+ ]
214
+
215
+ [[package]]
216
+ name = "quote"
217
+ version = "1.0.47"
218
+ source = "registry+https://github.com/rust-lang/crates.io-index"
219
+ checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
220
+ dependencies = [
221
+ "proc-macro2",
222
+ ]
223
+
224
+ [[package]]
225
+ name = "rawpointer"
226
+ version = "0.2.1"
227
+ source = "registry+https://github.com/rust-lang/crates.io-index"
228
+ checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
229
+
230
+ [[package]]
231
+ name = "rayon"
232
+ version = "1.12.0"
233
+ source = "registry+https://github.com/rust-lang/crates.io-index"
234
+ checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
235
+ dependencies = [
236
+ "either",
237
+ "rayon-core",
238
+ ]
239
+
240
+ [[package]]
241
+ name = "rayon-core"
242
+ version = "1.13.0"
243
+ source = "registry+https://github.com/rust-lang/crates.io-index"
244
+ checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
245
+ dependencies = [
246
+ "crossbeam-deque",
247
+ "crossbeam-utils",
248
+ ]
249
+
250
+ [[package]]
251
+ name = "rustc-hash"
252
+ version = "2.1.3"
253
+ source = "registry+https://github.com/rust-lang/crates.io-index"
254
+ checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d"
255
+
256
+ [[package]]
257
+ name = "syn"
258
+ version = "2.0.119"
259
+ source = "registry+https://github.com/rust-lang/crates.io-index"
260
+ checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
261
+ dependencies = [
262
+ "proc-macro2",
263
+ "quote",
264
+ "unicode-ident",
265
+ ]
266
+
267
+ [[package]]
268
+ name = "target-lexicon"
269
+ version = "0.13.5"
270
+ source = "registry+https://github.com/rust-lang/crates.io-index"
271
+ checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
272
+
273
+ [[package]]
274
+ name = "torchsx-native"
275
+ version = "0.1.0"
276
+ dependencies = [
277
+ "memchr",
278
+ "numpy",
279
+ "pyo3",
280
+ "rayon",
281
+ ]
282
+
283
+ [[package]]
284
+ name = "unicode-ident"
285
+ version = "1.0.26"
286
+ source = "registry+https://github.com/rust-lang/crates.io-index"
287
+ checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954"
@@ -0,0 +1,25 @@
1
+ [package]
2
+ name = "torchsx-native"
3
+ version = "0.1.0"
4
+ edition = "2021"
5
+ license = "MIT"
6
+ description = "Native CrystFEL .stream reader/writer for torchsx"
7
+ authors = ["Hans Peter Seidel"]
8
+ repository = "https://github.com/HatPdotS/TorchSX"
9
+ readme = "README.md"
10
+ # Released to PyPI as a wheel (see ../.github/workflows/publish.yml), not to crates.io.
11
+ publish = false
12
+
13
+ [lib]
14
+ name = "torchsx_native"
15
+ crate-type = ["cdylib", "rlib"]
16
+
17
+ [dependencies]
18
+ pyo3 = { version = "0.29", features = ["abi3-py310"] }
19
+ numpy = "0.29"
20
+ rayon = "1.12"
21
+ memchr = "2.8"
22
+
23
+ [profile.release]
24
+ lto = "fat"
25
+ codegen-units = 1
@@ -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.
@@ -0,0 +1,78 @@
1
+ Metadata-Version: 2.4
2
+ Name: torchsx-native
3
+ Version: 0.1.0
4
+ Classifier: Programming Language :: Python :: 3
5
+ Classifier: Programming Language :: Rust
6
+ Classifier: Intended Audience :: Science/Research
7
+ Classifier: Topic :: Scientific/Engineering :: Physics
8
+ Requires-Dist: numpy>=2
9
+ License-File: LICENSE
10
+ Summary: Native CrystFEL .stream reader/writer for torchsx
11
+ Author: Hans Peter Seidel
12
+ License-Expression: MIT
13
+ Requires-Python: >=3.10
14
+ Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
15
+ Project-URL: Repository, https://github.com/HatPdotS/TorchSX
16
+
17
+ # torchsx-native
18
+
19
+ The CrystFEL `.stream` indexer, parser and writer for torchsx, in Rust
20
+ ([PyO3](https://pyo3.rs) and [maturin](https://www.maturin.rs)). It is a
21
+ separate wheel so torchsx itself stays pure Python. Arrays cross the boundary
22
+ as numpy arrays, and the module does not link torch.
23
+
24
+ | Function | Does |
25
+ |----------|------|
26
+ | `index_chunks(path, n_threads=1, max_crystals=None)` | `(n, 4)` int64 rows `(start, end, n_crystals, n_obs_upper)` per chunk; `max_crystals` stops early (head mode) |
27
+ | `parse_chunks(path, chunks, panel_names, n_threads=1)` | Parses the given rows into `(columns, filenames, events, report)`; the columns use the `_allocate_tensors` names and dtypes |
28
+ | `write_stream(path, arrays, filenames, events, panel_names, lattice_type, centering, unique_axis, header, n_threads=1, images_per_batch=1024)` | Writes a stream byte-identical to the Python writer |
29
+ | `PROTOCOL` | Interface version; `torchsx/io/_native.py` refuses a build whose value differs |
30
+
31
+ Every call releases the GIL and runs on its own pool of `n_threads`. The
32
+ output does not depend on the thread count.
33
+
34
+ ## Behaviour
35
+
36
+ - On well-formed streams the parse is bit-identical to the pure-Python parser
37
+ (`torchsx/io/_parse/records.py`). Both follow the same rules, and
38
+ `tests/test_native_parse.py` compares them.
39
+ - Malformed content never raises. It is excluded at the smallest unit that
40
+ keeps the rest valid, and `report` counts it by reason, with the byte offset
41
+ of the first occurrence. `records.py` lists the rules.
42
+ - Python's `str.split()`, `str.strip()`, `int()` and `float()` are reproduced,
43
+ including PEP 515 underscores and `inf`/`nan`. Unicode decimal digits other
44
+ than ASCII, which Python's `int()` also accepts, are not numbers here.
45
+
46
+ ## Build
47
+
48
+ ```bash
49
+ maturin build --release # → target/wheels/torchsx_native-*-cp310-abi3-*.whl
50
+ pip install target/wheels/*.whl # or, from the repo root: pip install ./native
51
+ cargo test --release # Rust unit tests (needs libpython on LD_LIBRARY_PATH)
52
+ ```
53
+
54
+ The wheel uses the stable ABI (`abi3`, Python ≥ 3.10), so one build serves every
55
+ Python version on a platform. A wheel built on a host carries that host's glibc
56
+ floor (for example `manylinux_2_34` on SLES 15 SP6). For older systems, build
57
+ with `maturin build --release --zig --compatibility manylinux_2_28`.
58
+
59
+ Bump `PROTOCOL` in `src/lib.rs` and in `torchsx/io/_native.py` together
60
+ whenever a signature or a returned layout changes.
61
+
62
+ ## Release
63
+
64
+ torchsx-native is released from this repository together with torchsx
65
+ (`.github/workflows/publish.yml`, run by publishing a GitHub release `v<torchsx
66
+ version>`, or by hand against TestPyPI). The version is `Cargo.toml`'s. It changes
67
+ only when the native build changes, and torchsx pins its minor version
68
+ (`torchsx-native==0.1.*`):
69
+
70
+ - Rust change, same interface: bump the patch in `Cargo.toml`.
71
+ - Interface change: bump `PROTOCOL` (both files), the minor in `Cargo.toml`
72
+ and the pin in the root `pyproject.toml`.
73
+
74
+ The workflow checks all of this (`.github/scripts/check_release.py`). It builds
75
+ abi3 wheels for manylinux_2_28 x86_64/aarch64, macOS universal2 and Windows x64,
76
+ smoke-tests each (`tests/smoke.py`) on its platform, and uploads torchsx-native
77
+ before torchsx.
78
+
@@ -0,0 +1,61 @@
1
+ # torchsx-native
2
+
3
+ The CrystFEL `.stream` indexer, parser and writer for torchsx, in Rust
4
+ ([PyO3](https://pyo3.rs) and [maturin](https://www.maturin.rs)). It is a
5
+ separate wheel so torchsx itself stays pure Python. Arrays cross the boundary
6
+ as numpy arrays, and the module does not link torch.
7
+
8
+ | Function | Does |
9
+ |----------|------|
10
+ | `index_chunks(path, n_threads=1, max_crystals=None)` | `(n, 4)` int64 rows `(start, end, n_crystals, n_obs_upper)` per chunk; `max_crystals` stops early (head mode) |
11
+ | `parse_chunks(path, chunks, panel_names, n_threads=1)` | Parses the given rows into `(columns, filenames, events, report)`; the columns use the `_allocate_tensors` names and dtypes |
12
+ | `write_stream(path, arrays, filenames, events, panel_names, lattice_type, centering, unique_axis, header, n_threads=1, images_per_batch=1024)` | Writes a stream byte-identical to the Python writer |
13
+ | `PROTOCOL` | Interface version; `torchsx/io/_native.py` refuses a build whose value differs |
14
+
15
+ Every call releases the GIL and runs on its own pool of `n_threads`. The
16
+ output does not depend on the thread count.
17
+
18
+ ## Behaviour
19
+
20
+ - On well-formed streams the parse is bit-identical to the pure-Python parser
21
+ (`torchsx/io/_parse/records.py`). Both follow the same rules, and
22
+ `tests/test_native_parse.py` compares them.
23
+ - Malformed content never raises. It is excluded at the smallest unit that
24
+ keeps the rest valid, and `report` counts it by reason, with the byte offset
25
+ of the first occurrence. `records.py` lists the rules.
26
+ - Python's `str.split()`, `str.strip()`, `int()` and `float()` are reproduced,
27
+ including PEP 515 underscores and `inf`/`nan`. Unicode decimal digits other
28
+ than ASCII, which Python's `int()` also accepts, are not numbers here.
29
+
30
+ ## Build
31
+
32
+ ```bash
33
+ maturin build --release # → target/wheels/torchsx_native-*-cp310-abi3-*.whl
34
+ pip install target/wheels/*.whl # or, from the repo root: pip install ./native
35
+ cargo test --release # Rust unit tests (needs libpython on LD_LIBRARY_PATH)
36
+ ```
37
+
38
+ The wheel uses the stable ABI (`abi3`, Python ≥ 3.10), so one build serves every
39
+ Python version on a platform. A wheel built on a host carries that host's glibc
40
+ floor (for example `manylinux_2_34` on SLES 15 SP6). For older systems, build
41
+ with `maturin build --release --zig --compatibility manylinux_2_28`.
42
+
43
+ Bump `PROTOCOL` in `src/lib.rs` and in `torchsx/io/_native.py` together
44
+ whenever a signature or a returned layout changes.
45
+
46
+ ## Release
47
+
48
+ torchsx-native is released from this repository together with torchsx
49
+ (`.github/workflows/publish.yml`, run by publishing a GitHub release `v<torchsx
50
+ version>`, or by hand against TestPyPI). The version is `Cargo.toml`'s. It changes
51
+ only when the native build changes, and torchsx pins its minor version
52
+ (`torchsx-native==0.1.*`):
53
+
54
+ - Rust change, same interface: bump the patch in `Cargo.toml`.
55
+ - Interface change: bump `PROTOCOL` (both files), the minor in `Cargo.toml`
56
+ and the pin in the root `pyproject.toml`.
57
+
58
+ The workflow checks all of this (`.github/scripts/check_release.py`). It builds
59
+ abi3 wheels for manylinux_2_28 x86_64/aarch64, macOS universal2 and Windows x64,
60
+ smoke-tests each (`tests/smoke.py`) on its platform, and uploads torchsx-native
61
+ before torchsx.
@@ -0,0 +1,26 @@
1
+ [build-system]
2
+ requires = ["maturin>=1.15,<2"]
3
+ build-backend = "maturin"
4
+
5
+ [project]
6
+ name = "torchsx-native"
7
+ # The version is Cargo.toml's (maturin fills it in).
8
+ dynamic = ["version"]
9
+ description = "Native CrystFEL .stream reader/writer for torchsx"
10
+ readme = "README.md"
11
+ license = "MIT"
12
+ authors = [{ name = "Hans Peter Seidel" }]
13
+ requires-python = ">=3.10"
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "Programming Language :: Rust",
17
+ "Intended Audience :: Science/Research",
18
+ "Topic :: Scientific/Engineering :: Physics",
19
+ ]
20
+ dependencies = ["numpy>=2"]
21
+
22
+ [project.urls]
23
+ Repository = "https://github.com/HatPdotS/TorchSX"
24
+
25
+ [tool.maturin]
26
+ module-name = "torchsx_native"
@@ -0,0 +1,189 @@
1
+ //! Chunk index: `(start, end, n_crystals, n_obs_upper)` per
2
+ //! ``----- Begin chunk -----`` block, matching the Python ``_index_chunks``.
3
+ //!
4
+ //! A chunk marker always ends a reflection table, so the scan state at every
5
+ //! ``----- Begin chunk -----`` line is the same whatever came before. That is
6
+ //! what lets the threaded scan split the file at arbitrary chunk starts and
7
+ //! still return exactly the serial result, also for truncated tables.
8
+
9
+ use std::io;
10
+ use std::path::Path;
11
+
12
+ use rayon::prelude::*;
13
+
14
+ use crate::lines::LineReader;
15
+ use crate::pystr::bytes_space;
16
+
17
+ pub const BEGIN_CHUNK: &[u8] = b"----- Begin chunk -----";
18
+ pub const END_CHUNK: &[u8] = b"----- End chunk -----";
19
+ const BEGIN_CRYSTAL: &[u8] = b"--- Begin crystal";
20
+ const REFL_BEGIN: &[u8] = b"Reflections measured after indexing";
21
+ const REFL_END: &[u8] = b"End of reflections";
22
+
23
+ #[derive(Clone, Copy, Debug, PartialEq, Eq)]
24
+ pub struct Chunk {
25
+ pub start: u64,
26
+ pub end: u64,
27
+ pub n_crystals: i64,
28
+ pub n_obs: i64,
29
+ }
30
+
31
+ #[inline]
32
+ pub fn is_chunk_marker(line: &[u8]) -> bool {
33
+ line.starts_with(BEGIN_CHUNK) || line.starts_with(END_CHUNK)
34
+ }
35
+
36
+ /// ``len(line.split()) >= 10`` on raw bytes.
37
+ #[inline]
38
+ fn has_ten_fields(line: &[u8]) -> bool {
39
+ let mut n = 0;
40
+ let mut in_tok = false;
41
+ for &b in line {
42
+ let sp = bytes_space(b);
43
+ if !sp && !in_tok {
44
+ n += 1;
45
+ if n == 10 {
46
+ return true;
47
+ }
48
+ }
49
+ in_tok = !sp;
50
+ }
51
+ false
52
+ }
53
+
54
+ struct Scan {
55
+ chunks: Vec<Chunk>,
56
+ /// The last chunk of the range, whose end is the caller's to set.
57
+ open: Option<Chunk>,
58
+ }
59
+
60
+ /// Scan `[start, end)`. With `max_crystals`, stop once the flushed chunks hold
61
+ /// at least that many crystals (the chunk that crosses it is included).
62
+ fn scan(path: &Path, start: u64, end: u64, max_crystals: Option<i64>) -> io::Result<Scan> {
63
+ let mut r = LineReader::open(path)?;
64
+ r.set_range(start, end)?;
65
+ let mut chunks = Vec::new();
66
+ let mut open: Option<Chunk> = None;
67
+ let mut cum = 0i64;
68
+ let mut in_refl = false;
69
+ let mut header_pending = false;
70
+
71
+ while let Some((pos, line)) = r.next_line()? {
72
+ if in_refl {
73
+ if header_pending {
74
+ header_pending = false;
75
+ if !is_chunk_marker(line) {
76
+ continue;
77
+ }
78
+ }
79
+ if line.starts_with(REFL_END) {
80
+ in_refl = false;
81
+ continue;
82
+ }
83
+ if !is_chunk_marker(line) {
84
+ if has_ten_fields(line) {
85
+ if let Some(c) = open.as_mut() {
86
+ c.n_obs += 1;
87
+ }
88
+ }
89
+ continue;
90
+ }
91
+ in_refl = false;
92
+ }
93
+ match line.first() {
94
+ Some(b'-') => {
95
+ if line.starts_with(BEGIN_CHUNK) {
96
+ if let Some(mut c) = open.take() {
97
+ c.end = pos;
98
+ cum += c.n_crystals;
99
+ chunks.push(c);
100
+ if max_crystals.is_some_and(|m| cum >= m) {
101
+ return Ok(Scan { chunks, open: None });
102
+ }
103
+ }
104
+ open = Some(Chunk { start: pos, end: 0, n_crystals: 0, n_obs: 0 });
105
+ } else if line.starts_with(BEGIN_CRYSTAL) {
106
+ // Crystals before the first chunk (header region) are not counted.
107
+ if let Some(c) = open.as_mut() {
108
+ c.n_crystals += 1;
109
+ }
110
+ }
111
+ }
112
+ Some(b'R') if line.starts_with(REFL_BEGIN) => {
113
+ in_refl = true;
114
+ header_pending = true;
115
+ }
116
+ _ => {}
117
+ }
118
+ }
119
+ Ok(Scan { chunks, open })
120
+ }
121
+
122
+ /// Offset of the first ``----- Begin chunk -----`` line starting at or after
123
+ /// the line that contains `target` (the partial line at `target` is skipped),
124
+ /// or the file size.
125
+ fn find_chunk_at_or_after(path: &Path, target: u64, size: u64) -> io::Result<u64> {
126
+ let mut r = LineReader::open(path)?;
127
+ r.set_range(target.min(size), size)?;
128
+ if target > 0 && r.next_line()?.is_none() {
129
+ return Ok(size);
130
+ }
131
+ while let Some((pos, line)) = r.next_line()? {
132
+ if line.starts_with(BEGIN_CHUNK) {
133
+ return Ok(pos);
134
+ }
135
+ }
136
+ Ok(size)
137
+ }
138
+
139
+ pub fn index_chunks(path: &Path, n_threads: usize, max_crystals: Option<i64>) -> io::Result<Vec<Chunk>> {
140
+ let size = std::fs::metadata(path)?.len();
141
+ let finish = |s: Scan, end: u64| {
142
+ let mut v = s.chunks;
143
+ if let Some(mut c) = s.open {
144
+ c.end = end;
145
+ v.push(c);
146
+ }
147
+ v
148
+ };
149
+ if max_crystals.is_some() || n_threads <= 1 {
150
+ return Ok(finish(scan(path, 0, size, max_crystals)?, size));
151
+ }
152
+
153
+ let first = find_chunk_at_or_after(path, 0, size)?;
154
+ if first >= size {
155
+ return Ok(Vec::new());
156
+ }
157
+ let mut starts = vec![first];
158
+ for i in 1..n_threads as u64 {
159
+ let target = first + (size - first) * i / n_threads as u64;
160
+ let s = find_chunk_at_or_after(path, target, size)?;
161
+ if s > *starts.last().unwrap() && s < size {
162
+ starts.push(s);
163
+ }
164
+ }
165
+ let mut bounds: Vec<(u64, u64)> = starts.windows(2).map(|w| (w[0], w[1])).collect();
166
+ bounds.push((*starts.last().unwrap(), size));
167
+
168
+ let scans: Vec<Scan> = bounds
169
+ .par_iter()
170
+ .map(|&(s, e)| scan(path, s, e, None))
171
+ .collect::<io::Result<_>>()?;
172
+ let mut out = Vec::new();
173
+ for (scan, &(_, e)) in scans.into_iter().zip(&bounds) {
174
+ out.extend(finish(scan, e));
175
+ }
176
+ Ok(out)
177
+ }
178
+
179
+ #[cfg(test)]
180
+ mod tests {
181
+ use super::*;
182
+
183
+ #[test]
184
+ fn ten_fields() {
185
+ assert!(has_ten_fields(b" 1 2 3 1.0 2.0 3.0 4.0 5.0 6.0 p0"));
186
+ assert!(!has_ten_fields(b" 1 2 3 1.0 2.0 3.0 4.0 5.0 6.0"));
187
+ assert!(!has_ten_fields(b""));
188
+ }
189
+ }