diffct 2.0.1__tar.gz → 2.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 (63) hide show
  1. {diffct-2.0.1 → diffct-2.1.0}/CHANGELOG.md +51 -0
  2. {diffct-2.0.1 → diffct-2.1.0}/CITATION.cff +2 -2
  3. {diffct-2.0.1 → diffct-2.1.0}/PKG-INFO +54 -28
  4. {diffct-2.0.1 → diffct-2.1.0}/README.md +53 -27
  5. {diffct-2.0.1 → diffct-2.1.0}/diffct/__init__.py +1 -1
  6. diffct-2.1.0/diffct/chunking.py +73 -0
  7. {diffct-2.0.1 → diffct-2.1.0}/diffct/kernels/cone_beam.py +46 -36
  8. {diffct-2.0.1 → diffct-2.1.0}/diffct/kernels/fan_beam.py +30 -22
  9. {diffct-2.0.1 → diffct-2.1.0}/diffct/kernels/geometry_vjp.py +77 -48
  10. {diffct-2.0.1 → diffct-2.1.0}/diffct/kernels/parallel_beam.py +32 -22
  11. diffct-2.1.0/diffct/operators.py +1033 -0
  12. {diffct-2.0.1 → diffct-2.1.0}/diffct/projectors.py +143 -18
  13. {diffct-2.0.1 → diffct-2.1.0}/diffct/utils.py +17 -9
  14. diffct-2.1.0/docs/video/README.md +119 -0
  15. {diffct-2.0.1 → diffct-2.1.0}/examples/README.md +23 -6
  16. {diffct-2.0.1 → diffct-2.1.0}/formal/README.md +41 -9
  17. {diffct-2.0.1 → diffct-2.1.0}/pyproject.toml +1 -1
  18. diffct-2.1.0/tests/test_chunked_projector.py +1202 -0
  19. diffct-2.1.0/tests/test_cuda_exit_deallocations.py +83 -0
  20. diffct-2.1.0/tests/test_detector_surface_dtypes.py +80 -0
  21. diffct-2.1.0/tests/test_detector_surface_mapping.py +79 -0
  22. diffct-2.1.0/tests/test_detector_surface_units.py +62 -0
  23. diffct-2.1.0/tests/test_detector_surfaces.py +705 -0
  24. {diffct-2.0.1 → diffct-2.1.0}/tests/test_multi_gpu_exit.py +5 -0
  25. {diffct-2.0.1 → diffct-2.1.0}/tests/test_projector_api.py +9 -3
  26. diffct-2.0.1/diffct/operators.py +0 -555
  27. diffct-2.0.1/docs/video/README.md +0 -23
  28. {diffct-2.0.1 → diffct-2.1.0}/.gitignore +0 -0
  29. {diffct-2.0.1 → diffct-2.1.0}/LICENSE +0 -0
  30. {diffct-2.0.1 → diffct-2.1.0}/NOTICE +0 -0
  31. {diffct-2.0.1 → diffct-2.1.0}/diffct/analytical.py +0 -0
  32. {diffct-2.0.1 → diffct-2.1.0}/diffct/constants.py +0 -0
  33. {diffct-2.0.1 → diffct-2.1.0}/diffct/differentiable.py +0 -0
  34. {diffct-2.0.1 → diffct-2.1.0}/diffct/geometry.py +0 -0
  35. {diffct-2.0.1 → diffct-2.1.0}/diffct/kernels/__init__.py +0 -0
  36. {diffct-2.0.1 → diffct-2.1.0}/examples/data/NOTICE +0 -0
  37. {diffct-2.0.1 → diffct-2.1.0}/pytest.ini +0 -0
  38. {diffct-2.0.1 → diffct-2.1.0}/tests/__init__.py +0 -0
  39. {diffct-2.0.1 → diffct-2.1.0}/tests/_phantoms.py +0 -0
  40. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/README.md +0 -0
  41. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/__init__.py +0 -0
  42. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/_common.py +0 -0
  43. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/test_bench_cone.py +0 -0
  44. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/test_bench_fan.py +0 -0
  45. {diffct-2.0.1 → diffct-2.1.0}/tests/benchmarks/test_bench_parallel.py +0 -0
  46. {diffct-2.0.1 → diffct-2.1.0}/tests/distributed_projector_check.py +0 -0
  47. {diffct-2.0.1 → diffct-2.1.0}/tests/test_adjoint_inner_product.py +0 -0
  48. {diffct-2.0.1 → diffct-2.1.0}/tests/test_analytical_regressions.py +0 -0
  49. {diffct-2.0.1 → diffct-2.1.0}/tests/test_cone_projector_autograd.py +0 -0
  50. {diffct-2.0.1 → diffct-2.1.0}/tests/test_cuda_smoke.py +0 -0
  51. {diffct-2.0.1 → diffct-2.1.0}/tests/test_fbp_fan_accuracy.py +0 -0
  52. {diffct-2.0.1 → diffct-2.1.0}/tests/test_fbp_fan_offsets.py +0 -0
  53. {diffct-2.0.1 → diffct-2.1.0}/tests/test_fbp_parallel_accuracy.py +0 -0
  54. {diffct-2.0.1 → diffct-2.1.0}/tests/test_fdk_cone_accuracy.py +0 -0
  55. {diffct-2.0.1 → diffct-2.1.0}/tests/test_fdk_cone_offsets.py +0 -0
  56. {diffct-2.0.1 → diffct-2.1.0}/tests/test_geometry_autograd.py +0 -0
  57. {diffct-2.0.1 → diffct-2.1.0}/tests/test_geometry_gradients.py +0 -0
  58. {diffct-2.0.1 → diffct-2.1.0}/tests/test_gradcheck.py +0 -0
  59. {diffct-2.0.1 → diffct-2.1.0}/tests/test_projector_degenerate_geometry.py +0 -0
  60. {diffct-2.0.1 → diffct-2.1.0}/tests/test_projector_geometry_devices.py +0 -0
  61. {diffct-2.0.1 → diffct-2.1.0}/tests/test_ramp_filter_windows.py +0 -0
  62. {diffct-2.0.1 → diffct-2.1.0}/tests/test_release_2_0_1.py +0 -0
  63. {diffct-2.0.1 → diffct-2.1.0}/tests/test_weights.py +0 -0
