diff --git a/CHANGES.rst b/CHANGES.rst index 3f3af49b..e473b9b8 100644 --- a/CHANGES.rst +++ b/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. diff --git a/examples/scripts/ct_3d_tv_padmm.py b/examples/scripts/ct_3d_tv_padmm.py new file mode 100644 index 00000000..74682db7 --- /dev/null +++ b/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") diff --git a/examples/scripts/ct_astra_3d_tv_padmm.py b/examples/scripts/ct_astra_3d_tv_padmm.py index 6797d7ab..c4da3599 100644 --- a/examples/scripts/ct_astra_3d_tv_padmm.py +++ b/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 diff --git a/examples/scripts/index.rst b/examples/scripts/index.rst index 58416b95..01840739 100644 --- a/examples/scripts/index.rst +++ b/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 diff --git a/scico/linop/xray/_xray.py b/scico/linop/xray/_xray.py index 53457928..28f7f2ad 100644 --- a/scico/linop/xray/_xray.py +++ b/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: