Stochastic Primal Dual Hybrid Gradient Algorithm: Tomography¶

In this demo, we learn how to use the Stochastic Primal-Dual Hybrid (SPDHG) algorithm introduced in Chambolle et al., see as well Ehrhardt et al., Papoutsellis et al., Gutierrez et al.. This algorithm is a stochastic version of the Primal-Dual Hybrid Gradient (PDHG) Algorithm by Chambolle & Pock. We focus on Tomography Reconstruction under an edge-preserving prior, i.e., the Total variation regularisation, see ROF.

Why SPDHG?¶

In the previous demo, we presented PDHG for tomography reconstruction.

SPDHG is a stochastic version of PDHG. As PDHG, SPDHG can solve convex, non necessarily smooth optimization problems and is provenly convergent.

The central improvement with respect to PDHG is that SPDHG can deal with subsets of data, making it computationally more efficient, so that it can deal with larger datasets.

For CT and PET for example, the data can be divided into a collection of partial sinograms, each of them corresponding to a subset of the total set of angles. Accordingly, the projection operator $A: X \rightarrow Y$ can be decomposed as $A=(A_1, \dots, A_n): X \rightarrow Y=Y_1\times \dots \times Y_n$ where $n$ is the number of subsets and each $A_i:X \rightarrow Y_i$ is a partial projection.

At each iteration, SPDHG computes only one partial forward and back-projection $A_i,\,A_i^*$ for a randomly chosen index $i$, whereas PDHG computes $A$ and $A^*$. Hence SPDHG requires only a fraction of the computations of PDHG per iteration, which makes it faster per iteration.

Learning objectives¶

1. Implement the subsetting
2. Implement default set-up
3. Compare PDHG and SPDHG: notion of epoch, visual comparison, measures (primal objective function)
4. Tweak the default set-up: explicit / implicit, stepsizes

Prerequisites¶

Pre-requisites:

1. ImageData, AcquisitionData, Block Framework
2. Projection operator
3. Total Variation (TV)
4. PDHG implicit set-up for tomographic reconstruction. Optional: explicit set-up

Tomography reconstruction¶

The minimisation problem for Total Variation Tomography reconstruction with TV is

$$ \underset{u}{\operatorname{argmin}} \mathcal{O}(u)= \frac{1}{2}\| A u - g\|^{2} + \alpha\,\mathrm{TV}(u) + \mathbb{I}_{\{u>0\}}(u) \tag{1}$$

where,

1. TV is defined as $$\mathrm{TV}(u) = \|\nabla u \|_{2,1} = \sum \sqrt{ (\partial_{y}u)^{2} + (\partial_{x}u)^{2} }$$
2. g is the Acquisition data obtained from the detector.
3. $A$ is the projection operator ( Radon transform ) that maps from an image-space to an acquisition space, i.e., $A : X \rightarrow Y, $ where X is an ImageGeometry and Y is a BlockDataContainer of AcquisitionGeometry objects.
4. $\alpha$ is a regularising parameter that measures a trade-off between the fidelity and the regulariser terms.
5. The functional $\mathbb{I}_{{u\geq 0}}(u) : =

\begin{cases} 0, & \text{ if } u\geq 0\\ \infty , & \text{ otherwise} \quad \end{cases}

$, $\quad$ is a positivity constraint for the minimiser $u$.

SPDHG form¶

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

$$\min_{u} \sum_{i=1}^n F_i(K_i u) + G(u) \tag{2}$$

where we assume that:

1. $F_i$, $G$ are convex functionals

   • $F_i: Y_i \rightarrow \mathbb{R}$

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

2. $K_i$ is a continuous linear operator acting from a space X to another space Y_i :

$$K_i : X \rightarrow Y_i \quad $$

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

Note: We use the Power Method to approximate the norm of the $K_i$'s and if needed of $\nabla$.

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 filename variable below.

How to setup and run SPDHG¶

In order to setup and run SPDHG, we need to define the following:

• The operator $K=(K_1,\dots,K_n)$.
• The functions $F=(F_1,\dots,F_N)$ and $G$.
• The maximum number of iterations

The setup and run of SPDHG:

spdhg = SPDHG(f = F, g = G, operator = K)

spdhg.run(number_of_iterations)

We can express our problem in the form, if we let