@@ -5,6 +5,54 @@ 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.1.0] - 2026-10-10
11
+
12
+ ### Added
13
+
14
+ - `Projector(..., detector_surface=surface)` samples local parameterized surfaces
15
+ into per-pixel world positions for native parallel, fan and cone projection
16
+ and matched backprojection. Shared or per-view surfaces can contain trainable
17
+ PyTorch parameters. Flat detectors remain the default; analytical helpers
18
+ retain their flat-detector assumptions.
19
+ - CUDA regression tests and CPU formal verification for detector
20
+ surfaces, including bounded production-kernel checks against cell intersections.
21
+ - Automatic CPU-backed spatial tiles and view batches for native `Projector`
22
+ projection, matched backprojection, data autograd/Hessians and first-order
23
+ geometry gradients. CUDA inputs retain the full-volume path when its estimated
24
+ working set fits. Optional `volume_chunk_shape` and `view_chunk_size` override
25
+ automatic limits; full CPU arrays are not staged on CUDA. Several local
26
+ devices compute their view batches concurrently.
27
+ - A CPU-backed CGLS example and chunk-boundary, CUDA memory and
28
+ formal partition/translation/adjoint checks.
29
+
30
+ ### Fixed
31
+
32
+ - The exit-time cleanup hook returns early under the CUDA simulator
33
+ (``NUMBA_ENABLE_CUDASIM=1``). Before, each simulator process printed an
34
+ ``ImportError`` for ``_PendingDeallocs`` at exit. The simulator loads no
35
+ CUDA modules, so it needs no cleanup.
36
+
37
+ ## [2.0.2] - 2026-10-09
38
+
39
+ ### Fixed
40
+
41
+ - The exit-time cleanup from 2.0.1 set numba-cuda's pending-deallocation
42
+ ratio to infinity. numba-cuda converts ``capacity * ratio`` to ``int``, so
43
+ each deallocation queued during shutdown raised ``ValueError`` or
44
+ ``OverflowError``. The exit code stayed 0, but stderr showed the error. The
45
+ exit hook now disables numba-cuda's pending-deallocation queues instead of
46
+ changing the limits, so no queue is flushed at exit and no error is raised.
47
+ The hook no longer imports ``numba.cuda.core.config``, which older
48
+ numba-cuda releases do not have.
49
+
50
+ ### Documentation
51
+
52
+ - The install instructions state that ``pip install diffct`` without a CUDA
53
+ extra needs a system CUDA Toolkit. Use ``diffct[cu12]`` or ``diffct[cu13]``.
54
+ - README images and links use absolute URLs, so they also work on PyPI.
55
+
8
56
  ## [2.0.1] - 2026-10-09
