Primal Dual Hybrid Gradient Algorithm¶

In this demo, we learn how to use the Primal Dual Hybrid Algorithm (PDHG) introduced by Chambolle & Pock for Tomography Reconstruction. We will solve the following minimisation problem under three different regularisation terms, i.e.,

• $\|\cdot\|_{1}$ or
• Tikhonov regularisation or
• with $L=\nabla$ and Total variation:

$$ u^{*} =\underset{u}{\operatorname{argmin}} \frac{1}{2} \| \mathcal{A} u - g\|^{2} + \underbrace{ \begin{cases} \alpha\,\|u\|_{1}, & \\[10pt] \alpha\,\|\nabla u\|_{2}^{2}, & \\[10pt] \alpha\,\mathrm{TV}(u) + \mathbb{I}_{\{u\geq 0\}}(u). \end{cases}}_{Regularisers} \tag{all reg} $$

where,

1. $g$ is the Acquisition data obtained from the detector.

2. $\mathcal{A}$ is the projection operator ( Radon transform ) that maps from an image-space to an acquisition space, i.e., $\mathcal{A} : \mathbb{X} \rightarrow \mathbb{Y}, $ where $\mathbb{X}$ is an ImageGeometry and $\mathbb{Y}$ is an AcquisitionGeometry.

3. $\alpha$: regularising parameter that measures the trade-off between the fidelity and the regulariser terms.

4. The total variation (isotropic) is defined as $$\mathrm{TV}(u) = \|\nabla u \|_{2,1} = \sum \sqrt{ (\partial_{y}u)^{2} + (\partial_{x}u)^{2} }$$

5. $\mathbb{I}_{{u\geq 0}}(u) : =

$, $\quad$ a non-negativity constraint for the minimiser $u$.

Learning objectives¶

• Load the data using the CIL reader: ZEISSDataReader.
• Preprocess the data using the CIL processors: Binner, TransmissionAbsorptionConverter.
• Run FBP and SIRT reconstructions.
• Setup PDHG for 3 different regularisers: $L^{1}$, Tikhonov and Total variation.

We first import all the necessary libraries for this notebook.

Data information¶

In this demo, we use the Walnut found in Jørgensen_et_all. In total, there are 6 individual micro Computed Tomography datasets in the native Zeiss TXRM/TXM format. The six datasets were acquired at the 3D Imaging Center at Technical University of Denmark in 2014 (HDTomo3D in 2016) as part of the ERC-funded project High-Definition Tomography (HDTomo) headed by Prof. Per Christian Hansen.

This example requires the dataset walnut.zip from https://zenodo.org/record/4822516 :

• https://zenodo.org/record/4822516/files/walnut.zip

If running locally please download the data and update the path variable below.

Load walnut data¶

Acquisition and Image geometry information¶

Show Acquisition geometry and full 3D sinogram.¶

Slice through projections¶

For demonstration purposes, we extract the central slice and select only 160 angles from the total 1601 angles.¶

1. We use the Slicer processor with step size of 10.
2. We use the Binner processor to crop and bin the acquisition data in order to reduce the field of view.
3. We use the TransmissionAbsorptionConverter to convert from transmission measurements to absorption based on the Beer-Lambert law.

Note: To avoid circular artifacts in the reconstruction space, we subtract the mean value of a background Region of interest (ROI), i.e., ROI that does not contain the walnut.

Define Projection Operator¶

We can define our projection operator using our astra plugin that wraps the Astra-Toolbox library.

FBP and SIRT reconstructions¶

Now, let's perform simple reconstructions using the Filtered Back Projection (FBP) and Simultaneous Iterative Reconstruction Technique (see the appendix for more details).

Recall, for FBP we type

fbp_recon = FBP(ig, ag, device = 'gpu')(absorption_data)

For SIRT, we type

x_init = ig.allocate()
    sirt = SIRT(initial = x_init, operator = A, data=absorption_data, 
            max_iteration = 50, update_objective_interval=10)
    sirt.run(verbose=1)
    sirt_recon = sirt.solution

