diff --git a/examples/notebooks/tgv_minimization_reconstruction_pdhg.ipynb b/examples/notebooks/tgv_minimization_reconstruction_pdhg.ipynb new file mode 100644 index 000000000..fe8ba51a5 --- /dev/null +++ b/examples/notebooks/tgv_minimization_reconstruction_pdhg.ipynb @@ -0,0 +1,872 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PTB-MR/mrpro/blob/main/examples/notebooks/tgv_minimization_reconstruction_pdhg.ipynb)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": { + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "import importlib\n", + "\n", + "if not importlib.util.find_spec('mrpro'):\n", + " %pip install mrpro[notebooks]" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "# Total-generalized-variation (TGV)-minimization reconstruction" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "In this notebook we demonstrate TGV-based variational reconstruction with a primal-dual hybrid gradient (PDHG) solver\n", + "from **mrpro**.\n", + "\n", + "We work through two examples:\n", + "1. **Denoising** of a piecewise-constant “square” image — to show how TGV suppresses the staircasing artifact common\n", + "with TV.\n", + "2. **MRI reconstruction** from retrospectively undersampled 2-D **radial** spokes on a Cartesian grid using\n", + "$A = M \\circ \\mathcal{F}$." + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "## First Example: Denoising" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "### Load data\n", + "We use the classic square test image (grayscale, scaled to $[0,1]$) and a noisy version with additive white Gaussian\n", + "noise of standard deviation $\\sigma \\approx 0.05$:\n", + "- `square.png` (clean)\n", + "- `square_noisy_0_05.png` (noisy)\n", + "\n", + "The goal is to recover a clean image from the noisy observation using TV and then TGV for comparison." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": { + "mystnb": { + "code_prompt_show": "Show download details" + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "# Download raw data from Zenodo\n", + "import tempfile\n", + "from pathlib import Path\n", + "\n", + "import torch\n", + "\n", + "if torch.cuda.is_available():\n", + " torch.set_default_device('cuda')\n", + "\n", + "import zenodo_get\n", + "\n", + "tmp = tempfile.TemporaryDirectory() # RAII, automatically cleaned up\n", + "data_folder = Path(tmp.name)\n", + "zenodo_get.download(record='16811276', retry_attempts=5, output_dir=data_folder)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": { + "mystnb": { + "code_prompt_show": "Show plotting details" + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "def show_images(\n", + " *images: torch.Tensor, titles: list[str] | None = None, clim: tuple[float, float] | None = None\n", + ") -> None:\n", + " \"\"\"Plot images.\"\"\"\n", + " n_images = len(images)\n", + " _, axes = plt.subplots(1, n_images, squeeze=False, figsize=(n_images * 3, 3))\n", + " for i in range(n_images):\n", + " axes[0][i].imshow(images[i].cpu().squeeze(), cmap='gray', clim=clim)\n", + " axes[0][i].axis('off')\n", + " if titles:\n", + " axes[0][i].set_title(titles[i])\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Display the noisy input and the clean reference for visual comparison." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "import mrpro\n", + "from PIL import Image\n", + "from torchvision.transforms.functional import to_tensor\n", + "\n", + "square_clean = to_tensor(Image.open(data_folder / 'square.png').convert('L')).to(device=torch.get_default_device())\n", + "square_noisy = to_tensor(Image.open(data_folder / 'square_noisy_0_05.png').convert('L')).to(\n", + " device=torch.get_default_device()\n", + ")\n", + "\n", + "show_images(square_noisy, square_clean, titles=['Noisy', 'Clean'], clim=(0, 1))" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Before introducing TGV, we first solve a TV-regularized problem as a baseline\n", + "(see [tv_minimization_reconstruction_pdhg.ipynb](examples/notebooks/tv_minimization_reconstruction_pdhg.ipynb)\n", + "for details).\n", + "\n", + "With an observation $y$ and acquisition operator $A$, the TV model reads\n", + "$$\n", + "\\min_x\\; \\tfrac12\\,\\lVert Ax - y\\rVert_2^2\\; +\\; \\lambda\\,\\lVert \\nabla x\\rVert_1,\n", + "$$\n", + "where $\\nabla$ is the (forward) finite-difference gradient. For denoising we take $A = I$,\n", + "so the first term penalizes the pixel-wise squared error against $y$." + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "### Set up the operator $A$\n", + "For denoising, the acquisition is the identity operator $A = I$ (implemented with `mrpro.operators.IdentityOp`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "identity_op = mrpro.operators.IdentityOp() # acquisition operator A here is simply the identity operator" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Here we define a small helper that runs PDHG for TV. The three key components are:\n", + "- Data term: $\\tfrac12\\lVert Ax - y\\rVert_2^2$\n", + "- Gradient term (Regularization term): $\\lambda\\lVert \\nabla x\\rVert_1$\n", + "- Operator: $Kx = \\begin{bmatrix}Ax\\\\ \\nabla x\\end{bmatrix}$, handled with `LinearOperatorMatrix`\n", + "\n", + "We use the adjoint reconstruction as the initial point." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "from collections.abc import Sequence\n", + "\n", + "from mrpro.operators import (\n", + " FiniteDifferenceOp,\n", + " IdentityOp,\n", + " LinearOperator,\n", + " LinearOperatorMatrix,\n", + " ProximableFunctionalSeparableSum,\n", + " RearrangeOp,\n", + " ZeroOp,\n", + ")\n", + "from mrpro.operators.functionals import L1NormViewAsReal, L2NormSquared\n", + "\n", + "\n", + "def tv_minimization_reconstruction(\n", + " measurement: torch.Tensor,\n", + " acquisition_op: LinearOperator,\n", + " grad_term_weight: torch.Tensor,\n", + " dim: Sequence[int] = (-2, -1),\n", + " **pdhg_kwargs,\n", + ") -> torch.Tensor:\n", + " \"\"\"Perform TV-minimization reconstruction.\"\"\"\n", + " # 1. Compute initial values using the adjoint of the acquisition operator\n", + " adjoint_recon = acquisition_op.adjoint(measurement)[0]\n", + " initial_values = (adjoint_recon,)\n", + "\n", + " # 2. Define the objective functional\n", + " data_term = 0.5 * L2NormSquared(target=measurement)\n", + " grad_term = L1NormViewAsReal(weight=grad_term_weight)\n", + " minimization_sum = ProximableFunctionalSeparableSum(data_term, grad_term)\n", + "\n", + " # 3. Define the operator matrix K\n", + " data_term_row = (acquisition_op,)\n", + " nabla = FiniteDifferenceOp(dim=dim, mode='forward')\n", + " grad_term_row = (nabla,)\n", + " operator_matrix = LinearOperatorMatrix((data_term_row, grad_term_row))\n", + "\n", + " return mrpro.algorithms.optimizers.pdhg(\n", + " f=minimization_sum, g=None, operator=operator_matrix, initial_values=initial_values, **pdhg_kwargs\n", + " )[0]" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "### Run PDHG with TV" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Let us run PDHG for a number of iterations. The regularization parameter $\\lambda = 0.05$ was retrospectively picked\n", + "to achieve a good reconstruction for this example. We can later tune the regularization parameter $\\lambda$ and the\n", + "number of iterations to balance smoothing with detail preservation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": { + "mystnb": { + "code_prompt_show": "Show callback details" + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "# This is a \"callback\" function that will be called after each iteration of the PDHG algorithm.\n", + "# We use it here to print progress information.\n", + "\n", + "from mrpro.algorithms.optimizers.pdhg import PDHGStatus\n", + "\n", + "\n", + "def callback(optimizer_status: PDHGStatus) -> None:\n", + " \"\"\"Print the value of the objective functional every 16th iteration.\"\"\"\n", + " iteration = optimizer_status['iteration_number']\n", + " solution = optimizer_status['solution']\n", + " if iteration % 16 == 0:\n", + " print(f'Iteration {iteration: >3}: Objective = {optimizer_status[\"objective\"](*solution).item():.3e}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "square_tv_denoised = tv_minimization_reconstruction(\n", + " measurement=square_noisy,\n", + " acquisition_op=identity_op,\n", + " grad_term_weight=torch.tensor(0.05),\n", + " max_iterations=257,\n", + " callback=callback,\n", + ")\n", + "show_images(square_noisy, square_tv_denoised, square_clean, titles=['Noisy', 'TV Denoised', 'Clean'], clim=(0, 1))" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "TV reduces noise but exhibits staircasing (piecewise-constant plateaus with sharp steps). Next we apply TGV and show\n", + "that it can mitigate this artifact." + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "### Total-generalized-variation (TGV)" + ] + }, + { + "cell_type": "markdown", + "id": "21", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Similar to the TV case, let $y$ denote the measured data and $A$ the acquisition model. We assume\n", + "$$\n", + "y = Ax_{\\mathrm{true}} + n,\n", + "$$\n", + "with (complex) Gaussian noise $n$.\n", + "\n", + "In second-order TGV we introduce an auxiliary field $v$ (with the same spatial shape as $\\nabla x$) and minimize\n", + "$$\n", + "\\min_{x,\\,v}\\; \\tfrac12\\,\\lVert Ax - y\\rVert_2^2\n", + "\\; +\\; \\lambda_1\\,\\lVert \\nabla x - v\\rVert_1\n", + "\\; +\\; \\lambda_0\\,\\Big\\lVert \\mathcal{E}v\\Big\\rVert_1, \\tag{1}\n", + "$$\n", + "where $\\lambda_1$ and $\\lambda_0$ are regularization parameters, and $\\mathcal{E}$ is the symmetrized gradient\n", + "$$\n", + "\\mathcal{E}v \\;=\\; \\tfrac12\\,(\\nabla v + (\\nabla v)^{\\top})\n", + "$$" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "#### Recast for PDHG\n", + "To use PDHG we write (1) in the form\n", + "$$\n", + "\\min_{x,v}\\; g(x,v) + f\\big(K(x,v)\\big).\n", + "$$" + ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "#### Operator $K$\n", + "$$\n", + "K(x,v) \\;=\\;\n", + "\\begin{bmatrix}\n", + "A & 0\\\\\n", + "\\nabla & -I\\\\\n", + "0 & \\mathcal{E}\n", + "\\end{bmatrix}\n", + "\\begin{bmatrix}\n", + "x\\\\ v\n", + "\\end{bmatrix}\n", + "\\;=\\;\n", + "\\begin{bmatrix}\n", + "Ax\\\\\n", + "\\nabla x - v\\\\\n", + "\\mathcal{E}v\n", + "\\end{bmatrix}.\n", + "$$\n", + "\n", + "We build this as a `LinearOperatorMatrix` using `FiniteDifferenceOp` (forward/backward modes), a `RearrangeOp` to\n", + "align dimensions, and the simple `ZeroOp`." + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "#### Functionals $g$ and $f$\n", + "We take $g \\equiv 0$ (implemented implicitly by passing `g=None`), and\n", + "$$\n", + "f(z_1, z_2, z_3) \\,=\\, \\tfrac12\\,\\lVert z_1 - y\\rVert_2^2\n", + "\\; +\\; \\lambda_1\\,\\lVert z_2\\rVert_1\n", + "\\; +\\; \\lambda_0\\,\\lVert z_3\\rVert_1.\n", + "$$\n", + "In code this is\n", + "`ProximableFunctionalSeparableSum(L2NormSquared(y)/2, L1NormViewAsReal(weight=λ1), L1NormViewAsReal(weight=λ0))`.\n", + "We call $\\lambda_1$ the gradient term's weight and $\\lambda_0$ the symmetrized gradient term's weight.\n", + "\n", + "With this choice of $g$ and $f$, the proximal operators are closed-form." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "def tgv_minimization_reconstruction(\n", + " measurement: torch.Tensor,\n", + " acquisition_op: LinearOperator,\n", + " grad_term_weight: torch.Tensor,\n", + " sym_grad_term_weight: torch.Tensor,\n", + " dim: Sequence[int] = (-2, -1),\n", + " **pdhg_kwargs,\n", + ") -> torch.Tensor:\n", + " \"\"\"Perform TGV-minimization reconstruction.\"\"\"\n", + " # 1. Compute initial values using the adjoint of the acquisition operator\n", + " adjoint_recon = acquisition_op.adjoint(measurement)[0]\n", + " # Auxiliary tensor v which is in the gradient domain.\n", + " auxiliary_v_tensor = adjoint_recon.new_zeros((len(dim), *adjoint_recon.shape))\n", + " # Increase the number of dimensions of initial image by one.\n", + " # Must make the number of dimensions of the elements in initial values list match one another,\n", + " # because (currently) pdhg function concatenates the individual norms of each operator\n", + " # in a matrix's row and the norm has the same shape as the input.\n", + " # If the number of dimensions don't match, we get a runtime error like this:\n", + " # RuntimeError: stack expects each tensor to be equal size, but got ... at entry 0 and ... at entry 1\n", + " adjoint_recon = adjoint_recon.unsqueeze(0)\n", + " initial_values = (adjoint_recon, auxiliary_v_tensor)\n", + "\n", + " # 2. Define the objective functional\n", + " data_term = 0.5 * L2NormSquared(target=measurement)\n", + " grad_term = L1NormViewAsReal(weight=grad_term_weight)\n", + " sym_grad_term = L1NormViewAsReal(weight=sym_grad_term_weight)\n", + " minimization_sum = ProximableFunctionalSeparableSum(data_term, grad_term, sym_grad_term)\n", + "\n", + " # 3. Define the operator matrix K\n", + " # 3.1. First row (corresponding to the L2-norm data term): Ax + 0\n", + " data_term_row = (acquisition_op, ZeroOp())\n", + "\n", + " # 3.2. Second row (corresponding to the first L1-norm regularization term): \\nabla x - v\n", + " # Reduce the number of dimensions of the image by one before applying the finite difference operator\n", + " # to make the output of the finite difference operator match the auxiliary tensor v.\n", + " squeeze_op = RearrangeOp('1 ... -> ...')\n", + " forward_nabla = FiniteDifferenceOp(dim=dim, mode='forward')\n", + " grad_term_row = (forward_nabla @ squeeze_op, -1 * IdentityOp())\n", + "\n", + " # 3.3. Third row (corresponding to the second L1-norm regularization term): 0 + \\mathcal{E} v\n", + " backward_nabla = FiniteDifferenceOp(dim=dim, mode='backward')\n", + " transpose_op = RearrangeOp('sym_grad_dim grad_dim ... -> grad_dim sym_grad_dim ...')\n", + " symmetric_gradient_op = 0.5 * (1 + transpose_op) @ backward_nabla\n", + " sym_grad_term_row = (ZeroOp(), symmetric_gradient_op)\n", + "\n", + " operator_matrix = LinearOperatorMatrix((data_term_row, grad_term_row, sym_grad_term_row))\n", + "\n", + " return mrpro.algorithms.optimizers.pdhg(\n", + " f=minimization_sum,\n", + " g=None, # automatically converted to zero functionals\n", + " operator=operator_matrix,\n", + " initial_values=initial_values,\n", + " **pdhg_kwargs,\n", + " )[0]" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "### Run PDHG with TGV" + ] + }, + { + "cell_type": "markdown", + "id": "27", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Let us run TGV with weights $(\\lambda_1,\\lambda_0) = (0.05, 0.1)$; a common heuristic is\n", + "$\\lambda_0 \\approx 2\\lambda_1$. We can later adjust these and the number of iterations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "square_tgv_denoised = tgv_minimization_reconstruction(\n", + " measurement=square_noisy,\n", + " acquisition_op=identity_op,\n", + " grad_term_weight=torch.tensor(0.05),\n", + " sym_grad_term_weight=torch.tensor(0.1),\n", + " max_iterations=257,\n", + " callback=callback,\n", + ")\n", + "show_images(\n", + " square_noisy,\n", + " square_tv_denoised,\n", + " square_tgv_denoised,\n", + " square_clean,\n", + " titles=['Noisy', 'TV Denoised', 'TGV Denoised', 'Clean'],\n", + " clim=(0, 1),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Compared to TV, TGV can still preserve edges and details while avoiding staircasing; slowly varying ramps are\n", + "reconstructed more faithfully." + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "## Second example: Radial undersampling" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "We next reconstruct a multi-coil brain MRI from retrospectively undersampled radial spokes drawn on a Cartesian grid.\n", + "We use 4-coil k-space data (`1_rawdata_brainT2_4ch.mat`) as a reference dataset\n", + "(fully sampled in Cartesian coordinates)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": { + "mystnb": { + "code_prompt_show": "Show download details" + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "# Download ground-truth k-space data from Zenodo\n", + "zenodo_get.download(record='800525', retry_attempts=5, output_dir=data_folder)" + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "To begin, we load the fully sampled k-space, form the image with an FFT operator, and compute a\n", + "root-sum-of-squares (RSS) reference." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "from einops import rearrange\n", + "from mrpro.data import SpatialDimension\n", + "from scipy.io import loadmat\n", + "\n", + "file_name = '1_rawdata_brainT2_4ch.mat'\n", + "kdata_true = torch.tensor(loadmat(data_folder / file_name)['rawdata'])\n", + "kdata_true = rearrange(kdata_true, 'k1 k0 coils -> 1 coils 1 k1 k0')\n", + "# kdata = torch.flip(kdata, dims=(-2, -1))\n", + "kdata_true = kdata_true.to(dtype=torch.complex64) # If default dtype is torch.float32, use complex64\n", + "\n", + "recon_matrix = SpatialDimension(z=kdata_true.shape[-3], y=kdata_true.shape[-2], x=kdata_true.shape[-1])\n", + "encoding_matrix = SpatialDimension(z=kdata_true.shape[-3], y=kdata_true.shape[-2], x=kdata_true.shape[-1])\n", + "\n", + "fourier_op = mrpro.operators.FastFourierOp(\n", + " dim=(-2, -1),\n", + " recon_matrix=recon_matrix,\n", + " encoding_matrix=encoding_matrix,\n", + ")\n", + "x_true = fourier_op(kdata_true)[0]\n", + "\n", + "x_true_rss = x_true.abs().square().sum(dim=-4).sqrt().squeeze()\n", + "show_images(x_true_rss, titles=['Ground Truth'], clim=(0, 7e-4))" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "### Set up the operator $A$" + ] + }, + { + "cell_type": "markdown", + "id": "36", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "To set up the acquisition operator $A$, we first create a radial mask with a chosen number of spokes (here 48) on the\n", + "Cartesian grid and wrap it as a `CartesianMaskingOp`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "37", + "metadata": {}, + "outputs": [], + "source": [ + "from torchvision.transforms.functional import rotate\n", + "\n", + "\n", + "def radial_mask(ny: int, nx: int, num_spokes: int) -> torch.Tensor:\n", + " \"\"\"Generate a radial mask with the specified number of spokes.\"\"\"\n", + " theta = 180 * (3.0 - 5**0.5) # golden angle ~137.508°\n", + " mask = torch.zeros((1, 1, ny, nx), dtype=torch.bool)\n", + "\n", + " # prototype spoke: horizontal line through center\n", + " base_spoke = torch.zeros((1, 1, ny, nx))\n", + " base_spoke[0, 0, ny // 2, :] = 1.0\n", + "\n", + " for i_spoke in range(num_spokes):\n", + " spoke = rotate(base_spoke, angle=theta * i_spoke, fill=0)\n", + " mask |= spoke > 0.5\n", + " return mask\n", + "\n", + "\n", + "Nx, Ny = 256, 256\n", + "num_spokes = 48\n", + "mask_op = mrpro.operators.CartesianMaskingOp(radial_mask(Ny, Nx, num_spokes))\n", + "show_images(torch.tensor(mask_op.mask), titles=[f'Mask ({num_spokes} spokes)'])" + ] + }, + { + "cell_type": "markdown", + "id": "38", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "Then we can define the acquisition operator as\n", + "$$\n", + "A \\,=\\, M \\circ \\mathcal{F},\n", + "$$\n", + "where $\\mathcal{F}$ is the (multi-dimensional) FFT mapping images to k-space, and $M$ applies the binary mask.\n", + "We simulate undersampling by computing $y = A \\ x_{\\mathrm{true}}$, and display the adjoint reconstruction $A^* \\ y$\n", + "(i.e. zero-filled inverse FFT)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "39", + "metadata": {}, + "outputs": [], + "source": [ + "acquisition_op = mask_op @ fourier_op\n", + "# Apply A(x_true) to get undersampled k-space data\n", + "kdata_undersampled = acquisition_op(x_true)[0]\n", + "x_adjoint_recon = acquisition_op.adjoint(kdata_undersampled)[0]\n", + "x_adjoint_recon_rss = torch.sum(x_adjoint_recon.abs() ** 2, dim=1).sqrt()\n", + "show_images(x_adjoint_recon_rss, x_true_rss, titles=['Adjoint', 'Ground Truth'], clim=(0, 7e-4))" + ] + }, + { + "cell_type": "markdown", + "id": "40", + "metadata": {}, + "source": [ + "### Run PDHG with TV" + ] + }, + { + "cell_type": "markdown", + "id": "41", + "metadata": {}, + "source": [ + "Again, let us start with the TV-regularized reconstruction method. The regularization parameter\n", + "$\\lambda = 5 \\times 10^{-5}$ was chosen to achieve a good reconstruction for this example." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "42", + "metadata": {}, + "outputs": [], + "source": [ + "x_tv_recon = tv_minimization_reconstruction(\n", + " measurement=kdata_undersampled,\n", + " acquisition_op=acquisition_op,\n", + " grad_term_weight=torch.tensor(5e-6),\n", + " max_iterations=257,\n", + " callback=callback,\n", + ")\n", + "x_tv_recon_rss = torch.sum(x_tv_recon.abs() ** 2, dim=-4).sqrt()\n", + "show_images(x_adjoint_recon_rss, x_tv_recon_rss, x_true_rss, titles=['Adjoint', 'TV', 'Ground Truth'], clim=(0, 7e-4))" + ] + }, + { + "cell_type": "markdown", + "id": "43", + "metadata": {}, + "source": [ + "### Run PDHG with TGV" + ] + }, + { + "cell_type": "markdown", + "id": "44", + "metadata": {}, + "source": [ + "Now we can run the TGV-regularized reconstruction with $(\\lambda_1, \\lambda_0) = (5 \\times 10^{-6}, 10^{-5})$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45", + "metadata": {}, + "outputs": [], + "source": [ + "x_tgv_recon = tgv_minimization_reconstruction(\n", + " measurement=kdata_undersampled,\n", + " acquisition_op=acquisition_op,\n", + " grad_term_weight=torch.tensor(5e-6),\n", + " sym_grad_term_weight=torch.tensor(10e-6),\n", + " max_iterations=257,\n", + " callback=callback,\n", + ")\n", + "x_tgv_recon_rss = torch.sum(x_tgv_recon.abs() ** 2, dim=-4).sqrt()\n", + "show_images(\n", + " x_adjoint_recon_rss,\n", + " x_tv_recon_rss,\n", + " x_tgv_recon_rss,\n", + " x_true_rss,\n", + " titles=['Adjoint', 'TV', 'TGV', 'Ground Truth'],\n", + " clim=(0, 7e-4),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "46", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "### Compare the results\n", + "We compare adjoint (zero-filled), TV, TGV, and the ground truth (RSS) on a zoomed-in portion.\n", + "As expected, adjoint reconstruction shows aliasing from undersampling.\n", + "TV reduces aliasing but introduces some staircasing,\n", + "while TGV is able to reduce aliasing without the staircasing effect." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "47", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "show_images(\n", + " x_adjoint_recon_rss[..., :128, :128],\n", + " x_tv_recon_rss[..., :128, :128],\n", + " x_tgv_recon_rss[..., :128, :128],\n", + " x_true_rss[..., :128, :128],\n", + " titles=['Adjoint', 'TV', 'TGV', 'Ground Truth'],\n", + " clim=(0, 7e-4),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "48", + "metadata": {}, + "source": [ + "That`s it — we performed denoising and MRI reconstruction via TGV-minimization solved with PDHG.\n", + "We can later try different $(\\lambda_1,\\lambda_2)$ value combinations and adjust the iteration count\n", + "to see how the reconstruction quality changes." + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "jupytext": { + "cell_metadata_filter": "mystnb,tags,-all" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/scripts/tgv_minimization_reconstruction_pdhg.py b/examples/scripts/tgv_minimization_reconstruction_pdhg.py new file mode 100644 index 000000000..bec62f6de --- /dev/null +++ b/examples/scripts/tgv_minimization_reconstruction_pdhg.py @@ -0,0 +1,463 @@ +# %% [markdown] +# # Total-generalized-variation (TGV)-minimization reconstruction + +# %% [markdown] +# In this notebook we demonstrate TGV-based variational reconstruction with a primal-dual hybrid gradient (PDHG) solver +# from **mrpro**. +# +# We work through two examples: +# 1. **Denoising** of a piecewise-constant “square” image — to show how TGV suppresses the staircasing artifact common +# with TV. +# 2. **MRI reconstruction** from retrospectively undersampled 2-D **radial** spokes on a Cartesian grid using +# $A = M \circ \mathcal{F}$. +# %% [markdown] +# ## First Example: Denoising +# %% [markdown] +# ### Load data +# We use the classic square test image (grayscale, scaled to $[0,1]$) and a noisy version with additive white Gaussian +# noise of standard deviation $\sigma \approx 0.05$: +# - `square.png` (clean) +# - `square_noisy_0_05.png` (noisy) +# +# The goal is to recover a clean image from the noisy observation using TV and then TGV for comparison. +# %% mystnb={"code_prompt_show": "Show download details"} tags=["hide-cell"] +# Download raw data from Zenodo +import tempfile +from pathlib import Path + +import torch + +if torch.cuda.is_available(): + torch.set_default_device('cuda') + +import zenodo_get + +tmp = tempfile.TemporaryDirectory() # RAII, automatically cleaned up +data_folder = Path(tmp.name) +zenodo_get.download(record='16811276', retry_attempts=5, output_dir=data_folder) + +# %% mystnb={"code_prompt_show": "Show plotting details"} tags=["hide-cell"] +import matplotlib.pyplot as plt + + +def show_images( + *images: torch.Tensor, titles: list[str] | None = None, clim: tuple[float, float] | None = None +) -> None: + """Plot images.""" + n_images = len(images) + _, axes = plt.subplots(1, n_images, squeeze=False, figsize=(n_images * 3, 3)) + for i in range(n_images): + axes[0][i].imshow(images[i].cpu().squeeze(), cmap='gray', clim=clim) + axes[0][i].axis('off') + if titles: + axes[0][i].set_title(titles[i]) + plt.show() + + +# %% [markdown] +# Display the noisy input and the clean reference for visual comparison. +# %% +import mrpro +from PIL import Image +from torchvision.transforms.functional import to_tensor + +square_clean = to_tensor(Image.open(data_folder / 'square.png').convert('L')).to(device=torch.get_default_device()) +square_noisy = to_tensor(Image.open(data_folder / 'square_noisy_0_05.png').convert('L')).to( + device=torch.get_default_device() +) + +show_images(square_noisy, square_clean, titles=['Noisy', 'Clean'], clim=(0, 1)) + +# %% [markdown] +# Before introducing TGV, we first solve a TV-regularized problem as a baseline +# (see [tv_minimization_reconstruction_pdhg.ipynb](examples/notebooks/tv_minimization_reconstruction_pdhg.ipynb) +# for details). +# +# With an observation $y$ and acquisition operator $A$, the TV model reads +# $$ +# \min_x\; \tfrac12\,\lVert Ax - y\rVert_2^2\; +\; \lambda\,\lVert \nabla x\rVert_1, +# $$ +# where $\nabla$ is the (forward) finite-difference gradient. For denoising we take $A = I$, +# so the first term penalizes the pixel-wise squared error against $y$. +# %% [markdown] +# ### Set up the operator $A$ +# For denoising, the acquisition is the identity operator $A = I$ (implemented with `mrpro.operators.IdentityOp`). +# %% +identity_op = mrpro.operators.IdentityOp() # acquisition operator A here is simply the identity operator + +# %% [markdown] +# Here we define a small helper that runs PDHG for TV. The three key components are: +# - Data term: $\tfrac12\lVert Ax - y\rVert_2^2$ +# - Gradient term (Regularization term): $\lambda\lVert \nabla x\rVert_1$ +# - Operator: $Kx = \begin{bmatrix}Ax\\ \nabla x\end{bmatrix}$, handled with `LinearOperatorMatrix` +# +# We use the adjoint reconstruction as the initial point. +# %% +from collections.abc import Sequence + +from mrpro.operators import ( + FiniteDifferenceOp, + IdentityOp, + LinearOperator, + LinearOperatorMatrix, + ProximableFunctionalSeparableSum, + RearrangeOp, + ZeroOp, +) +from mrpro.operators.functionals import L1NormViewAsReal, L2NormSquared + + +def tv_minimization_reconstruction( + measurement: torch.Tensor, + acquisition_op: LinearOperator, + grad_term_weight: torch.Tensor, + dim: Sequence[int] = (-2, -1), + **pdhg_kwargs, +) -> torch.Tensor: + """Perform TV-minimization reconstruction.""" + # 1. Compute initial values using the adjoint of the acquisition operator + adjoint_recon = acquisition_op.adjoint(measurement)[0] + initial_values = (adjoint_recon,) + + # 2. Define the objective functional + data_term = 0.5 * L2NormSquared(target=measurement) + grad_term = L1NormViewAsReal(weight=grad_term_weight) + minimization_sum = ProximableFunctionalSeparableSum(data_term, grad_term) + + # 3. Define the operator matrix K + data_term_row = (acquisition_op,) + nabla = FiniteDifferenceOp(dim=dim, mode='forward') + grad_term_row = (nabla,) + operator_matrix = LinearOperatorMatrix((data_term_row, grad_term_row)) + + return mrpro.algorithms.optimizers.pdhg( + f=minimization_sum, g=None, operator=operator_matrix, initial_values=initial_values, **pdhg_kwargs + )[0] + + +# %% [markdown] +# ### Run PDHG with TV + +# %% [markdown] +# Let us run PDHG for a number of iterations. The regularization parameter $\lambda = 0.05$ was retrospectively picked +# to achieve a good reconstruction for this example. We can later tune the regularization parameter $\lambda$ and the +# number of iterations to balance smoothing with detail preservation. +# %% mystnb={"code_prompt_show": "Show callback details"} tags=["hide-cell"] +# This is a "callback" function that will be called after each iteration of the PDHG algorithm. +# We use it here to print progress information. + +from mrpro.algorithms.optimizers.pdhg import PDHGStatus + + +def callback(optimizer_status: PDHGStatus) -> None: + """Print the value of the objective functional every 16th iteration.""" + iteration = optimizer_status['iteration_number'] + solution = optimizer_status['solution'] + if iteration % 16 == 0: + print(f'Iteration {iteration: >3}: Objective = {optimizer_status["objective"](*solution).item():.3e}') + + +# %% +square_tv_denoised = tv_minimization_reconstruction( + measurement=square_noisy, + acquisition_op=identity_op, + grad_term_weight=torch.tensor(0.05), + max_iterations=257, + callback=callback, +) +show_images(square_noisy, square_tv_denoised, square_clean, titles=['Noisy', 'TV Denoised', 'Clean'], clim=(0, 1)) + + +# %% [markdown] +# TV reduces noise but exhibits staircasing (piecewise-constant plateaus with sharp steps). Next we apply TGV and show +# that it can mitigate this artifact. +# %% [markdown] +# ### Total-generalized-variation (TGV) +# %% [markdown] +# Similar to the TV case, let $y$ denote the measured data and $A$ the acquisition model. We assume +# $$ +# y = Ax_{\mathrm{true}} + n, +# $$ +# with (complex) Gaussian noise $n$. +# +# In second-order TGV we introduce an auxiliary field $v$ (with the same spatial shape as $\nabla x$) and minimize +# $$ +# \min_{x,\,v}\; \tfrac12\,\lVert Ax - y\rVert_2^2 +# \; +\; \lambda_1\,\lVert \nabla x - v\rVert_1 +# \; +\; \lambda_0\,\Big\lVert \mathcal{E}v\Big\rVert_1, \tag{1} +# $$ +# where $\lambda_1$ and $\lambda_0$ are regularization parameters, and $\mathcal{E}$ is the symmetrized gradient +# $$ +# \mathcal{E}v \;=\; \tfrac12\,(\nabla v + (\nabla v)^{\top}) +# $$ +# %% [markdown] +# #### Recast for PDHG +# To use PDHG we write (1) in the form +# $$ +# \min_{x,v}\; g(x,v) + f\big(K(x,v)\big). +# $$ +# %% [markdown] +# #### Operator $K$ +# $$ +# K(x,v) \;=\; +# \begin{bmatrix} +# A & 0\\ +# \nabla & -I\\ +# 0 & \mathcal{E} +# \end{bmatrix} +# \begin{bmatrix} +# x\\ v +# \end{bmatrix} +# \;=\; +# \begin{bmatrix} +# Ax\\ +# \nabla x - v\\ +# \mathcal{E}v +# \end{bmatrix}. +# $$ +# +# We build this as a `LinearOperatorMatrix` using `FiniteDifferenceOp` (forward/backward modes), a `RearrangeOp` to +# align dimensions, and the simple `ZeroOp`. +# %% [markdown] +# #### Functionals $g$ and $f$ +# We take $g \equiv 0$ (implemented implicitly by passing `g=None`), and +# $$ +# f(z_1, z_2, z_3) \,=\, \tfrac12\,\lVert z_1 - y\rVert_2^2 +# \; +\; \lambda_1\,\lVert z_2\rVert_1 +# \; +\; \lambda_0\,\lVert z_3\rVert_1. +# $$ +# In code this is +# `ProximableFunctionalSeparableSum(L2NormSquared(y)/2, L1NormViewAsReal(weight=λ1), L1NormViewAsReal(weight=λ0))`. +# We call $\lambda_1$ the gradient term's weight and $\lambda_0$ the symmetrized gradient term's weight. +# +# With this choice of $g$ and $f$, the proximal operators are closed-form. +# %% +def tgv_minimization_reconstruction( + measurement: torch.Tensor, + acquisition_op: LinearOperator, + grad_term_weight: torch.Tensor, + sym_grad_term_weight: torch.Tensor, + dim: Sequence[int] = (-2, -1), + **pdhg_kwargs, +) -> torch.Tensor: + """Perform TGV-minimization reconstruction.""" + # 1. Compute initial values using the adjoint of the acquisition operator + adjoint_recon = acquisition_op.adjoint(measurement)[0] + # Auxiliary tensor v which is in the gradient domain. + auxiliary_v_tensor = adjoint_recon.new_zeros((len(dim), *adjoint_recon.shape)) + # Increase the number of dimensions of initial image by one. + # Must make the number of dimensions of the elements in initial values list match one another, + # because (currently) pdhg function concatenates the individual norms of each operator + # in a matrix's row and the norm has the same shape as the input. + # If the number of dimensions don't match, we get a runtime error like this: + # RuntimeError: stack expects each tensor to be equal size, but got ... at entry 0 and ... at entry 1 + adjoint_recon = adjoint_recon.unsqueeze(0) + initial_values = (adjoint_recon, auxiliary_v_tensor) + + # 2. Define the objective functional + data_term = 0.5 * L2NormSquared(target=measurement) + grad_term = L1NormViewAsReal(weight=grad_term_weight) + sym_grad_term = L1NormViewAsReal(weight=sym_grad_term_weight) + minimization_sum = ProximableFunctionalSeparableSum(data_term, grad_term, sym_grad_term) + + # 3. Define the operator matrix K + # 3.1. First row (corresponding to the L2-norm data term): Ax + 0 + data_term_row = (acquisition_op, ZeroOp()) + + # 3.2. Second row (corresponding to the first L1-norm regularization term): \nabla x - v + # Reduce the number of dimensions of the image by one before applying the finite difference operator + # to make the output of the finite difference operator match the auxiliary tensor v. + squeeze_op = RearrangeOp('1 ... -> ...') + forward_nabla = FiniteDifferenceOp(dim=dim, mode='forward') + grad_term_row = (forward_nabla @ squeeze_op, -1 * IdentityOp()) + + # 3.3. Third row (corresponding to the second L1-norm regularization term): 0 + \mathcal{E} v + backward_nabla = FiniteDifferenceOp(dim=dim, mode='backward') + transpose_op = RearrangeOp('sym_grad_dim grad_dim ... -> grad_dim sym_grad_dim ...') + symmetric_gradient_op = 0.5 * (1 + transpose_op) @ backward_nabla + sym_grad_term_row = (ZeroOp(), symmetric_gradient_op) + + operator_matrix = LinearOperatorMatrix((data_term_row, grad_term_row, sym_grad_term_row)) + + return mrpro.algorithms.optimizers.pdhg( + f=minimization_sum, + g=None, # automatically converted to zero functionals + operator=operator_matrix, + initial_values=initial_values, + **pdhg_kwargs, + )[0] + + +# %% [markdown] +# ### Run PDHG with TGV + +# %% [markdown] +# Let us run TGV with weights $(\lambda_1,\lambda_0) = (0.05, 0.1)$; a common heuristic is +# $\lambda_0 \approx 2\lambda_1$. We can later adjust these and the number of iterations. +# %% +square_tgv_denoised = tgv_minimization_reconstruction( + measurement=square_noisy, + acquisition_op=identity_op, + grad_term_weight=torch.tensor(0.05), + sym_grad_term_weight=torch.tensor(0.1), + max_iterations=257, + callback=callback, +) +show_images( + square_noisy, + square_tv_denoised, + square_tgv_denoised, + square_clean, + titles=['Noisy', 'TV Denoised', 'TGV Denoised', 'Clean'], + clim=(0, 1), +) + +# %% [markdown] +# Compared to TV, TGV can still preserve edges and details while avoiding staircasing; slowly varying ramps are +# reconstructed more faithfully. +# %% [markdown] +# ## Second example: Radial undersampling +# %% [markdown] +# We next reconstruct a multi-coil brain MRI from retrospectively undersampled radial spokes drawn on a Cartesian grid. +# We use 4-coil k-space data (`1_rawdata_brainT2_4ch.mat`) as a reference dataset +# (fully sampled in Cartesian coordinates). +# %% mystnb={"code_prompt_show": "Show download details"} tags=["hide-cell"] +# Download ground-truth k-space data from Zenodo +zenodo_get.download(record='800525', retry_attempts=5, output_dir=data_folder) + +# %% [markdown] +# To begin, we load the fully sampled k-space, form the image with an FFT operator, and compute a +# root-sum-of-squares (RSS) reference. +# %% +from einops import rearrange +from mrpro.data import SpatialDimension +from scipy.io import loadmat + +file_name = '1_rawdata_brainT2_4ch.mat' +kdata_true = torch.tensor(loadmat(data_folder / file_name)['rawdata']) +kdata_true = rearrange(kdata_true, 'k1 k0 coils -> 1 coils 1 k1 k0') +# kdata = torch.flip(kdata, dims=(-2, -1)) +kdata_true = kdata_true.to(dtype=torch.complex64) # If default dtype is torch.float32, use complex64 + +recon_matrix = SpatialDimension(z=kdata_true.shape[-3], y=kdata_true.shape[-2], x=kdata_true.shape[-1]) +encoding_matrix = SpatialDimension(z=kdata_true.shape[-3], y=kdata_true.shape[-2], x=kdata_true.shape[-1]) + +fourier_op = mrpro.operators.FastFourierOp( + dim=(-2, -1), + recon_matrix=recon_matrix, + encoding_matrix=encoding_matrix, +) +x_true = fourier_op(kdata_true)[0] + +x_true_rss = x_true.abs().square().sum(dim=-4).sqrt().squeeze() +show_images(x_true_rss, titles=['Ground Truth'], clim=(0, 7e-4)) + +# %% [markdown] +# ### Set up the operator $A$ +# %% [markdown] +# To set up the acquisition operator $A$, we first create a radial mask with a chosen number of spokes (here 48) on the +# Cartesian grid and wrap it as a `CartesianMaskingOp`. +# %% +from torchvision.transforms.functional import rotate + + +def radial_mask(ny: int, nx: int, num_spokes: int) -> torch.Tensor: + """Generate a radial mask with the specified number of spokes.""" + theta = 180 * (3.0 - 5**0.5) # golden angle ~137.508° + mask = torch.zeros((1, 1, ny, nx), dtype=torch.bool) + + # prototype spoke: horizontal line through center + base_spoke = torch.zeros((1, 1, ny, nx)) + base_spoke[0, 0, ny // 2, :] = 1.0 + + for i_spoke in range(num_spokes): + spoke = rotate(base_spoke, angle=theta * i_spoke, fill=0) + mask |= spoke > 0.5 + return mask + + +Nx, Ny = 256, 256 +num_spokes = 48 +mask_op = mrpro.operators.CartesianMaskingOp(radial_mask(Ny, Nx, num_spokes)) +show_images(torch.tensor(mask_op.mask), titles=[f'Mask ({num_spokes} spokes)']) + +# %% [markdown] +# Then we can define the acquisition operator as +# $$ +# A \,=\, M \circ \mathcal{F}, +# $$ +# where $\mathcal{F}$ is the (multi-dimensional) FFT mapping images to k-space, and $M$ applies the binary mask. +# We simulate undersampling by computing $y = A \ x_{\mathrm{true}}$, and display the adjoint reconstruction $A^* \ y$ +# (i.e. zero-filled inverse FFT). +# %% +acquisition_op = mask_op @ fourier_op +# Apply A(x_true) to get undersampled k-space data +kdata_undersampled = acquisition_op(x_true)[0] +x_adjoint_recon = acquisition_op.adjoint(kdata_undersampled)[0] +x_adjoint_recon_rss = torch.sum(x_adjoint_recon.abs() ** 2, dim=1).sqrt() +show_images(x_adjoint_recon_rss, x_true_rss, titles=['Adjoint', 'Ground Truth'], clim=(0, 7e-4)) + +# %% [markdown] +# ### Run PDHG with TV + +# %% [markdown] +# Again, let us start with the TV-regularized reconstruction method. The regularization parameter +# $\lambda = 5 \times 10^{-5}$ was chosen to achieve a good reconstruction for this example. + +# %% +x_tv_recon = tv_minimization_reconstruction( + measurement=kdata_undersampled, + acquisition_op=acquisition_op, + grad_term_weight=torch.tensor(5e-6), + max_iterations=257, + callback=callback, +) +x_tv_recon_rss = torch.sum(x_tv_recon.abs() ** 2, dim=-4).sqrt() +show_images(x_adjoint_recon_rss, x_tv_recon_rss, x_true_rss, titles=['Adjoint', 'TV', 'Ground Truth'], clim=(0, 7e-4)) + +# %% [markdown] +# ### Run PDHG with TGV + +# %% [markdown] +# Now we can run the TGV-regularized reconstruction with $(\lambda_1, \lambda_0) = (5 \times 10^{-6}, 10^{-5})$. + +# %% +x_tgv_recon = tgv_minimization_reconstruction( + measurement=kdata_undersampled, + acquisition_op=acquisition_op, + grad_term_weight=torch.tensor(5e-6), + sym_grad_term_weight=torch.tensor(10e-6), + max_iterations=257, + callback=callback, +) +x_tgv_recon_rss = torch.sum(x_tgv_recon.abs() ** 2, dim=-4).sqrt() +show_images( + x_adjoint_recon_rss, + x_tv_recon_rss, + x_tgv_recon_rss, + x_true_rss, + titles=['Adjoint', 'TV', 'TGV', 'Ground Truth'], + clim=(0, 7e-4), +) + +# %% [markdown] +# ### Compare the results +# We compare adjoint (zero-filled), TV, TGV, and the ground truth (RSS) on a zoomed-in portion. +# As expected, adjoint reconstruction shows aliasing from undersampling. +# TV reduces aliasing but introduces some staircasing, +# while TGV is able to reduce aliasing without the staircasing effect. +# %% +show_images( + x_adjoint_recon_rss[..., :128, :128], + x_tv_recon_rss[..., :128, :128], + x_tgv_recon_rss[..., :128, :128], + x_true_rss[..., :128, :128], + titles=['Adjoint', 'TV', 'TGV', 'Ground Truth'], + clim=(0, 7e-4), +) + + +# %% [markdown] +# That`s it — we performed denoising and MRI reconstruction via TGV-minimization solved with PDHG. +# We can later try different $(\lambda_1,\lambda_2)$ value combinations and adjust the iteration count +# to see how the reconstruction quality changes.