9
57
 
10
58
  ### Fixed
@@ -13,6 +61,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
13
61
  weighted backprojections run in the CUDA context of their input. Before,
14
62
  inputs on another GPU than the current device could use the wrong context
15
63
  and stream.
64
+ - A process that ran diffct kernels on several GPUs no longer crashes at
65
+ interpreter exit: numba-cuda's flush of pending CUDA module unloads during
66
+ shutdown is skipped, and process exit releases the modules.
16
67
  - ``custom_trajectory_3d`` and ``custom_trajectory_2d_fan`` raise
17
68
  ``ValueError`` for a source at the isocenter or a non-finite source,
18
69
  instead of returning NaN geometry. The figure-8 example in the
@@ -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.1
17
- date-released: "2026-10-09"
16
+ version: 2.1.0
17
+ date-released: "2026-10-10"
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.1
3
+ Version: 2.1.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/
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
31
31
 
32
32
  <h1 align="center">diffct</h1>
33
33
 
34
- <p align="center">English · <a href="README.zh.md">简体中文</a></p>
34
+ <p align="center">English · <a href="https://github.com/sypsyp97/diffct/blob/main/README.zh.md">简体中文</a></p>
35
35
 
36
36
  <p align="center">
37
37
  Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU, multi-node, geometry gradients.
@@ -43,15 +43,15 @@ Description-Content-Type: text/markdown
43
43
  <a href="https://pypi.org/project/diffct/"><img src="https://img.shields.io/pypi/v/diffct.svg?style=flat-square&logo=pypi&logoColor=white" alt="PyPI version"></a>
44
44
  <a href="https://sypsyp97.github.io/diffct/"><img src="https://img.shields.io/github/actions/workflow/status/sypsyp97/diffct/docs.yml?branch=main&label=docs&style=flat-square" alt="Documentation"></a>
45
45
  <a href="https://github.com/sypsyp97/diffct/actions"><img src="https://img.shields.io/github/actions/workflow/status/sypsyp97/diffct/ci.yml?branch=main&label=CI&style=flat-square" alt="CI/CD"></a>
46
- <a href="https://deepwiki.com/sypsyp97/diffct"><img src="docs/assets/deepwiki-badge.svg" alt="Ask DeepWiki"></a>
46
+ <a href="https://deepwiki.com/sypsyp97/diffct"><img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/deepwiki-badge.svg" alt="Ask DeepWiki"></a>
47
47
  </p>
48
48
 
49
49
  <p align="center">
50
- <img src="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="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,7 +65,9 @@ Description-Content-Type: text/markdown
65
65
  <a href="#citation">Citation</a>
66
66
  </p>
67
67
 
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](docs/MIGRATION.md).
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
+ > **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).
69
71
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
70
72
 
71
73
  ## Why diffct
@@ -73,20 +75,25 @@ Description-Content-Type: text/markdown
73
75
  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.
74
76
 
75
77
  - **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.
78
+ - **Curved detectors.** A `detector_surface(u, v)` callback defines each pixel's position for arc, cylindrical or other parameterized surfaces. 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).
76
79
  - **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.
77
- - **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.
80
+ - **Large volumes.** CPU volumes and sinograms automatically stream spatial tiles and view batches through CUDA. Full arrays and data gradients stay on CPU; tile sizes follow the available device memory and can be overridden. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
81
+ - **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. CPU-backed volumes stay in host RAM while GPUs compute tiles; distributed ranks retain replicated-volume semantics. Speedup depends on workload and communication costs.
78
82
  - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
79
83
 
80
- Capabilities, limits and isocenter rules: [docs/REFERENCE.md](docs/REFERENCE.md#capabilities-and-limits).
84
+ Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
81
85
 
82
86
  ## Install
83
87
 
84
- You need a CUDA GPU and PyTorch. Install PyTorch for your CUDA version first.
88
+ You need an NVIDIA GPU. Install PyTorch for your CUDA version first. Then install diffct with the extra for the same CUDA major version:
85
89
 
86
90
  ```bash
87
- pip install "diffct[cu12]" # [cu13] for CUDA 13; install PyTorch for your CUDA first
91
+ pip install "diffct[cu12]" # PyTorch built for CUDA 12
92
+ pip install "diffct[cu13]" # PyTorch built for CUDA 13
88
93
  ```
89
94
 
95
+ The extra installs the CUDA compiler libraries that Numba CUDA needs to compile the kernels. A plain `pip install diffct` does not install them. Use it only if a CUDA Toolkit is already installed on the system.
96
+
90
97
  From source, with the examples:
91
98
 
92
99
  ```bash
@@ -129,22 +136,39 @@ loss = 0.5 * (A.project(x) - sinogram).square().sum()
129
136
  loss.backward() # x.grad = A^T (A x - y)
130
137
  ```
131
138
 
132
- **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](docs/REFERENCE.md#custom-or-calibrated-trajectories).
139
+ **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
133
140
 
134
- **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](docs/REFERENCE.md#geometry-gradients).
141
+ **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).
135
142
 
136
- **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](docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](docs/DISTRIBUTED.md).
143
+ **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).
137
144
 
138
145
  ```python
139
146
  A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
140
147
  ```
141
148
 
149
+ **Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
150
+ CUDA tiles automatically; CUDA inputs retain the full-volume path when its
151
+ estimated working set fits. Explicit limits are optional:
152
+
153
+ ```python
154
+ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
155
+ volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
156
+ sinogram = A.project(volume) # CPU output, CUDA computation
157
+ sinogram.square().mean().backward() # CPU volume.grad
158
+ # Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
159
+ ```
160
+
161
+ Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
162
+ CUDA allocation. Chunking preserves the Siddon model, with float32 summation
163
+ differences; smaller blocks can increase transfer and launch overhead. Full
164
+ arrays and iterative state still need enough host RAM.
165
+
142
166
  ## Performance
143
167
 
144
168
  <p align="center">
145
169
  <picture>
146
- <source media="(prefers-color-scheme: dark)" srcset="docs/assets/scaling_dark.png">
147
- <img src="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">
170
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
171
+ <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">
148
172
  </picture>
149
173
  </p>
150
174
 
@@ -154,7 +178,7 @@ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, dev
154
178
  | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
155
179
  | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
156
180
 
157
- 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](docs/assets/scaling.json).
181
+ 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).
158
182
 
159
183
  To measure one against several GPUs on your machine (default scan of the script, not the table setup):
160
184
 
@@ -164,10 +188,12 @@ python examples/benchmark_projector.py --devices 0 1
164
188
 
165
189
  ## Gallery
166
190
 
191
+ TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
192
+
167
193
  **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.
168
194
 
169
195
  <p align="center">
170
- <img src="docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
196
+ <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">
171
197
  </p>
172
198
 
173
199
  ```bash
@@ -177,7 +203,7 @@ python examples/walnut_reconstruction.py --figure out.png
177
203
  **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.
178
204
 
179
205
  <p align="center">
180
- <img src="docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
206
+ <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">
181
207
  </p>
182
208
 
183
209
  ```bash
