diffct 2.0.0__tar.gz → 2.0.1__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 (54) hide show
  1. {diffct-2.0.0 → diffct-2.0.1}/CHANGELOG.md +29 -1
  2. {diffct-2.0.0 → diffct-2.0.1}/CITATION.cff +1 -1
  3. {diffct-2.0.0 → diffct-2.0.1}/PKG-INFO +2 -4
  4. {diffct-2.0.0 → diffct-2.0.1}/README.md +1 -3
  5. {diffct-2.0.0 → diffct-2.0.1}/diffct/__init__.py +1 -1
  6. {diffct-2.0.0 → diffct-2.0.1}/diffct/analytical.py +4 -0
  7. {diffct-2.0.0 → diffct-2.0.1}/diffct/geometry.py +15 -2
  8. {diffct-2.0.0 → diffct-2.0.1}/diffct/operators.py +9 -0
  9. {diffct-2.0.0 → diffct-2.0.1}/diffct/projectors.py +49 -48
  10. {diffct-2.0.0 → diffct-2.0.1}/diffct/utils.py +66 -0
  11. {diffct-2.0.0 → diffct-2.0.1}/docs/video/README.md +2 -1
  12. {diffct-2.0.0 → diffct-2.0.1}/examples/README.md +0 -68
  13. {diffct-2.0.0 → diffct-2.0.1}/examples/data/NOTICE +4 -4
  14. {diffct-2.0.0 → diffct-2.0.1}/formal/README.md +1 -1
  15. {diffct-2.0.0 → diffct-2.0.1}/pyproject.toml +1 -1
  16. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/README.md +1 -1
  17. diffct-2.0.1/tests/test_multi_gpu_exit.py +83 -0
  18. diffct-2.0.1/tests/test_release_2_0_1.py +326 -0
  19. {diffct-2.0.0 → diffct-2.0.1}/.gitignore +0 -0
  20. {diffct-2.0.0 → diffct-2.0.1}/LICENSE +0 -0
  21. {diffct-2.0.0 → diffct-2.0.1}/NOTICE +0 -0
  22. {diffct-2.0.0 → diffct-2.0.1}/diffct/constants.py +0 -0
  23. {diffct-2.0.0 → diffct-2.0.1}/diffct/differentiable.py +0 -0
  24. {diffct-2.0.0 → diffct-2.0.1}/diffct/kernels/__init__.py +0 -0
  25. {diffct-2.0.0 → diffct-2.0.1}/diffct/kernels/cone_beam.py +0 -0
  26. {diffct-2.0.0 → diffct-2.0.1}/diffct/kernels/fan_beam.py +0 -0
  27. {diffct-2.0.0 → diffct-2.0.1}/diffct/kernels/geometry_vjp.py +0 -0
  28. {diffct-2.0.0 → diffct-2.0.1}/diffct/kernels/parallel_beam.py +0 -0
  29. {diffct-2.0.0 → diffct-2.0.1}/pytest.ini +0 -0
  30. {diffct-2.0.0 → diffct-2.0.1}/tests/__init__.py +0 -0
  31. {diffct-2.0.0 → diffct-2.0.1}/tests/_phantoms.py +0 -0
  32. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/__init__.py +0 -0
  33. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/_common.py +0 -0
  34. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/test_bench_cone.py +0 -0
  35. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/test_bench_fan.py +0 -0
  36. {diffct-2.0.0 → diffct-2.0.1}/tests/benchmarks/test_bench_parallel.py +0 -0
  37. {diffct-2.0.0 → diffct-2.0.1}/tests/distributed_projector_check.py +0 -0
  38. {diffct-2.0.0 → diffct-2.0.1}/tests/test_adjoint_inner_product.py +0 -0
  39. {diffct-2.0.0 → diffct-2.0.1}/tests/test_analytical_regressions.py +0 -0
  40. {diffct-2.0.0 → diffct-2.0.1}/tests/test_cone_projector_autograd.py +0 -0
  41. {diffct-2.0.0 → diffct-2.0.1}/tests/test_cuda_smoke.py +0 -0
  42. {diffct-2.0.0 → diffct-2.0.1}/tests/test_fbp_fan_accuracy.py +0 -0
  43. {diffct-2.0.0 → diffct-2.0.1}/tests/test_fbp_fan_offsets.py +0 -0
  44. {diffct-2.0.0 → diffct-2.0.1}/tests/test_fbp_parallel_accuracy.py +0 -0
  45. {diffct-2.0.0 → diffct-2.0.1}/tests/test_fdk_cone_accuracy.py +0 -0
  46. {diffct-2.0.0 → diffct-2.0.1}/tests/test_fdk_cone_offsets.py +0 -0
  47. {diffct-2.0.0 → diffct-2.0.1}/tests/test_geometry_autograd.py +0 -0
  48. {diffct-2.0.0 → diffct-2.0.1}/tests/test_geometry_gradients.py +0 -0
  49. {diffct-2.0.0 → diffct-2.0.1}/tests/test_gradcheck.py +0 -0
  50. {diffct-2.0.0 → diffct-2.0.1}/tests/test_projector_api.py +0 -0
  51. {diffct-2.0.0 → diffct-2.0.1}/tests/test_projector_degenerate_geometry.py +0 -0
  52. {diffct-2.0.0 → diffct-2.0.1}/tests/test_projector_geometry_devices.py +0 -0
  53. {diffct-2.0.0 → diffct-2.0.1}/tests/test_ramp_filter_windows.py +0 -0
  54. {diffct-2.0.0 → diffct-2.0.1}/tests/test_weights.py +0 -0
@@ -5,6 +5,34 @@ 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
+ ## [2.0.1] - 2026-10-09
9
+
10
+ ### Fixed
11
+
12
+ - The low-level projector and backprojector Functions and the analytical
13
+ weighted backprojections run in the CUDA context of their input. Before,
14
+ inputs on another GPU than the current device could use the wrong context
15
+ and stream.
16
+ - ``custom_trajectory_3d`` and ``custom_trajectory_2d_fan`` raise
17
+ ``ValueError`` for a source at the isocenter or a non-finite source,
18
+ instead of returning NaN geometry. The figure-8 example in the
19
+ ``custom_trajectory_3d`` docstring no longer puts a source at the origin.
20
+ - The class docstring examples of the six projector Functions use the
21
+ per-view trajectory API.
22
+ - ``_trig_tables`` accepts NumPy dtypes for tensor input.
23
+ - Walnut example: ``preprocess_walnut.py`` raises when projection archives or
24
+ views are missing or an air level is not positive, instead of writing NaN
25
+ data. A non-default crop keeps the detector centre offset, and
26
+ ``walnut_reconstruction.py`` applies it to the geometry and the FDK
27
+ weights. The default crop has no offset.
28
+ - Examples: ``psnr`` returns infinity for identical images; distributed
29
+ reductions in CGLS, TV and the noise peak use the operator's process
30
+ group.
31
+
32
+ ### Added
33
+
34
+ - ``Projector.process_group``: the process group of a distributed operator.
35
+
8
36
  ## [2.0.0] - 2026-10-09
9
37
 
10
38
  Version 2.0.0 replaces the circular-orbit API of 1.x with arbitrary
@@ -28,7 +56,7 @@ per-view trajectories. It breaks the 1.x interface; see
28
56
  - Install extras ``diffct[cu12]`` and ``diffct[cu13]`` pull in the matching
29
57
  numba-cuda CUDA bindings, so one ``pip install`` gives a working import.
30
58
  - Measured walnut example (``examples/walnut_reconstruction.py``), intro video
31
- and README figures with provenance records in ``docs/assets``.
59
+ and README figures.
32
60
  - Geometry gradients. Trajectory tensors may require gradients; ``Projector``
33
61
  and the low-level Function classes return gradients for source, detector