Note: In SIRT, a non-negative constraint can be used with

constraint=IndicatorBox(lower=0)

Exercise 1: Run FBP and SIRT reconstructions¶

Use the code blocks described above and run FBP (fbp_recon) and SIRT (sirt_recon) reconstructions.

Note: To display the results, use

show2D([fbp_recon,sirt_recon], title = ['FBP reconstruction','SIRT reconstruction'], cmap = 'inferno')

Exercise 1: Solution¶

Uncomment the following line to see the solution:

Why PDHG?¶

In the previous notebook, we presented the Tikhonov regularisation for tomography reconstruction, i.e.,

where we can use either the GradientOperator ($L = \nabla) $ or the IdentityOperator ($L = \mathbb{I}$). Due to the $\|\cdot\|^{2}_{2}$ terms, one can observe that the above objective function is differentiable. As shown in the previous notebook, we can use the standard GradientDescent algorithm namely

f1 = LeastSquares(A, absorption_data)
    D = GradientOperator(ig2D)
    f2 = OperatorCompositionFunction(L2NormSquared(),D)
    f = f1 + alpha_tikhonov*f2

    gd = GD(initial=ig2D.allocate(), objective_function=f, step_size=None, 
            max_iteration=1000, update_objective_interval = 10)
    gd.run(100, verbose=1)

However, this is not always the case. Consider for example an $L^{1}$ norm for the fidelity, i.e., $\|\mathcal{A} u - g\|_{1}$ or an $L^{1}$ norm of the regulariser i.e., $\|u\|_{1}$ or a non-negativity constraint $\mathbb{I}_{\{u>0\}}(u)$. An alternative is to use Proximal Gradient Methods, discussed in the previous notebook, e.g., the FISTA algorithm, where we require one of the functions to be differentiable and the other to have a simple proximal method, i.e., "easy to solve". For more information, we refer to Parikh_Boyd.

Using the PDHG algorithm, we can solve minimisation problems where the objective is not differentiable, and the only required assumption is convexity with simple proximal problems.

$L^{1}$ regularisation¶

Let $L=$IdentityOperator in equation (1) and replace the

which results to a non-differentiable objective function. Hence, we have

$$ u^{*} =\underset{u}{\operatorname{argmin}} \frac{1}{2} \| \mathcal{A} u - g\|^{2} + \alpha\|u\|_{1} \tag{$L^{2}-L^{1}$} $$

How to setup and run PDHG?¶

In order to use the PDHG algorithm for the problem above, we need to express our minimisation problem into the following form:

$$ \min_{u\in\mathbb{X}} \mathcal{F}(K u) + \mathcal{G}(u) \tag{PDHG form} $$

where we assume that:

1. $\mathcal{F}$, $\mathcal{G}$ are convex functionals:

   • $\mathcal{F}: \mathbb{Y} \rightarrow \mathbb{R}$

   • $\mathcal{G}: \mathbb{X} \rightarrow \mathbb{R}$

2. $K$ is a continuous linear operator acting from a space $\mathbb{X}$ to another space $\mathbb{Y}$ :

   $$K : \mathbb{X} \rightarrow \mathbb{Y} \quad $$

   with operator norm defined as $$\| K \| = \max\{ \|K x\|_{\mathbb{Y}} : \|x\|_{\mathbb{X}}\leq 1 \}.$$

We can write the problem $L^{2}-L^{1}$ into PDHG_form, if we let

1. $K = \mathcal{A} \quad \Longleftrightarrow \quad $ K = A

2. $\mathcal{F}: Y \rightarrow \mathbb{R}, \text{ with } \mathcal{F}(z) := \frac{1}{2}\| z - g \|^{2}, \quad \Longleftrightarrow \quad$ F = 0.5 * L2NormSquared(absorption_data)

3. $\mathcal{G}: X \rightarrow \mathbb{R}, \text{ with } \mathcal{G}(z) := \alpha\|z\|_{1}, \quad \Longleftrightarrow \quad$ G = alpha * L1Norm()

