diffct 2.0.2__tar.gz → 2.2.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 (83) hide show
  1. {diffct-2.0.2 → diffct-2.2.0}/CHANGELOG.md +72 -0
  2. {diffct-2.0.2 → diffct-2.2.0}/CITATION.cff +2 -2
  3. {diffct-2.0.2 → diffct-2.2.0}/PKG-INFO +45 -19
  4. {diffct-2.0.2 → diffct-2.2.0}/README.md +44 -18
  5. {diffct-2.0.2 → diffct-2.2.0}/diffct/__init__.py +6 -1
  6. diffct-2.2.0/diffct/chunking.py +201 -0
  7. {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/cone_beam.py +46 -36
  8. {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/fan_beam.py +30 -22
  9. {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/geometry_vjp.py +77 -48
  10. {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/parallel_beam.py +32 -22
  11. diffct-2.2.0/diffct/operators.py +2047 -0
  12. diffct-2.2.0/diffct/pipeline.py +258 -0
  13. {diffct-2.0.2 → diffct-2.2.0}/diffct/projectors.py +155 -125
  14. diffct-2.2.0/diffct/storage.py +125 -0
  15. diffct-2.2.0/diffct/streaming.py +48 -0
  16. diffct-2.2.0/diffct/surfaces.py +103 -0
  17. {diffct-2.0.2 → diffct-2.2.0}/diffct/utils.py +5 -0
  18. diffct-2.2.0/docs/video/README.md +119 -0
  19. {diffct-2.0.2 → diffct-2.2.0}/examples/README.md +55 -2
  20. diffct-2.2.0/formal/README.md +151 -0
  21. {diffct-2.0.2 → diffct-2.2.0}/pyproject.toml +1 -1
  22. diffct-2.2.0/tests/test_allocator_capacity.py +349 -0
  23. diffct-2.2.0/tests/test_axis_aligned_chords.py +118 -0
  24. diffct-2.2.0/tests/test_batched_surfaces.py +705 -0
  25. diffct-2.2.0/tests/test_block_cgls.py +996 -0
  26. diffct-2.2.0/tests/test_block_storage.py +539 -0
  27. diffct-2.2.0/tests/test_cgls_state_admission.py +173 -0
  28. diffct-2.2.0/tests/test_chunked_projector.py +1298 -0
  29. diffct-2.2.0/tests/test_completed_transfer_reuse.py +92 -0
  30. diffct-2.2.0/tests/test_cuda_sink_lifetimes.py +125 -0
  31. diffct-2.2.0/tests/test_detector_surface_dtypes.py +80 -0
  32. diffct-2.2.0/tests/test_detector_surface_mapping.py +79 -0
  33. diffct-2.2.0/tests/test_detector_surface_units.py +62 -0
  34. diffct-2.2.0/tests/test_detector_surfaces.py +705 -0
  35. diffct-2.2.0/tests/test_example_psnr.py +75 -0
  36. diffct-2.2.0/tests/test_execution_buffers.py +613 -0
  37. diffct-2.2.0/tests/test_execution_schedules.py +582 -0
  38. diffct-2.2.0/tests/test_output_pipeline_overlap.py +130 -0
  39. {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_api.py +9 -3
  40. diffct-2.2.0/tests/test_spatial_partition.py +433 -0
  41. diffct-2.2.0/tests/test_transfer_pipeline.py +646 -0
  42. diffct-2.2.0/tests/test_transfer_retirement_lifetimes.py +166 -0
  43. diffct-2.0.2/diffct/operators.py +0 -555
  44. diffct-2.0.2/docs/video/README.md +0 -47
  45. diffct-2.0.2/formal/README.md +0 -82
  46. {diffct-2.0.2 → diffct-2.2.0}/.gitignore +0 -0
  47. {diffct-2.0.2 → diffct-2.2.0}/LICENSE +0 -0
  48. {diffct-2.0.2 → diffct-2.2.0}/NOTICE +0 -0
  49. {diffct-2.0.2 → diffct-2.2.0}/diffct/analytical.py +0 -0
  50. {diffct-2.0.2 → diffct-2.2.0}/diffct/constants.py +0 -0
  51. {diffct-2.0.2 → diffct-2.2.0}/diffct/differentiable.py +0 -0
  52. {diffct-2.0.2 → diffct-2.2.0}/diffct/geometry.py +0 -0
  53. {diffct-2.0.2 → diffct-2.2.0}/diffct/kernels/__init__.py +0 -0
  54. {diffct-2.0.2 → diffct-2.2.0}/examples/data/NOTICE +0 -0
  55. {diffct-2.0.2 → diffct-2.2.0}/pytest.ini +0 -0
  56. {diffct-2.0.2 → diffct-2.2.0}/tests/__init__.py +0 -0
  57. {diffct-2.0.2 → diffct-2.2.0}/tests/_phantoms.py +0 -0
  58. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/README.md +0 -0
  59. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/__init__.py +0 -0
  60. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/_common.py +0 -0
  61. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_cone.py +0 -0
  62. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_fan.py +0 -0
  63. {diffct-2.0.2 → diffct-2.2.0}/tests/benchmarks/test_bench_parallel.py +0 -0
  64. {diffct-2.0.2 → diffct-2.2.0}/tests/distributed_projector_check.py +0 -0
  65. {diffct-2.0.2 → diffct-2.2.0}/tests/test_adjoint_inner_product.py +0 -0
  66. {diffct-2.0.2 → diffct-2.2.0}/tests/test_analytical_regressions.py +0 -0
  67. {diffct-2.0.2 → diffct-2.2.0}/tests/test_cone_projector_autograd.py +0 -0
  68. {diffct-2.0.2 → diffct-2.2.0}/tests/test_cuda_exit_deallocations.py +0 -0
  69. {diffct-2.0.2 → diffct-2.2.0}/tests/test_cuda_smoke.py +0 -0
  70. {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_fan_accuracy.py +0 -0
  71. {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_fan_offsets.py +0 -0
  72. {diffct-2.0.2 → diffct-2.2.0}/tests/test_fbp_parallel_accuracy.py +0 -0
  73. {diffct-2.0.2 → diffct-2.2.0}/tests/test_fdk_cone_accuracy.py +0 -0
  74. {diffct-2.0.2 → diffct-2.2.0}/tests/test_fdk_cone_offsets.py +0 -0
  75. {diffct-2.0.2 → diffct-2.2.0}/tests/test_geometry_autograd.py +0 -0
  76. {diffct-2.0.2 → diffct-2.2.0}/tests/test_geometry_gradients.py +0 -0
  77. {diffct-2.0.2 → diffct-2.2.0}/tests/test_gradcheck.py +0 -0
  78. {diffct-2.0.2 → diffct-2.2.0}/tests/test_multi_gpu_exit.py +0 -0
  79. {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_degenerate_geometry.py +0 -0
  80. {diffct-2.0.2 → diffct-2.2.0}/tests/test_projector_geometry_devices.py +0 -0
  81. {diffct-2.0.2 → diffct-2.2.0}/tests/test_ramp_filter_windows.py +0 -0
  82. {diffct-2.0.2 → diffct-2.2.0}/tests/test_release_2_0_1.py +0 -0
  83. {diffct-2.0.2 → diffct-2.2.0}/tests/test_weights.py +0 -0
@@ -5,6 +5,78 @@ All notable changes to this project are documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [Unreleased]
9
+
10
+ ## [2.2.0] - 2026-10-11
11
+
12
+ ### Added
13
+
14
+ - Operation-specific automatic schedules and slab/block comparisons for
15
+ native projection and matched backprojection, with bounded CUDA pilots
16
+ and per-call execution/copy metadata.
17
+ - Detector-pixel tiling and explicit `ParameterizedSurface` batches using
18
+ global pixel/view coordinates, saved parameter/frame values and local VJP
19
+ reduction. Legacy two-argument surface callbacks remain compatible.
20
+ - Caller-owned numerical `project_into`/`backproject_into` outputs, tensor
21
+ block access and mapped `.npy` stores for data exceeding host RAM.
22
+ - Distributed `partition="space"` with fixed volume slabs, ownership-aware
23
+ autograd/norms and disk CGLS state/checkpoints without full-volume gathering.
24
+ - Independent CUDA tests and bounded formal protocol/ownership verification.
25
+
26
+ ### Changed
27
+
28
+ - GPU backprojection accumulators now span all contributing view batches,
29
+ with final transfer after tile completion. Contiguous native volume layouts
30
+ and ray buffers are reused across batches.
31
+ - Transfer execution uses bounded pinned slots, separate streams and events;
32
+ all extra staging/cached GPU buffers are included in capacity estimates.
33
+ Local GPU and NCCL reductions occur before completed-result downloads.
34
+ - Example CGLS norms accumulate FP64 blocks and vector updates reuse buffers.
35
+ The disk example keeps all three volume states plus residual/projection state
36
+ blockwise and retains `q` until the global `alpha` decision.
37
+
38
+ ### Fixed
39
+
40
+ - Tensor CGLS ranks agree on CPU fallback before jointly replanning unequal
41
+ spatial shards, preserving collective order.
42
+ - Parameterized detector surfaces use a common arithmetic dtype of at least
43
+ float32, supporting mixed frame dtypes and finite half-precision inputs.
44
+ - CUDA store uploads wait for newly generated blocks on their source stream,
45
+ including transfers within one GPU.
46
+ - Single-slot execution retires prior native computation before allocating
47
+ the next tile, keeping deferred CUDA allocations within its admission bound.
48
+ - Mapped input and output stores reject hard links to the same file before writes.
49
+ - Example PSNR supports CPU reconstruction results with CUDA reference volumes.
50
+ - Deferred CPU output writes retain their submission order across spatial tiles.
51
+ - Coronal slice displays preserve the measured detector row orientation.
52
+
53
+ ## [2.1.0] - 2026-10-10
54
+
55
+ ### Added
56
+
57
+ - `Projector(..., detector_surface=surface)` samples local parameterized surfaces
58
+ into per-pixel world positions for native parallel, fan and cone projection
59
+ and matched backprojection. Shared or per-view surfaces can contain trainable
60
+ PyTorch parameters. Flat detectors remain the default; analytical helpers
61
+ retain their flat-detector assumptions.
62
+ - CUDA regression tests and CPU formal verification for detector
63
+ surfaces, including bounded production-kernel checks against cell intersections.
64
+ - Automatic CPU-backed spatial tiles and view batches for native `Projector`
65
+ projection, matched backprojection, data autograd/Hessians and first-order
66
+ geometry gradients. CUDA inputs retain the full-volume path when its estimated
67
+ working set fits. Optional `volume_chunk_shape` and `view_chunk_size` override
68
+ automatic limits; full CPU arrays are not staged on CUDA. Several local
69
+ devices compute their view batches concurrently.
70
+ - A CPU-backed CGLS example and chunk-boundary, CUDA memory and
71
+ formal partition/translation/adjoint checks.
72
+
73
+ ### Fixed
74
+
75
+ - The exit-time cleanup hook returns early under the CUDA simulator
76
+ (``NUMBA_ENABLE_CUDASIM=1``). Before, each simulator process printed an
77
+ ``ImportError`` for ``_PendingDeallocs`` at exit. The simulator loads no
78
+ CUDA modules, so it needs no cleanup.
79
+
8
80
  ## [2.0.2] - 2026-10-09
9
81
 
10
82
  ### Fixed
@@ -13,8 +13,8 @@ authors:
13
13
  given-names: Yipeng
14
14
  email: yipeng.sun@fau.de
15
15
  alias: sypsyp97
16
- version: 2.0.2
17
- date-released: "2026-10-09"
16
+ version: 2.2.0
17
+ date-released: "2026-10-11"
18
18
  license: Apache-2.0
19
19
  repository-code: "https://github.com/sypsyp97/diffct"
20
20
  url: "https://sypsyp97.github.io/diffct/"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: diffct
3
- Version: 2.0.2
3
+ Version: 2.2.0
4
4
  Summary: Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU and multi-node execution, geometry gradients
5
5
  Project-URL: Homepage, https://github.com/sypsyp97/diffct
6
6
  Project-URL: Documentation, https://sypsyp97.github.io/diffct/
@@ -47,11 +47,11 @@ Description-Content-Type: text/markdown
47
47
  </p>
48
48
 
49
49
  <p align="center">
50
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
50
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
51
51
  </p>
52
52
 
53
53
  <p align="center">
54
- <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
54
+ <a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
55
55
  <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
56
56
  <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
57
57
  </p>
@@ -65,8 +65,6 @@ Description-Content-Type: text/markdown
65
65
  <a href="#citation">Citation</a>
66
66
  </p>
67
67
 
68
- > **Install** with the CUDA extra, `pip install "diffct[cu12]"` or `pip install "diffct[cu13]"`, after PyTorch. A plain `pip install diffct` needs a system CUDA Toolkit.
69
- >
70
68
  > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
71
69
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
72
70
 
@@ -75,8 +73,10 @@ Description-Content-Type: text/markdown
75
73
  To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
76
74
 
77
75
  - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
76
+ - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
78
77
  - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
79
- - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun and NCCL. Views are partitioned and the volume is replicated. Speedup depends on workload and communication costs.
78
+ - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
79
+ - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
80
80
  - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
81
81
 
82
82
  Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
@@ -138,41 +138,67 @@ loss.backward() # x.grad = A^T (A x - y)
138
138
 
139
139
  **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
140
140
 
141
- **Several GPUs and nodes.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun for one or more nodes. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
141
+ **Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
142
142
 
143
143
  ```python
144
144
  A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
145
145
  ```
146
146
 
147
+ **Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
148
+ CUDA tiles automatically; CUDA inputs retain the full-volume path when its
149
+ estimated working set fits. Explicit limits are optional:
150
+
151
+ ```python
152
+ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
153
+ volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
154
+ sinogram = A.project(volume) # CPU output, CUDA computation
155
+ sinogram.square().mean().backward() # CPU volume.grad
156
+ # Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
157
+ ```
158
+
159
+ Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
160
+ CUDA allocation. Chunking preserves the Siddon model, with float32 summation
161
+ differences; smaller blocks can increase transfer and launch overhead. Full
162
+ tensor arrays and iterative state still need enough host RAM. For disk-backed
163
+ state and sharded checkpoints, run
164
+ `python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
165
+ volume slabs use `--partition space` under torchrun. Run
166
+ `python examples/benchmark_execution.py` to compare schedules and copy payloads.
167
+
147
168
  ## Performance
148
169
 
149
170
  <p align="center">
150
171
  <picture>
151
- <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
152
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 64^3 and 128^3">
172
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
173
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
153
174
  </picture>
154
175
  </p>
155
176
 
156
- | 128³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
177
+ | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
157
178
  |---|---:|---:|---:|
158
- | Forward projection | 8.45 ms | 2.66 ms (3.2×) | 1.68 ms (5.0×) |
159
- | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
160
- | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
179
+ | Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
180
+ | Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
181
+ | CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
182
+
183
+ diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
161
184
 
162
- Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB GPUs. Small volumes scale less: at 64³, one CGLS iteration takes 4.63 ms on 1 GPU, 1.57 ms on 4 and 1.56 ms on 8. Raw data: [docs/assets/scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
185
+ Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
163
186
 
164
- To measure one against several GPUs on your machine (default scan of the script, not the table setup):
187
+ Run the same workload:
165
188
 
166
189
  ```bash
167
- python examples/benchmark_projector.py --devices 0 1
190
+ python examples/benchmark_scaling.py --output one.json
191
+ torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
168
192
  ```
169
193
 
170
194
  ## Gallery
171
195
 
196
+ TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
197
+
172
198
  **Measured walnut.** 240 measured views, circular cone beam, 256³: FDK (Hann window), SIRT (200 iterations), CGLS (20), TV (300, weight 0.3); axial and coronal centre slices.
173
199
 
174
200
  <p align="center">
175
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
201
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png?v=e0092c4" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
176
202
  </p>
177
203
 
178
204
  ```bash
@@ -182,7 +208,7 @@ python examples/walnut_reconstruction.py --figure out.png
182
208
  **Simulated helical scan of the walnut.** The FDK walnut volume as ground truth, 720 views, 1% noise, 256³. PSNR: FDK 26.67 dB, CGLS (30 iterations) 34.13 dB, SIRT (200) 34.52 dB, TV (200, weight 1.0) 37.77 dB.
183
209
 
184
210
  <p align="center">
185
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
211
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png?v=e0092c4" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
186
212
  </p>
187
213
 
188
214
  ```bash
@@ -194,7 +220,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
194
220
 
195
221
  ## Examples
196
222
 
197
- `quickstart.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
223
+ `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
198
224
 
199
225
  ## Documentation
200
226
 
@@ -16,11 +16,11 @@
16
16
  </p>
17
17
 
18
18
  <p align="center">
19
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
19
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.gif" width="100%" alt="diffct intro: projection, sinogram, trajectories, multi-GPU split, reconstruction and geometry calibration">
20
20
  </p>
21
21
 
22
22
  <p align="center">
23
- <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
23
+ <a href="https://github.com/sypsyp97/diffct/blob/edf7c52227ec2cc14c988918faa7313f8ee0fe03/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
24
24
  <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
25
25
  <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
26
26
  </p>
@@ -34,8 +34,6 @@
34
34
  <a href="#citation">Citation</a>
35
35
  </p>
36
36
 
37
- > **Install** with the CUDA extra, `pip install "diffct[cu12]"` or `pip install "diffct[cu13]"`, after PyTorch. A plain `pip install diffct` needs a system CUDA Toolkit.
38
- >
39
37
  > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
40
38
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
41
39
 
@@ -44,8 +42,10 @@
44
42
  To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
45
43
 
46
44
  - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
45
+ - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
47
46
  - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
48
- - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun and NCCL. Views are partitioned and the volume is replicated. Speedup depends on workload and communication costs.
47
+ - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
48
+ - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Default `partition="views"` shards projections and replicates volumes. `partition="space"` gives each rank an owned volume slab and combines ray contributions. Speedup depends on workload and communication costs.
49
49
  - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
50
50
 
51
51
  Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
@@ -107,41 +107,67 @@ loss.backward() # x.grad = A^T (A x - y)
107
107
 
108
108
  **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
109
109
 
110
- **Several GPUs and nodes.** Use `devices=[0, 1, 2, 3]` in one process, or one process per GPU with torchrun for one or more nodes. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
110
+ **Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
111
111
 
112
112
  ```python
113
113
  A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
114
114
  ```
115
115
 
116
+ **Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
117
+ CUDA tiles automatically; CUDA inputs retain the full-volume path when its
118
+ estimated working set fits. Explicit limits are optional:
119
+
120
+ ```python
121
+ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
122
+ volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
123
+ sinogram = A.project(volume) # CPU output, CUDA computation
124
+ sinogram.square().mean().backward() # CPU volume.grad
125
+ # Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
126
+ ```
127
+
128
+ Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
129
+ CUDA allocation. Chunking preserves the Siddon model, with float32 summation
130
+ differences; smaller blocks can increase transfer and launch overhead. Full
131
+ tensor arrays and iterative state still need enough host RAM. For disk-backed
132
+ state and sharded checkpoints, run
133
+ `python examples/disk_reconstruction.py --output disk-run`; for fixed rank-local
134
+ volume slabs use `--partition space` under torchrun. Run
135
+ `python examples/benchmark_execution.py` to compare schedules and copy payloads.
136
+
116
137
  ## Performance
117
138
 
118
139
  <p align="center">
119
140
  <picture>
120
- <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
121
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 64^3 and 128^3">
141
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
142
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
122
143
  </picture>
123
144
  </p>
124
145
 
125
- | 128³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
146
+ | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
126
147
  |---|---:|---:|---:|
127
- | Forward projection | 8.45 ms | 2.66 ms (3.2×) | 1.68 ms (5.0×) |
128
- | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
129
- | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
148
+ | Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
149
+ | Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
150
+ | CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
151
+
152
+ diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
130
153
 
131
- Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB GPUs. Small volumes scale less: at 64³, one CGLS iteration takes 4.63 ms on 1 GPU, 1.57 ms on 4 and 1.56 ms on 8. Raw data: [docs/assets/scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
154
+ Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
132
155
 
133
- To measure one against several GPUs on your machine (default scan of the script, not the table setup):
156
+ Run the same workload:
134
157
 
135
158
  ```bash
136
- python examples/benchmark_projector.py --devices 0 1
159
+ python examples/benchmark_scaling.py --output one.json
160
+ torchrun --standalone --nproc-per-node=4 examples/benchmark_scaling.py --expected-nodes 1 --output four.json
137
161
  ```
138
162
 
139
163
  ## Gallery
140
164
 
165
+ TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
166
+
141
167
  **Measured walnut.** 240 measured views, circular cone beam, 256³: FDK (Hann window), SIRT (200 iterations), CGLS (20), TV (300, weight 0.3); axial and coronal centre slices.
142
168
 
143
169
  <p align="center">
144
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
170
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_measured.png?v=e0092c4" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
145
171
  </p>
146
172
 
147
173
  ```bash
@@ -151,7 +177,7 @@ python examples/walnut_reconstruction.py --figure out.png
151
177
  **Simulated helical scan of the walnut.** The FDK walnut volume as ground truth, 720 views, 1% noise, 256³. PSNR: FDK 26.67 dB, CGLS (30 iterations) 34.13 dB, SIRT (200) 34.52 dB, TV (200, weight 1.0) 37.77 dB.
152
178
 
153
179
  <p align="center">
154
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
180
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/walnut_helical.png?v=e0092c4" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
155
181
  </p>
156
182
 
157
183
  ```bash
@@ -163,7 +189,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
163
189
 
164
190
  ## Examples
165
191
 
166
- `quickstart.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
192
+ `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `analytical_reconstruction.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
167
193
 
168
194
  ## Documentation
169
195
 
@@ -15,6 +15,8 @@ from .projectors import (
15
15
  )
16
16
 
17
17
  from .operators import Projector
18
+ from .storage import TensorStore, NpyStore
19
+ from .surfaces import ParameterizedSurface
18
20
 
19
21
  from .geometry import (
20
22
  circular_trajectory_3d,
@@ -49,10 +51,13 @@ try:
49
51
  except ImportError:
50
52
  pass
51
53
 
52
- __version__ = '2.0.2'
54
+ __version__ = '2.2.0'
53
55
 
54
56
  __all__ = [
55
57
  'Projector',
58
+ 'TensorStore',
59
+ 'NpyStore',
60
+ 'ParameterizedSurface',
56
61
  'ParallelProjectorFunction',
57
62
  'ParallelBackprojectorFunction',
58
63
  'FanProjectorFunction',
@@ -0,0 +1,201 @@
1
+ """Spatial slices and immutable execution limits for streamed projectors."""
2
+
3
+ from dataclasses import dataclass
4
+ from itertools import product
5
+ import math
6
+
7
+ import torch
8
+
9
+
10
+ @dataclass(frozen=True)
11
+ class _ChunkPlan:
12
+ chunk_shape: tuple
13
+ view_chunk_size: int
14
+ devices: tuple
15
+ detector_chunk_shape: tuple = None
16
+ gpu_budget_bytes: int = 0
17
+ allocation_limit: int = (1 << 63) - 1
18
+ schedule: str = "spatial"
19
+ window_size: int = 1
20
+ slots: int = 1
21
+ estimated_gpu_bytes: int = 0
22
+ fallback_reason: str = ""
23
+ foreign_output: bool = False
24
+ reservation_budget_bytes: int = (1 << 63) - 1
25
+ estimated_reserved_bytes: int = 0
26
+
27
+
28
+ def _pixel_tiles(shape, chunk_shape):
29
+ for starts in product(*(range(0, size, part) for size, part in zip(shape, chunk_shape))):
30
+ clipped = tuple(min(part, size - start) for size, start, part in zip(shape, starts, chunk_shape))
31
+ yield tuple(slice(start, start + size) for start, size in zip(starts, clipped)), clipped
32
+
33
+
34
+ def _memory_budget(device):
35
+ """Usable allocator bytes, retaining 25% headroom for allocation overhead.
36
+
37
+ Reusable reserved blocks count as available, but never bypass the process
38
+ fraction ceiling. Older PyTorch without its public fraction getter can
39
+ only observe driver/allocator memory; explicit limits remain available.
40
+ """
41
+ free, total = torch.cuda.mem_get_info(device)
42
+ allocated = torch.cuda.memory_allocated(device)
43
+ reusable = max(0, torch.cuda.memory_reserved(device) - allocated)
44
+ getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
45
+ fraction = getter(device) if getter is not None else 1.0
46
+ allocator_room = max(0, int(total * fraction) - allocated)
47
+ backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
48
+ foreign_splits = 0
49
+ if backend == "native":
50
+ caller = torch.cuda.current_stream(device).cuda_stream
51
+ for segment in _native_segments(device):
52
+ blocks = segment["blocks"]
53
+ if segment["stream"] != caller and any(block["state"] != "inactive" for block in blocks):
54
+ foreign_splits += sum(block["size"] for block in blocks if block["state"] == "inactive")
55
+ available = max(0, min(free + reusable, allocator_room) - foreign_splits)
56
+ budget = available * 3 // 4
57
+ allocation_limit = (1 << 63) - 1 # Unbounded limit in scalar int64 agreement.
58
+ # Native requests >1 MiB and <10 MiB reserve a 20 MiB segment. A smaller
59
+ # current available room must use the <=1 MiB small pool, whose segments are 2 MiB.
60
+ # Source: pytorch/c10/cuda/CUDACachingAllocator.cpp (Native allocator notes).
61
+ if backend == "native" and available < 20 * 1024 ** 2:
62
+ allocation_limit = 1024 ** 2
63
+ return budget, allocation_limit
64
+
65
+
66
+ def _native_segments(device):
67
+ index = torch.device(device).index if not isinstance(device, int) else device
68
+ index = torch.cuda.current_device() if index is None else index
69
+ return [segment for segment in torch.cuda.memory_snapshot() if segment["device"] == index]
70
+
71
+
72
+ def _reservation_budget(device):
73
+ """Room for fresh segments, conservatively excluding every pinned segment.
74
+
75
+ Free splits in a live segment may serve the caller, but cannot guarantee
76
+ placement of the complete new inventory. Wholly inactive segments can be
77
+ released by the native allocator, including those owned by other streams.
78
+ """
79
+ if torch.cuda.get_allocator_backend() != "native":
80
+ return (1 << 63) - 1
81
+ free, total = torch.cuda.mem_get_info(device)
82
+ segments = _native_segments(device)
83
+ pinned = sum(segment["total_size"] for segment in segments
84
+ if any(block["state"] != "inactive" for block in segment["blocks"]))
85
+ reclaimable = sum(segment["total_size"] for segment in segments
86
+ if all(block["state"] == "inactive" for block in segment["blocks"]))
87
+ getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
88
+ fraction = getter(device) if getter is not None else 1.0
89
+ return max(0, min(free + reclaimable, int(total * fraction) - pinned))
90
+
91
+
92
+ def _working_set_bytes(beam, shape, views, detector_shape, surface):
93
+ """Conservative native staging/layout, cotangent, geometry and VJP scratch."""
94
+ rank = len(shape)
95
+ voxels = math.prod(shape)
96
+ rays = views * math.prod(detector_shape)
97
+ geometry = (rank * (views + rays) if surface
98
+ else rank * (4 if beam == "cone" else 3) * views)
99
+ # A resident native tile, staging and completed cone layout can coexist.
100
+ # Two ray buffers cover reusable output/cotangent and transfer scratch; two
101
+ # geometry buffers cover staged values and their first-order VJP/reduction.
102
+ tile_buffers = 3 if beam == "cone" else 2
103
+ return 4 * (tile_buffers * voxels + 2 * rays + 2 * geometry)
104
+
105
+
106
+ def _largest_buffer_bytes(beam, shape, views, detector_shape, surface):
107
+ rank = len(shape)
108
+ rays = views * math.prod(detector_shape)
109
+ geometry = (rank * (views + rays) if surface
110
+ else rank * (4 if beam == "cone" else 3) * views)
111
+ return 4 * max(math.prod(shape), rays, geometry)
112
+
113
+
114
+ def _allocation_bytes(elements):
115
+ size = 4 * elements
116
+ if torch.cuda.get_allocator_backend() == "native":
117
+ return (size + 511) // 512 * 512
118
+ return size
119
+
120
+
121
+ def _reservation_bytes(buffer_sizes):
122
+ """Fresh, monotone fixed-buffer inventory bound for standard native pools.
123
+
124
+ Arbitrary free holes and allocator overrides are outside this model. Large
125
+ and middle requests share one pool, so mixed inventories use separate fresh
126
+ segments per request instead of assuming their placement is independent.
127
+ """
128
+ sizes = [size for size in buffer_sizes if size > 0]
129
+ if torch.cuda.get_allocator_backend() != "native":
130
+ return sum(sizes)
131
+ mib = 1024 ** 2
132
+ sizes = [(size + 511) // 512 * 512 for size in sizes]
133
+ large = [size for size in sizes if size >= 10 * mib]
134
+ reservation = sum((size + 2 * mib - 1) // (2 * mib) * (2 * mib) for size in large)
135
+ middle = [size for size in sizes if mib < size < 10 * mib]
136
+ if large:
137
+ reservation += len(middle) * 20 * mib
138
+ middle = []
139
+ for segment_size, requests in (
140
+ (2 * mib, [size for size in sizes if size <= mib]),
141
+ (20 * mib, middle),
142
+ ):
143
+ if not requests:
144
+ continue
145
+ total = sum(requests)
146
+ # At a new segment's birth every older segment holds more than
147
+ # segment_size - max(requests); the new one holds at least min(requests).
148
+ # This bounds online best-fit placement in any allocation order.
149
+ count = (1 if total <= segment_size else
150
+ (total - min(requests)) // (segment_size - max(requests)) + 1)
151
+ reservation += count * segment_size
152
+ return reservation
153
+
154
+
155
+ def _pipeline_buffers(beam, shape, views, pixels, points, *, slots=1, retained=1,
156
+ geometry_grad=False, peers=False, ray_accumulators=0):
157
+ """Live typed allocation inventory, with native rounding per buffer."""
158
+ rank, rays = len(shape), views * math.prod(pixels)
159
+ native = _allocation_bytes(math.prod(shape))
160
+ # Two producer epochs permit the preceding native/layout/output storage
161
+ # to stay in flight while the next retained tile window queues work.
162
+ tiles = slots * (retained + int(beam == "cone")) + int(peers)
163
+ sizes = ((views * rank, rays * rank) if points
164
+ else (views * rank,) * (4 if beam == "cone" else 3))
165
+ geometry = tuple(_allocation_bytes(size) for size in sizes)
166
+ return ((native,) * tiles + (_allocation_bytes(rays),) * (slots + ray_accumulators)
167
+ + geometry * (slots * (1 + int(geometry_grad))))
168
+
169
+
170
+ def _pipeline_bytes(beam, shape, views, pixels, points, *, slots=1, retained=1,
171
+ geometry_grad=False, peers=False, ray_accumulators=0):
172
+ """Live typed bytes; segment reservation is admitted separately."""
173
+ rank, rays = len(shape), views * math.prod(pixels)
174
+ native = _allocation_bytes(math.prod(shape))
175
+ tile = native * (slots * (retained + int(beam == "cone")) + int(peers))
176
+ sizes = ((views * rank, rays * rank) if points
177
+ else (views * rank,) * (4 if beam == "cone" else 3))
178
+ geometry = sum(_allocation_bytes(size) for size in sizes)
179
+ return (tile + slots * (_allocation_bytes(rays) + geometry)
180
+ + ray_accumulators * _allocation_bytes(rays) + (slots * geometry if geometry_grad else 0))
181
+
182
+
183
+ def _host_geometry_bytes(beam, views, pixels, points):
184
+ rank = 3 if beam == "cone" else 2
185
+ if not points:
186
+ return 8 * 2 * rank * views * (4 if beam == "cone" else 3)
187
+ # Grids/offsets, world expressions and first-order cotangent/graph work.
188
+ # Arbitrary sampler internals and explicit parameter snapshots are separate.
189
+ return 8 * rank * (math.prod(pixels) * (1 + 6 * views) + 4 * views)
190
+
191
+
192
+ def _tiles(shape, chunk_shape, spacing):
193
+ """Yield tensor slices, clipped shapes and physical centres in xyz order."""
194
+ starts = (range(0, size, chunk) for size, chunk in zip(shape, chunk_shape))
195
+ for start in product(*starts):
196
+ tile_shape = tuple(min(chunk, size - first)
197
+ for size, chunk, first in zip(shape, chunk_shape, start))
198
+ slices = tuple(slice(first, first + size) for first, size in zip(start, tile_shape))
199
+ centre = tuple((first + size / 2 - full / 2) * spacing
200
+ for first, size, full in zip(start, tile_shape, shape))
201
+ yield slices, tile_shape, centre[::-1]