1. $K = (A_1,\dots,A_n) \quad \Longleftrightarrow \quad $ K = BlockOperator(A_i)

2. $F_i: Y_i \rightarrow \mathbb{R}, \text{ with } F_i(y_i) := \frac{1}{2}\| y_i - g_i \|^{2}, \quad \Longleftrightarrow \quad$ for all $i$'s, F_i = 0.5*L2NormSquared(data[i]) and F = BlockFunction(F_i)

3. $G: X \rightarrow \mathbb{R}, \text{ with } G(x) := \alpha\|\nabla x\|_{1} + \mathbb{I}_{\{x>0\}}(x), \quad \Longleftrightarrow \quad$ G = alpha * FGP_TV().

Define operator $K$, functions $F$ and $G$¶

Setup and run SPDHG¶

Note: in this example, there are some parameters which we don't pass, so that SPDHG will use default, pre-defined parameters. These parameters consist of:

1. The probabilities $p_i$ of choosing each subset $i$. By default these probabilities are uniform: if there are $n$ subsets, $p_i=1/n$. We'll see another choice in the part about explicit set-up.
2. The step-sizes of the algorithm. To know more about step-sizes setting, refer to the last part of the tutorial.

Show reconstruction result¶

Comparison to PDHG¶

PDHG and SPDHG converge to the same set of solutions, but SPDHG is faster per iteration. How can we compare PDHG and SPDHG? We want to compare PDHG and SPDHG for the same amount of computational effort. Remember that at each iteration, PDHG computes one entire forward and back-projection, while SPDHG computes only a partial one. For example, if we have 10 subsets, $10$ iterations of SPDHG have the same computational cost than one iteration of PDHG.

In order to have a fair comparison between both algorithms, we then define the notion of epoch, where one epoch corresponds to a full forward and back-projection. Hence, one epoch of PDHG corresponds to one iteration, while one epoch of SPDHG corresponds in our case to $10$ iteration.

Let's set-up and run PDHG and SPDHG for the same number of epochs and compare the results.

After 5 epochs, SPDHG reconstruction seems visually more converged than PDHG reconstruction.

To further compare PDHG and SPDHG speeds, let's use a quantitative measure of convergence, which is given by the value of the objective function $\mathcal{O}(u)$.

Explicit set-up¶

As for PDHG, there are various ways to put our reconstruction problem: $$ \underset{u}{\operatorname{argmin}} \mathcal{O}(u)= \frac{1}{2}\| A u - g\|^{2} + \alpha\,\mathrm{TV}(u) + \mathbb{I}_{\{u>0\}}(u) $$ into SPDHG-appropriate form: $$\min_{u} \sum_{i=1}^n F_i(K_i u) + G(u).$$

Implicit set-up

In the first part of the tutorial, we chose the so-called implicit set-up, where

1. $K = (A_1,\dots,A_n) \quad \Longleftrightarrow \quad $ K = BlockOperator(A_i)

2. $F_i: Y_i \rightarrow \mathbb{R}, \text{ with } F_i(y_i) := \frac{1}{2}\| y_i - g_i \|^{2}, \quad \Longleftrightarrow \quad$ for all $i$'s, F_i = 0.5*L2NormSquared(data[i]) and F = BlockFunction(F_i)

3. $G: X \rightarrow \mathbb{R}, \text{ with } G(x) := \alpha\|\nabla x\|_{1} + \mathbb{I}_{\{x>0\}}(x), \quad \Longleftrightarrow \quad$ G = alpha * FGP_TV().

Explicit set-up

As for PDHG, another way is to include the gradient of $\mathrm{TV}(u)=\alpha \|\nabla u\|_{2,1}$ into the operator $K$:

1. $K = (A_1,\dots,A_n, \nabla) \quad \Longleftrightarrow \quad $ K = BlockOperator(A_1,...,A_n, Grad)

2. For all $1\leq i \leq n, F_i: Y_i \rightarrow \mathbb{R}, \text{ with } F_i(y_i) := \frac{1}{2}\| y_i - g_i \|^{2}$ and $F_{n+1}(z)=\|z\|_{1,2}$, $\quad \Longleftrightarrow \quad$

