Improve linop.xray.XRayTransform3D

CHANGES.rst

@@ -5,6 +5,8 @@ SCICO Release Notes
 Version 0.0.8   (unreleased)
 ----------------------------
 
+• Substantial computational improvement in ``linop.xray.XRayTransform3D``,
+  which is now faster than ``linop.xray.astra.XRayTransform3D``.
 • Enable certain parameters of array creation functions to trigger
  ``BlockArray`` creation when they receive lists (currently ``device``).
 • New interface to ASTRA Toolbox cone beam X-ray projectors.

examples/scripts/ct_3d_tv_padmm.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# This file is part of the SCICO package. Details of the copyright
+# and user license can be found in the 'LICENSE.txt' file distributed
+# with the package.
+
+r"""
+3D TV-Regularized Sparse-View CT Reconstruction (Proximal ADMM Solver)
+======================================================================
+
+This example demonstrates solution of a sparse-view, 3D CT
+reconstruction problem with isotropic total variation (TV)
+regularization
+
+  $$\mathrm{argmin}_{\mathbf{x}} \; (1/2) \| \mathbf{y} - C \mathbf{x}
+  \|_2^2 + \lambda \| D \mathbf{x} \|_{2,1} \;,$$
+
+where $C$ is the X-ray transform (the CT forward projection operator),
+$\mathbf{y}$ is the sinogram, $D$ is a 3D finite difference operator,
+and $\mathbf{x}$ is the reconstructed image.
+
+This example uses the native scico 3d X-Ray projector, while the
+[companion example](ct_astra_3d_tv_padmm.rst) uses the astra projector.
+"""
+
+import numpy as np
+
+import komplot as kplt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+import scico.numpy as snp
+from scico import functional, linop, loss, metric
+from scico.examples import create_tangle_phantom
+from scico.linop.xray import XRayTransform3D
+from scico.linop.xray.astra import angle_to_vector, convert_to_scico_geometry
+from scico.optimize import ProximalADMM
+from scico.util import device_info
+
+"""
+Create a ground truth image and projector.
+"""
+Nx = 128
+Ny = 256
+Nz = 64
+
+tangle = snp.array(create_tangle_phantom(Nx, Ny, Nz))
+
+n_projection = 10  # number of projections
+angles = np.linspace(0, np.pi, n_projection, endpoint=False)  # evenly spaced projection angles
+det_spacing = [1.0, 1.0]
+det_count = (Nz, max(Nx, Ny))
+vectors = angle_to_vector(det_spacing, angles)
+
+# It would have been more straightforward to use the det_spacing and angles keywords
+# in this case (since vectors is just computed directly from these two quantities), but
+# the more general form is used here as a demonstration.
+matrices = convert_to_scico_geometry(input_shape=tangle.shape, det_count=det_count, vectors=vectors)
+C = XRayTransform3D(tangle.shape, matrices, det_count)  # CT projection operator
+y = C @ tangle  # sinogram
+
+
+r"""
+Set up problem and solver. We want to minimize the functional
+
+  $$\mathrm{argmin}_{\mathbf{x}} \; (1/2) \| \mathbf{y} - C \mathbf{x}
+  \|_2^2 + \lambda \| D \mathbf{x} \|_{2,1} \;,$$
+
+where $C$ is the X-ray transform and $D$ is a finite difference
+operator. This problem can be expressed as
+
+  $$\mathrm{argmin}_{\mathbf{x}, \mathbf{z}} \; (1/2) \| \mathbf{y} -
+  \mathbf{z}_0 \|_2^2 + \lambda \| \mathbf{z}_1 \|_{2,1} \;\;
+  \text{such that} \;\; \mathbf{z}_0 = C \mathbf{x} \;\; \text{and} \;\;
+  \mathbf{z}_1 = D \mathbf{x} \;,$$
+
+which can be written in the form of a standard ADMM problem
+
+  $$\mathrm{argmin}_{\mathbf{x}, \mathbf{z}} \; f(\mathbf{x}) + g(\mathbf{z})
+  \;\; \text{such that} \;\; A \mathbf{x} + B \mathbf{z} = \mathbf{c}$$
+
+with
+
+  $$f = 0 \qquad g = g_0 + g_1$$
+  $$g_0(\mathbf{z}_0) = (1/2) \| \mathbf{y} - \mathbf{z}_0 \|_2^2 \qquad
+  g_1(\mathbf{z}_1) = \lambda \| \mathbf{z}_1 \|_{2,1}$$
+  $$A = \left( \begin{array}{c} C \\ D \end{array} \right) \qquad
+  B = \left( \begin{array}{cc} -I & 0 \\ 0 & -I \end{array} \right) \qquad
+  \mathbf{c} = \left( \begin{array}{c} 0 \\ 0 \end{array} \right) \;.$$
+
+This is a more complex splitting than that used in the
+[companion example](ct_astra_3d_tv_admm.rst), but it allows the use of a
+proximal ADMM solver in a way that avoids the need for the conjugate
+gradient sub-iterations used by the ADMM solver in the
+[companion example](ct_astra_3d_tv_admm.rst).
+"""
+𝛼 = 1e2  # improve problem conditioning by balancing C and D components of A
+λ = 2e0  # ℓ2,1 norm regularization parameter
+ρ = 5e-3  # ADMM penalty parameter
+maxiter = 1000  # number of ADMM iterations
+
+f = functional.ZeroFunctional()
+g0 = loss.SquaredL2Loss(y=y)
+g1 = (λ / 𝛼) * functional.L21Norm()
+g = functional.SeparableFunctional((g0, g1))
+D = linop.FiniteDifference(input_shape=tangle.shape, append=0)
+
+A = linop.VerticalStack((C, 𝛼 * D))
+mu, nu = ProximalADMM.estimate_parameters(A)
+
+solver = ProximalADMM(
+    f=f,
+    g=g,
+    A=A,
+    B=None,
+    rho=ρ,
+    mu=mu,
+    nu=nu,
+    maxiter=maxiter,
+    itstat_options={"display": True, "period": 50},
+)
+
+"""
+Run the solver.
+"""
+print(f"Solving on {device_info()}\n")
+tangle_recon = solver.solve()
+
+print(
+    "TV Restruction\nSNR: %.2f (dB), MAE: %.3f"
+    % (metric.snr(tangle, tangle_recon), metric.mae(tangle, tangle_recon))
+)
+
+
+"""
+Show the recovered volume.
+"""
+fig, ax = kplt.subplots(nrows=1, ncols=2, sharex=True, sharey=True, figsize=(7, 6))
+kplt.imview(
+    tangle[32],
+    title="Ground truth",
+    cmap=kplt.cm.viridis,
+    show_cbar=None,
+    ax=ax[0],
+)
+kplt.imview(
+    tangle_recon[32],
+    title="TV Reconstruction\nSNR: %.2f (dB), MAE: %.3f"
+    % (metric.snr(tangle, tangle_recon), metric.mae(tangle, tangle_recon)),
+    cmap=kplt.cm.viridis,
+    ax=ax[1],
+)
+divider = make_axes_locatable(ax[1])
+cax = divider.append_axes("right", size="5%", pad=0.2)
+fig.colorbar(ax[1].get_images()[0], cax=cax, label="arbitrary units")
+fig.suptitle("Central slice on $z$ axis (axis 0)")
+fig.tight_layout()
+fig.show()
+
+fig, ax = kplt.subplots(
+    nrows=1,
+    ncols=2,
+    sharex=True,
+    sharey=True,
+    gridspec_kw={"width_ratios": [1, 1.08]},
+    figsize=(13, 4),
+)
+kplt.imview(
+    tangle[:, 128],
+    title="Ground truth",
+    cmap=kplt.cm.viridis,
+    ax=ax[0],
+)
+kplt.imview(
+    tangle_recon[:, 128],
+    title="TV Reconstruction\nSNR: %.2f (dB), MAE: %.3f"
+    % (metric.snr(tangle, tangle_recon), metric.mae(tangle, tangle_recon)),
+    cmap=kplt.cm.viridis,
+    ax=ax[1],
+)
+divider = make_axes_locatable(ax[1])
+cax = divider.append_axes("right", size="5%", pad=0.2)
+fig.colorbar(ax[1].get_images()[0], ax=ax[1], cax=cax, label="arbitrary units")
+fig.suptitle("Central slice on $y$ axis (axis 1)")
+fig.tight_layout()
+fig.show()
+
+input("\nWaiting for input to close figures and exit")

