PnP_ULA_notebook.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5f612dcb",
   "metadata": {},
   "source": [
    "# Inverse Problems : The Bayesian Approach with Plug and Play (PnP) priors\n",
    "\n",
    "This numerical tour follows a Bayesian approach with Plug and Play (PnP) priors for solving inverse problems. These data-driven approach make use of learned implicit representation of the prior density while keeping an explicit likelihood density. Our aim is to perform image deconvolution and uncertainty quantification using [Langevin algorithms](https://epubs.siam.org/doi/abs/10.1137/21M1406349?journalCode=sjisbi) for sampling under PnP priors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f6bfc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import time as time\n",
    "from tqdm.auto import tqdm\n",
    "from PIL import Image\n",
    "\n",
    "from sampling_tools import *\n",
    "\n",
    "# Check if there's a GPU available and run on GPU, otherwise run on CPU\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "38c5335e",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "For many problems in imaging, we are interested in recovering an image  $x \\in \\mathbb{R}^d$, from a set of measurements $y \\in \\mathbb{R}^d$. We consider measurement models of the form \n",
    "\n",
    "$$y = Ax + w,$$ \n",
    "\n",
    "where $A \\in \\mathbb{R}^{d \\times d}$ is a linear operator and $w \\sim N(0, \\sigma^2 I_d)$ is a noise term.  Here, we study the case where both $A$ and $\\sigma>0$ are known.  The recovery of $x$ from $y$ is often ill posed or ill conditioned, so we regularize with prior knowledge.\n",
    "\n",
    "The measurement model defines a likelihood function $p(y|x)\\propto e^{-f_{y}(x)}$, where $f_{y}$ (negative log-likelihood) is convex and continuously differentiable with $L_{y}$-Lipschitz gradient. The prior knowledge about $x$ is encoded by a prior distribution $p(x)$. Then, from the Bayes Theorem, the posterior distribution $p(x|y)$ is defined as\n",
    "\n",
    "$$p(x|y) = \\dfrac{p(y|x)p(x)}{\\int_{\\mathbb{R}^d}p(y|\\tilde{x})p(\\tilde{x})d\\tilde{x}} ~~~,$$\n",
    "\n",
    "which underpins all inference about $x$ given the observation $y$. \n",
    "\n",
    "Many works in the Bayesian imaging literature consider relatively simple handcrafted priors promoting sparsity in transformed domains or piecewise regularity (e.g., involving the $\\ell_1$ norm or the total variation pseudonorm). Special  attention  is  given  in  the  literature  to  models that are log-concave, as this enables the use of Bayesian computation  algorithms  that  scale  efficiently  to  high  dimensions  and  which  have  detailed convergence guarantees. For a relevant implementation of such models check the notebooks for [MYULA](https://github.com/MI2G/sampling-tutorials) and [SKROCK](https://github.com/MI2G/sampling-tutorials).\n",
    "\n",
    "Our aim here is to calculate the Minimum Mean Square Error (MMSE) Bayesian estimator defined as\n",
    "\n",
    "$$\\hat{x}_{MMSE} = \\operatorname{argmin}_{u\\in\\mathbb{R}^{d}}\\mathbb{E}\\left[||x-u||^{2}|y\\right]= \\mathbb{E}(x|y) = \\int_{\\mathbb{R}^{d}}\\tilde{x}p(\\tilde{x}|y)d\\tilde{x}.$$\n",
    "\n",
    "Moreover, as an illustration of an uncertainty visualisation analysis, we also calculate the posterior variance for each image pixel $x_{i}$, for $i = 1,\\ldots,d$, defined as\n",
    "\n",
    "$$Var(x_i|y) = \\mathbb{E}(x_i^{2}|y) - (\\mathbb{E}(x_i|y))^{2},$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\\mathbb{E}(x_i^{2}|y)=\\int_{\\mathbb{R}^{d}}\\tilde{x_i}^{2}p(\\tilde{x_i}|y)d\\tilde{x_i}.$$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e8310c03",
   "metadata": {},
   "source": [
    "## Setting up a deconvolution problem\n",
    "In this example, we have a deconvolution problem where the operator $A$ corresponds to a blurring operator. Deconvolution corresponds to removing the blur from the image. First we load the image to be used for the rest of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1160daf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load grayscale image\n",
    "im = np.array(Image.open(\"cman.png\")) \n",
    "plot_im(im, \"ground truth image\")\n",
    "\n",
    "# Convert to torch tensor\n",
    "x = torch.Tensor(im/255.).to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "57f67eda",
   "metadata": {},
   "source": [
    "We define the convolution kernel $A$ in Fourier space. Here a $5\\times 5$ uniform blur is used. The `blur_operators()` function returns the functions handles $A$ and $A^T$ that implement the forward operations $Ax$ and $A^Tx$ respectively (computed by calling $A(x)$ and $AT(x)$) as well as the spectral norm $||AA^T||_{2}^{2}$.\n",
    "\n",
    "The calculation of $||AA^T||_{2}^{2}$ is necessary since it appears in the Lipschitz constant of the model and plays a role for the stepsize of the presented algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2639a612",
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel_len = [5,5]\n",
    "size = [im.shape[0],im.shape[1]]\n",
    "type_blur = \"uniform\"\n",
    "A, AT, AAT_norm = blur_operators(kernel_len, size, type_blur, device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bb1f66c7",
   "metadata": {},
   "source": [
    "We apply the blur operator and add some noise to obtain the measurements $y = Ax + w$ and achieve a blurred signal-to-noise ratio (BSNR) of $40$ dB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfa0a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "y0 = A(x)\n",
    "\n",
    "BSNRdb = 40\n",
    "sigma = torch.linalg.matrix_norm(A(x)-torch.mean(A(x)), ord='fro')/math.sqrt(torch.numel(x)*10**(BSNRdb/10))\n",
    "\n",
    "y = y0 + sigma * torch.randn_like(x)\n",
    "plot_im(y, \"noisy and blurry observation y\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "98e90cab",
   "metadata": {},
   "source": [
    "# Markov Chain Monte Carlo (MCMC) methods in imaging inverse problems\n",
    "\n",
    "It is clear that the calculation of the integrals defined above is a computationally infeasible task in high dimensions. Stochastic simulation algorithms, namely Markov Chain Monte Carlo (MCMC) algorithms, estimate these integrals by sampling. For example, if we could get i.i.d samples $\\{X_{i}\\}_{i=1}^{N}$ from $p(x|y)$ the posterior mean is approximated by\n",
    "\n",
    "$$\\hat{x}_{MMSE}\\approx\\dfrac{1}{N}\\sum_{i=1}^{N}X_{i},$$\n",
    "\n",
    "with the Monte Carlo approximation error vanishing as $N → ∞$.\n",
    "\n",
    "The question that arises here is how can we sample from the posterior. Under a smooth log-concave prior, a common approach is to solve the overdamped Langevin stochastic differential equation (SDE)\n",
    "\n",
    "$$dX_{t} = \\nabla \\log p(X_{t}|y) dt + \\sqrt{2}dW_{t} = \\nabla \\log p(y|X_{t}) dt+ \\nabla \\log p(X_{t}) dt+ \\sqrt{2}dW_{t} ~~,$$\n",
    "\n",
    "where $(W_{t})_{t\\geq0}$ is a $d$-dimensional Brownian motion. Under mild conditions, the above SDE has a unique strong solution $(X_{t})_{t>0}$ that admits the posterior of interest $p(x|y)$ as its  unique stationary density.\n",
    "\n",
    "In general, one cannot solve the overdamped Langevin SDE analytically and needs to resort to the use of numerical methods. The simplest numerical method one can employ is the Euler-Maruyama method, which gives rise to the unadjusted Langevin algorithm (ULA), which defines a Markov Chain $\\{ X_{k},  k \\in \\mathbb{N}\\}$ using the following one step recursion \n",
    "\n",
    "\n",
    "$$X_{k+1} = X_{k} + \\delta\\nabla \\log p(y|X_{k}) + \\delta\\nabla \\log p(X_{k})+ \\sqrt{2\\delta}Z_{k+1} ~~,$$\n",
    "\n",
    "where $\\{Z_{k} : k\\in \\mathbb{N}\\}$ is a family of i.i.d $d$-dimensional Gaussian random variables with zero mean and identity covariance matrix and $\\delta$ is the time-step. When the prior $p(x)$ is log-concave but not smooth, one can approximate the gradient of $U(x) = -\\log p(x)$ by the gradient of the Moreau Yosida (MY) envelope of $U_\\lambda(x)$ given by $\\nabla U_\\lambda(x) = \\frac{1}{\\lambda}(x-\\mathrm{prox}_U^\\lambda(x))$. Then, the Moreau Yosida ULA (MYULA) can be used as\n",
    "\n",
    "$$X_{k+1} = X_{k} + \\delta\\nabla \\log p(y|X_{k}) + \\frac{\\delta}{\\lambda}(\\mathrm{prox}_U^\\lambda(X_{k})-X_{k})+ \\sqrt{2\\delta}Z_{k+1} ~~.$$\n",
    "\n",
    "For implementation of MY samplers under log-concave priors check the notebooks here and here.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cd29e8a8",
   "metadata": {},
   "source": [
    "## Likelihood\n",
    "\n",
    "Since, we assumed Gaussian noise, the likelihood $p(y|x)$ is Gaussian given from the forward model $y = Ax + w$. Here we keep track of the negative log-likelihood $$-\\log p(y|x) = f_y(x) =\\dfrac{1}{2\\sigma^2}||y-Ax||^{2}_{2}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679d2b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x,A : (torch.linalg.matrix_norm(y-A(x), ord='fro')**2.0)/(2.0*sigma**2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "652e4b2d",
   "metadata": {},
   "source": [
    "Define the gradient of the negative log-likelihood $f_{y}$.\n",
    "$$\\nabla f_y(x) = A^T(Ax - y).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6d2ae09",
   "metadata": {},
   "outputs": [],
   "source": [
    "gradf = lambda x,A,AT : AT(A(x)-y)/sigma**2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "05b2eacc",
   "metadata": {},
   "source": [
    "The Lipschitz constant $L_{y}$ related to the specific likelihood's gradient (useful for the stepsize choice) is $\\dfrac{||AA^T||_2}{\\sigma^2}$. We calculated $||AA^T||_2$ earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd432501",
   "metadata": {},
   "outputs": [],
   "source": [
    "L_y = AAT_norm/(sigma**2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2162c4a3",
   "metadata": {},
   "source": [
    "## Introduce a Plug and Play prior\n",
    "\n",
    "We are now ready to outline the way PnP priors can be used for Bayesian computation. For more details look the work of [Remi and al.](https://epubs.siam.org/doi/abs/10.1137/21M1406349?journalCode=sjisbi). Suppose the existence of an optimal and intractable prior distribution over $x$, namely $p^*(x)$, and an optimal posterior distribution over $x|y$, namely $p^*(x|y)$. The posterior distribution associated with the optimal $p^*(x)$ is given as \n",
    "\n",
    "$$p^*(x|y) = \\dfrac{p(y|x)p^*(x)}{\\int_{\\mathbb{R}^d}p(y|\\tilde{x})p^*(\\tilde{x})d\\tilde{x}} ~~~,$$\n",
    "\n",
    "Since $p^*$ and its respective Lebesque measure are unknown, it is not straightforward to verify that $p^*(x|y)$ is proper and differentiable or that $\\nabla \\log p^*(x|y)$ is Lipschitz continuous - an important feature for convergent gradient-based algorithms. Thus, we approximate $p^*$ by a smooth density $p_\\epsilon^*=G_\\epsilon p^*$ where $G_\\epsilon$ is a Gaussian smoothing kernel of bandwidth $\\epsilon>0$. Then, by definition\n",
    "\n",
    "$$p_\\epsilon^*(x) = (2\\pi\\epsilon)^{-d/2} \\int_{\\mathbb{R}^d} \\exp\\left[-\\frac{||x-\\tilde{x}||_2^2}{2\\epsilon}\\right]p^*(\\tilde{x})d\\tilde{x} ~~~,$$\n",
    "\n",
    "$$p_\\epsilon^*(x|y) = \\dfrac{p(y|x)p_\\epsilon^*(x)}{\\int_{\\mathbb{R}^d}p(y|\\tilde{x})p_\\epsilon^*(\\tilde{x})d\\tilde{x}} ~~~.$$\n",
    "\n",
    "\n",
    "The core idea of Bayesian computation with PnP priors is to perform approximate inference w.r.t. $p_\\epsilon^*(x|y)$ and consequently make use of the $\\nabla \\log p_\\epsilon^*(x|y)$. Specifically, we introduce the oracle MMSE denoiser $D^{*}_{\\epsilon}$ defined as \n",
    "\n",
    "$$D^{*}_{\\epsilon}(x) = (2\\pi\\epsilon)^{-d/2} \\int_\\mathbb{R}  \\tilde{x}\\exp\\left[-\\frac{||x-\\tilde{x}||_2^2}{2\\epsilon}\\right]p^*(\\tilde{x})d\\tilde{x} ~~~.$$\n",
    "\n",
    "It can be seen that $D^{*}_{\\epsilon}$ is the MMSE estimator to recover an image $x \\sim p^*(x)$ from a noisy observation $x_\\epsilon\\sim \\mathcal{N}(x, \\epsilon I)$. To associate $D^{*}_{\\epsilon}$ with the gradient $x \\rightarrow \\nabla \\log p_\\epsilon^*(x)$, we make use of the Tweedie's identity which states that\n",
    "\n",
    "$$\\nabla \\log p_\\epsilon^*(x) = \\frac{1}{\\epsilon}(D^{*}_{\\epsilon}(x) - x) ~~~.$$\n",
    "\n",
    "Then, we can  express  a  ULA  recursion  for  sampling approximately from $p^*_\\epsilon(x|y)$ as\n",
    "\n",
    "$$X_{k+1} = X_{k} + \\delta\\nabla \\log p^*_\\epsilon(X_{k}|y ) + \\sqrt{2\\delta}Z_{k+1}  = X_{k} + \\delta\\nabla \\log p(y|X_{k}) + \\frac{\\delta}{\\epsilon}(D_\\epsilon^*(X_{k})-X_{k}) + \\sqrt{2\\delta}Z_{k+1} ~~.$$\n",
    "\n",
    "Since $D_\\epsilon^*(\\cdot)$ is intractable, we use the above equation as a starting point and we replace $D_\\epsilon^*(\\cdot)$ by a surrogate denoiser $D_\\epsilon(\\cdot)$ to approximate $ p_\\epsilon^*(x|y)$ by the posterior density $p_\\epsilon(x|y)$ related to $D_\\epsilon(\\cdot)$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "be811a2b",
   "metadata": {},
   "source": [
    " ### Load the denoiser"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9ddb393f",
   "metadata": {},
   "source": [
    "Below, we load the denoiser $D_\\epsilon(\\cdot)$ which is a pretrained residual network presented in [Ryu and al.](https://proceedings.mlr.press/v97/ryu19a) with controlled Lipschitz constant $L = 1$. The Lipschitz continuity of the denoiser is important since this guarantees that $\\nabla \\log p_\\epsilon(x|y)$ is also Lipschitz. Since the network is residual, $D_\\epsilon(\\cdot) = x - D'_\\epsilon(\\cdot)$, where $D'_\\epsilon(\\cdot)$ is the loaded network, denoted as `denoise()` and `model()` respectively. Finally, we assume that the denoiser $D'_\\epsilon(\\cdot)$ is pretrained at noise level of $\\epsilon=5/255$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c737ff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "L_net = 1.0\n",
    "# keep in mind that we've got a residual network\n",
    "model = load_model(\"RealSN_DnCNN\", 5, device)\n",
    "denoise = lambda x: (x - model(x[None][None])[0][0]).detach() "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c22421f5",
   "metadata": {},
   "source": [
    "## PnP Algorithms"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "844d5e25",
   "metadata": {},
   "source": [
    "\n",
    "For convergence verification reasons, instead of targeting the posterior $p_{\\epsilon}(X_{t}|y)$, we will target $p_{\\epsilon,C}(X_{t}|y)$ where\n",
    "\n",
    "$$p_{\\epsilon,C}(X_{t}|y)=p_{\\epsilon}(X_{t}|y)e^{-d^{2}(x,C)/\\lambda}\\,\\,,$$\n",
    "\n",
    "where $C\\subset \\mathbb{R}^{d}$ is some large compact convex set that controls the far tail behaviour of $p_{\\epsilon,C}$ and $d(x, C)$ is the minimum Euclidean distance between $x$ and the set $C$, and $\\lambda>0$ is a tail regularisation parameter. It is easy to check that\n",
    "\n",
    "$$\\nabla \\log p_{\\epsilon,C}(X_{t}|y) = \\nabla \\log p_{\\epsilon}(X_{t}|y) + \\dfrac{1}{\\lambda}(\\Pi_{C}(X_{t})-X_{t})\\,\\,,$$\n",
    "\n",
    "where $\\Pi_{C}(\\cdot)$ denotes the projection function to the convex set $C$. We define **PnP-ULA** as the following recursion:  \n",
    "\n",
    "$$X_{k+1} = X_{k} + \\delta\\nabla \\log p_{\\epsilon, C}(X_{k}|y) + \\sqrt{2\\delta}Z_{k+1}  = X_{k} + \\delta\\nabla \\log p(y|X_{k}) + \\frac{\\alpha\\delta}{\\epsilon}(D_\\epsilon(X_{k})-X_{k}) + \\dfrac{\\delta}{\\lambda}(\\Pi_{C}(X_{t})-X_{t}) + \\sqrt{2\\delta}Z_{k+1} ~~.$$\n",
    "\n",
    "Alternative strategy is to modify PnP-ULA to include a hard projection onto $C$, and define the **projected  PnP-ULA (PPnP-ULA)** as the following recursion:\n",
    "\n",
    "$$X_{k+1} = \\Pi_{C}(X_{k} + \\delta\\nabla \\log p_{\\epsilon}(X_{k}|y) + \\sqrt{2\\delta}Z_{k+1})  = \\Pi_{C}\\left(X_{k} + \\delta\\nabla \\log p(y|X_{k}) + \\frac{\\alpha\\delta}{\\epsilon}(D_\\epsilon(X_{k})-X_{k}) + \\sqrt{2\\delta}Z_{k+1}\\right) ~~.$$\n",
    "\n",
    "It should be noted that the hard projection guarantees geometric convergence with weaker restrictions on $\\delta$ and hence PPnP-ULA can converge faster than PnP-ULA, with the expense of more bias. Finally, note the presence of a regularisation parameter $\\alpha$ in these algorithms, which permits to balance the weights between the prior and data terms."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3eb6ef94",
   "metadata": {},
   "source": [
    "### Setting parameters"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c3fa7483",
   "metadata": {},
   "source": [
    "For this notebook, we set the noise level $\\epsilon$ of the denoiser to be $5/255$. We also set $\\lambda = (2L/\\epsilon + 4L_y)^{-1}$, and $C = [0,1] $. These are guidelines suggested [here](https://epubs.siam.org/doi/abs/10.1137/21M1406349?journalCode=sjisbi)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4514af15",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1\n",
    "eps =  (5/255)**2\n",
    "max_lambd = 1.0/((2.0*alpha*L_net)/eps+4.0*L_y)\n",
    "lambd_frac = 0.99\n",
    "lambd = max_lambd*lambd_frac\n",
    "\n",
    "C_upper_lim = torch.tensor(1).to(device)\n",
    "C_lower_lim = torch.tensor(0).to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3e0d4798",
   "metadata": {},
   "source": [
    "### PnP-ULA and PPnP-ULA kernel updates\n",
    "\n",
    "Below, we define the kernel updates based on the algorithms presented above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa8e0d13",
   "metadata": {},
   "outputs": [],
   "source": [
    "projbox = lambda x: torch.clamp(x, min = C_lower_lim, max = C_upper_lim)\n",
    "\n",
    "def Markov_kernel(X, delta, projected):\n",
    "    if projected:\n",
    "        return projbox(X - delta * gradf(X,A,AT) + alpha*delta/eps*(denoise(X)-X) + math.sqrt(2*delta) * torch.randn_like(X))\n",
    "    else:\n",
    "        return X - delta * gradf(X,A,AT) + alpha*delta/eps*(denoise(X)-X) + delta/lambd*(projbox(X)-X) + math.sqrt(2*delta) * torch.randn_like(X)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b105070e",
   "metadata": {},
   "source": [
    "### Setting the stepsize\n",
    "\n",
    "It should be noted that $\\delta$ is subject to a numerical stability constraint related to the inverse of the Lipschitz constant of $\\nabla \\log p_\\epsilon(x|y)$. For **PnP-ULA** we require $\\delta < (1/3) L_{total}^{-1}$, where $ L_{total} = (\\alpha L)/\\epsilon+L_y+1/\\lambda$. For **PPnP-ULA**, we require $\\delta < (L/\\epsilon+L_y)^{-1}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54f9c363",
   "metadata": {},
   "outputs": [],
   "source": [
    "projected = True\n",
    "\n",
    "if projected:\n",
    "    delta_max = (1.0)/(L_net/eps+L_y)\n",
    "else:\n",
    "    delta_max = (1.0/3.0)/((alpha*L_net)/eps+L_y+1/lambd)\n",
    "delta_frac = 0.99\n",
    "delta = delta_max*delta_frac"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "afaef368",
   "metadata": {},
   "source": [
    "### Bayesian computations\n",
    "\n",
    "As mentioned before our aim is to calculate the posterior mean as\n",
    "$$\\hat{x}_{MMSE} = \\mathbb{E}(x|y) = \\int_{\\mathbb{R}^{d}}\\tilde{x}p_\\epsilon(\\tilde{x}|y)d\\tilde{x}\\approx \\dfrac{1}{N}\\sum_{i=1}^{N}X_{i}$$\n",
    "\n",
    "and the posterior variance for each image pixel $x_{i}$, for $i = 1,\\ldots,d$, as\n",
    "\n",
    "$$\\operatorname{Var}(x_i|y) = \\mathbb{E}(x_i^{2}|y) - (\\mathbb{E}(x_i|y))^{2} \\approx \\dfrac{1}{N}\\sum_{i=1}^{N}X_{i}^2 - \\left(\\dfrac{1}{N}\\sum_{i=1}^{N}X_{i}\\right)^2\\,\\,.$$\n",
    "\n",
    "Additionally, we are keeping track of the reconstruction quality of the MMSE estimate by calculating the Normalized Mean Square Error (NRMSE), the Peak-Signal-to-Noise Ratio (PSNR) and the Structural Similarity Index Measure (SSIM) w.r.t. the ground truth image."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "94eb2c4d",
   "metadata": {},
   "source": [
    "### Algorithm choices and initializations\n",
    "\n",
    "Choose the number of sampling iterations, here it is $10^{4}$. A burn-in period of $5\\%$ is also included.  The Markov chain is initialized with the observation $y$. In total, we would like to save a trace of 2000 samples, so we set the thinning variable accordingly. In practice, you will need to run the sampling for longer (500k-1m iterations) for uncertainty quantification (UQ) tasks, but this is enough to illustrate how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777d586a",
   "metadata": {},
   "outputs": [],
   "source": [
    "maxit = 10000\n",
    "burnin = np.int64(maxit*0.05)\n",
    "n_samples = np.int64(2000)\n",
    "X = y.clone()\n",
    "MC_X = []\n",
    "thinned_trace_counter = 0\n",
    "thinning_step = np.int64(maxit/n_samples)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "933354d5",
   "metadata": {},
   "source": [
    "### Quality metrics\n",
    "Keep track of the NRMSE, PSNR, SSIM w.r.t. to the ground truth image and the log-posterior on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cda1b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "nrmse_values = []\n",
    "psnr_values = []\n",
    "ssim_values = []"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d761010a",
   "metadata": {},
   "source": [
    "### The Algorithm in Practice\n",
    "\n",
    "1. Each iteration is calculated using the MYULA kernel.\n",
    "\n",
    "2. After the burning-period we start calculating the posterior sample mean and variance on the fly by updating them in each iteration. We are using the `welford` class for this. The `update()` function of the class updates the sample mean and variance in each iteration.\n",
    "\n",
    "3. We also keeping track of the samples in the Fourier domain calculating respective the posterior mean and variance on the fly in the Fourier domain. This can be accomplished by just giving as input in the `welford` class the absolute value of the FT of the samples.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63fbc2b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "for i_x in tqdm(range(maxit)):\n",
    "\n",
    "    # Update X\n",
    "    X = Markov_kernel(X, delta, projected=projected)\n",
    "\n",
    "    if i_x == burnin:\n",
    "        # Initialise recording of sample summary statistics after burnin period\n",
    "        post_meanvar = welford(X)\n",
    "        absfouriercoeff = welford(torch.fft.fft2(X).abs())\n",
    "        count=0\n",
    "    elif i_x > burnin:\n",
    "        # update the sample summary statistics\n",
    "        post_meanvar.update(X)\n",
    "        absfouriercoeff.update(torch.fft.fft2(X).abs())\n",
    "\n",
    "        # collect quality measurements\n",
    "        current_mean = post_meanvar.get_mean()\n",
    "        nrmse_values.append(NRMSE(x, current_mean))\n",
    "        psnr_values.append(PSNR(x, current_mean))\n",
    "        ssim_values.append(SSIM(x, current_mean))\n",
    "\n",
    "        # collect thinned trace\n",
    "        if count == thinning_step-1:\n",
    "            MC_X.append(X.detach().cpu().numpy())\n",
    "            count = 0\n",
    "        else:\n",
    "            count += 1\n",
    "\n",
    "end_time = time.time()\n",
    "elapsed = end_time - start_time       "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4197c370",
   "metadata": {},
   "source": [
    "Evaluate metrics of noisy image $y$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48d4610",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Initial NRMSE: {NRMSE(x,y):.4f}\")\n",
    "print(f\"Initial PSNR: {PSNR(x,y):.2f} dB\")\n",
    "print(f\"Initial SSIM: {SSIM(x,y):.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "50936d1f",
   "metadata": {},
   "source": [
    "Evaluate metrics of posterior mean.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef272967",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Result NRMSE: {NRMSE(post_meanvar.get_mean(),x):.4f}\" )\n",
    "print(f\"Result PSNR: {PSNR(post_meanvar.get_mean(),x):.2f} dB\")\n",
    "print(f\"Result SSIM: {SSIM(post_meanvar.get_mean(),x):.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cae0887a",
   "metadata": {},
   "source": [
    "## Plotting Results\n",
    "We display below the ground truth and noisy and blurry image, the posterior mean (**MMSE** solution of the inverse problem), the posterior variance (which is useful for quantifying the uncertainties) and further statistics in the second row. More specifically, we display ratios of the posterior mean over the posterior standard deviation (SD, taking the root of the variance) as well as the ratio of the SD over the posterior mean in the spatial domain (so called coefficients of variation), and the mean and variance in log scale in Fourier domain. The following row shows the image quality metrics (NMRSE, PSNR, SSIM) for the cumulative mean, which tells us about the convergence of the algorithm (when the curves are completely flattened out)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdc11221",
   "metadata": {},
   "outputs": [],
   "source": [
    "plots(x, y, post_meanvar, absfouriercoeff, nrmse_values, psnr_values, ssim_values)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3b9722a6",
   "metadata": {},
   "source": [
    "## Autocorrelation function of the Markov chain\n",
    "\n",
    "To analyse the convergence properties of the Markov chain generated with `ULA`, we display below the autocorrelation function of the slowest, median and fastest component of the resulting chain. We obtained the slowest component by computing the pixel index resulting in the smallest variance, and extracting a trace from the Markov chain for this pixel. Similarly, the median and fastest pixels are obtained by extracting traces from the Markov chain with the median and largest variance respectively. From where the autocorrelation drops to zero, we can see how many samples we need to generate another independent sample. Therefore we would like to see fast decay.\n",
    "\n",
    "In addition, in the legend we give information about the effective sample size (ESS). This is also an important quantity that tells us about the quality of the chain. This number gives an indication how many independent samples are generated by the present Markov chain, and a higher number is better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a57c0fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "MC_X = np.stack(MC_X)\n",
    "autocor_plots(MC_X, \"PnP-ULA\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "879d313b",
   "metadata": {},
   "source": [
    "## Uncertainty quantification\n",
    "\n",
    "We would like to use the Markov chain to visualize the variance at different scales, giving us an indication how sure the algorithm is about the solution superpixels. Dark means lower variance and depict areas where the algorithm is sure about the solution, bright areas indicate high variance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac7f34af",
   "metadata": {},
   "outputs": [],
   "source": [
    "downsampling_variance(MC_X)"
   ]
  }
 ],
 "metadata": {
  "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.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "a6c315f1f155508bf54fbe27284b29881dcf6b46de61a5c58cc7ce542ce59b45"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}