{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "d304d2b3", "metadata": {}, "outputs": [], "source": [ "# -*- coding: utf-8 -*-\n", "# Copyright 2019 - 2022 United Kingdom Research and Innovation\n", "# Copyright 2019 - 2022 The University of Manchester\n", "#\n", "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# http://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License.\n", "#\n", "# Authored by: Evangelos Papoutsellis (UKRI-STFC)\n", "# Gemma Fardell (UKRI-STFC)\n", "# Laura Murgatroyd (UKRI-STFC) " ] }, { "attachments": {}, "cell_type": "markdown", "id": "c06fd185-0bbb-48c1-ba6c-ac61cad8266b", "metadata": {}, "source": [ "# Primal Dual Hybrid Gradient Algorithm\n", "\n", "In this demo, we learn how to use the **Primal Dual Hybrid Algorithm (PDHG)** introduced by [Chambolle & Pock](https://hal.archives-ouvertes.fr/hal-00490826/document) for Tomography Reconstruction. We will solve the following minimisation problem under three different regularisation terms, i.e., \n", "\n", "* $\\|\\cdot\\|_{1}$ or\n", "* Tikhonov regularisation or\n", "* with $L=\\nabla$ and Total variation:\n", "\n", "\n", "\n", "$$\n", "u^{*} =\\underset{u}{\\operatorname{argmin}} \\frac{1}{2} \\| \\mathcal{A} u - g\\|^{2} +\n", "\\underbrace{\n", "\\begin{cases}\n", "\\alpha\\,\\|u\\|_{1}, & \\\\[10pt]\n", "\\alpha\\,\\|\\nabla u\\|_{2}^{2}, & \\\\[10pt]\n", "\\alpha\\,\\mathrm{TV}(u) + \\mathbb{I}_{\\{u\\geq 0\\}}(u).\n", "\\end{cases}}_{Regularisers}\n", "\\tag{all reg}\n", "$$\n", "\n", "where,\n", "\n", "1. $g$ is the Acquisition data obtained from the detector.\n", "\n", "1. $\\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__.\n", "\n", "1. $\\alpha$: regularising parameter that measures the trade-off between the fidelity and the regulariser terms.\n", "\n", "1. The total variation (isotropic) is defined as $$\\mathrm{TV}(u) = \\|\\nabla u \\|_{2,1} = \\sum \\sqrt{ (\\partial_{y}u)^{2} + (\\partial_{x}u)^{2} }$$\n", "\n", "1. $\\mathbb{I}_{\\{u\\geq 0\\}}(u) : = \n", "\\begin{cases}\n", "0, & \\text{ if } u\\geq 0\\\\\n", "\\infty , & \\text{ otherwise}\n", "\\,\n", "\\end{cases}\n", "$, $\\quad$ a non-negativity constraint for the minimiser $u$." ] }, { "attachments": {}, "cell_type": "markdown", "id": "385c848a-9d9a-48ff-b029-a42e90e4cb2c", "metadata": {}, "source": [ "# Learning objectives\n", "\n", "- Load the data using the CIL reader: `ZEISSDataReader`.\n", "- Preprocess the data using the CIL processors: `Binner`, `TransmissionAbsorptionConverter`.\n", "- Run FBP and SIRT reconstructions.\n", "- Setup PDHG for 3 different regularisers: $L^{1}$, Tikhonov and Total variation.\n", "\n", "\n", "\n", "\n", "\n", "We first import all the necessary libraries for this notebook.\n", "\n", "" ] }, { "cell_type": "code", "execution_count": null, "id": "f731d970-d6ce-489a-bf93-7888cc855a60", "metadata": {}, "outputs": [], "source": [ "# Import libraries\n", "\n", "from cil.framework import BlockDataContainer\n", "\n", "from cil.optimisation.functions import L2NormSquared, L1Norm, BlockFunction, MixedL21Norm, IndicatorBox, TotalVariation\n", "from cil.optimisation.operators import GradientOperator, BlockOperator\n", "from cil.optimisation.algorithms import PDHG, SIRT\n", "\n", "from cil.plugins.astra.operators import ProjectionOperator\n", "from cil.plugins.astra.processors import FBP\n", "\n", "from cil.plugins.ccpi_regularisation.functions import FGP_TV\n", "\n", "from cil.utilities.display import show2D, show1D, show_geometry\n", "from cil.utilities.jupyter import islicer\n", "\n", "from cil.io import ZEISSDataReader\n", "\n", "from cil.processors import Binner, TransmissionAbsorptionConverter, Slicer\n", "\n", "import matplotlib.pyplot as plt\n", "\n", "import numpy as np\n", "\n", "import os" ] }, { "cell_type": "markdown", "id": "36dc5a6d-c4e9-4ece-9f37-cde82dc93964", "metadata": {}, "source": [ "# Data information\n", "\n", "In this demo, we use the **Walnut** found in [Jørgensen_et_all](https://zenodo.org/record/4822516#.YLXyAJMzZp8). 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. \n", "\n", "This example requires the dataset walnut.zip from https://zenodo.org/record/4822516 :\n", "\n", " - https://zenodo.org/record/4822516/files/walnut.zip\n", "\n", "If running locally please download the data and update the `path` variable below." ] }, { "cell_type": "markdown", "id": "732c5f6b-6fd4-43ea-b5f6-796632f62528", "metadata": {}, "source": [ "## Load walnut data" ] }, { "cell_type": "code", "execution_count": null, "id": "ce2c36aa-1231-4669-9486-197f9332fe2e", "metadata": {}, "outputs": [], "source": [ "path = '/mnt/materials/SIRF/Fully3D/CIL/Walnut'" ] }, { "cell_type": "code", "execution_count": null, "id": "d64a1b3a", "metadata": {}, "outputs": [], "source": [ "reader = ZEISSDataReader()\n", "filename = os.path.join(path, \"valnut_2014-03-21_643_28\",\"tomo-A\",\"valnut_tomo-A.txrm\")\n", "data3D = ZEISSDataReader(file_name=filename).read()" ] }, { "cell_type": "code", "execution_count": null, "id": "0fd50459", "metadata": {}, "outputs": [], "source": [ "# reorder data to match default order for Astra/Tigre operator\n", "data3D.reorder('astra')\n", "\n", "# Get Image and Acquisition geometries\n", "ag3D = data3D.geometry\n", "ig3D = ag3D.get_ImageGeometry()" ] }, { "cell_type": "markdown", "id": "2f97e39e-12db-4eae-9de5-cc31e667490e", "metadata": {}, "source": [ "### Acquisition and Image geometry information" ] }, { "cell_type": "code", "execution_count": null, "id": "ece32bb6-5066-4f7c-b6b6-8edc62ecb817", "metadata": {}, "outputs": [], "source": [ "print(ag3D)" ] }, { "cell_type": "code", "execution_count": null, "id": "14083f74-3490-4d18-a784-c659f2c5984b", "metadata": {}, "outputs": [], "source": [ "print(ig3D)" ] }, { "cell_type": "markdown", "id": "bab54a03-3ab6-4100-a82a-45a8e807e5f1", "metadata": {}, "source": [ "### Show Acquisition geometry and full 3D sinogram." ] }, { "cell_type": "code", "execution_count": null, "id": "d6b6a162-08d5-4c28-9a2b-a06091747f0a", "metadata": {}, "outputs": [], "source": [ "show_geometry(ag3D)" ] }, { "cell_type": "code", "execution_count": null, "id": "fc29dd0a-2f27-4a89-8a49-9733d9453b86", "metadata": {}, "outputs": [], "source": [ "show2D(data3D, slice_list = [('vertical',512), ('angle',800), ('horizontal',512)], cmap=\"inferno\", num_cols=3, size=(15,15))" ] }, { "cell_type": "markdown", "id": "dc945eb5-17d9-476e-b4dd-998b9b90d51e", "metadata": {}, "source": [ "### Slice through projections" ] }, { "cell_type": "code", "execution_count": null, "id": "e6e19c5d-0724-46a2-907c-d9e86b24ee91", "metadata": {}, "outputs": [], "source": [ "islicer(data3D, direction=1, cmap=\"inferno\")" ] }, { "cell_type": "markdown", "id": "5afa53ea-6723-44ef-a803-587630568ec4", "metadata": {}, "source": [ "### For demonstration purposes, we extract the central slice and select only 160 angles from the total 1601 angles.\n", "\n", "1. We use the `Slicer` processor with step size of 10.\n", "1. We use the `Binner` processor to crop and bin the acquisition data in order to reduce the field of view.\n", "1. We use the `TransmissionAbsorptionConverter` to convert from transmission measurements to absorption based on the Beer-Lambert law.\n", "\n", "**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." ] }, { "cell_type": "code", "execution_count": null, "id": "3dae73ff-165f-4ac3-990a-e5584a627e9f", "metadata": {}, "outputs": [], "source": [ "# Extract vertical slice\n", "data2D = data3D.get_slice(vertical='centre')\n", "\n", "# Select every 10 angles\n", "sliced_data = Slicer(roi={'angle':(0,1600,10)})(data2D)\n", "\n", "# Reduce background regions\n", "binned_data = Binner(roi={'horizontal':(120,-120,2)})(sliced_data)\n", "\n", "# Create absorption data \n", "absorption_data = TransmissionAbsorptionConverter()(binned_data) \n", "\n", "# Remove circular artifacts\n", "absorption_data -= np.mean(absorption_data.as_array()[80:100,0:30])" ] }, { "cell_type": "code", "execution_count": null, "id": "4165bcfc-1e28-44ea-bfe3-f85648bc4dc8", "metadata": {}, "outputs": [], "source": [ "# Get Image and Acquisition geometries for one slice\n", "ag2D = absorption_data.geometry\n", "ag2D.set_angles(ag2D.angles, initial_angle=0.2, angle_unit='radian')\n", "ig2D = ag2D.get_ImageGeometry()" ] }, { "cell_type": "code", "execution_count": null, "id": "5fdafdab-19f2-499a-8b6b-acfdaa83bf95", "metadata": {}, "outputs": [], "source": [ "print(\" Acquisition Geometry 2D: {} with labels {}\".format(ag2D.shape, ag2D.dimension_labels))\n", "print(\" Image Geometry 2D: {} with labels {}\".format(ig2D.shape, ig2D.dimension_labels))" ] }, { "cell_type": "markdown", "id": "3156de74-773a-4ec9-aadb-1f7d31a435b5", "metadata": {}, "source": [ "### Define Projection Operator \n", "We can define our projection operator using our __astra__ __plugin__ that wraps the Astra-Toolbox library." ] }, { "cell_type": "code", "execution_count": null, "id": "e4919986-1aad-4e2d-9af3-68d7b1d685ea", "metadata": {}, "outputs": [], "source": [ "A = ProjectionOperator(ig2D, ag2D, device = \"gpu\")" ] }, { "cell_type": "markdown", "id": "b3e9e994-a8f8-4877-b366-789989760f8a", "metadata": {}, "source": [ "## FBP and SIRT reconstructions\n", "\n", "Now, let's perform simple reconstructions using the **Filtered Back Projection (FBP)** and **Simultaneous Iterative Reconstruction Technique** (see the [appendix](../appendix.ipynb) for more details).\n", "\n", "Recall, for FBP we type\n", "\n", "\n", "```python\n", "\n", " fbp_recon = FBP(ig, ag, device = 'gpu')(absorption_data)\n", " \n", "```\n", "\n" ] }, { "cell_type": "markdown", "id": "b912854b-d7bc-49d4-8aa0-fa6f563c5d72", "metadata": {}, "source": [ "For SIRT, we type\n", "\n", " \n", "```python\n", " \n", " x_init = ig.allocate()\n", " sirt = SIRT(initial = x_init, operator = A, data=absorption_data, \n", " max_iteration = 50, update_objective_interval=10)\n", " sirt.run(verbose=1)\n", " sirt_recon = sirt.solution \n", " \n", "```\n", "\n", "**Note**: In SIRT, a non-negative constraint can be used with\n", "\n", "\n", "```python \n", " \n", " constraint=IndicatorBox(lower=0) \n", " \n", "```\n", " " ] }, { "cell_type": "markdown", "id": "54097efa-1fec-44cb-bc12-18fbf6d8d797", "metadata": {}, "source": [ "### **Exercise 1: Run FBP and SIRT reconstructions**\n", "\n", "Use the code blocks described above and run FBP (`fbp_recon`) and SIRT (`sirt_recon`) reconstructions.\n", "\n", "**Note**: To display the results, use \n", "\n", " \n", "```python \n", " \n", " show2D([fbp_recon,sirt_recon], title = ['FBP reconstruction','SIRT reconstruction'], cmap = 'inferno') \n", " \n", "```\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "bcb5e5c1-21d2-4210-982d-3ee87da6892c", "metadata": {}, "outputs": [], "source": [ "# Setup and run the FBP algorithm\n", "fbp_recon = FBP(..., ..., device = 'gpu')(absorption_data)\n", "\n", "# Setup and run the SIRT algorithm, with non-negative constraint\n", "x_init = ig2D.allocate() \n", "sirt = SIRT(initial = x_init, \n", " operator = ..., \n", " data= ..., \n", " constraint = ...,\n", " max_iteration = 300, \n", " update_objective_interval=100)\n", "sirt.run(verbose=1)\n", "sirt_recon = sirt.solution\n", "\n", "# Show reconstructions\n", "show2D([fbp_recon,sirt_recon], \n", " title = ['FBP reconstruction','SIRT reconstruction'], \n", " cmap = 'inferno', fix_range=(0,0.05))" ] }, { "cell_type": "markdown", "id": "62c95c5d", "metadata": {}, "source": [ "### **Exercise 1: Solution**\n", "\n", "Uncomment the following line to see the solution:" ] }, { "cell_type": "code", "execution_count": null, "id": "2e348820", "metadata": {}, "outputs": [], "source": [ "# %load ./snippets/03_ex1.py" ] }, { "attachments": {}, "cell_type": "markdown", "id": "024626bc-bc07-4f65-9b7e-7641c7c19528", "metadata": {}, "source": [ "## Why PDHG?\n", "\n", "In the previous notebook, we presented the __Tikhonov regularisation__ for tomography reconstruction, i.e.,\n", "\n", "$$\n", "u^{*} =\\underset{u}{\\operatorname{argmin}} \\frac{1}{2} \\| \\mathcal{A} u - g\\|^{2} + \\alpha\\|L u\\|^{2}_{2}\n", "\\tag{1}\n", "$$\n", "\n", "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\n", " \n", "\n", "```python\n", " \n", " f1 = LeastSquares(A, absorption_data)\n", " D = GradientOperator(ig2D)\n", " f2 = OperatorCompositionFunction(L2NormSquared(),D)\n", " f = f1 + alpha_tikhonov*f2\n", "\n", " gd = GD(initial=ig2D.allocate(), objective_function=f, step_size=None, \n", " max_iteration=1000, update_objective_interval = 10)\n", " gd.run(100, verbose=1)\n", " \n", "```\n", "\n", "\n", "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](https://web.stanford.edu/~boyd/papers/pdf/prox_algs.pdf#page=30).\n", "\n", "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." ] }, { "cell_type": "markdown", "id": "29699f88-da86-4be7-9145-9d9ff51198fc", "metadata": {}, "source": [ "### $L^{1}$ regularisation\n", "\n", "Let $L=$`IdentityOperator` in equation (1) and replace the\n", "\n", "$$\\alpha^{2}\\|L u\\|^{2}_{2}\\quad\\text{ with }\\quad \\alpha\\|u\\|_{1}, $$ \n", "\n", "which results to a non-differentiable objective function. Hence, we have \n", "\n", "\n", "$$\n", "u^{*} =\\underset{u}{\\operatorname{argmin}} \\frac{1}{2} \\| \\mathcal{A} u - g\\|^{2} + \\alpha\\|u\\|_{1} \n", "\\tag{$L^{2}-L^{1}$}\n", "$$" ] }, { "cell_type": "markdown", "id": "a5a99548-d786-4dd6-87b8-6301d27c0989", "metadata": {}, "source": [ "## How to setup and run PDHG?\n", "\n", "In order to use the PDHG algorithm for the problem above, we need to express our minimisation problem into the following form:\n", "\n", "\n", "$$\n", "\\min_{u\\in\\mathbb{X}} \\mathcal{F}(K u) + \\mathcal{G}(u)\n", "\\tag{PDHG form}\n", "$$\n", "\n", "where we assume that:\n", "\n", "1. $\\mathcal{F}$, $\\mathcal{G}$ are __convex__ functionals:\n", " \n", " - $\\mathcal{F}: \\mathbb{Y} \\rightarrow \\mathbb{R}$ \n", " \n", " - $\\mathcal{G}: \\mathbb{X} \\rightarrow \\mathbb{R}$\n", " \n", " \n", "1. $K$ is a continuous linear operator acting from a space $\\mathbb{X}$ to another space $\\mathbb{Y}$ :\n", "\n", " $$K : \\mathbb{X} \\rightarrow \\mathbb{Y} \\quad $$ \n", "\n", " with operator norm defined as $$\\| K \\| = \\max\\{ \\|K x\\|_{\\mathbb{Y}} : \\|x\\|_{\\mathbb{X}}\\leq 1 \\}.$$ \n", "\n", "\n", "We can write the problem [$L^{2}-L^{1}$](#Lasso) into [PDHG_form](#PDHG_form), if we let\n", "\n", "1. $K = \\mathcal{A} \\quad \\Longleftrightarrow \\quad $ `K = A` \n", "\n", "1. $\\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)`\n", "\n", "1. $\\mathcal{G}: X \\rightarrow \\mathbb{R}, \\text{ with } \\mathcal{G}(z) := \\alpha\\|z\\|_{1}, \\quad \\Longleftrightarrow \\quad$ ` G = alpha * L1Norm()`\n", "\n", "Hence, we can verify that with the above setting we have that [$L^{2}-L^{1}$](#Lasso)$\\Rightarrow$[PDHG_form](#PDHG_form) for $x=u$, $$\\underset{u}{\\operatorname{argmin}} \\frac{1}{2}\\|\\mathcal{A} u - g\\|^{2}_{2} + \\alpha\\|u\\|_{1} = \n", "\\underset{u}{\\operatorname{argmin}} \\mathcal{F}(\\mathcal{A}u) + \\mathcal{G}(u) = \\underset{x}{\\operatorname{argmin}} \\mathcal{F}(Kx) + \\mathcal{G}(x) $$" ] }, { "cell_type": "markdown", "id": "2e3a6b63-c855-4094-9e7d-2a20c06d762f", "metadata": {}, "source": [ "The algorithm is described in the [Appendix](../appendix.ipynb) and for every iteration, we solve two (proximal-type) subproblems, i.e., __primal & dual problems__ where \n", "$\\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.,\n", "\n", "$$\n", "\\text{prox}_{\\lambda \\mathcal{F}}(x) = \\underset{z}{\\operatorname{argmin}} \\frac{1}{2}\\|z - x \\|^{2} + \\lambda \n", "\\mathcal{F}(z)\n", "$$\n", "\n", "One application of the proximal operator is similar to a gradient step but is defined for convex and not necessarily differentiable functions.\n", "\n", "\n", "To setup and run PDHG in CIL:\n", "\n", "\n", "```python\n", " \n", " pdhg = PDHG(f = F, g = G, operator = K, \n", " max_iterations = 500, update_objective_interval = 100)\n", " pdhg.run(verbose=1)\n", " \n", "```\n", "\n", "**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`." ] }, { "cell_type": "markdown", "id": "d18293aa-19d0-435f-8aca-1f9575131d03", "metadata": {}, "source": [ "\n", "### Define operator $K$, functions $\\mathcal{F}$ and $\\mathcal{G}$" ] }, { "cell_type": "code", "execution_count": null, "id": "b2879663-a0fe-4a3f-900f-c2cd8ba31226", "metadata": {}, "outputs": [], "source": [ "K = A\n", "F = 0.5 * L2NormSquared(b=absorption_data)\n", "alpha = 0.01\n", "G = alpha * L1Norm()" ] }, { "cell_type": "markdown", "id": "96234d76-14d3-4f7c-961c-8c717fb531d7", "metadata": {}, "source": [ "### Setup and run PDHG" ] }, { "cell_type": "code", "execution_count": null, "id": "85ff3e81-2ed1-4aee-a97f-051fe4ad5580", "metadata": {}, "outputs": [], "source": [ "# Setup and run PDHG\n", "pdhg_l1 = PDHG(f = F, g = G, operator = K, \n", " max_iteration = 500,\n", " update_objective_interval = 100)\n", "pdhg_l1.run(verbose=1)" ] }, { "cell_type": "code", "execution_count": null, "id": "b0d969cd-225a-47cf-af38-0a2be75965c6", "metadata": {}, "outputs": [], "source": [ "# Show reconstruction and ground truth\n", "show2D([pdhg_l1.solution,fbp_recon], fix_range=(0,0.05), title = ['L1 regularisation', 'FBP'], cmap = 'inferno')\n", "\n", "# Plot middle line profile\n", "show1D([fbp_recon,pdhg_l1.solution], slice_list=[('horizontal_y',int(ig2D.voxel_num_y/2))],\n", " label = ['FBP','L1 regularisation'], title='Middle Line Profiles')\n" ] }, { "cell_type": "markdown", "id": "c695cd3f-a612-4bbe-99d2-ee2e9dd033b8", "metadata": {}, "source": [ "# PDHG for Total Variation Regularisation\n", "\n", "Now, we continue with the setup of the PDHG algorithm using the Total variation regulariser from the [introduction](#Primal-Dual-Hybrid-Gradient-Algorithm).\n", "\n", "Similarly, to the [($L^{2}-L^{1}$)](#Lasso) problem, we need to express [($L^{2}-TV$)](#all_reg) in the general [PDHG form](#PDHG_form). This can be done using two different formulations:\n", "\n", "1. Explicit formulation: All the subproblems in the PDHG algorithm have a closed form solution.\n", "1. Implicit formulation: One of the subproblems in the PDHG algorithm is not solved explicitly but an inner solver is used.\n", "\n", "---\n", "## ($L^{2}-TV$) with Explicit PDHG\n", "\n", "For the setup of the **($L^{2}-TV$) Explicit PDHG**, we let\n", "\n", "$$\\begin{align}\n", "& f_{1}: \\mathbb{Y} \\rightarrow \\mathbb{R}, \\quad f_{1}(z_{1}) = \\alpha\\,\\|z_{1}\\|_{2,1}, \\text{ ( the TV term ) }\\\\\n", "& f_{2}: \\mathbb{X} \\rightarrow \\mathbb{R}, \\quad f_{2}(z_{2}) = \\frac{1}{2}\\|z_{2} - g\\|_{2}^{2}, \\text{ ( the data-fitting term ). }\n", "\\tag{1}\n", "\\end{align}$$\n", "\n", "```python\n", "\n", " f1 = alpha * MixedL21Norm()\n", " f2 = 0.5 * L2NormSquared(b=absorption_data)\n", "\n", "```\n", "\n", "\n", "\n", "For $z = (z_{1}, z_{2})\\in \\mathbb{Y}\\times \\mathbb{X}$, we define a separable function, e.g., BlockFunction, see the [appendix](../appendix.ipynb).\n", "\n", "$$\\mathcal{F}(z) : = \\mathcal{F}(z_{1},z_{2}) = f_{1}(z_{1}) + f_{2}(z_{2})$$\n", "\n", "\n", "\n", "```python\n", " \n", " F = BlockFunction(f1, f2)\n", " \n", "```\n", "\n", "\n", "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$](#TomoTV), i.e., the `GradientOperator` $\\nabla$ and the `ProjectionOperator` $\\mathcal{A}$.\n", "\n", "$$ \\mathcal{K} = \n", "\\begin{bmatrix}\n", "\\nabla\\\\\n", "\\mathcal{A}\n", "\\end{bmatrix}\n", "$$\n", "\n", "\n", "\n", "```python\n", " \n", " Grad = GradientOperator(ig)\n", " K = BlockOperator(Grad, A)\n", " \n", "```\n", "\n", "\n", "\n", "Finally, we enforce a non-negativity constraint by letting $\\mathcal{G} = \\mathbb{I}_{\\{u>0\\}}(u)$ $\\Longleftrightarrow$ `G = IndicatorBox(lower=0)`\n", " \n", "Again, we can verify that with the above setting we can express our problem into [PDHG form](#PDHG_form), for $x=u$\n", "\n", "$$\n", "\\begin{align}\n", "\\underset{u}{\\operatorname{argmin}}\\alpha\\|\\nabla u\\|_{2,1} + \\frac{1}{2}\\|\\mathcal{A} u - g\\|^{2}_{2} + \\mathbb{I}_{\\{u>0\\}}(u) = \\underset{u}{\\operatorname{argmin}} f_{1}(\\nabla u) + f_{2}(\\mathcal{A}u) + \\mathbb{I}_{\\{u>0\\}}(u) \\\\ = \\underset{u}{\\operatorname{argmin}} F(\n", "\\begin{bmatrix}\n", "\\nabla \\\\\n", "\\mathcal{A}\n", "\\end{bmatrix}u) + \\mathbb{I}_{\\{u>0\\}}(u) = \n", "\\underset{u}{\\operatorname{argmin}} \\mathcal{F}(Ku) + \\mathcal{G}(u) = \\underset{x}{\\operatorname{argmin}} \\mathcal{F}(Kx) + \\mathcal{G}(x) \n", "\\end{align}\n", "\\tag{2}\n", "$$\n" ] }, { "cell_type": "code", "execution_count": null, "id": "33dc3542-e984-4eeb-9e12-1abadedfc9a4", "metadata": {}, "outputs": [], "source": [ "# Define BlockFunction F\n", "alpha_tv = 0.0003\n", "f1 = alpha_tv * MixedL21Norm()\n", "f2 = 0.5 * L2NormSquared(b=absorption_data)\n", "F = BlockFunction(f1, f2)\n", "\n", "# Define BlockOperator K\n", "Grad = GradientOperator(ig2D)\n", "K = BlockOperator(Grad, A)\n", "\n", "# Define Function G\n", "G = IndicatorBox(lower=0)\n", "\n", "\n", "# Setup and run PDHG\n", "pdhg_tv_explicit = PDHG(f = F, g = G, operator = K,\n", " max_iteration = 1000,\n", " update_objective_interval = 200)\n", "pdhg_tv_explicit.run(verbose=1)" ] }, { "cell_type": "code", "execution_count": null, "id": "af31caee-4684-49ad-bcf0-0ff05221a5da", "metadata": {}, "outputs": [], "source": [ "# Show reconstruction and ground truth\n", "show2D([pdhg_tv_explicit.solution,fbp_recon], fix_range=(0,0.055), title = ['TV regularisation','FBP'], cmap = 'inferno')\n", "\n", "# Plot middle line profile\n", "show1D([fbp_recon,pdhg_tv_explicit.solution], slice_list=[('horizontal_y',int(ig2D.voxel_num_y/2))],\n", " label = ['FBP','TV regularisation'], title='Middle Line Profiles')" ] }, { "cell_type": "markdown", "id": "26da6bf7", "metadata": {}, "source": [ "## Speed of PDHG convergence" ] }, { "cell_type": "markdown", "id": "1154a706-45f5-406c-b9b5-080e0083cb8c", "metadata": {}, "source": [ "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:\n", "\n", "- $\\sigma=\\frac{1}{\\|K\\|}$\n", "- $\\tau =\\frac{1}{\\|K\\|}$\n", "\n", "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:\n", "\n", "\n", "- $\\sigma=1.0$\n", "- $\\tau = \\frac{1.0}{\\sigma\\|K\\|^{2}}$,\n", "\n", "\n", "where $\\|K\\|$ is the operator norm of $K$. \n", "\n", "```python\n", "\n", "normK = K.norm()\n", "sigma = 1.\n", "tau = 1./(sigma*normK**2)\n", "\n", "\n", "PDHG(f = F, g = G, operator = K, sigma=sigma, tau=tau,\n", " max_iteration = 2000,\n", " update_objective_interval = 500)\n", "\n", "```\n", "\n", "The operator norm is computed using the [Power Method](https://en.wikipedia.org/wiki/Power_iteration) to approximate the greatest eigenvalue of $K$.\n", "\n", "\n" ] }, { "cell_type": "markdown", "id": "fb5d3a68-4c1c-45c6-9304-072f5d1e55ad", "metadata": {}, "source": [ "## **Exercise 2: Setup and run PDHG algorithm for Tikhonov regularisation**\n", "\n", "Use exactly the same code as above and replace:\n", "\n", "$$f_{1}(z_{1}) = \\alpha\\,\\|z_{1}\\|_{2,1} \\text{ with } f_{1}(z_{1}) = \\alpha\\,\\|z_{1}\\|_{2}^{2}.$$\n" ] }, { "cell_type": "code", "execution_count": null, "id": "90129b41-48eb-4487-b2a0-cdc2df1c70ee", "metadata": {}, "outputs": [], "source": [ "# Define BlockFunction F\n", "alpha_tikhonov = 0.05\n", "f1 = ... \n", "F = BlockFunction(f1, f2)\n", "\n", "# Setup and run PDHG\n", "pdhg_tikhonov_explicit = PDHG(f = F, g = G, operator = K,\n", " max_iteration = 500,\n", " update_objective_interval = 100)\n", "pdhg_tikhonov_explicit.run(verbose=1)" ] }, { "cell_type": "markdown", "id": "9bf82bca", "metadata": {}, "source": [ "## **Exercise 2: Solution**\n", "\n", "Uncomment the following line to see the solution:" ] }, { "cell_type": "code", "execution_count": null, "id": "7340eadf", "metadata": {}, "outputs": [], "source": [ "# %load ./snippets/03_ex2.py" ] }, { "cell_type": "code", "execution_count": null, "id": "960b7f91-b6ae-419f-b5ed-40f369a496db", "metadata": {}, "outputs": [], "source": [ "# Show reconstruction and ground truth\n", "show2D([pdhg_tikhonov_explicit.solution,fbp_recon], fix_range=(0,0.055), title = ['Tikhonov regularisation','FBP'], cmap = 'inferno')\n", "\n", "# Plot middle line profile\n", "show1D([fbp_recon,pdhg_tikhonov_explicit.solution], slice_list=[('horizontal_y',int(ig2D.voxel_num_y/2))],\n", " label = ['FBP','Tikhonov regularisation'], title='Middle Line Profiles')" ] }, { "cell_type": "markdown", "id": "78c4bcbc-f3a5-4cb6-b388-8c3cb294070a", "metadata": {}, "source": [ "---\n", "## ($L^{2}-TV$) with Implicit PDHG\n", "\n", "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\n", "\n", "$$\\begin{align}\n", "& \\mathcal{F}: \\mathbb{Y} \\rightarrow \\mathbb{R}, \\quad \\mathcal{F}(z_{1}) = \\frac{1}{2}\\|z_{1} - g\\|_{2}^{2}\\\\\n", "& \\mathcal{G}: \\mathbb{X} \\rightarrow \\mathbb{R}, \\quad \\mathcal{G}(z_{2}) = \\alpha\\, \\mathrm{TV}(z_{2}) = \\|\\nabla z_{2}\\|_{2,1}\n", "\\end{align}$$\n", "\n", "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](https://github.com/vais-ral/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\n", "\n", "$$\\begin{equation}\n", "\\mathrm{prox}_{\\tau G}(u) = \\underset{z}{\\operatorname{argmin}} \\frac{1}{2} \\| u - z\\|^{2} + \\tau\\,\\alpha\\,\\mathrm{TV}(z) + \\mathbb{I}_{\\{z>0\\}}(z),\n", "\\end{equation}\n", "$$\n", "\n", "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.\n", "\n", "\n", "```python\n", "\n", " G = alpha * FGP_TV(max_iteration=100, nonnegativity = True, device = 'gpu')\n", " \n", " G = alpha * TotalVariation(max_iteration=100, lower=0.)\n", " \n", "```\n" ] }, { "cell_type": "markdown", "id": "1d32911d-b3d5-495f-a787-9aac95492efb", "metadata": {}, "source": [ "## **Exercise 3: Setup and run implicit PDHG algorithm with the Total variation regulariser**\n", "\n", "- Using the TotalVariation class, from CIL. This solves the TV denoising problem (using the FGP algorithm) in CPU.\n", "\n", "\n", "- Using the FGP_TV class from the CCPi regularisation plugin.\n", "\n", " **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}}$$\n", " \n", "\n" ] }, { "cell_type": "markdown", "id": "b0c70da8-a069-4df2-b946-6a5900ad8794", "metadata": {}, "source": [ "## $(L^{2}-TV)$ Implicit PDHG: using FGP_TV" ] }, { "cell_type": "code", "execution_count": null, "id": "5ccf2d6d-bbbe-4f55-aabc-13c453bf16da", "metadata": {}, "outputs": [], "source": [ "F = 0.5 * L2NormSquared(b=absorption_data)\n", "G = (alpha_tv/ig2D.voxel_size_y) * ... \n", "K = A\n", "\n", "# Setup and run PDHG\n", "pdhg_tv_implicit_regtk = PDHG(f = F, g = G, operator = K,\n", " max_iteration = 1000,\n", " update_objective_interval = 200)\n", "pdhg_tv_implicit_regtk.run(verbose=1)" ] }, { "cell_type": "markdown", "id": "31f6d280", "metadata": {}, "source": [ "## **Exercise 3: Solution**\n", "\n", "Uncomment the following line to see the solution" ] }, { "cell_type": "code", "execution_count": null, "id": "98231fdf", "metadata": {}, "outputs": [], "source": [ "# %load ./snippets/03_ex3.py" ] }, { "cell_type": "code", "execution_count": null, "id": "796bf380-81c6-424c-9ab1-9af37b863745", "metadata": {}, "outputs": [], "source": [ "# Show reconstruction and ground truth\n", "show2D([pdhg_tv_implicit_regtk.solution,pdhg_tv_explicit.solution, \n", " (pdhg_tv_explicit.solution-pdhg_tv_implicit_regtk.solution).abs()], \n", " fix_range=[(0,0.055),(0,0.055),(0,1e-3)],\n", " title = ['TV (Implicit CCPi-RegTk)','TV (Explicit)', 'Absolute Difference'], \n", " cmap = 'inferno', num_cols=3)\n", "\n", "\n", "# Plot middle line profile\n", "show1D([pdhg_tv_explicit.solution,pdhg_tv_implicit_regtk.solution], slice_list=[('horizontal_y',int(ig2D.voxel_num_y/2))],\n", " label = ['TV (explicit)','TV (implicit)'], title='Middle Line Profiles')\n" ] }, { "cell_type": "markdown", "id": "a0f699e0", "metadata": {}, "source": [ "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." ] }, { "cell_type": "markdown", "id": "a2232e53-a927-461b-9e95-2c9f322257b7", "metadata": {}, "source": [ "## $(L^{2}-TV)$ Implicit PDHG: using TotalVariation" ] }, { "cell_type": "code", "execution_count": null, "id": "35b0e04b-4393-46d5-b8a3-bbf6161feeb8", "metadata": {}, "outputs": [], "source": [ "G = alpha_tv * TotalVariation(max_iteration=100, lower=0.)\n", "\n", "# Setup and run PDHG\n", "pdhg_tv_implicit_cil = PDHG(f = F, g = G, operator = K,\n", " max_iteration = 500,\n", " update_objective_interval = 100)\n", "pdhg_tv_implicit_cil.run(verbose=1)" ] }, { "cell_type": "code", "execution_count": null, "id": "b78c87bb-9dcf-4d40-b41e-b7118c3ac120", "metadata": {}, "outputs": [], "source": [ "# Show reconstruction and ground truth\n", "show2D([pdhg_tv_implicit_regtk.solution,\n", " pdhg_tv_implicit_cil.solution,\n", " (pdhg_tv_implicit_cil.solution-pdhg_tv_implicit_regtk.solution).abs()], \n", " fix_range=[(0,0.055),(0,0.055),(0,1e-3)], num_cols=3,\n", " title = ['TV (CIL)','TV (CCPI-RegTk)', 'Absolute Difference'], \n", " cmap = 'inferno')\n", "\n", "# Plot middle line profile\n", "show1D([pdhg_tv_implicit_regtk.solution,pdhg_tv_implicit_cil.solution], slice_list=[('horizontal_y',int(ig2D.voxel_num_y/2))],\n", " label = ['TV (CCPi-RegTk)','TV (CIL)'], title='Middle Line Profiles')\n" ] }, { "cell_type": "markdown", "id": "5f0e9178-110e-4a03-84ff-ef5391e25461", "metadata": {}, "source": [ "# FBP reconstruction with all the projection angles." ] }, { "cell_type": "code", "execution_count": null, "id": "24afa098-da30-4e33-b605-c875b3288767", "metadata": {}, "outputs": [], "source": [ "binned_data3D = Binner(roi={'horizontal':(120,-120,2)})(data3D)\n", "absorption_data3D = TransmissionAbsorptionConverter()(binned_data3D.get_slice(vertical=512))\n", "absorption_data3D -= np.mean(absorption_data3D.as_array()[80:100,0:30])\n", "ag3D = absorption_data3D.geometry\n", "ag3D.set_angles(ag3D.angles, initial_angle=0.2, angle_unit='radian')\n", "ig3D = ag3D.get_ImageGeometry()\n", "\n", "fbp_recon3D = FBP(ig3D, ag3D)(absorption_data3D)" ] }, { "cell_type": "markdown", "id": "40c57ee9-a7a6-4d9c-963b-9757f1086ee8", "metadata": {}, "source": [ "# Show all reconstructions \n", "\n", "- FBP (1601 projections)\n", "- FBP (160 projections)\n", "- SIRT (160 projections)\n", "- $L^{1}$ regularisation (160 projections)\n", "- Tikhonov regularisation (160 projections)\n", "- Total variation regularisation (160 projections)" ] }, { "cell_type": "code", "execution_count": null, "id": "980d19c1-32d3-4f7f-ba2a-b2260e7be27a", "metadata": {}, "outputs": [], "source": [ "show2D([fbp_recon3D, \n", " fbp_recon, \n", " sirt_recon, \n", " pdhg_l1.solution, \n", " pdhg_tikhonov_explicit.solution,\n", " pdhg_tv_explicit.solution],\n", " title=['FBP 1601 projections', 'FBP', 'SIRT','$L^{1}$','Tikhonov','TV'],\n", " cmap=\"inferno\",num_cols=3, size=(25,20), fix_range=(0,0.05))" ] }, { "cell_type": "markdown", "id": "ae6a6825-09ca-45ba-80f2-edf6bb459f4e", "metadata": {}, "source": [ "## Zoom ROIs" ] }, { "cell_type": "code", "execution_count": null, "id": "4e20863a-f0d8-49ba-be86-ee041acccbad", "metadata": {}, "outputs": [], "source": [ "show2D([fbp_recon3D.as_array()[175:225,150:250], \n", " fbp_recon.as_array()[175:225,150:250], \n", " sirt_recon.as_array()[175:225,150:250], \n", " pdhg_l1.solution.as_array()[175:225,150:250], \n", " pdhg_tikhonov_explicit.solution.as_array()[175:225,150:250],\n", " pdhg_tv_implicit_regtk.solution.as_array()[175:225,150:250]],\n", " title=['FBP 1601 projections', 'FBP', 'SIRT','$L^{1}$','Tikhonov','TV'],\n", " cmap=\"inferno\",num_cols=3, size=(25,20), fix_range=(0,0.05))" ] }, { "attachments": {}, "cell_type": "markdown", "id": "fb10bb8e-bdd7-4945-b291-aa98a0018292", "metadata": {}, "source": [ "# Conclusions\n", "\n", "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:\n", "\n", "* $\\sigma = \\tau = \\frac{1.0}{\\|K\\|}$\n", "\n", "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](https://hal.archives-ouvertes.fr/hal-00490826/document), [Chambolle_Pock2011](https://ieeexplore.ieee.org/document/6126441), [Goldstein et al](https://arxiv.org/pdf/1305.0546.pdf), [Malitsky_Pock](https://arxiv.org/pdf/1608.08883.pdf).\n", "\n", "In the following notebook, we are going to present a stochastic version of PDHG, namely **SPDHG** introduced in [Chambolle et al](https://arxiv.org/pdf/1706.04957.pdf) 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\n", "over the PDHG algorithm, see [Ehrhardt et al](https://arxiv.org/abs/1808.07150) and [Papoutsellis et al](https://arxiv.org/pdf/2102.06126.pdf)." ] }, { "cell_type": "code", "execution_count": null, "id": "ad6242be-de86-4ae4-9e9a-a8c33dacce58", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "interpreter": { "hash": "43cbf82c2f716cd564b762322e13d4dbd881fd8a341d231fe608abc3118da208" }, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.16" } }, "nbformat": 4, "nbformat_minor": 5 }