Hence, we can verify that with the above setting we have that $L^{2}-L^{1}$$\Rightarrow$PDHG_form for $x=u$, $$\underset{u}{\operatorname{argmin}} \frac{1}{2}\|\mathcal{A} u - g\|^{2}_{2} + \alpha\|u\|_{1} = \underset{u}{\operatorname{argmin}} \mathcal{F}(\mathcal{A}u) + \mathcal{G}(u) = \underset{x}{\operatorname{argmin}} \mathcal{F}(Kx) + \mathcal{G}(x) $$

The algorithm is described in the Appendix and for every iteration, we solve two (proximal-type) subproblems, i.e., primal & dual problems where $\text{prox}_{\tau \mathcal{G}}(x)$ and $\text{prox}_{\sigma \mathcal{F^{*}}}(x)$ are the proximal operators of $\mathcal{G}$ and $\mathcal{F}^{*}$ (convex conjugate of $\mathcal{F}$), i.e.,

One application of the proximal operator is similar to a gradient step but is defined for convex and not necessarily differentiable functions.

To setup and run PDHG in CIL:

pdhg = PDHG(f = F, g = G, operator = K, 
            max_iterations = 500, update_objective_interval = 100)
    pdhg.run(verbose=1)

Note: To monitor convergence, we use pdhg.run(verbose=1) that prints the objective value of the primal problem, or pdhg.run(verbose=2) that prints the objective value of the primal and dual problems, as well as the primal dual gap. Nothing is printed with verbose=0.

Define operator $K$, functions $\mathcal{F}$ and $\mathcal{G}$¶

Setup and run PDHG¶

PDHG for Total Variation Regularisation¶

Now, we continue with the setup of the PDHG algorithm using the Total variation regulariser from the introduction.

Similarly, to the ($L^{2}-L^{1}$) problem, we need to express ($L^{2}-TV$) in the general PDHG form. This can be done using two different formulations:

1. Explicit formulation: All the subproblems in the PDHG algorithm have a closed form solution.
2. Implicit formulation: One of the subproblems in the PDHG algorithm is not solved explicitly but an inner solver is used.

($L^{2}-TV$) with Explicit PDHG¶

For the setup of the ($L^{2}-TV$) Explicit PDHG, we let

f1 = alpha * MixedL21Norm()
                                f2 = 0.5 * L2NormSquared(b=absorption_data)

For $z = (z_{1}, z_{2})\in \mathbb{Y}\times \mathbb{X}$, we define a separable function, e.g., BlockFunction, see the appendix.

F = BlockFunction(f1, f2)

In order to obtain an element $z = (z_{1}, z_{2})\in \mathbb{Y}\times \mathbb{X}$, we need to define a BlockOperator $K$, using the two operators involved in $L^{2}-TV$, i.e., the GradientOperator $\nabla$ and the ProjectionOperator $\mathcal{A}$.

Grad = GradientOperator(ig)
                                        K = BlockOperator(Grad, A)

Finally, we enforce a non-negativity constraint by letting $\mathcal{G} = \mathbb{I}_{\{u>0\}}(u)$ $\Longleftrightarrow$ G = IndicatorBox(lower=0)

Again, we can verify that with the above setting we can express our problem into PDHG form, for $x=u$

Speed of PDHG convergence¶

The PDHG algorithm converges when $\sigma\tau\|K\|^{2}<1$, where the variable $\sigma$, $\tau$ are called the primal and dual stepsizes. When we setup the PDHG algorithm, the default values of $\sigma$ and $\tau$ are used:

• $\sigma=\frac{1}{\|K\|}$
• $\tau =\frac{1}{\|K\|}$

and are not passed as arguments in the setup of PDHG. However, the speed of the algorithm depends heavily on the choice of these stepsizes. For the following, we encourage you to use different values, such as:

• $\sigma=1.0$
• $\tau = \frac{1.0}{\sigma\|K\|^{2}}$,

where $\|K\|$ is the operator norm of $K$.

normK = K.norm()
sigma = 1.
tau = 1./(sigma*normK**2)


PDHG(f = F, g = G, operator = K, sigma=sigma, tau=tau,
            max_iteration = 2000,
            update_objective_interval = 500)

The operator norm is computed using the Power Method to approximate the greatest eigenvalue of $K$.

Exercise 2: Setup and run PDHG algorithm for Tikhonov regularisation¶

Use exactly the same code as above and replace:

Exercise 2: Solution¶

Uncomment the following line to see the solution:

($L^{2}-TV$) with Implicit PDHG¶

In the implicit PDHG, one of the proximal subproblems, i.e., $\mathrm{prox}_{\tau\mathcal{F}^{*}}$ or $\mathrm{prox}_{\sigma\mathcal{G}}$ are not solved exactly and an iterative solver is used. For the setup of the Implicit PDHG, we let

For the function $\mathcal{G}$, we can use the TotalVariation Function class from CIL. Alternatively, we can use the FGP_TV Function class from our cil.plugins.ccpi_regularisation that wraps regularisation routines from the CCPi-Regularisation Toolkit. For these functions, the proximal method implements an iterative solver, namely the Fast Gradient Projection (FGP) algorithm that solves the dual problem of

for every PDHG iteration. Hence, we need to specify the number of iterations for the FGP algorithm. In addition, we can enforce a non-negativity constraint using lower=0.0. For the FGP_TV class, we can either use device=cpu or device=gpu to speed up this inner solver.

G = alpha * FGP_TV(max_iteration=100, nonnegativity = True, device = 'gpu')

    G = alpha * TotalVariation(max_iteration=100, lower=0.)

Exercise 3: Setup and run implicit PDHG algorithm with the Total variation regulariser¶

• Using the TotalVariation class, from CIL. This solves the TV denoising problem (using the FGP algorithm) in CPU.

• Using the FGP_TV class from the CCPi regularisation plugin.

  Note: In the FGP_TV implementation no pixel size information is included when in the forward and backward of the finite difference operator. Hence, we need to divide our regularisation parameter by the pixel size, e.g., $$\frac{\alpha}{\mathrm{ig2D.voxel\_size\_y}}$$

$(L^{2}-TV)$ Implicit PDHG: using FGP_TV¶

Exercise 3: Solution¶

Uncomment the following line to see the solution

In the above comparison between explicit and implicit TV reconstructions, we may observe some differences in the reconstructions and in the middle line profiles. This is due to a) the number of iterations and b) $\sigma, \tau$ values used in both the explicit and implicit setup of the PDHG algorithm. You can try more iterations with different values of $\sigma$ and $\tau$ for both cases in order to be sure that converge to the same solution.

$(L^{2}-TV)$ Implicit PDHG: using TotalVariation¶

FBP reconstruction with all the projection angles.¶

Show all reconstructions¶

• FBP (1601 projections)
• FBP (160 projections)
• SIRT (160 projections)
• $L^{1}$ regularisation (160 projections)
• Tikhonov regularisation (160 projections)
• Total variation regularisation (160 projections)

Zoom ROIs¶

Conclusions¶

In the PDHG algorithm, the step-sizes $\sigma, \tau$ play a significant role in terms of the convergence speed. In the above problems, we used the default values:

• $\sigma = \tau = \frac{1.0}{\|K\|}$

and we encourage you to try different values provided that $\sigma\tau\|K\|^{2}<1$ is satisfied. Certainly, these values are not the optimal ones and there are several acceleration methods in the literature to tune these parameters appropriately, see for instance Chambolle_Pock2010, Chambolle_Pock2011, Goldstein et al, Malitsky_Pock.

In the following notebook, we are going to present a stochastic version of PDHG, namely SPDHG introduced in Chambolle et al which is extremely useful to reconstruct large datasets, e.g., 3D walnut data. The idea behind SPDHG is to split our initial dataset into smaller chunks and apply forward and backward operations to these randomly selected subsets of the data. SPDHG has been used for different imaging applications and produces significant computational improvements over the PDHG algorithm, see Ehrhardt et al and Papoutsellis et al.