@@ -185,22 +211,22 @@ python examples/walnut_reconstruction.py --algorithms --save-volume walnut256.np
185
211
  python examples/iterative_reconstruction.py --size 256 --views 720 --trajectory helical --noise 0.01 --phantom walnut256.npy --figure out.png
186
212
  ```
187
213
 
188
- Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](examples/data/NOTICE).
214
+ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](https://github.com/sypsyp97/diffct/blob/main/examples/data/NOTICE).
189
215
 
190
216
  ## Examples
191
217
 
192
- `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, distributed-loss rules and measured results are in [examples/README.md](examples/README.md).
218
+ `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).
193
219
 
194
220
  ## Documentation
195
221
 
196
222
  | Document | Contents |
197
223
  |---|---|
198
- | [Trajectory guide](docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
199
- | [docs/REFERENCE.md](docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
200
- | [docs/DISTRIBUTED.md](docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
201
- | [docs/MIGRATION.md](docs/MIGRATION.md) | Moving from the circular-only API |
202
- | [docs/video/README.md](docs/video/README.md) | How the intro video is rendered |
203
- | [CHANGELOG.md](CHANGELOG.md) | Release history |
224
+ | [Trajectory guide](https://github.com/sypsyp97/diffct/blob/main/docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
225
+ | [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
226
+ | [docs/DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
227
+ | [docs/MIGRATION.md](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md) | Moving from the circular-only API |
228
+ | [docs/video/README.md](https://github.com/sypsyp97/diffct/blob/main/docs/video/README.md) | How the intro video is rendered |
229
+ | [CHANGELOG.md](https://github.com/sypsyp97/diffct/blob/main/CHANGELOG.md) | Release history |
204
230
 
205
231
  Run the test suite on a CUDA host with `python -m pytest tests/ -q`.
206
232
 
@@ -232,7 +258,7 @@ If you use this library in your research, please cite the software and the techn
232
258
 
233
259
  ## License and acknowledgements
234
260
 
235
- Apache 2.0, see [LICENSE](LICENSE). The project draws on
261
+ Apache 2.0, see [LICENSE](https://github.com/sypsyp97/diffct/blob/main/LICENSE). The project draws on
236
262
  [PYRO-NN](https://github.com/csyben/PYRO-NN) and
237
263
  [geometry_gradients_CT](https://github.com/mareikethies/geometry_gradients_CT).
238
- Issues and pull requests are welcome. The walnut data are CC BY 4.0 (Meaney 2022); see [NOTICE](NOTICE).
264
+ Issues and pull requests are welcome. The walnut data are CC BY 4.0 (Meaney 2022); see [NOTICE](https://github.com/sypsyp97/diffct/blob/main/NOTICE).
@@ -1,6 +1,6 @@
1
1
  <h1 align="center">diffct</h1>
2
2
 
3
- <p align="center">English · <a href="README.zh.md">简体中文</a></p>
3
+ <p align="center">English · <a href="https://github.com/sypsyp97/diffct/blob/main/README.zh.md">简体中文</a></p>
4
4
 
5
5
  <p align="center">
6
6
  Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU, multi-node, geometry gradients.
@@ -12,15 +12,15 @@
12
12
  <a href="https://pypi.org/project/diffct/"><img src="https://img.shields.io/pypi/v/diffct.svg?style=flat-square&logo=pypi&logoColor=white" alt="PyPI version"></a>
13
13
  <a href="https://sypsyp97.github.io/diffct/"><img src="https://img.shields.io/github/actions/workflow/status/sypsyp97/diffct/docs.yml?branch=main&label=docs&style=flat-square" alt="Documentation"></a>
14
14
  <a href="https://github.com/sypsyp97/diffct/actions"><img src="https://img.shields.io/github/actions/workflow/status/sypsyp97/diffct/ci.yml?branch=main&label=CI&style=flat-square" alt="CI/CD"></a>
15
- <a href="https://deepwiki.com/sypsyp97/diffct"><img src="docs/assets/deepwiki-badge.svg" alt="Ask DeepWiki"></a>
15
+ <a href="https://deepwiki.com/sypsyp97/diffct"><img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/deepwiki-badge.svg" alt="Ask DeepWiki"></a>
16
16
  </p>
17
17
 
18
18
  <p align="center">
19
- <img src="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="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,7 +34,9 @@
34
34
  <a href="#citation">Citation</a>
35
35
  </p>
36
36
 
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](docs/MIGRATION.md).
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
+ > **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).
38
40
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
39
41
 
40
42
  ## Why diffct
@@ -42,20 +44,25 @@
42
44
  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.
43
45
 
44
46
  - **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.
47
+ - **Curved detectors.** A `detector_surface(u, v)` callback defines each pixel's position for arc, cylindrical or other parameterized surfaces. 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).
45
48
  - **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.
46
- - **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.
49
+ - **Large volumes.** CPU volumes and sinograms automatically stream spatial tiles and view batches through CUDA. Full arrays and data gradients stay on CPU; tile sizes follow the available device memory and can be overridden. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
50
+ - **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. CPU-backed volumes stay in host RAM while GPUs compute tiles; distributed ranks retain replicated-volume semantics. Speedup depends on workload and communication costs.
47
51
  - **Analytical helpers.** `diffct.analytical` provides ramp filters (ram-lak, shepp-logan, cosine, hamming, hann), fan, cone and Parker weights, and FBP and FDK backprojection.
48
52
 
49
- Capabilities, limits and isocenter rules: [docs/REFERENCE.md](docs/REFERENCE.md#capabilities-and-limits).
53
+ Capabilities, limits and isocenter rules: [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#capabilities-and-limits).
50
54
 
51
55
  ## Install
52
56
 
53
- You need a CUDA GPU and PyTorch. Install PyTorch for your CUDA version first.
57
+ You need an NVIDIA GPU. Install PyTorch for your CUDA version first. Then install diffct with the extra for the same CUDA major version:
54
58
 
55
59
  ```bash
56
- pip install "diffct[cu12]" # [cu13] for CUDA 13; install PyTorch for your CUDA first
60
+ pip install "diffct[cu12]" # PyTorch built for CUDA 12
61
+ pip install "diffct[cu13]" # PyTorch built for CUDA 13
57
62
  ```
58
63
 
64
+ The extra installs the CUDA compiler libraries that Numba CUDA needs to compile the kernels. A plain `pip install diffct` does not install them. Use it only if a CUDA Toolkit is already installed on the system.
65
+
59
66
  From source, with the examples:
60
67
 
61
68
  ```bash
@@ -98,22 +105,39 @@ loss = 0.5 * (A.project(x) - sinogram).square().sum()
98
105
  loss.backward() # x.grad = A^T (A x - y)
99
106
  ```
100
107
 
101
- **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](docs/REFERENCE.md#custom-or-calibrated-trajectories).
108
+ **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
102
109
 
103
- **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](docs/REFERENCE.md#geometry-gradients).
110
+ **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).
104
111
 
105
- **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](docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](docs/DISTRIBUTED.md).
112
+ **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).
106
113
 
107
114
  ```python
108
115
  A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, devices=[0, 1, 2, 3])
109
116
  ```
110
117
 
118
+ **Volumes larger than GPU memory.** Keep the arrays on CPU. The same calls use
119
+ CUDA tiles automatically; CUDA inputs retain the full-volume path when its
120
+ estimated working set fits. Explicit limits are optional:
121
+
122
+ ```python
123
+ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8)
124
+ volume = torch.ones(A.volume_shape, device="cpu", requires_grad=True)
125
+ sinogram = A.project(volume) # CPU output, CUDA computation
126
+ sinogram.square().mean().backward() # CPU volume.grad
127
+ # Optional: volume_chunk_shape=(32, 64, 64), view_chunk_size=16
128
+ ```
129
+
130
+ Run `python examples/chunked_reconstruction.py` for CPU-backed CGLS and measured
131
+ CUDA allocation. Chunking preserves the Siddon model, with float32 summation
132
+ differences; smaller blocks can increase transfer and launch overhead. Full
133
+ arrays and iterative state still need enough host RAM.
134
+
111
135
  ## Performance
112
136
 
113
137
  <p align="center">
114
138
  <picture>
115
- <source media="(prefers-color-scheme: dark)" srcset="docs/assets/scaling_dark.png">
116
- <img src="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">
139
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png">
140
+ <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">
117
141
  </picture>
118
142
  </p>
119
143
 
@@ -123,7 +147,7 @@ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, dev
123
147
  | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
124
148
  | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
125
149
 
126
- 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](docs/assets/scaling.json).
150
+ 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).
127
151
 
128
152
  To measure one against several GPUs on your machine (default scan of the script, not the table setup):
129
153
 
@@ -133,10 +157,12 @@ python examples/benchmark_projector.py --devices 0 1
133
157
 
134
158
  ## Gallery
135
159
 
160
+ TV denotes TV-regularized iterative reconstruction; Adam is the optimizer. Figure labels distinguish the method from its optimizer.
161
+
136
162
  **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.
137
163
 
138
164
  <p align="center">
139
- <img src="docs/assets/walnut_measured.png" width="100%" alt="Measured walnut: FDK, SIRT, CGLS and TV reconstructions">
165
+ <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">
140
166
  </p>
141
167
 
142
168
  ```bash
@@ -146,7 +172,7 @@ python examples/walnut_reconstruction.py --figure out.png
146
172
  **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.
147
173
 
148
174
  <p align="center">
149
- <img src="docs/assets/walnut_helical.png" width="100%" alt="Simulated helical scan of the walnut: FDK, CGLS, SIRT and TV reconstructions">
175
+ <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">
150
176
  </p>
151
177
 
152
178
  ```bash
@@ -154,22 +180,22 @@ python examples/walnut_reconstruction.py --algorithms --save-volume walnut256.np
154
180
  python examples/iterative_reconstruction.py --size 256 --views 720 --trajectory helical --noise 0.01 --phantom walnut256.npy --figure out.png
155
181
  ```
156
182
 
157
- Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](examples/data/NOTICE).
183
+ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](https://github.com/sypsyp97/diffct/blob/main/examples/data/NOTICE).
158
184
 
159
185
  ## Examples
160
186
 
161
- `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, distributed-loss rules and measured results are in [examples/README.md](examples/README.md).
187
+ `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).
162
188
 
163
189
  ## Documentation
164
190
 
165
191
  | Document | Contents |
166
192
  |---|---|
167
- | [Trajectory guide](docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
168
- | [docs/REFERENCE.md](docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
169
- | [docs/DISTRIBUTED.md](docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
170
- | [docs/MIGRATION.md](docs/MIGRATION.md) | Moving from the circular-only API |
171
- | [docs/video/README.md](docs/video/README.md) | How the intro video is rendered |
172
- | [CHANGELOG.md](CHANGELOG.md) | Release history |
193
+ | [Trajectory guide](https://github.com/sypsyp97/diffct/blob/main/docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
194
+ | [docs/REFERENCE.md](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
195
+ | [docs/DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
196
+ | [docs/MIGRATION.md](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md) | Moving from the circular-only API |
197
+ | [docs/video/README.md](https://github.com/sypsyp97/diffct/blob/main/docs/video/README.md) | How the intro video is rendered |
198
+ | [CHANGELOG.md](https://github.com/sypsyp97/diffct/blob/main/CHANGELOG.md) | Release history |
173
199
 
174
200
  Run the test suite on a CUDA host with `python -m pytest tests/ -q`.
175
201
 
@@ -201,7 +227,7 @@ If you use this library in your research, please cite the software and the techn
201
227
 
202
228
  ## License and acknowledgements
203
229
 
204
- Apache 2.0, see [LICENSE](LICENSE). The project draws on
230
+ Apache 2.0, see [LICENSE](https://github.com/sypsyp97/diffct/blob/main/LICENSE). The project draws on
205
231
  [PYRO-NN](https://github.com/csyben/PYRO-NN) and
206
232
  [geometry_gradients_CT](https://github.com/mareikethies/geometry_gradients_CT).
207
- Issues and pull requests are welcome. The walnut data are CC BY 4.0 (Meaney 2022); see [NOTICE](NOTICE).
233
+ Issues and pull requests are welcome. The walnut data are CC BY 4.0 (Meaney 2022); see [NOTICE](https://github.com/sypsyp97/diffct/blob/main/NOTICE).
@@ -49,7 +49,7 @@ try:
49
49
  except ImportError:
50
50
  pass
51
51
 
52
- __version__ = '2.0.1'
52
+ __version__ = '2.1.0'
53
53
 
54
54
  __all__ = [
55
55
  'Projector',
@@ -0,0 +1,73 @@
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
+
16
+
17
+ def _memory_budget(device):
18
+ """Usable allocator bytes, retaining 25% headroom for allocation overhead.
19
+
20
+ Reusable reserved blocks count as available, but never bypass the process
21
+ fraction ceiling. Older PyTorch without its public fraction getter can
22
+ only observe driver/allocator memory; explicit limits remain available.
23
+ """
24
+ free, total = torch.cuda.mem_get_info(device)
25
+ allocated = torch.cuda.memory_allocated(device)
26
+ reusable = max(0, torch.cuda.memory_reserved(device) - allocated)
27
+ getter = getattr(torch.cuda, "get_per_process_memory_fraction", None)
28
+ fraction = getter(device) if getter is not None else 1.0
29
+ allocator_room = max(0, int(total * fraction) - allocated)
30
+ available = max(0, min(free + reusable, allocator_room))
31
+ budget = available * 3 // 4
32
+ allocation_limit = (1 << 63) - 1 # Unbounded limit in scalar int64 agreement.
33
+ backend = getattr(torch.cuda, "get_allocator_backend", lambda: "native")()
34
+ # Native requests >1 MiB and <10 MiB reserve a 20 MiB segment. A smaller
35
+ # current available room must use the <=1 MiB small pool, whose segments are 2 MiB.
36
+ # Source: pytorch/c10/cuda/CUDACachingAllocator.cpp (Native allocator notes).
37
+ if backend == "native" and available < 20 * 1024 ** 2:
38
+ allocation_limit = 1024 ** 2
39
+ return budget, allocation_limit
40
+
41
+
42
+ def _working_set_bytes(beam, shape, views, detector_shape, surface):
43
+ """Conservative native staging/layout, cotangent, geometry and VJP scratch."""
44
+ rank = len(shape)
45
+ voxels = math.prod(shape)
46
+ rays = views * math.prod(detector_shape)
47
+ geometry = (rank * (views + rays) if surface
48
+ else rank * (4 if beam == "cone" else 3) * views)
49
+ # Cone permuted layout and its returned tensor can coexist with staging.
50
+ # Two ray buffers cover native output/cotangent and transfer scratch; two
51
+ # geometry buffers cover staged values and their first-order VJP/reduction.
52
+ tile_buffers = 3 if beam == "cone" else 2
53
+ return 4 * (tile_buffers * voxels + 2 * rays + 2 * geometry)
54
+
55
+
56
+ def _largest_buffer_bytes(beam, shape, views, detector_shape, surface):
57
+ rank = len(shape)
58
+ rays = views * math.prod(detector_shape)
59
+ geometry = (rank * (views + rays) if surface
60
+ else rank * (4 if beam == "cone" else 3) * views)
61
+ return 4 * max(math.prod(shape), rays, geometry)
62
+
63
+
64
+ def _tiles(shape, chunk_shape, spacing):
65
+ """Yield tensor slices, clipped shapes and physical centres in xyz order."""
66
+ starts = (range(0, size, chunk) for size, chunk in zip(shape, chunk_shape))
67
+ for start in product(*starts):
68
+ tile_shape = tuple(min(chunk, size - first)
69
+ for size, chunk, first in zip(shape, chunk_shape, start))
70
+ slices = tuple(slice(first, first + size) for first, size in zip(start, tile_shape))
71
+ centre = tuple((first + size / 2 - full / 2) * spacing
72
+ for first, size, full in zip(start, tile_shape, shape))
73
+ yield slices, tile_shape, centre[::-1]