Add Abel linear operator

CHANGES.rst

@@ -7,12 +7,15 @@ Version 0.0.2   (unreleased)
 ----------------------------
 
 • Additional optimization algorithms: Linearized ADMM and PDHG.
+• Additional Abel transform and array slicing linear operators.
+• Additional nuclear norm functional.
+• New module ``scico.ray.tune`` providing a simplified interface to Ray Tune.
 • Move optimization algorithms into ``optimize`` subpackage.
 • Additional iteration stats columns for iterative ADMM subproblem solvers.
 • Renamed "Primal Rsdl" to "Prml Rsdl" in displayed iteration stats.
 • Move some functions from ``util`` and ``math`` modules to new ``array``
   module.
-• Bump pinned `jaxlib` and `jax` versions to 0.1.70 and 0.2.19 respectively.
+• Bump pinned `jaxlib` and `jax` versions to 0.1.75 and 0.2.26 respectively.
 
 
 Version 0.0.1   (2021-11-24)

docs/source/examples.rst

@@ -22,6 +22,7 @@ Computed Tomography
 .. toctree::
    :maxdepth: 1
 
+   examples/ct_abel_tv_admm
    examples/ct_astra_pcg
    examples/ct_astra_tv_admm
    examples/ct_astra_weighted_tv_admm
@@ -97,6 +98,7 @@ Total Variation
 .. toctree::
    :maxdepth: 1
 
+   examples/ct_abel_tv_admm
    examples/ct_astra_tv_admm
    examples/ct_astra_weighted_tv_admm
    examples/ct_svmbir_tv_multi
@@ -135,6 +137,7 @@ ADMM
 .. toctree::
    :maxdepth: 1
 
+   examples/ct_abel_tv_admm
    examples/ct_astra_tv_admm
    examples/ct_astra_weighted_tv_admm
    examples/ct_svmbir_ppp_bm3d_admm_cg

docs/source/references.bib

@@ -173,7 +173,7 @@ @Article {esser-2010-general
                   Primal-Dual Algorithms for Convex Optimization in
                   Imaging Science},
   journal =	 {SIAM Journal on Imaging Sciences},
-  doi =	 {10.1137/09076934x},
+  doi =	         {10.1137/09076934x},
   year =	 2010,
   month =	 Jan,
   volume =	 3,
@@ -369,6 +369,15 @@ @Misc {svmbir-2020
   year =	 2020
 }
 
