Cone Beam Iterative Reconstruction#

examples/iterative_reconstruction.py reconstructs a 3D cone-beam phantom with diffct.Projector. It runs on any supported trajectory, on one GPU, on several GPUs in one process, or on several processes and nodes.

Axial and coronal slices of a simulated helical scan

Simulated helical walnut scan from Measured Walnut: 720 views, 256³, 1% noise.#

Run it#

Run the commands from the root of a source checkout; pip install diffct does not install the examples/ directory. See the “From source” tab in Installation. --figure needs Matplotlib: pip install matplotlib.

python examples/iterative_reconstruction.py --trajectory helical
python examples/iterative_reconstruction.py --trajectory circular --devices 0 1 2 3
torchrun --nproc-per-node=4 examples/iterative_reconstruction.py

--devices and torchrun cannot be used together. For multiple nodes, see Multiple GPUs and Nodes.

Options#

Option

Meaning

--size

Volume edge length in voxels (default 128).

--views

Number of views (default 360).

--trajectory

circular, helical, saddle or sinusoidal (default helical).

--algorithms

One or more of cgls, sirt and tv (default: all).

--iterations

Override the per-algorithm defaults (CGLS 30, SIRT 200, TV 200).

--tv-weight

TV penalty weight for tv (default 1.0).

--noise

Gaussian noise, relative to the largest projection (default 0.0).

--phantom

(size, size, size) .npy volume with values in [0, 1], used instead of the Shepp-Logan phantom.

--figure

Save central slices to this PNG. Requires Matplotlib.

Pipeline#

  1. Project the phantom with Projector (cone beam, on the selected trajectory).

  2. Reconstruct with CGLS and SIRT. Both use the projector and its adjoint.

  3. Reconstruct with TV-regularized nonnegative least squares. The script uses Adam and autograd for this solver.

  4. When the full sinogram is on one GPU, compute an FDK baseline. On non-circular trajectories, FDK is a heuristic baseline, not an exact reconstruction.

Notes#

  • The iteration counts are demonstration settings, not a general convergence schedule. Choose a stopping rule for your data.

  • The TV weight is an example setting. Choose it for your geometry, resolution, voxel spacing and data scale. The solver uses half the mean squared residual per ray and a mean TV term. Its voxel differences are not divided by voxel spacing.

  • The script keeps complete input/output tensors on the process’s CUDA device. Native working buffers may stream; CGLS state can fall back to CPU.

  • To adapt the loss for several ranks, see Multiple GPUs and Nodes.

Source#