for all $1\leq i \leq n$, F_i = 0.5*L2NormSquared(data[i]), F_{n+1}=MixedL21Norm() and F = BlockFunction(F_i)

1. $G: X \rightarrow \mathbb{R}, \text{ with } G(x) := \mathbb{I}_{\{x>0\}}(x), \quad \Longleftrightarrow \quad$ G = IndicatorBox(lower=0).

Non-uniform sampling¶

This time, we also use non-uniform sampling: half of the time, we update the dual variable corresponding to TV, so that $p_{n+1}=1/2$, and half of the time we update one of the dual variable corresponding to one subset, so that $p_i=1/(2n)$ for $1\leq i \leq n$.

Comment 1: if you don't want to use the default uniform sampling, the probabilities you specify need sum to $1$.

Comment 2: the setting and sampling have an influence on the value of one epoch in terms of iterations. With the sampling defined above, we use the partial forward operator only half of the time in mean, so that $1$ epoch (that is, the unit of one entire forward and back-projection) corresponds to twice the number of subsets.

Improve the speed of explicit reconstruction by re-normalizing¶

You can observe that the explicit reconstruction converges slower than the implicit one. However, the block operator $K$ used in the explicit framework is now formed by partial operators $A_{\text{subsets}}[i]$ and $\nabla$ with potentially very different norms. Empirically it hence makes sense to renormalize the partial operators so that they all have norm one. To leave the objective function unchanged, we'd also need to redefine the function $f_i$'s according to the following equations: $$ \| K u - g\|_2^2 = \|K\|^2 \cdot \left\| \frac{K}{\|K\|} u - \frac{g}{\|K\|}\right\|_2^2$$ $$ \| \nabla u\|_{2,1} = \|\nabla \| \cdot \left\| \frac{\nabla}{\|\nabla\|} u\right\|_{2,1}$$

When to use explicit or implicit framework?¶

1. Set-up: implicit framework is generally easier to set-up

2. Time-per-iteration: the time-per-iteration is greater in the implicit framework; indeed, at each iteration one needs to compute the proximal operator of the TV, which is done by in a (hidden) inner loop. This does not occur in the explicit framework, where all the operators have explicit expressions (hence its name.)

3. Speed of convergence: there is no theoretical result as to which framework is the fastest. In the example above, one can observe that renormalized explicit SPDHG seems to converge at roughly the same speed than implicit SPDHG.

Theoretical background¶

Algorithm

SPDHG is described as follows:

1. Input: Step-sizes $\tau, (\sigma_i)_{1\leq i \leq n}$, $\tau$; probabilities $p_i$

2. Initialize: $x=0$, $y=0$, $z = \bar{z} = A^T y = 0$

3. Iterate for $m=1, \dots$:

   1. Update primal variable: $$x = \text{prox}_{\tau g}(x - \tau \bar{z}) $$
   2. Pick index $i$ at random with probability $p_i$ $$ $$
   3. Update dual variable with index $i$: $$ y_i^+ = \text{prox}_{\sigma_i f_i^*}(y_i + \sigma_i A_i x) $$
   4. Update secondary variables: $$ d z = A_i^T(y_i^+ - y_i), \quad z = z + d z, \quad y_i = y_i^+ $$
   5. Perform extrapolation step:
      $$ \bar{z}= z + (1/p_i) \cdot d z $$

Admissible step-sizes

SPDHG converges to the set of solutions of the saddle-point problem (in terms of Bregman distance [ref] and almost surely [ref1, ref2]) if the following conditions on step-sizes $\tau, (\sigma_i)_{1\leq i \leq n}$ is verified: $$ \tau \, \sigma_i \, \|K_i\|^2 < p_i^2, \quad 1 \leq i \leq n$$

CIL default implementation is $\sigma_i = \rho/\|K_i\|$ for all $i$'s and $\tau =\rho\, \min_i p_i^2/\|K_i\|$ with $\rho=0.99$.

Other admissible choices would be $\sigma_i = \gamma \rho/\|K_i\|$ for all $i$'s and $\tau =\rho\, \min_i p_i^2/\|K_i\| / \gamma$ with $\rho=0.99$. Changing the value of $\gamma$ allows to change the balance between primal and dual convergence. It can be changed in the following way:

spdhg = SPDHG(f = F, g = G, operator = K, gamma = gamma)