34
62
  centre and detector axes (parallel: ray direction, detector origin and
@@ -13,7 +13,7 @@ authors:
13
13
  given-names: Yipeng
14
14
  email: yipeng.sun@fau.de
15
15
  alias: sypsyp97
16
- version: 2.0.0
16
+ version: 2.0.1
17
17
  date-released: "2026-10-09"
18
18
  license: Apache-2.0
19
19
  repository-code: "https://github.com/sypsyp97/diffct"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: diffct
3
- Version: 2.0.0
3
+ Version: 2.0.1
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/
@@ -76,7 +76,6 @@ To our knowledge, diffct is the only open-source GPU CT library that combines ar
76
76
  - **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
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.
78
78
  - **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
- - **Validated.** FDK matches ASTRA 2.5.0 `FDK_CUDA`: PSNR within 0.3 dB. Geometry gradients match an independent float64 reference to about 1e-6. The [recorded A100 validation](docs/VALIDATION.md) reports 229 passing pytest tests. Full notes: [REFERENCE](docs/REFERENCE.md#validation-summary).
80
79
 
81
80
  Capabilities, limits and isocenter rules: [docs/REFERENCE.md](docs/REFERENCE.md#capabilities-and-limits).
82
81
 
@@ -155,7 +154,7 @@ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, dev
155
154
  | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
156
155
  | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
157
156
 
158
- Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB on Leonardo Booster. 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).
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).
159
158
 
160
159
  To measure one against several GPUs on your machine (default scan of the script, not the table setup):
161
160
 
@@ -199,7 +198,6 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
199
198
  | [Trajectory guide](docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
200
199
  | [docs/REFERENCE.md](docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
201
200
  | [docs/DISTRIBUTED.md](docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
202
- | [docs/VALIDATION.md](docs/VALIDATION.md) | Test commands, validation details and the limits of each check |
203
201
  | [docs/MIGRATION.md](docs/MIGRATION.md) | Moving from the circular-only API |
204
202
  | [docs/video/README.md](docs/video/README.md) | How the intro video is rendered |
205
203
  | [CHANGELOG.md](CHANGELOG.md) | Release history |
@@ -45,7 +45,6 @@ To our knowledge, diffct is the only open-source GPU CT library that combines ar
45
45
  - **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
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.
47
47
  - **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
- - **Validated.** FDK matches ASTRA 2.5.0 `FDK_CUDA`: PSNR within 0.3 dB. Geometry gradients match an independent float64 reference to about 1e-6. The [recorded A100 validation](docs/VALIDATION.md) reports 229 passing pytest tests. Full notes: [REFERENCE](docs/REFERENCE.md#validation-summary).
49
48
 
50
49
  Capabilities, limits and isocenter rules: [docs/REFERENCE.md](docs/REFERENCE.md#capabilities-and-limits).
51
50
 
@@ -124,7 +123,7 @@ A = Projector(trajectory, (128, 128, 128), (384, 256), detector_spacing=0.8, dev
124
123
  | Adjoint (backprojection) | 20.66 ms | 5.89 ms (3.5×) | 4.62 ms (4.5×) |
125
124
  | CGLS iteration | 30.08 ms | 8.08 ms (3.7×) | 5.59 ms (5.4×) |
126
125
 
127
- Circular trajectory, 360 views, detector (2n, 1.5n) cells at pitch 1.25, A100 64 GB on Leonardo Booster. 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).
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).
128
127
 
129
128
  To measure one against several GPUs on your machine (default scan of the script, not the table setup):
130
129
 
@@ -168,7 +167,6 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
168
167
  | [Trajectory guide](docs/source/trajectories.rst) | Trajectory tuples and `Projector` usage |
169
168
  | [docs/REFERENCE.md](docs/REFERENCE.md) | Capabilities and limits, custom trajectories, geometry gradients, multi-GPU details, geometry and units |
170
169
  | [docs/DISTRIBUTED.md](docs/DISTRIBUTED.md) | Execution/memory choices, distributed loss rules, Slurm and cross-node checks |
171
- | [docs/VALIDATION.md](docs/VALIDATION.md) | Test commands, validation details and the limits of each check |
172
170
  | [docs/MIGRATION.md](docs/MIGRATION.md) | Moving from the circular-only API |
173
171
  | [docs/video/README.md](docs/video/README.md) | How the intro video is rendered |
174
172
  | [CHANGELOG.md](CHANGELOG.md) | Release history |
@@ -49,7 +49,7 @@ try:
49
49
  except ImportError:
50
50
  pass
51
51
 
52
- __version__ = '2.0.0'
52
+ __version__ = '2.0.1'
53
53
 
54
54
  __all__ = [
55
55
  'Projector',
@@ -34,6 +34,7 @@ from .utils import (
34
34
  _get_numba_external_stream_for,
35
35
  _grid_2d,
36
36
  _grid_3d,
37
+ _on_device_of,
37
38
  )
38
39
  from .kernels import (
39
40
  _parallel_2d_fbp_backproject_kernel,
@@ -365,6 +366,7 @@ def _as_contig_f32(t, device):
365
366
  return DeviceManager.ensure_device(t, device).to(dtype=torch.float32).contiguous()
366
367
 
367
368
 
369
+ @_on_device_of("sinogram")
368
370
  def parallel_weighted_backproject(sinogram, ray_dir, det_origin, det_u_vec,
369
371
  detector_spacing, H, W, voxel_spacing=1.0):
370
372
  """Voxel-driven parallel-beam FBP backprojection with analytical constant.
@@ -448,6 +450,7 @@ def _analytical_geometry(src_pos, det_center, n, isocenter):
448
450
  return isocenter, float(sid_n.mean().item()), float(sdd_n.mean().item())
449
451
 
450
452
 
453
+ @_on_device_of("sinogram")
451
454
  def fan_weighted_backproject(sinogram, src_pos, det_center, det_u_vec,
452
455
  detector_spacing, H, W, voxel_spacing=1.0, *, isocenter=None):
453
456
  """Voxel-driven fan-beam FBP backprojection with analytical constant.
@@ -503,6 +506,7 @@ def fan_weighted_backproject(sinogram, src_pos, det_center, det_u_vec,
503
506
  return reco * scale
504
507
 
505
508
 
509
+ @_on_device_of("sinogram")
506
510
  def cone_weighted_backproject(sinogram, src_pos, det_center, det_u_vec, det_v_vec,
507
511
  D, H, W, du, dv, voxel_spacing=1.0, *, isocenter=None):
508
512
  """Voxel-driven cone-beam FDK backprojection with analytical constant.
@@ -13,6 +13,17 @@ from typing import Callable, Tuple
13
13
  # Trajectory Generation Functions
14
14
  # ============================================================================
15
15
 
16
+ def _check_source_positions(src_pos):
17
+ """Reject non-finite sources and sources at the isocenter, which have no viewing direction."""
18
+ if not torch.isfinite(src_pos).all():
19
+ raise ValueError("source_path_fn returned non-finite source positions")
20
+ if (torch.linalg.vector_norm(src_pos.detach(), dim=1) == 0).any():
21
+ raise ValueError(
22
+ "source_path_fn returned a source at the isocenter; every source needs a "
23
+ "nonzero distance from the origin"
24
+ )
25
+
26
+
16
27
  def circular_trajectory_3d(n_views, sid, sdd, start_angle=0.0, end_angle=None, device='cuda', dtype=torch.float32):
17
28
  """Generate circular trajectory geometry for cone-beam CT.
18
29
 
@@ -500,11 +511,11 @@ def custom_trajectory_3d(n_views, sid, sdd,
500
511
 
501
512
  Examples
502
513
  --------
503
- >>> # Define a custom figure-8 trajectory
514
+ >>> # Circular orbit with a two-lobe (figure-8) vertical motion
504
515
  >>> def figure8_path(angles, sid):
505
516
  ... src_pos = torch.zeros((len(angles), 3), device=angles.device, dtype=angles.dtype)
506
517
  ... src_pos[:, 0] = -sid * torch.sin(angles)
507
- ... src_pos[:, 1] = sid * torch.cos(angles) * torch.sin(angles)
518
+ ... src_pos[:, 1] = sid * torch.cos(angles)
508
519
  ... src_pos[:, 2] = 50 * torch.sin(2 * angles)
509
520
  ... return src_pos
510
521
  >>>
@@ -523,6 +534,7 @@ def custom_trajectory_3d(n_views, sid, sdd,
523
534
  if src_pos.shape != (n_views, 3):
524
535
  raise ValueError(f"source_path_fn must return tensor of shape ({n_views}, 3), "
525
536
  f"got {src_pos.shape}")
537
+ _check_source_positions(src_pos)
526
538
 
527
539
  # Keep the entire frame in the source-path autograd graph. Constructing
528
540
  # tensors from scalar tensors detaches them; rewriting view slices in
@@ -750,6 +762,7 @@ def custom_trajectory_2d_fan(n_views, sid, sdd,
750
762
  if src_pos.shape != (n_views, 2):
751
763
  raise ValueError(f"source_path_fn must return tensor of shape ({n_views}, 2), "
752
764
  f"got {src_pos.shape}")
765
+ _check_source_positions(src_pos)
753
766
 
754
767
  # Preallocate remaining position matrices
755
768
  det_center = torch.zeros((n_views, 2), device=device, dtype=dtype)
@@ -523,6 +523,15 @@ class Projector:
523
523
  *self.detector_spacing, self.voxel_spacing,
524
524
  )
525
525
 
526
+ @property
527
+ def process_group(self):
528
+ """Process group of a distributed operator, or None for the default group.
529
+
530
+ Pass it as ``group=`` to collectives that combine per-rank results of this
531
+ operator, so they use the same ranks as the operator.
532
+ """
533
+ return self._process_group
534
+
526
535
  def project(self, volume):
527
536
  """Project a ``(H, W)`` or ``(D, H, W)`` CUDA tensor to float32.
528
537
 
@@ -12,6 +12,7 @@ from .utils import (
12
12
  DeviceManager,
13
13
  TorchCUDABridge,
14
14
  _get_numba_external_stream_for,
15
+ _on_device_of,
15
16
  _trig_tables,
16
17
  _validate_3d_memory_layout,
17
18
  _grid_2d,
@@ -77,6 +78,7 @@ class _GeometryVJPFunction(torch.autograd.Function):
77
78
  """First-order geometry gradient; differentiating it again raises."""
78
79
 
79
80
  @staticmethod
81
+ @_on_device_of("volume")
80
82
  def forward(ctx, beam, volume, grad_sino, detector_spacing, voxel_spacing, *geometry):
81
83
  return _geometry_vjp(beam, volume, grad_sino, geometry, detector_spacing, voxel_spacing)
82
84
 
@@ -119,23 +121,17 @@ class ParallelProjectorFunction(torch.autograd.Function):
119
121
  Examples
120
122
  --------
121
123
  >>> import torch
122
- >>> from diffct.differentiable import ParallelProjectorFunction
124
+ >>> from diffct import ParallelProjectorFunction, circular_trajectory_2d_parallel
123
125
  >>>
124
- >>> # Create a 2D image with gradient tracking
125
126
  >>> image = torch.randn(128, 128, device='cuda', requires_grad=True)
126
- >>> # Define projection parameters
127
- >>> angles = torch.linspace(0, torch.pi, 180, device='cuda')
128
- >>> num_detectors = 128
129
- >>> detector_spacing = 1.0
130
- >>> # Compute forward projection
131
- >>> projector = ParallelProjectorFunction.apply
132
- >>> sinogram = projector(image, angles, num_detectors, detector_spacing)
133
- >>> # Compute loss and gradients
134
- >>> loss = sinogram.sum()
135
- >>> loss.backward()
136
- >>> print(f"Gradient shape: {image.grad.shape}") # (128, 128)
127
+ >>> ray_dir, det_origin, det_u_vec = circular_trajectory_2d_parallel(180, device='cuda')
128
+ >>> sinogram = ParallelProjectorFunction.apply(image, ray_dir, det_origin, det_u_vec, 192, 1.0)
129
+ >>> sinogram.sum().backward()
130
+ >>> image.grad.shape
131
+ torch.Size([128, 128])
137
132
  """
138
133
  @staticmethod
134
+ @_on_device_of("image")
139
135
  def forward(ctx, image, ray_dir, det_origin, det_u_vec, num_detectors, detector_spacing=1.0, voxel_spacing=1.0):
140
136
  """Compute the 2D parallel beam forward projection with arbitrary trajectories using CUDA acceleration.
141
137
 
@@ -254,16 +250,17 @@ class ParallelBackprojectorFunction(torch.autograd.Function):
254
250
  Examples
255
251
  --------
256
252
  >>> import torch
257
- >>> from diffct.differentiable import ParallelBackprojectorFunction
253
+ >>> from diffct import ParallelBackprojectorFunction, circular_trajectory_2d_parallel
258
254
  >>>
259
- >>> sinogram = torch.randn(180, 128, device='cuda', requires_grad=True)
260
- >>> angles = torch.linspace(0, torch.pi, 180, device='cuda')
261
- >>> recon = ParallelBackprojectorFunction.apply(sinogram, angles, 1.0, 128, 128)
262
- >>> loss = recon.sum()
263
- >>> loss.backward()
264
- >>> print(sinogram.grad.shape) # (180, 128)
255
+ >>> sinogram = torch.randn(180, 192, device='cuda', requires_grad=True)
256
+ >>> ray_dir, det_origin, det_u_vec = circular_trajectory_2d_parallel(180, device='cuda')
257
+ >>> image = ParallelBackprojectorFunction.apply(sinogram, ray_dir, det_origin, det_u_vec, 1.0, 128, 128)
258
+ >>> image.sum().backward()
259
+ >>> sinogram.grad.shape
260
+ torch.Size([180, 192])
265
261
  """
266
262
  @staticmethod
263
+ @_on_device_of("sinogram")
267
264
  def forward(ctx, sinogram, ray_dir, det_origin, det_u_vec, detector_spacing=1.0, H=128, W=128, voxel_spacing=1.0):
268
265
  """Compute the 2D parallel beam backprojection with arbitrary trajectories using CUDA acceleration.
269
266
 
@@ -384,16 +381,17 @@ class FanProjectorFunction(torch.autograd.Function):
384
381
  Examples
385
382
  --------
386
383
  >>> import torch
387
- >>> from diffct.differentiable import FanProjectorFunction
384
+ >>> from diffct import FanProjectorFunction, circular_trajectory_2d_fan
388
385
  >>>
389
386
  >>> image = torch.randn(256, 256, device='cuda', requires_grad=True)
390
- >>> angles = torch.linspace(0, 2 * torch.pi, 360, device='cuda')
391
- >>> sinogram = FanProjectorFunction.apply(image, angles, 512, 1.0, 1500.0, 1000.0)
392
- >>> loss = sinogram.sum()
393
- >>> loss.backward()
394
- >>> print(image.grad.shape) # (256, 256)
387
+ >>> src_pos, det_center, det_u_vec = circular_trajectory_2d_fan(360, sid=1000.0, sdd=1500.0, device='cuda')
388
+ >>> sinogram = FanProjectorFunction.apply(image, src_pos, det_center, det_u_vec, 512, 1.0)
389
+ >>> sinogram.sum().backward()
390
+ >>> image.grad.shape
391
+ torch.Size([256, 256])
395
392
  """
396
393
  @staticmethod
394
+ @_on_device_of("image")
397
395
  def forward(ctx, image, src_pos, det_center, det_u_vec, num_detectors, detector_spacing, voxel_spacing=1.0):
398
396
  """Compute the 2D fan beam forward projection with arbitrary trajectories using CUDA acceleration.
399
397
 
@@ -511,16 +509,17 @@ class FanBackprojectorFunction(torch.autograd.Function):
511
509
  Examples
512
510
  --------
513
511
  >>> import torch
514
- >>> from diffct.differentiable import FanBackprojectorFunction
512
+ >>> from diffct import FanBackprojectorFunction, circular_trajectory_2d_fan
515
513
  >>>
516
514
  >>> sinogram = torch.randn(360, 512, device='cuda', requires_grad=True)
517
- >>> angles = torch.linspace(0, 2 * torch.pi, 360, device='cuda')
518
- >>> recon = FanBackprojectorFunction.apply(sinogram, angles, 1.0, 256, 256, 1500.0, 1000.0)
519
- >>> loss = recon.sum()
520
- >>> loss.backward()
521
- >>> print(sinogram.grad.shape) # (360, 512)
515
+ >>> src_pos, det_center, det_u_vec = circular_trajectory_2d_fan(360, sid=1000.0, sdd=1500.0, device='cuda')
516
+ >>> image = FanBackprojectorFunction.apply(sinogram, src_pos, det_center, det_u_vec, 1.0, 256, 256)
517
+ >>> image.sum().backward()
518
+ >>> sinogram.grad.shape
519
+ torch.Size([360, 512])
522
520
  """
523
521
  @staticmethod
522
+ @_on_device_of("sinogram")
524
523
  def forward(ctx, sinogram, src_pos, det_center, det_u_vec, detector_spacing, H, W, voxel_spacing=1.0):
525
524
  """Compute the 2D fan beam backprojection with arbitrary trajectories using CUDA acceleration.
526
525
 
@@ -639,16 +638,19 @@ class ConeProjectorFunction(torch.autograd.Function):
639
638
  Examples
640
639
  --------
641
640
  >>> import torch
642
- >>> from diffct.differentiable import ConeProjectorFunction
641
+ >>> from diffct import ConeProjectorFunction, circular_trajectory_3d
643
642
  >>>
644
643
  >>> volume = torch.randn(128, 128, 128, device='cuda', requires_grad=True)
645
- >>> angles = torch.linspace(0, 2 * torch.pi, 360, device='cuda')
646
- >>> projections = ConeProjectorFunction.apply(volume, angles, 256, 256, 1.0, 1.0, 1500.0, 1000.0)
647
- >>> loss = projections.sum()
648
- >>> loss.backward()
649
- >>> print(volume.grad.shape) # (128, 128, 128)
644
+ >>> src_pos, det_center, det_u_vec, det_v_vec = circular_trajectory_3d(360, sid=1000.0, sdd=1500.0, device='cuda')
645
+ >>> projections = ConeProjectorFunction.apply(
646
+ ... volume, src_pos, det_center, det_u_vec, det_v_vec, 256, 256, 1.0, 1.0
647
+ ... )
648
+ >>> projections.sum().backward()
649
+ >>> volume.grad.shape
650
+ torch.Size([128, 128, 128])
650
651
  """
651
652
  @staticmethod
653
+ @_on_device_of("volume")
652
654
  def forward(ctx, volume, src_pos, det_center, det_u_vec, det_v_vec, det_u, det_v, du, dv, voxel_spacing=1.0):
653
655
  """Compute the 3D cone beam forward projection with arbitrary trajectories using CUDA acceleration.
654
656
 
@@ -784,20 +786,19 @@ class ConeBackprojectorFunction(torch.autograd.Function):
784
786
  Examples
785
787
  --------
786
788
  >>> import torch
787
- >>> from diffct.differentiable import ConeBackprojectorFunction
789
+ >>> from diffct import ConeBackprojectorFunction, circular_trajectory_3d
788
790
  >>>
789
791
  >>> projections = torch.randn(360, 256, 256, device='cuda', requires_grad=True)
790
- >>> angles = torch.linspace(0, 2 * torch.pi, 360, device='cuda')
791
- >>> D, H, W = 128, 128, 128
792
- >>> du, dv = 1.0, 1.0
793
- >>> sdd, sid = 1500.0, 1000.0
794
- >>> backprojector = ConeBackprojectorFunction.apply
795
- >>> volume = backprojector(projections, angles, D, H, W, du, dv, sdd, sid)
796
- >>> loss = volume.sum()
797
- >>> loss.backward()
798
- >>> print(f"Projection gradient shape: {projections.grad.shape}") # (360, 256, 256)
792
+ >>> src_pos, det_center, det_u_vec, det_v_vec = circular_trajectory_3d(360, sid=1000.0, sdd=1500.0, device='cuda')
793
+ >>> volume = ConeBackprojectorFunction.apply(
794
+ ... projections, src_pos, det_center, det_u_vec, det_v_vec, 128, 128, 128, 1.0, 1.0
795
+ ... )
796
+ >>> volume.sum().backward()
797
+ >>> projections.grad.shape
798
+ torch.Size([360, 256, 256])
799
799
  """
800
800
  @staticmethod
801
+ @_on_device_of("sinogram")
801
802
  def forward(ctx, sinogram, src_pos, det_center, det_u_vec, det_v_vec, D, H, W, du, dv, voxel_spacing=1.0):
802
803
  """Compute the 3D cone beam backprojection with arbitrary trajectories using CUDA acceleration.
803
804
 
@@ -5,7 +5,13 @@ PyTorch-CUDA bridging, stream caching, trigonometric table generation,
5
5
  memory layout validation, and CUDA grid computation.
6
6
  """
7
7
 
8
+ import atexit
9
+ import contextlib
10
+ import functools
11
+ import inspect
8
12
  import math
13
+ import sys
14
+ import weakref
9
15
  import numpy as np
10
16
  import torch
11
17
  from numba import cuda
@@ -150,6 +156,64 @@ def _get_numba_external_stream_for(pt_stream=None):
150
156
  return numba_stream
151
157
 
152
158
 
159
+ @contextlib.contextmanager
160
+ def _cuda_context(device):
161
+ """Make ``device`` the current device for PyTorch and Numba.
162
+
163
+ Kernel launches then use the CUDA context and the current stream of the
164
+ device that holds the data, not of the device that happens to be current.
165
+ """
166
+ device = torch.device(device)
167
+ with torch.cuda.device(device), cuda.gpus[device.index]:
168
+ yield
169
+
170
+
171
+ def _on_device_of(name):
172
+ """Run the decorated function in the CUDA context of its argument ``name``.
173
+
174
+ The argument may be passed by position or by keyword.
175
+ """
176
+ def decorate(fn):
177
+ signature = inspect.signature(fn)
178
+
179
+ @functools.wraps(fn)
180
+ def wrapper(*args, **kwargs):
181
+ bound = signature.bind_partial(*args, **kwargs).arguments
182
+ tensor = bound[name] if name in bound else None
183
+ if not (isinstance(tensor, torch.Tensor) and tensor.is_cuda):
184
+ return fn(*args, **kwargs)
185
+ with _cuda_context(tensor.device):
186
+ return fn(*args, **kwargs)
187
+ return wrapper
188
+ return decorate
189
+
190
+
191
+ def _keep_cuda_modules_at_exit():
192
+ """Stop numba-cuda from unloading CUDA modules while the interpreter exits.
193
+
194
+ numba-cuda queues module unloads and flushes the queue once it holds more
195
+ than ``CUDA_DEALLOCS_COUNT`` items. After kernels were loaded on several
196
+ GPUs, a flush during interpreter shutdown can unload the same library twice
197
+ and corrupt the heap, so the process crashes after all work is done. The
198
+ process exit releases these resources anyway.
199
+ """
200
+ from numba.cuda.core import config as cuda_config
201
+
202
+ cuda_config.CUDA_DEALLOCS_COUNT = sys.maxsize
203
+ cuda_config.CUDA_DEALLOCS_RATIO = float("inf")
204
+
205
+
206
+ class _ExitOrderAnchor:
207
+ pass
208
+
209
+
210
+ # weakref.finalize registers its atexit hook with its first finalizer. Create one
211
+ # now, so that the handler below is registered later and runs before it.
212
+ _EXIT_ORDER_ANCHOR = _ExitOrderAnchor()
213
+ weakref.finalize(_EXIT_ORDER_ANCHOR, lambda: None)
214
+ atexit.register(_keep_cuda_modules_at_exit)
215
+
216
+
153
217
  # ============================================================================
154
218
  # GPU-aware Trigonometric Table Generation
155
219
  # ============================================================================
@@ -185,6 +249,8 @@ def _trig_tables(angles, dtype=_DTYPE, device=None):
185
249
  """
186
250
  if isinstance(angles, torch.Tensor):
187
251
  device = angles.device if device is None else device
252
+ if not isinstance(dtype, torch.dtype):
253
+ dtype = torch.from_numpy(np.zeros(0, dtype=dtype)).dtype
188
254
  # Compute both cos and sin in one call to avoid redundant kernel launches
189
255
  angles_device = angles.to(dtype=dtype, device=device)
190
256
  cos = torch.cos(angles_device)
@@ -6,6 +6,7 @@ and `docs/assets/diffct_intro.mp4`. Make the inputs on a CUDA GPU, then render a
6
6
  ```bash
7
7
  python docs/video/make_inputs.py --gpu # walnut reconstructions; without --gpu only data2d.npz
8
8
  cd docs/video
9
+ rm -f scenes.txt
9
10
  for s in S00Intro S01Trajectory S02Autograd S03Geometry S04MultiGPU S05Measured S06End; do
10
11
  manim -qh diffct_intro.py $s
11
12
  echo "file 'media/videos/diffct_intro/1080p60/$s.mp4'" >> scenes.txt
@@ -19,4 +20,4 @@ ffmpeg -i ../assets/diffct_intro.mp4 -i palette.png \
19
20
 
20
21
  The reconstructions come from the measured walnut in `examples/data/walnut_cone.npz`
21
22
  (Meaney 2022, CC BY 4.0; see `examples/data/NOTICE`). The timings in the multi-GPU scene
22
- are measured diffct results on A100 64 GB GPUs (Leonardo Booster); see `docs/assets/scaling.json`.
23
+ are measured diffct results on A100 64 GB GPUs; see `docs/assets/scaling.json`.
@@ -130,74 +130,6 @@ These rules apply when you write your own loss with `Projector(..., distributed=
130
130
 
131
131
  `iterative_reconstruction.py` and `geometry_calibration.py` show these rules in use.
132
132
 
133
- ## Results
134
-
135
- The following are recorded measurements from Leonardo Booster (A100 64 GB,
136
- PyTorch 2.10 cu126, numba-cuda 0.30.4) on 2026-10-09, not expected results for
137
- every GPU. See the [validation record](https://github.com/sypsyp97/diffct/blob/main/docs/VALIDATION.md) for related
138
- correctness and cross-node results.
139
-
140
- ### Iterative reconstruction
141
-
142
- 128³ Shepp-Logan, 360 views, one GPU, noise-free. Values are PSNR (dB).
143
-
144
- | trajectory | FDK | CGLS 30 it | SIRT 200 it | TV 200 it |
145
- | --- | ---: | ---: | ---: | ---: |
146
- | circular | 34.10 | 38.58 | 32.73 | 48.37 |
147
- | helical | 32.13 | 38.49 | 32.80 | 48.37 |
148
- | saddle | 29.83 | 38.82 | 32.63 | 48.27 |
149
- | sinusoidal | 33.84 | 38.49 | 32.76 | 48.50 |
150
-
151
- FDK is designed for the circular orbit and is approximate away from its central plane. For the other orbits the approximation is coarser.
152
-
153
- With 1% Gaussian noise (helical): FDK 31.19, CGLS 35.83, SIRT 32.68, TV 47.46.
154
-
155
- Time per iteration, same helical run (CGLS / SIRT / TV):
156
-
157
- | launch mode | CGLS | SIRT | TV |
158
- | --- | ---: | ---: | ---: |
159
- | 1 GPU | 69.2 ms | 66.2 ms | 67.5 ms |
160
- | one process, 4 GPUs | 22.7 ms | 20.3 ms | 21.7 ms |
161
- | torchrun, 4 GPUs | 19.3 ms | 17.5 ms | 18.5 ms |
162
- | torchrun, 8 GPUs on 2 nodes | 12.0 ms | 10.6 ms | 11.5 ms |
163
-
164
- PSNR is identical in every launch mode.
165
-
166
- ### Projector scaling
167
-
168
- Circular scan, 360 views, detector (2n, 1.5n) with pitch 1.25, 128³ volume. Times are for 1 GPU, then 4 GPUs (1 node), then 8 GPUs (2 nodes). Speedup is relative to 1 GPU.
169
-
170
- | operation | 1 GPU | 4 GPUs | 8 GPUs (2 nodes) |
171
- | --- | ---: | ---: | ---: |
172
- | forward | 8.46 ms | 2.67 ms (3.2×) | 1.69 ms (5.0×) |
173
- | adjoint | 20.69 ms | 5.91 ms (3.5×) | 4.53 ms (4.6×) |
174
- | CGLS iteration | 30.07 ms | 8.09 ms (3.7×) | 5.54 ms (5.4×) |
175
-
176
- 64³ CGLS iteration: 4.75 ms on 1 GPU, 1.57 ms on 4 GPUs (3.0×), 1.56 ms on 8 GPUs (3.0×). Small problems are dominated by transfers and kernel launches.
177
-
178
- ### Analytical reconstruction
179
-
180
- 128² and 128³ Shepp-Logan. PSNR (dB):
181
-
182
- | window | parallel FBP | fan FBP | cone FDK |
183
- | --- | ---: | ---: | ---: |
184
- | shepp-logan | 38.10 | 34.92 | 34.10 |
185
- | hann | 27.66 | 28.12 | 29.22 |
186
-
187
- The window and the detector sampling set the sharpness.
188
-
189
- ### Geometry calibration
190
-
191
- 64³ volume, 360 views. The true per-view angle error is 0.54° RMS. The true detector shift is (1.5, -1.0) voxels. Adam runs 300 steps.
192
-
193
- - Angle error after calibration: 0.0001° to 0.0025° RMS.
194
- - Detector shift: recovered to less than 0.001 voxel in every launch mode.
195
- - Time per step: 19.4 ms on 1 GPU, 8.2 ms on torchrun with 4 GPUs, 6.7 ms on 8 GPUs (2 nodes).
196
-
197
- ### FDK check against ASTRA
198
-
199
- FDK from diffct matches ASTRA 2.5.0 `FDK_CUDA` on the same projections. The PSNR difference is within 0.3 dB. The maximum pixel difference is below 1% of the phantom maximum.
200
-
201
133
  ## Example scan
202
134
 
203
135
  The examples use this scan for an n³ volume, unit voxels:
@@ -28,10 +28,10 @@ been:
28
28
  - spatially binned 8x (2368x2240 -> 296x280 detector pixels),
29
29
  - cropped to the central 256x256 detector region,
30
30
  - angularly subsampled every 3rd view (721 -> 241 projections),
31
- - flat-field normalised per-projection (using the 99.5th percentile
32
- of each raw uint16 image as the "air" reference; no separate
33
- dark scan is available in the source record, so we assume zero
34
- dark field) and converted to line integrals via ``-log(I / I0)``,
31
+ - divided per projection by an air level I0 (the 99.5th percentile
32
+ of each image; the published projections are already dark- and
33
+ flat-field corrected) and converted to line integrals via
34
+ ``-log(I / I0)``,
35
35
  - stored as float16 in a compressed .npz together with the
36
36
  physical geometry (sdd, sid, binned pixel pitch, angles in
37
37
  radians in diffct handedness) plus regeneration metadata such
@@ -44,7 +44,7 @@ All finite domains and numerical tolerances are in the check files.
44
44
  Numba's [simulator documentation](https://nvidia.github.io/numba-cuda/user/simulator.html) describes its execution and limitations.
45
45
  The simulator executes Python kernel bodies—its arithmetic and scheduling do not establish the behavior of compiled GPU code.
46
46
  This suite does not verify CUDA type inference, LLVM fastmath, streams, multi-GPU execution, geometry gradients, or PyTorch autograd wrappers.
47
- The existing GPU validation remains necessary for those paths; see [VALIDATION.md](../docs/VALIDATION.md).
47
+ The GPU test suite in `tests/` covers those paths on a CUDA machine.
48
48
 
49
49
  ## Installation
50
50
 
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "diffct"
7
- version = "2.0.0"
7
+ version = "2.0.1"
8
8
  description = "Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU and multi-node execution, geometry gradients"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -11,7 +11,7 @@ test suite stays fast.
11
11
  Run everything:
12
12
 
13
13
  ```bash
14
- /c/Users/sun/miniconda3/envs/cuda12/python.exe -m pytest tests/benchmarks/ \
14
+ python -m pytest tests/benchmarks/ \
15
15
  --benchmark-only --benchmark-columns=min,mean,median,stddev,rounds
16
16
  ```
17
17
 
@@ -0,0 +1,83 @@
1
+ """Interpreter shutdown after CUDA kernels run on two devices."""
2
+
3
+ import os
4
+ from pathlib import Path
5
+ import subprocess
6
+ import sys
7
+ import textwrap
8
+
9
+ import pytest
10
+ import torch
11
+
12
+
13
+ @pytest.mark.cuda
14
+ def test_multi_gpu_exit():
15
+ if torch.cuda.device_count() < 2:
16
+ pytest.skip("At least two CUDA devices are required")
17
+
18
+ workload = textwrap.dedent("""\
19
+ import torch
20
+ import diffct
21
+
22
+ for index in (0, 1):
23
+ device = torch.device(f"cuda:{index}")
24
+ with torch.cuda.device(device):
25
+ parallel = diffct.circular_trajectory_2d_parallel(4, device=device)
26
+ fan = diffct.circular_trajectory_2d_fan(
27
+ 4, sid=20.0, sdd=30.0, device=device
28
+ )
29
+ for geometry, projector, backprojector, weighted in (
30
+ (parallel, diffct.ParallelProjectorFunction,
31
+ diffct.ParallelBackprojectorFunction,
32
+ diffct.parallel_weighted_backproject),
33
+ (fan, diffct.FanProjectorFunction,
34
+ diffct.FanBackprojectorFunction,
35
+ diffct.fan_weighted_backproject),
36
+ ):
37
+ image = torch.ones(
38
+ (8, 8), device=device, dtype=torch.float32,
39
+ requires_grad=True,
40
+ )
41
+ projected = projector.apply(image, *geometry, 12, 1.0)
42
+ projected.sum().backward()
43
+ sinogram = projected.detach().requires_grad_(True)
44
+ backprojector.apply(sinogram, *geometry, 1.0, 8, 8).sum().backward()
45
+ weighted(sinogram.detach(), *geometry, 1.0, 8, 8)
46
+
47
+ cone = diffct.circular_trajectory_3d(
48
+ 4, sid=20.0, sdd=30.0, device=device
49
+ )
50
+ volume = torch.ones(
51
+ (4, 4, 4), device=device, dtype=torch.float32,
52
+ requires_grad=True,
53
+ )
54
+ projected = diffct.ConeProjectorFunction.apply(
55
+ volume, *cone, 8, 8, 1.0, 1.0,
56
+ )
57
+ projected.sum().backward()
58
+ sinogram = projected.detach().requires_grad_(True)
59
+ diffct.ConeBackprojectorFunction.apply(
60
+ sinogram, *cone, 4, 4, 4, 1.0, 1.0,
61
+ ).sum().backward()
62
+ diffct.cone_weighted_backproject(
63
+ sinogram.detach(), *cone, 4, 4, 4, 1.0, 1.0,
64
+ )
65
+ torch.cuda.synchronize(device)
66
+ """)
67
+ env = os.environ.copy()
68
+ repository = Path(__file__).resolve().parents[1]
69
+ env["PYTHONPATH"] = os.pathsep.join([
70
+ str(repository),
71
+ *(str(Path(entry or os.getcwd()).resolve()) for entry in sys.path),
72
+ ])
73
+ child = subprocess.run(
74
+ [sys.executable, "-c", workload],
75
+ env=env,
76
+ capture_output=True,
77
+ text=True,
78
+ timeout=300,
79
+ )
80
+ assert child.returncode == 0, (
81
+ f"CUDA subprocess exited with return code {child.returncode}\n"
82
+ f"Child stderr tail:\n{child.stderr[-8000:]}"
83
+ )
@@ -0,0 +1,326 @@
1
+ """Release contracts for device selection, geometry, and walnut examples."""
2
+
3
+ import doctest
4
+ import importlib
5
+ import io
6
+ import math
7
+ import sys
8
+ import zipfile
9
+ from functools import partial
10
+ from pathlib import Path
11
+
12
+ import numpy as np
13
+ import pytest
14
+ import torch
15
+
16
+
17
+ ROOT = Path(__file__).resolve().parents[1]
18
+ FUNCTIONS = [
19
+ "ParallelProjectorFunction", "ParallelBackprojectorFunction",
20
+ "FanProjectorFunction", "FanBackprojectorFunction",
21
+ "ConeProjectorFunction", "ConeBackprojectorFunction",
22
+ ]
23
+ BEAMS = ["parallel", "fan", "cone"]
24
+ requires_cuda = pytest.mark.skipif(
25
+ not torch.cuda.is_available(), reason="requires CUDA"
26
+ )
27
+ requires_two_gpus = pytest.mark.skipif(
28
+ torch.cuda.device_count() < 2, reason="requires two CUDA devices"
29
+ )
30
+
31
+
32
+ @pytest.fixture
33
+ def examples(monkeypatch):
34
+ monkeypatch.syspath_prepend(str(ROOT / "examples"))
35
+ return importlib.import_module("data.preprocess_walnut")
36
+
37
+
38
+ def trajectory(beam, device):
39
+ geometry = importlib.import_module("diffct.geometry")
40
+ if beam == "parallel":
41
+ return geometry.circular_trajectory_2d_parallel(5, start_angle=0.13, device=device)
42
+ if beam == "fan":
43
+ return geometry.circular_trajectory_2d_fan(5, 20.0, 35.0, start_angle=0.13, device=device)
44
+ return geometry.circular_trajectory_3d(5, 20.0, 35.0, start_angle=0.13, device=device)
45
+
46
+
47
+ def function_arguments(name, geometry):
48
+ if name == "ConeProjectorFunction":
49
+ return (*geometry, 9, 7, 0.9, 1.1, 0.8)
50
+ if name == "ConeBackprojectorFunction":
51
+ return (*geometry, 4, 6, 8, 0.9, 1.1, 0.8)
52
+ if "Backprojector" in name:
53
+ return (*geometry, 0.9, 6, 8, 0.8)
54
+ return (*geometry, 9, 0.9, 0.8)
55
+
56
+
57
+ @pytest.mark.cuda
58
+ @requires_two_gpus
59
+ @pytest.mark.parametrize("name", FUNCTIONS)
60
+ def test_function_uses_input_cuda_device_for_output_and_gradients(name):
61
+ projectors = importlib.import_module("diffct.projectors")
62
+ beam = next(beam for beam in BEAMS if name.lower().startswith(beam))
63
+ geometry = trajectory(beam, "cuda:1")
64
+ if "Backprojector" in name:
65
+ shape = (5, 9, 7) if beam == "cone" else (5, 9)
66
+ else:
67
+ shape = (4, 6, 8) if beam == "cone" else (6, 8)
68
+ data = torch.linspace(0.1, 0.9, math.prod(shape), device="cuda:1").reshape(shape)
69
+
70
+ def evaluate(current_device):
71
+ with torch.cuda.device(current_device):
72
+ inputs = [value.detach().clone().requires_grad_() for value in (data, *geometry)]
73
+ result = getattr(projectors, name).apply(
74
+ inputs[0], *function_arguments(name, inputs[1:])
75
+ )
76
+ assert result.device == torch.device("cuda:1")
77
+ assert torch.isfinite(result).all()
78
+ weights = torch.linspace(0.2, 1.0, result.numel(), device="cuda:1").reshape(result.shape)
79
+ gradients = torch.autograd.grad(result, inputs, grad_outputs=weights)
80
+ torch.cuda.synchronize(1)
81
+ assert torch.cuda.current_device() == current_device
82
+ for gradient in gradients:
83
+ assert gradient.device == torch.device("cuda:1")
84
+ assert torch.isfinite(gradient).all()
85
+ return result.detach(), tuple(gradient.detach() for gradient in gradients)
86
+
87
+ expected, expected_gradients = evaluate(1)
88
+ actual, actual_gradients = evaluate(0)
89
+ torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
90
+ for actual_gradient, expected_gradient in zip(actual_gradients, expected_gradients):
91
+ torch.testing.assert_close(actual_gradient, expected_gradient, rtol=1e-5, atol=1e-6)
92
+
93
+
94
+ @pytest.mark.cuda
95
+ @requires_two_gpus
96
+ @pytest.mark.parametrize("beam", BEAMS)
97
+ def test_analytical_backprojection_uses_input_cuda_device(beam):
98
+ analytical = importlib.import_module("diffct.analytical")
99
+ geometry = trajectory(beam, "cuda:1")
100
+ shape = (5, 9, 7) if beam == "cone" else (5, 9)
101
+ sinogram = torch.linspace(0.1, 0.9, math.prod(shape), device="cuda:1").reshape(shape)
102
+ sizes = (4, 6, 8, 0.9, 1.1) if beam == "cone" else (0.9, 6, 8)
103
+ function = getattr(analytical, f"{beam}_weighted_backproject")
104
+
105
+ def evaluate(current_device):
106
+ with torch.cuda.device(current_device):
107
+ result = function(sinogram, *geometry, *sizes, voxel_spacing=0.8)
108
+ assert result.device == torch.device("cuda:1")
109
+ torch.cuda.synchronize(1)
110
+ assert torch.isfinite(result).all()
111
+ assert torch.cuda.current_device() == current_device
112
+ return result
113
+
114
+ torch.testing.assert_close(evaluate(0), evaluate(1), rtol=1e-5, atol=1e-6)
115
+
116
+
117
+ def weighted_backprojection_arguments(beam, device):
118
+ geometry = trajectory(beam, device)
119
+ shape = (5, 9, 7) if beam == "cone" else (5, 9)
120
+ sinogram = torch.linspace(0.1, 0.9, math.prod(shape), device=device).reshape(shape)
121
+ if beam == "cone":
122
+ names = ("sinogram", "src_pos", "det_center", "det_u_vec", "det_v_vec",
123
+ "D", "H", "W", "du", "dv", "voxel_spacing")
124
+ args = (sinogram, *geometry, 4, 6, 8, 0.9, 1.1, 0.8)
125
+ else:
126
+ geometry_names = ("ray_dir", "det_origin", "det_u_vec") if beam == "parallel" else (
127
+ "src_pos", "det_center", "det_u_vec"
128
+ )
129
+ names = ("sinogram", *geometry_names, "detector_spacing", "H", "W", "voxel_spacing")
130
+ args = (sinogram, *geometry, 0.9, 6, 8, 0.8)
131
+ keyword_only = {} if beam == "parallel" else {
132
+ "isocenter": torch.zeros(3 if beam == "cone" else 2, device=device)
133
+ }
134
+ return args, dict(zip(names, args), **keyword_only), keyword_only
135
+
136
+
137
+ @pytest.mark.cuda
138
+ @requires_cuda
139
+ @pytest.mark.parametrize("beam", BEAMS)
140
+ def test_analytical_backprojection_accepts_all_arguments_by_keyword(beam):
141
+ analytical = importlib.import_module("diffct.analytical")
142
+ function = getattr(analytical, f"{beam}_weighted_backproject")
143
+ args, kwargs, keyword_only = weighted_backprojection_arguments(beam, "cuda:0")
144
+ with torch.cuda.device(0):
145
+ expected = function(*args, **keyword_only)
146
+ actual = function(**kwargs)
147
+ assert actual.device == args[0].device
148
+ assert torch.isfinite(actual).all()
149
+ torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
150
+
151
+
152
+ @pytest.mark.cuda
153
+ @requires_two_gpus
154
+ @pytest.mark.parametrize("beam", BEAMS)
155
+ def test_analytical_keyword_backprojection_uses_input_cuda_device(beam):
156
+ analytical = importlib.import_module("diffct.analytical")
157
+ function = getattr(analytical, f"{beam}_weighted_backproject")
158
+ args, kwargs, keyword_only = weighted_backprojection_arguments(beam, "cuda:1")
159
+ with torch.cuda.device(1):
160
+ expected = function(*args, **keyword_only)
161
+ with torch.cuda.device(0):
162
+ actual = function(**kwargs)
163
+ assert actual.device == torch.device("cuda:1")
164
+ torch.cuda.synchronize(1)
165
+ assert torch.cuda.current_device() == 0
166
+ assert torch.isfinite(actual).all()
167
+ torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
168
+
169
+
170
+ @pytest.mark.parametrize("dimensions", [2, 3])
171
+ @pytest.mark.parametrize("bad_value", [0.0, math.nan, math.inf, -math.inf], ids=["origin", "nan", "inf", "negative-inf"])
172
+ @pytest.mark.parametrize("bad_view", [0, 2, 4])
173
+ def test_custom_trajectory_rejects_invalid_source(dimensions, bad_value, bad_view):
174
+ geometry = importlib.import_module("diffct.geometry")
175
+ function = geometry.custom_trajectory_3d if dimensions == 3 else geometry.custom_trajectory_2d_fan
176
+
177
+ def source_path(angles, sid):
178
+ positions = torch.zeros((len(angles), dimensions), device=angles.device, dtype=angles.dtype)
179
+ positions[:, 0] = -sid * torch.sin(angles)
180
+ positions[:, 1] = sid * torch.cos(angles)
181
+ positions[bad_view] = bad_value
182
+ return positions
183
+
184
+ with pytest.raises(ValueError):
185
+ function(5, 20.0, 35.0, source_path, device="cpu")
186
+
187
+
188
+ def test_figure8_docstring_example_is_finite_and_constructs_projector():
189
+ import diffct
190
+
191
+ function = diffct.custom_trajectory_3d
192
+ snippets = doctest.DocTestParser().get_examples(function.__doc__)
193
+ assert any("def figure8_path" in snippet.source for snippet in snippets)
194
+ namespace = {"torch": torch, "custom_trajectory_3d": partial(function, device="cpu")}
195
+ for snippet in snippets:
196
+ exec(compile(snippet.source, "<custom_trajectory_3d docstring>", "exec"), namespace)
197
+ geometry = tuple(namespace[name] for name in ("src_pos", "det_center", "det_u_vec", "det_v_vec"))
198
+ assert all(value.device.type == "cpu" for value in geometry)
199
+ assert all(torch.isfinite(value).all() for value in geometry)
200
+ assert (torch.linalg.vector_norm(geometry[0], dim=1) > 0).all()
201
+ operator = diffct.Projector(geometry, (8, 8, 8), (9, 7))
202
+ assert operator is not None
203
+
204
+
205
+ @pytest.mark.cuda
206
+ @requires_cuda
207
+ @pytest.mark.parametrize("name", FUNCTIONS)
208
+ def test_function_class_docstring_example(name):
209
+ projectors = importlib.import_module("diffct.projectors")
210
+ function = getattr(projectors, name)
211
+ example = doctest.DocTestParser().get_doctest(
212
+ function.__doc__, {}, name, str(ROOT / "diffct" / "projectors.py"), 0
213
+ )
214
+ assert example.examples
215
+ runner = doctest.DocTestRunner()
216
+ result = runner.run(example)
217
+ assert result.failed == 0
218
+ assert result.attempted == len(example.examples)
219
+
220
+
221
+ def kept_region_offset(raw_size, bin_factor, crop):
222
+ """Compute the midpoint of the kept raw-pixel interval from the contract."""
223
+ binned = raw_size // bin_factor
224
+ size = binned if crop == 0 else crop
225
+ start = (binned - size) // 2
226
+ left_edge = start * bin_factor
227
+ right_edge = (start + size) * bin_factor
228
+ return (left_edge + right_edge) / 2 - raw_size / 2
229
+
230
+
231
+ @pytest.mark.parametrize("raw_size,bin_factor,crop", [
232
+ (2240, 8, 256), (2368, 8, 256), (2240, 8, 255),
233
+ (2368, 8, 255), (11, 2, 0), (11, 2, 4), (12, 2, 3),
234
+ (13, 3, 3), (14, 3, 0), (9, 1, 4), (8, 1, 3), (9, 1, 3),
235
+ ])
236
+ def test_detector_center_offset(examples, raw_size, bin_factor, crop):
237
+ expected = kept_region_offset(raw_size, bin_factor, crop)
238
+ assert examples.detector_center_offset(raw_size, bin_factor, crop) == expected
239
+
240
+
241
+ def test_preprocess_without_projection_zips_raises(examples, tmp_path, monkeypatch):
242
+ output = tmp_path / "result.npz"
243
+ monkeypatch.setattr(sys, "argv", ["preprocess_walnut.py", str(output), str(tmp_path), "--view-stride", "721"])
244
+ with pytest.raises(FileNotFoundError):
245
+ examples.main()
246
+ assert not output.exists()
247
+
248
+
249
+ @pytest.mark.parametrize("projection_number", [1, 2, 721])
250
+ def test_preprocess_kept_projection_presence(examples, tmp_path, monkeypatch, projection_number):
251
+ tifffile = pytest.importorskip("tifffile")
252
+ image = np.full((2368, 2240), 1000, dtype=np.uint16)
253
+ buffer = io.BytesIO()
254
+ tifffile.imwrite(buffer, image)
255
+ with zipfile.ZipFile(tmp_path / "20201111_walnut_projections_1.zip", "w", compression=zipfile.ZIP_DEFLATED) as archive:
256
+ archive.writestr(f"nested/20201111_walnut_{projection_number:04d}.tif", buffer.getvalue())
257
+ output = tmp_path / "result.npz"
258
+ monkeypatch.setattr(sys, "argv", ["preprocess_walnut.py", str(output), str(tmp_path), "--view-stride", "721"])
259
+ if projection_number == 1:
260
+ examples.main()
261
+ with np.load(output) as data:
262
+ assert data["sinogram"].shape == (1, 256, 256)
263
+ np.testing.assert_array_equal(data["sinogram"], np.zeros((1, 256, 256)))
264
+ np.testing.assert_array_equal(data["view_indices"], [0])
265
+ else:
266
+ with pytest.raises(ValueError, match="missing"):
267
+ examples.main()
268
+ assert not output.exists()
269
+
270
+
271
+ @pytest.mark.parametrize("bin_factor,crop_h,crop_w,du,dv", [
272
+ (8, 255, 255, 0.4, 0.4), (7, 0, 0, 0.35, 0.35),
273
+ (8, 255, 256, 0.4, 0.4), (8, 255, 255, 0.4, 0.6),
274
+ ])
275
+ def test_load_scan_detector_offsets(examples, tmp_path, bin_factor, crop_h, crop_w, du, dv):
276
+ walnut = importlib.import_module("walnut_reconstruction")
277
+ path = tmp_path / "scan.npz"
278
+ stored = np.arange(24, dtype=np.float16).reshape(2, 3, 4)
279
+ angles = np.array([0.0, 0.7], dtype=np.float64)
280
+ np.savez(path, sinogram=stored, angles=angles, sid=20.0, sdd=35.0,
281
+ du=du, dv=dv, detector_bin=bin_factor, crop_h=crop_h, crop_w=crop_w)
282
+ result = walnut.load_scan(path)
283
+ assert len(result) == 8
284
+ sinogram, actual_angles, sid, sdd, actual_du, actual_dv, offset_u, offset_v = result
285
+ np.testing.assert_array_equal(sinogram, stored.astype(np.float32).transpose(0, 2, 1))
286
+ np.testing.assert_array_equal(actual_angles, angles)
287
+ assert (sid, sdd, actual_du, actual_dv) == (20.0, 35.0, du, dv)
288
+ assert offset_u == pytest.approx(kept_region_offset(2240, bin_factor, crop_w) * du / bin_factor)
289
+ assert offset_v == pytest.approx(kept_region_offset(2368, bin_factor, crop_h) * dv / bin_factor)
290
+
291
+
292
+ @pytest.mark.parametrize("offset_u,offset_v", [(0.0, 0.0), (1.25, 0.0), (0.0, -2.5), (-1.25, 2.5)])
293
+ def test_circular_geometry_moves_detector_along_axes(examples, offset_u, offset_v):
294
+ walnut = importlib.import_module("walnut_reconstruction")
295
+ angles = np.array([0.0, math.pi / 2, 0.37])
296
+ source, detector, u_axis, v_axis = walnut.circular_geometry(angles, 20.0, 35.0, offset_u, offset_v)
297
+ sin, cos = np.sin(angles), np.cos(angles)
298
+ zero = np.zeros_like(angles)
299
+ expected_u = np.column_stack((cos, sin, zero))
300
+ expected_v = np.column_stack((zero, zero, np.ones_like(angles)))
301
+ nominal_detector = np.column_stack((15 * sin, -15 * cos, zero))
302
+ np.testing.assert_allclose(source.numpy(), np.column_stack((-20 * sin, 20 * cos, zero)), rtol=1e-6, atol=1e-6)
303
+ np.testing.assert_allclose(u_axis.numpy(), expected_u, rtol=1e-6, atol=1e-6)
304
+ np.testing.assert_allclose(v_axis.numpy(), expected_v, rtol=1e-6, atol=1e-6)
305
+ np.testing.assert_allclose(detector.numpy(), nominal_detector + offset_u * expected_u + offset_v * expected_v,
306
+ rtol=1e-6, atol=1e-6)
307
+
308
+
309
+ @pytest.mark.parametrize("values", [[0.0, 0.0], [0.1, 0.7, -0.2]])
310
+ def test_psnr_identical_images_is_positive_infinity(examples, values):
311
+ common = importlib.import_module("_common")
312
+ image = torch.tensor(values, dtype=torch.float32)
313
+ assert common.psnr(image.clone(), image) == math.inf
314
+
315
+
316
+ @pytest.mark.parametrize("beam", BEAMS)
317
+ def test_projector_process_group_is_none_and_read_only(beam):
318
+ import diffct
319
+
320
+ shape = (4, 6, 8) if beam == "cone" else (6, 8)
321
+ detector = (9, 7) if beam == "cone" else 9
322
+ operator = diffct.Projector(trajectory(beam, "cpu"), shape, detector, beam=beam)
323
+ assert operator.process_group is None
324
+ with pytest.raises(AttributeError):
325
+ operator.process_group = object()
326
+ assert operator.process_group is None
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes