Add support for primal domain regularization and FISTA algorithm (#26)

Signed-off-by: Nicola VIGANÒ <nicola.vigano@cea.fr>

Author: Obi-Wan

README.md

@@ -19,6 +19,7 @@ This package provides the following functionality:
     * Various data fitting terms, including **Gaussian and Poisson noise** modelling.
     * Various optional regularization terms, including: **TV-min**, l1-min, laplacian, and **wavelet** l1-min.
     * Multi-channel (collaborative) regularization terms, like: **TNV** (Total Nuclear Variation).
+  - Fast Iterative Shrinkage-Thresholding Algorithm (FISTA).
   - Filtered Back-Projection (**FBP**), and its data-dependent filter learning variant
 (**[PyMR-FBP](https://github.com/dmpelt/pymrfbp)**).
 * Two projector backends, based on: [astra-toolbox](https://github.com/astra-toolbox/astra-toolbox) and

doc_sources/tutorial.md

@@ -124,12 +124,15 @@ filter can be selected, by passing the `fbp_filter` parameter at initialization:
 solver_fbp = cct.solvers.FBP(fbp_filter="shepp-logan")
 ```
 
-#### SIRT, PDHG, and MLEM
+#### SIRT, PDHG, MLEM, and FISTA
 
-The SIRT and PDHG algorithms, are algebraic (iterative) methods. They both
+The SIRT, PDHG and FISTA algorithms, are algebraic (iterative) methods. They all
 support regularization, and box constraints on the solution. The PDHG also
-supports various data fidelity terms.
-The MLEM algorithm is also an iterative algorithm to find the maximum likelihood estimation of the reconstructed signal. The MLEM does not currently support any regularization.
+supports various data fidelity terms. FISTA only supports regularizations that
+form a tight Parseval frame (like certain wavelet transforms).
+The MLEM algorithm is also an iterative algorithm to find the maximum likelihood
+estimation of the reconstructed signal. The MLEM does not currently support any
+regularization.
 
 The interface of the iterative methods is the same as for the FBP, with the only
 difference of requiring an iterations count:
@@ -146,7 +149,7 @@ with cct.projectors.ProjectorUncorrected(vol_shape_xy, angles_rad) as p:
     vol, _ = solver_sirt(p, sino, iterations=100)
 ```
 
-It is possible to specify an intial solution or box limits on the solutions like
+It is possible to specify an initial solution or box limits on the solutions like
 the following:
 ```python
 x0 = np.ones(vol_shape_xy)  # Initial solution

examples/example_09_synthetic_phantom_PDHG_vs_FISTA.py

@@ -0,0 +1,129 @@
+"""
+This example compares the fista solver against SIRT and weighted least-squares
+reconstructions (implemented with the PDHG algorithm) of the Shepp-Logan phantom.
+
+@author: Jérome Lesaint, ESRF - The European Synchrotron, Grenoble, France
+"""
+
+from collections.abc import Sequence
+
+import matplotlib.pyplot as plt
+import numpy as np
+import corrct as cct
+
+try:
+    import phantom
+except ImportError:
+    cct.testing.download_phantom()
+    import phantom
+
+
+def cm2inch(x: Sequence[float]) -> tuple[float, float]:
+    """Convert cm to inch.
+
+    Parameters
+    ----------
+    x : ArrayLike
+        Sizes in cm.
+
+    Returns
+    -------
+    tuple[float, float]
+        Sizes in inch.
+    """
+    return tuple(np.array(x) / 2.54)
+
+
+vol_shape = [256, 256, 3]
+data_type = np.float32
+
+ph_or = np.squeeze(phantom.modified_shepp_logan(vol_shape).astype(data_type))
+ph_or = ph_or[:, :, 1]
+
+ph, vol_att_in, vol_att_out = cct.testing.phantom_assign_concentration(ph_or)
+# Create sino with no background noise.
+sino, angles, expected_ph, background_avg = cct.testing.create_sino(
+    ph, 30, psf=None, add_poisson=True, dwell_time_s=1e-1, background_avg=1e-5, background_std=1e-6, dtype=np.float32
+)
+
+num_iterations = 100
+lower_limit = 0
+vol_mask = cct.processing.circular_mask(ph_or.shape)
+
+sino_variance = cct.processing.compute_variance_poisson(sino)
+sino_weights = cct.processing.compute_variance_weight(sino_variance, normalized=True)
+
+data_term_ls = cct.data_terms.DataFidelity_l2(background=background_avg)
+data_term_lsw = cct.data_terms.DataFidelity_wl2(sino_weights, background=background_avg)
+data_term_kl = cct.data_terms.DataFidelity_KL(background=background_avg)
+
+with cct.projectors.ProjectorUncorrected(ph.shape, angles) as A:
+    reg = cct.regularizers.Regularizer_l1dwl(1e-1, "Haar", 3)
+    solver_pdhg = cct.solvers.PDHG(verbose=True, data_term=data_term_kl, regularizer=reg)
+    rec_pdhg, _ = solver_pdhg(A, sino, num_iterations, x_mask=vol_mask)
+
+    solver_sirt = cct.solvers.SIRT(verbose=True, data_term=data_term_lsw)
+    rec_sirt, _ = solver_sirt(A, sino, num_iterations, x_mask=vol_mask)
+
+    reg = cct.regularizers.Regularizer_l1dwl(1e-2, "Haar", 3)
+    solver_fista = cct.solvers.FISTA(verbose=True, data_term=data_term_lsw, regularizer=reg)
+    rec_fista, _ = solver_fista(A, sino, num_iterations, x_mask=vol_mask)
+
+
+vminmax = dict(vmin=expected_ph.min(), vmax=expected_ph.max())
+
+# Reconstructions
+fig = plt.figure(figsize=cm2inch([36, 24]))
+gs = fig.add_gridspec(8, 3)
+ax_ph = fig.add_subplot(gs[:4, 0])
+im_ph = ax_ph.imshow(expected_ph, **vminmax)
+ax_ph.set_title("Phantom")
+fig.colorbar(im_ph, ax=ax_ph)
+
+ax_sino_clean = fig.add_subplot(gs[4, 0])
+with cct.projectors.ProjectorUncorrected(ph_or.shape, angles) as p:
+    sino_clean = p.fp(expected_ph)
+im_sino_clean = ax_sino_clean.imshow(sino_clean)
+ax_sino_clean.set_title("Clean sinogram")
+
+ax_sino_noise = fig.add_subplot(gs[5, 0])
+im_sino_noise = ax_sino_noise.imshow(sino - background_avg)
+ax_sino_noise.set_title("Noisy sinogram")
+
+ax_sino_lines = fig.add_subplot(gs[6:, 0])
+im_sino_lines = ax_sino_lines.plot(sino[9, :] - background_avg, label="Noisy")
+im_sino_lines = ax_sino_lines.plot(sino_clean[9, :], label="Clean")
+ax_sino_lines.set_title("Sinograms - angle: 10")
+ax_sino_lines.legend()
+ax_sino_lines.grid()
+
+ax_wls_l = fig.add_subplot(gs[:4, 1], sharex=ax_ph, sharey=ax_ph)
+im_wls_l = ax_wls_l.imshow(np.squeeze(rec_pdhg), **vminmax)
+ax_wls_l.set_title(solver_pdhg.info().upper())
+fig.colorbar(im_wls_l, ax=ax_wls_l)
+
+ax_sirt = fig.add_subplot(gs[4:, 1], sharex=ax_ph, sharey=ax_ph)
+im_sirt = ax_sirt.imshow(np.squeeze(rec_sirt), **vminmax)
+ax_sirt.set_title(solver_sirt.info().upper())
+fig.colorbar(im_sirt, ax=ax_sirt)
+
+ax_fista = fig.add_subplot(gs[:4, 2], sharex=ax_ph, sharey=ax_ph)
+im_fista = ax_fista.imshow(np.squeeze(rec_fista), **vminmax)
+ax_fista.set_title(solver_fista.info().upper())
+fig.colorbar(im_fista, ax=ax_fista)
+
+axs = fig.add_subplot(gs[4:, 2])
+axs.plot(np.squeeze(expected_ph[..., 172]), label="Phantom")
+axs.plot(np.squeeze(rec_pdhg[..., 172]), label=solver_pdhg.info().upper())
+axs.plot(np.squeeze(rec_sirt[..., 172]), label=solver_sirt.info().upper())
+axs.plot(np.squeeze(rec_fista[..., 172]), label=solver_fista.info().upper())
+axs.grid()
+axs.legend()
+fig.tight_layout()
+
+# Comparing FRCs for each reconstruction
+labels = [solver_pdhg.info().upper(), solver_sirt.info().upper(), solver_fista.info().upper()]
+vols = [rec_pdhg, rec_sirt, rec_fista]
+cct.processing.post.plot_frcs([(expected_ph, rec) for rec in vols], labels=labels, snrt=0.4142)
+
+plt.show(block=False)