Iterative Reconstruction Example (3D cone beam, any trajectory)#
  1"""Iterative cone-beam reconstruction of a 3D phantom on any trajectory.
  2
  3Runs CGLS, SIRT and TV-regularized nonnegative least squares (Adam through
  4autograd), and FDK as a baseline when the full sinogram is on one GPU. The same
  5script runs on one GPU, on several GPUs in one process, or on several processes
  6and nodes:
  7
  8    python examples/iterative_reconstruction.py --trajectory helical
  9    python examples/iterative_reconstruction.py --devices 0 1 2 3
 10    torchrun --nproc-per-node=4 examples/iterative_reconstruction.py
 11    (multiple nodes: see examples/slurm/multi_node.sbatch)
 12
 13The phantom is the 3D Shepp-Logan by default. ``--phantom`` loads a ``(size,)*3``
 14volume with values in [0, 1] instead, for example the walnut volume written by
 15``walnut_reconstruction.py --save-volume``.
 16"""
 17
 18import argparse
 19import time
 20from pathlib import Path
 21
 22import numpy as np
 23import torch
 24
 25from diffct import Projector
 26from _common import (TRAJECTORIES, Scan, cgls, fdk, init_process, psnr, save_slices,
 27                     shepp_logan_3d, sirt, synchronize, tv_reconstruction)
 28
 29ALGORITHMS = {"cgls": cgls, "sirt": sirt, "tv": tv_reconstruction}
 30# CGLS converges fastest; SIRT and Adam need more, cheaper-to-tune steps.
 31DEFAULT_ITERATIONS = {"cgls": 30, "sirt": 200, "tv": 200}
 32
 33
 34def main():
 35    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
 36    parser.add_argument("--size", type=int, default=128, help="volume edge length in voxels")
 37    parser.add_argument("--views", type=int, default=360)
 38    parser.add_argument("--trajectory", choices=TRAJECTORIES, default="helical")
 39    parser.add_argument("--algorithms", nargs="+", choices=list(ALGORITHMS), default=list(ALGORITHMS))
 40    parser.add_argument("--iterations", type=int, help="override the per-algorithm default "
 41                        + str(DEFAULT_ITERATIONS))
 42    parser.add_argument("--devices", type=int, nargs="+", help="GPUs for one process (not with torchrun)")
 43    parser.add_argument("--tv-weight", type=float, default=1.0, help="TV penalty weight for the tv algorithm")
 44    parser.add_argument("--noise", type=float, default=0.0, help="Gaussian noise, relative to the largest projection")
 45    parser.add_argument("--figure", type=Path, help="save central slices to this PNG")
 46    parser.add_argument("--phantom", type=Path, help="(size,)*3 .npy volume in [0, 1] instead of Shepp-Logan")
 47    args = parser.parse_args()
 48
 49    distributed, device = init_process()
 50    if distributed and args.devices:
 51        parser.error("use either torchrun or --devices, not both")
 52    scan = Scan(args.size, args.views)
 53    trajectory = scan.trajectory(args.trajectory)
 54    # One Projector serves every mode: devices=... splits views over local GPUs,
 55    # distributed=True splits them over torchrun ranks.
 56    operator = Projector(trajectory, (args.size,) * 3, scan.detector, detector_spacing=scan.pitch,
 57                         devices=args.devices, distributed=distributed)
 58    rank0 = operator.rank == 0
 59    if args.phantom:
 60        volume = np.load(args.phantom).astype(np.float32)
 61        if volume.shape != (args.size,) * 3 or not (np.isfinite(volume).all() and 0 <= volume.min() and volume.max() <= 1):
 62            raise ValueError(f"{args.phantom} must hold a {args.size}^3 volume with values in [0, 1]")
 63        truth = torch.from_numpy(volume).to(device)
 64    else:
 65        truth = torch.from_numpy(shepp_logan_3d(args.size)).to(device)
 66    phantom_name = args.phantom.stem if args.phantom else "Shepp-Logan"
 67
 68    measurements = operator.project(truth)  # this rank's views
 69    if args.noise > 0:
 70        # Same noise for every launch mode: global peak, noise drawn for all views.
 71        peak = measurements.max()
 72        if distributed:
 73            torch.distributed.all_reduce(peak, op=torch.distributed.ReduceOp.MAX, group=operator.process_group)
 74        generator = torch.Generator(device=device).manual_seed(1234)
 75        noise = torch.randn((args.views, *scan.detector), generator=generator, device=device)
 76        measurements = measurements + args.noise * peak * noise[operator.view_slice]
 77
 78    if rank0:
 79        print(f"{args.trajectory} trajectory, {args.size}^3 volume, {args.views} views, "
 80              f"{operator.world_size} process(es), devices per process: {args.devices or [device.index]}")
 81    panels = [("phantom\nground truth", truth)]
 82    if not distributed:
 83        fdk(scan, measurements, trajectory)  # compile kernels and warm up
 84        synchronize(distributed)
 85        start = time.perf_counter()
 86        reconstruction = fdk(scan, measurements, trajectory)
 87        synchronize(distributed)
 88        print(f"  fdk   {1e3 * (time.perf_counter() - start):8.1f} ms            PSNR {psnr(reconstruction, truth):6.2f} dB")
 89        panels.append((f"FDK\n{psnr(reconstruction, truth):.1f} dB", reconstruction))
 90
 91    for name in args.algorithms:
 92        iterations = args.iterations or DEFAULT_ITERATIONS[name]
 93        options = {"weight": args.tv_weight} if name == "tv" else {}
 94        ALGORITHMS[name](operator, measurements, 1, **options)  # compile kernels and warm up
 95        synchronize(distributed)
 96        start = time.perf_counter()
 97        reconstruction = ALGORITHMS[name](operator, measurements, iterations, **options)
 98        synchronize(distributed)
 99        elapsed = time.perf_counter() - start
100        if rank0:
101            print(f"  {name:5s} {1e3 * elapsed:8.1f} ms ({1e3 * elapsed / iterations:6.2f} ms/it, {iterations} it)"
102                  f"  PSNR {psnr(reconstruction, truth):6.2f} dB")
103        label = (f"TV-regularized\nAdam, {iterations} it; {psnr(reconstruction, truth):.1f} dB"
104                 if name == "tv" else f"{name.upper()}, {iterations} it\n{psnr(reconstruction, truth):.1f} dB")
105        panels.append((label, reconstruction))
106
107    if rank0 and args.figure:
108        noise = f", {args.noise:.0%} noise" if args.noise > 0 else ""
109        save_slices(args.figure, panels, f"{args.trajectory} cone-beam scan, {args.size}³ {phantom_name}{noise}",
110                    vmax=1.0 if args.phantom else 0.4, sections=("axial", "coronal"))
111    if distributed:
112        torch.distributed.destroy_process_group()
113
114
115if __name__ == "__main__":
116    main()