+@Misc {pyabel-2022,
+  author =	 {Stephen Gibson and Daniel Hickstein and Roman Yurchak,
+              Mikhail Ryazanov and Dhrubajyoti Das and Gilbert Shih},
+  title =	 {PyAbel},
+  howpublished = {PyAbel/PyAbel: v0.8.5},
+  year =	 2022,
+  doi =          {10.5281/zenodo.5888391}
+}
+
 @InProceedings {venkatakrishnan-2013-plugandplay2,
   author =	 {Singanallur V. Venkatakrishnan and Charles A. Bouman
                   and Brendt Wohlberg},
@@ -410,7 +419,7 @@ @article{yang-2012-linearized
   title =	 {Linearized augmented {L}agrangian and alternating
                   direction methods for nuclear norm minimization},
   journal =	 {Mathematics of Computation},
-  doi =	 {10.1090/s0025-5718-2012-02598-1},
+  doi =	         {10.1090/s0025-5718-2012-02598-1},
   year =	 2012,
   month =	 mar,
   volume =	 82,

docs/source/team.rst

@@ -31,6 +31,7 @@ Core Developers
 - `Thilo Balke <https://github.com/tbalke>`_
 - `Fernando Davis <https://github.com/FernandoDavis>`_
 - `Cristina Garcia Cardona <https://github.com/crstngc>`_
+- `Soumendu Majee <https://github.com/smajee>`_
 - `Michael McCann <https://github.com/Michael-T-McCann>`_
 - `Brendt Wohlberg <https://github.com/bwohlberg>`_
 
@@ -40,5 +41,4 @@ Contributors
 
 - `Oleg Korobkin <https://github.com/korobkin>`_ (BlockArray improvements)
 - `Yanpeng Yuan <https://github.com/yanpeng7>`_ (ASTRA interface improvements)
-- `Soumendu Majee <https://github.com/smajee>`_ (SVMBIR interface improvements)
 - `Saurav Maheshkar <https://github.com/SauravMaheshkar>`_ (Improvements to pre-commit configuration)

examples/scripts/ct_abel_tv_admm.py

@@ -0,0 +1,86 @@
+#!/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"""
+Regularized Abel Inversion
+==========================
+
+This example demonstrates a TV-regularized Abel inversion using
+an Abel projector based on PyAbel :cite:`pyabel-2022`
+"""
+
+import numpy as np
+
+import scico.numpy as snp
+from scico import functional, linop, loss, plot
+from scico.examples import create_circular_phantom
+from scico.linop.abel import AbelProjector
+from scico.optimize.admm import ADMM, LinearSubproblemSolver
+from scico.util import device_info
+
+"""
+Create a ground truth image.
+"""
+x_gt = create_circular_phantom((256, 254), [100, 50, 25], [1, 0, 0.5])
+
+"""
+Set up the forward operator and create a test measurement
+"""
+A = AbelProjector(x_gt.shape)
+y = A @ x_gt
+y = y + np.random.normal(size=y.shape).astype(np.float32)
+ATy = A.T @ y
+
+
+"""
+Set up ADMM solver object.
+"""
+λ = 1.71e01  # L1 norm regularization parameter
+ρ = 4.83e01  # ADMM penalty parameter
+maxiter = 100  # number of ADMM iterations
+cg_tol = 1e-4  # CG relative tolerance
+cg_maxiter = 25  # maximum CG iterations per ADMM iteration
+
+g = λ * functional.L1Norm()
+C = linop.FiniteDifference(input_shape=x_gt.shape)
+
+f = loss.SquaredL2Loss(y=y, A=A)
+
+x_inv = A.inverse(y)
+x0 = snp.clip(x_inv, 0, 1.0)
+
+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)
+x_tv = snp.clip(solver.x, 0, 1.0)
+
+norm = plot.matplotlib.colors.Normalize(vmin=-0.1, vmax=1.2)
+fig, ax = plot.subplots(nrows=2, ncols=2, figsize=(12, 12))
+plot.imview(x_gt, title="Ground Truth", cmap=plot.cm.Blues, fig=fig, ax=ax[0, 0], norm=norm)
+plot.imview(y, title="Measurements", cmap=plot.cm.Blues, fig=fig, ax=ax[0, 1])
+plot.imview(x_inv, title="Inverse Abel", cmap=plot.cm.Blues, fig=fig, ax=ax[1, 0], norm=norm)
+plot.imview(
+    x_tv, title="TV Regularized Inversion", cmap=plot.cm.Blues, fig=fig, ax=ax[1, 1], norm=norm
+)
+fig.show()
+
+
+input("\nWaiting for input to close figures and exit")

examples/scripts/ct_abel_tv_admm_tune.py

@@ -0,0 +1,148 @@
+#!/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"""
+Regularized Abel Inversion Tuning Demo
+==========================
+
+This example demonstrates the use of
+[scico.ray.tune](../_autosummary/scico.ray.tune.rst) to tune
+parameters for the companion [example script](ct_abel_tv_admm.rst).
+"""
+
+import numpy as np
+
+import jax
+
+import scico.ray as ray
+from scico import functional, linop, loss, metric, plot
+from scico.examples import create_circular_phantom
+from scico.linop.abel import AbelProjector
+from scico.optimize.admm import ADMM, LinearSubproblemSolver
+from scico.ray import tune
+
+"""
+Create a ground truth image.
+"""
+x_gt = create_circular_phantom((256, 256), [100, 50, 25], [1, 0, 0.5])
+
+
+"""
+Set up the forward operator and create a test measurement
+"""
+A = AbelProjector(x_gt.shape)
+y = A @ x_gt
+y = y + 1 * np.random.normal(size=y.shape)
+ATy = A.T @ y
+
+"""
+Put main arrays into ray object store.
+"""
+ray_x_gt, ray_y = ray.put(np.array(x_gt)), ray.put(np.array(y))
+
+
+"""
+Define performance evaluation function.
+"""
+
+
+def eval_params(config):
+    # Extract solver parameters from config dict.
+    λ, ρ = config["lambda"], config["rho"]
+    # Get main arrays from ray object store.
+    x_gt, y = ray.get([ray_x_gt, ray_y])
+    # Put main arrays on jax device.
+    x_gt, y = jax.device_put([x_gt, y])
+    # Set up problem to be solved.
+    A = AbelProjector(x_gt.shape)
+    f = loss.SquaredL2Loss(y=y, A=A)
+    g = λ * functional.L1Norm()
+    C = linop.FiniteDifference(input_shape=x_gt.shape)
+    # Define solver.
+    solver = ADMM(
+        f=f,
+        g_list=[g],
+        C_list=[C],
+        rho_list=[ρ],
+        x0=A.inverse(y),
+        maxiter=10,
+        subproblem_solver=LinearSubproblemSolver(),
+    )
+    # Perform 50 iterations, reporting performance to ray.tune every 10 iterations.
+    for step in range(10):
+        x_admm = solver.solve()
+        tune.report(psnr=float(metric.psnr(x_gt, x_admm)))
+
+
+"""
+Define parameter search space and resources per trial.
+"""
+config = {"lambda": tune.loguniform(1e-2, 1e3), "rho": tune.loguniform(1e-1, 1e3)}
+resources = {"gpu": 0, "cpu": 1}  # gpus per trial, cpus per trial
+
+
+"""
+Run parameter search.
+"""
+analysis = tune.run(
+    eval_params,
+    metric="psnr",
+    mode="max",
+    num_samples=100,
+    config=config,
+    resources_per_trial=resources,
+    hyperopt=True,
+    verbose=True,
+)
+
+"""
+Display best parameters and corresponding performance.
+"""
+best_config = analysis.get_best_config(metric="psnr", mode="max")
+print(f"Best PSNR: {analysis.get_best_trial().last_result['psnr']:.2f} dB")
+print("Best config: " + ", ".join([f"{k}: {v:.2e}" for k, v in best_config.items()]))
+
+
+"""
+Plot parameter values visited during parameter search. Marker sizes are
+proportional to number of iterations run at each parameter pair. The best
+point in the parameter space is indicated in red.
+"""
+fig = plot.figure(figsize=(8, 8))
+for t in analysis.trials:
+    n = t.metric_analysis["training_iteration"]["max"]
+    plot.plot(
+        t.config["lambda"],
+        t.config["rho"],
+        ptyp="loglog",
+        lw=0,
+        ms=(0.5 + 1.5 * n),
+        marker="o",
+        mfc="blue",
+        mec="blue",
+        fig=fig,
+    )
+plot.plot(
+    best_config["lambda"],
+    best_config["rho"],
+    ptyp="loglog",
+    title="Parameter search sampling locations\n(marker size proportional to number of iterations)",
+    xlbl=r"$\rho$",
+    ylbl=r"$\lambda$",
+    lw=0,
+    ms=5.0,
+    marker="o",
+    mfc="red",
+    mec="red",
+    fig=fig,
+)
+ax = fig.axes[0]
+ax.set_xlim([config["rho"].lower, config["rho"].upper])
+ax.set_ylim([config["lambda"].lower, config["lambda"].upper])
+fig.show()
+
+
+input("\nWaiting for input to close figures and exit")

examples/scripts/index.rst

@@ -9,6 +9,7 @@ Organized by Application
 Computed Tomography
 ^^^^^^^^^^^^^^^^^^^
 
+   - ct_abel_tv_admm.py
    - ct_astra_pcg.py
    - ct_astra_tv_admm.py
    - ct_astra_weighted_tv_admm.py
@@ -66,6 +67,7 @@ Plug and Play Priors
 Total Variation
 ^^^^^^^^^^^^^^^
 
+   - ct_abel_tv_admm.py
    - ct_astra_tv_admm.py
    - ct_astra_weighted_tv_admm.py
    - ct_svmbir_tv_multi.py

requirements.txt

@@ -7,3 +7,4 @@ jax==0.2.26
 flax
 bm3d
 svmbir>=0.2.7
+pyabel>=0.8.5