examples/scripts/ct_astra_3d_tv_padmm.py

@@ -47,7 +47,7 @@
 n_projection = 10  # number of projections
 angles = np.linspace(0, np.pi, n_projection, endpoint=False)  # evenly spaced projection angles
 det_spacing = [1.0, 1.0]
-det_count = [Nz, max(Nx, Ny)]
+det_count = (Nz, max(Nx, Ny))
 vectors = angle_to_vector(det_spacing, angles)
 
 # It would have been more straightforward to use the det_spacing and angles keywords

examples/scripts/index.rst

@@ -15,6 +15,7 @@ Computed Tomography
    - ct_astra_noreg_pcg.py
    - ct_astra_3d_tv_admm.py
    - ct_astra_3d_tv_padmm.py
+   - ct_3d_tv_padmm.py
    - ct_tv_admm.py
    - ct_astra_tv_admm.py
    - ct_multi_tv_admm.py
@@ -113,6 +114,7 @@ Total Variation
    - ct_astra_tv_admm.py
    - ct_astra_3d_tv_admm.py
    - ct_astra_3d_tv_padmm.py
+   - ct_3d_tv_padmm.py
    - ct_astra_weighted_tv_admm.py
    - ct_svmbir_tv_multi.py
    - deconv_circ_tv_admm.py
@@ -213,6 +215,7 @@ Proximal ADMM
 ^^^^^^^^^^^^^
 
     - ct_astra_3d_tv_padmm.py
+    - ct_3d_tv_padmm.py
     - deconv_tv_padmm.py
     - denoise_tv_multi.py
     - deconv_ppp_dncnn_padmm.py

scico/linop/xray/_xray.py

@@ -360,9 +360,6 @@ class XRayTransform3D(LinearOperator):
     adjoint of the forward projector. It is written purely in JAX,
     allowing it to run on either CPU or GPU and minimizing host copies.
 
-    Warning: This class is experimental and may be up to ten times slower
-    than :class:`scico.linop.xray.astra.XRayTransform3D`.
-
     For each view, the projection geometry is specified by an array
     with shape (2, 4) that specifies a :math:`2 \times 3` projection
     matrix and a :math:`2 \times 1` offset vector. Denoting the matrix
@@ -406,7 +403,7 @@ def __init__(
 
         self.input_shape: Shape = input_shape
         self.matrices = jnp.asarray(matrices, dtype=np.float32)
-        self.det_shape = det_shape
+        self.det_shape = tuple(det_shape)  # in case det_shape is a list
         self.output_shape = (len(matrices), *det_shape)
         self.input_device = input_device
         self.output_device = output_device
@@ -432,6 +429,7 @@ def back_project(self, proj: ArrayLike) -> snp.Array:
         )
 
     @staticmethod
+    @partial(jax.jit, static_argnames="det_shape")
     def _project(
         im: ArrayLike,
         matrices: ArrayLike,
@@ -447,25 +445,30 @@ def _project(
             device: (optional) :class:`~jax.Device` or :class:`~jax.sharding.Sharding`
                 to which the output will be committed.
         """
-        MAX_SLICE_LEN = 10
-        slice_offsets = list(range(0, im.shape[0], MAX_SLICE_LEN))
-
-        num_views = len(matrices)
-        proj = jnp.zeros((num_views,) + det_shape, dtype=im.dtype, device=device)
-        for view_ind, matrix in enumerate(matrices):
-            for slice_offset in slice_offsets:
-                proj = proj.at[view_ind].set(
-                    XRayTransform3D._project_single(
-                        im[slice_offset : slice_offset + MAX_SLICE_LEN],
-                        matrix,
-                        proj[view_ind],
-                        slice_offset=slice_offset,
-                    )
-                )
-        return proj
+        BATCH_SIZE = 8
+
+        # Apply gradient checkpointing to the underlying core operator
+        project_single = jax.remat(XRayTransform3D._project_single)
+
+        # Define projection behavior for a single matrix over the full image
+        def project_single_matrix(matrix):
+            # Start with an empty detector plane baseline
+            init_plane = jnp.zeros(det_shape, dtype=im.dtype, device=device)
+
+            # Call the rematerialized operator on the full image
+            return project_single(
+                im,
+                matrix,
+                init_plane,
+                slice_offset=0,  # No manual loops: processed as a whole
+            )
+
+        # Automatically chunk and execute views sequentially/parallelized via JAX.
+        # If len(matrices) is not divisible by BATCH_SIZE, JAX natively handles the
+        # remainder.
+        return jax.lax.map(project_single_matrix, matrices, batch_size=BATCH_SIZE)
 
     @staticmethod
-    @partial(jax.jit, donate_argnames="proj")
     def _project_single(
         im: ArrayLike, matrix: ArrayLike, proj: ArrayLike, slice_offset: int = 0
     ) -> snp.Array:
@@ -486,6 +489,7 @@ def _project_single(
         return proj
 
     @staticmethod
+    @partial(jax.jit, static_argnames="input_shape")
     def _back_project(
         proj: ArrayLike,
         matrices: ArrayLike,
@@ -494,32 +498,44 @@ def _back_project(
     ) -> snp.Array:
         r"""
         Args:
-            proj: Input (set of) projection(s).
+            proj: Input projection data of shape (num_views, *det_shape).
             matrix: (num_views, 2, 4) array of homogeneous projection matrices.
-            input_shape: Shape of desired back projection.
+            input_shape: Shape of back projection.
             device: (optional) :class:`~jax.Device` or :class:`~jax.sharding.Sharding`
                 to which the output will be committed.
         """
-        MAX_SLICE_LEN = 10
-        slice_offsets = list(range(0, input_shape[0], MAX_SLICE_LEN))
-
-        HTy = jnp.zeros(input_shape, dtype=proj.dtype, device=device)
-        for view_ind, matrix in enumerate(matrices):
-            for slice_offset in slice_offsets:
-                HTy = HTy.at[slice_offset : slice_offset + MAX_SLICE_LEN].set(
-                    XRayTransform3D._back_project_single(
-                        proj[view_ind],
-                        matrix,
-                        HTy[slice_offset : slice_offset + MAX_SLICE_LEN],
-                        slice_offset=slice_offset,
-                    )
-                )
-                HTy.block_until_ready()  # prevent OOM
+        BATCH_SIZE = 8
 
-        return HTy
+        # Wrap the single back-project function for gradient checkpointing
+        back_project_single = jax.remat(XRayTransform3D._back_project_single)
+
+        # Process an individual view slice-by-slice natively via map mapping
+        def back_project_single_view(packed_inputs):
+            # Unpack the active iteration variables provided by lax.map
+            single_proj, single_matrix = packed_inputs
+
+            # Initialize a full-sized target volume structure for this single projection
+            # contribution
+            init_volume = jnp.zeros(input_shape, dtype=proj.dtype, device=device)
+
+            return back_project_single(
+                single_proj,
+                single_matrix,
+                init_volume,
+                slice_offset=0,  # Let JAX optimize the internal execution structure
+            )
+
+        # Map across the zip-like structure of projections and matrices.
+        # lax.map accumulates a stacked array of individual volume reconstructions.
+        individual_volumes = jax.lax.map(
+            back_project_single_view, (proj, matrices), batch_size=BATCH_SIZE
+        )
+
+        # Collapse the mapped axis by summing the independent view contributions
+        # to finalize the reconstructed 3D output image array.
+        return jnp.sum(individual_volumes, axis=0)
 
     @staticmethod
-    @partial(jax.jit, donate_argnames="HTy")
     def _back_project_single(
         y: ArrayLike, matrix: ArrayLike, HTy: ArrayLike, slice_offset: int = 0
     ) -> snp.Array: