Add support for 3D tomographic projection with astra

.github/workflows/test_examples.yml

@@ -69,4 +69,4 @@ jobs:
       # Run example test
       - name: Run example test
         run: |
-          ${GITHUB_WORKSPACE}/examples/scriptcheck.sh -e -d -t
+          ${GITHUB_WORKSPACE}/examples/scriptcheck.sh -e -d -t -g

docs/source/examples.rst

@@ -26,6 +26,7 @@ Computed Tomography
    examples/ct_abel_tv_admm_tune
    examples/ct_astra_noreg_pcg
    examples/ct_astra_tv_admm
+   examples/ct_astra_3d_tv_admm
    examples/ct_astra_weighted_tv_admm
    examples/ct_svmbir_tv_multi
    examples/ct_svmbir_ppp_bm3d_admm_cg
@@ -126,6 +127,7 @@ Total Variation
    examples/ct_abel_tv_admm
    examples/ct_abel_tv_admm_tune
    examples/ct_astra_tv_admm
+   examples/ct_astra_3d_tv_admm
    examples/ct_astra_weighted_tv_admm
    examples/ct_svmbir_tv_multi
    examples/deconv_circ_tv_admm
@@ -193,6 +195,7 @@ ADMM
    examples/ct_abel_tv_admm
    examples/ct_abel_tv_admm_tune
    examples/ct_astra_tv_admm
+   examples/ct_astra_3d_tv_admm
    examples/ct_astra_weighted_tv_admm
    examples/ct_svmbir_tv_multi
    examples/ct_svmbir_ppp_bm3d_admm_cg

examples/scriptcheck.sh

@@ -12,19 +12,22 @@ Usage: $SCRIPT [-h] [-d]
           [-e] Display excerpt of error message on failure
           [-d] Skip tests involving additional data downloads
           [-t] Skip tests related to learned model training
+          [-g] Skip tests that need the GPU
 EOF
 )
 
 OPTIND=1
 DISPLAY_ERROR=0
 SKIP_DOWNLOAD=0
 SKIP_TRAINING=0
-while getopts ":hedt" opt; do
+SKIP_GPU=0
+while getopts ":hedtg" opt; do
     case $opt in
 	h) echo "$USAGE"; exit 0;;
 	e) DISPLAY_ERROR=1;;
 	d) SKIP_DOWNLOAD=1;;
 	t) SKIP_TRAINING=1;;
+    g) SKIP_GPU=1;;
 	\?) echo "Error: invalid option -$OPTARG" >&2
             echo "$USAGE" >&2
             exit 1
@@ -74,15 +77,19 @@ for f in $SCRIPTPATH/scripts/*.py; do
 
     # Skip problem cases.
     if [ $SKIP_DOWNLOAD -eq 1 ] && grep -q '_microscopy' <<< $f; then
-	printf "%s\n" skipped
-	continue
+	    printf "%s\n" skipped
+	    continue
     fi
     if [ $SKIP_TRAINING -eq 1 ]; then
 	if grep -q '_datagen' <<< $f || grep -q '_train' <<< $f; then
 	    printf "%s\n" skipped
 	    continue
         fi
     fi
+    if [ $SKIP_GPU -eq 1 ] && grep -q '_astra_3d' <<< $f; then
+	    printf "%s\n" skipped
+	    continue
+    fi
 
     # Create temporary copy of script with all algorithm maxiter values set
     # to small number and final input statements commented out.

examples/scripts/ct_astra_3d_tv_admm.py

@@ -0,0 +1,131 @@
+#!/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
+============================================
+
+This example demonstrates solution of a sparse-view CT reconstruction
+problem with isotropic total variation (TV) regularization
+
+  $$\mathrm{argmin}_{\mathbf{x}} \; (1/2) \| \mathbf{y} - A \mathbf{x}
+  \|_2^2 + \lambda \| C \mathbf{x} \|_{2,1} \;,$$
+
+where $A$ is the Radon transform, $\mathbf{y}$ is the sinogram, $C$ is
+a 3D finite difference operator, and $\mathbf{x}$ is the desired
+image.
+"""
+
+import numpy as np
+
+import jax
+
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from scico import functional, linop, loss, metric, plot
+from scico.linop.radon_astra import TomographicProjector
+from scico.optimize.admm import ADMM, LinearSubproblemSolver
+from scico.util import device_info
+
+
+def make_tangle(nx, ny, nz):
+    xs = 1.0 * np.linspace(-1.0, 1.0, nx)
+    ys = 1.0 * np.linspace(-1.0, 1.0, ny)
+    zs = 1.0 * np.linspace(-1.0, 1.0, nz)
+
+    # default ordering for meshgrid is `xy`, this makes inputs of length
+    # M, N, P will create a mesh of N, M, P. Thus we want ys, zs and xs.
+    xx, yy, zz = np.meshgrid(ys, zs, xs, copy=True)
+    xx = 3.0 * xx
+    yy = 3.0 * yy
+    zz = 3.0 * zz
+    values = (
+        xx * xx * xx * xx
+        - 5.0 * xx * xx
+        + yy * yy * yy * yy
+        - 5.0 * yy * yy
+        + zz * zz * zz * zz
+        - 5.0 * zz * zz
+        + 11.8
+    ) * 0.2 + 0.5
+    return values
+
+
+Nx = 128
+Ny = 256
+Nz = 64
+
+tangle = (make_tangle(Nx, Ny, Nz) < 2.0).astype(float)
+tangle = jax.device_put(tangle)
+
+n_projection = 10  # number of projections
+angles = np.linspace(0, np.pi, n_projection)  # evenly spaced projection angles
+A = TomographicProjector(
+    tangle.shape, [1.0, 1.0], [Nz, max(Nx, Ny)], angles
+)  # Radon transform operator
+y = A @ tangle  # sinogram
+
+
+"""
+Set up ADMM solver object.
+"""
+λ = 2e0  # L1 norm regularization parameter
+ρ = 5e0  # ADMM penalty parameter
+maxiter = 25  # number of ADMM iterations
+cg_tol = 1e-4  # CG relative tolerance
+cg_maxiter = 25  # maximum CG iterations per ADMM iteration
+
+# The append=0 option makes the results of horizontal and vertical
+# finite differences the same shape, which is required for the L21Norm,
+# which is used so that g(Cx) corresponds to isotropic TV.
+C = linop.FiniteDifference(input_shape=tangle.shape, append=0)
+g = λ * functional.L21Norm()
+
+f = loss.SquaredL2Loss(y=y, A=A)
+
+x0 = A.T(y)
+
+solver = ADMM(
+    f=f,
+    g_list=[g],
+    C_list=[C],
+    rho_list=[ρ],
+    x0=x0,
+    maxiter=maxiter,
+    subproblem_solver=LinearSubproblemSolver(cg_kwargs={"tol": cg_tol, "maxiter": cg_maxiter}),
+    itstat_options={"display": True, "period": 5},
+)
+
+"""
+Run the solver.
+"""
+print(f"Solving on {device_info()}\n")
+solver.solve()
+hist = solver.itstat_object.history(transpose=True)
+tangle_recon = solver.x
+
+print(
+    "TV Restruction\nSNR: %.2f (dB), MAE: %.3f"
+    % (metric.snr(tangle, tangle_recon), metric.mae(tangle, tangle_recon))
+)
+
+"""
+Show the recovered image.
+"""
+fig, ax = plot.subplots(nrows=1, ncols=2, figsize=(7, 5))
+plot.imview(tangle[32], title="Ground truth (central slice)", cbar=None, fig=fig, ax=ax[0])
+
+plot.imview(
+    tangle_recon[32],
+    title="TV Reconstruction (central slice)\nSNR: %.2f (dB), MAE: %.3f"
+    % (metric.snr(tangle, tangle_recon), metric.mae(tangle, tangle_recon)),
+    fig=fig,
+    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.show()

examples/scripts/index.rst

@@ -12,6 +12,7 @@ Computed Tomography
    - ct_abel_tv_admm.py
    - ct_abel_tv_admm_tune.py
    - ct_astra_noreg_pcg.py
+   - ct_astra_3d_tv_admm.py
    - ct_astra_tv_admm.py
    - ct_astra_weighted_tv_admm.py
    - ct_svmbir_tv_multi.py
@@ -95,6 +96,7 @@ Total Variation
    - ct_abel_tv_admm.py
    - ct_abel_tv_admm_tune.py
    - ct_astra_tv_admm.py
+   - ct_astra_3d_tv_admm.py
    - ct_astra_weighted_tv_admm.py
    - ct_svmbir_tv_multi.py
    - deconv_circ_tv_admm.py
@@ -150,6 +152,7 @@ ADMM
    - ct_abel_tv_admm.py
    - ct_abel_tv_admm_tune.py
    - ct_astra_tv_admm.py
+   - ct_astra_3d_tv_admm.py
    - ct_astra_weighted_tv_admm.py
    - ct_svmbir_tv_multi.py
    - ct_svmbir_ppp_bm3d_admm_cg.py