empirical_bayesian_estimation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "PC62m42-sIrB",
   "metadata": {
    "id": "PC62m42-sIrB"
   },
   "source": [
    "# Inverse Problems: The Bayesian Approach \n",
    "\n",
    "For many problems in imaging, we are interested in recovering an image of interest, $x \\in \\mathbb{R}^d$, from a set of measurements $y \\in \\mathbb{R}^d$. These two are related through a statistical model with likelihood function\n",
    "of the form\n",
    "\n",
    "$$p(y|x,\\sigma^2) = e^{-f^{\\sigma^2}_{y}(x)}\\,\\,,$$\n",
    "\n",
    "where $f^{\\sigma^2}_{y}$ is convex and continuously differentiable with $L_{y}$-Lipschitz gradient. This class includes important observation models, in particular Gaussian linear 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 the noise variance $\\sigma>0$ are known. \n",
    "\n",
    "However, the recovery of $x$ from $y$ is often ill posed or ill conditioned, so regularisation is required in order to deliver meaningful solutions. In the Bayesian framework, this is achieved by using prior knowledge about $x$. We consider prior distributions given for any $x\\in\\mathbb{R}^{d}$ and $\\theta\\in \\Theta_\\theta\\subset (0,+\\infty)$ by\n",
    "\n",
    "$$p(x|\\theta) = e^{-\\theta g(x)}/Z(\\theta)\\,\\,,$$\n",
    "\n",
    "for some function $g: \\mathbb{R}^{d} → \\mathbb{R}$ that is proper, convex but potentially not smooth. The normalising constant $Z(\\theta)$ is given by\n",
    "\n",
    "$$Z(\\theta)=\\int_{\\mathbb{R}^{d}}e^{-\\theta g(\\tilde{x})}d\\tilde{x}\\,\\,.$$\n",
    "\n",
    "The parameter $\\theta$ controls the amount of regularity enforced. This parameter is difficult to set a priori, ***so the purpose of this tutorial is to implement a stochastic approximation proximal gradient algorithm to estimate $\\theta$ and $\\sigma^2$ directly from $y$ by maximum marginal likelihood estimation.***"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "VcoKttKu81tz",
   "metadata": {
    "id": "VcoKttKu81tz"
   },
   "source": [
    "# Bayesian Inference\n",
    "\n",
    "Once the likelihood and prior $p(y|x,\\sigma^2)$ and $p(x|\\theta)$ are defined, we use Bayes’ theorem to derive the posterior for any $x\\in\\mathbb{R}^{d}$, $\\theta\\in\\Theta_\\theta$ and $\\sigma^2\\in\\Theta_{\\sigma^2}\\,\\,,$\n",
    "\n",
    "$$p(x|y,\\theta,\\sigma^2) = \\dfrac{p(y|x,\\sigma^2)p(x|\\theta)}{p(y|\\theta,\\sigma^2)} \\propto e^{-f^{\\sigma^2}_{y}(x) + \\theta g(x)}\\,\\,.$$\n",
    "\n",
    "The posterior distribution describe how $x$ is affected by the data $y$.\n",
    "Many imaging methods typically use the maximum-a-posteriori (MAP) estimator, given for any $\\theta\\in\\Theta_\\theta$ and $\\sigma^2\\in\\Theta_{\\sigma^2}$ by\n",
    "\n",
    "$$\\hat{x}_{\\theta,\\sigma^2,MAP} = \\mathrm{arg min}_{\\tilde{x}\\in\\mathbb{R}^{d}}\\{f^{\\sigma^2}_{y}(\\tilde{x})+\\theta g(\\tilde{x})\\}$$\n",
    "\n",
    "From a computation viewpoint, if the posterior $p(x|y,\\theta)$ is log-concave, the computation of $\\hat{x}_{\\theta,\\sigma^2,MAP}$ is a convex optimisation problem\n",
    "that can usually be efficiently solved using modern optimisation algorithms. However, $g$ might not be smooth. In this case, imaging MAP algorithms \n",
    "typically adopt a proximal splitting approach involving the gradient $\\nabla f^{\\sigma^2}_{y}$ and the proximal operator $\\mathrm{prox}_{g}^{\\lambda} : \\mathbb{R}^{d} \\rightarrow \\mathbb{R}^{d}$ of $g$. The operator is defined for any $\\lambda>0$ and $x\\in \\mathbb{R}^{d}$ as\n",
    "\n",
    "$$\\mathrm{prox}_{g}^{\\lambda}(x) = \\mathrm{arg min}_{\\tilde{x}\\in\\mathbb{R}^{d}}\\left\\{g(\\tilde{x}) + \\dfrac{1}{2\\lambda}||\\tilde{x}-x||^{2}\\right\\}\\,\\,.$$\n",
    "\n",
    "where the smoothness parameter $\\lambda>0$ controls the regularity properties of the proximal operator. In some cases, the proximal operator is given in closed-form. Otherwise, it is calculated by using specialized algorithms.\n",
    "\n",
    "In problems that are ill-posed or ill-conditioned, the value of $\\theta$ can significantly impact our inferences about $x$. Here we adopt an empirical Bayesian strategy that seeks to estimate $\\theta$ directly from $y$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "u32-tmX7waY0",
   "metadata": {
    "id": "u32-tmX7waY0"
   },
   "source": [
    "# Empirical Bayes estimation\n",
    "\n",
    "Adopting an empirical Bayesian approach, the regularisation parameter $\\theta$ is estimated directly from the observed data $y$ by maximum marginal likelihood estimation, i.e.,\n",
    "\n",
    "$$\\theta^{*}, \\sigma^{*2} = \\mathrm{argmax}_{\\theta\\in\\Theta_\\theta,\\sigma^2\\in\\Theta_{\\sigma^2}}\\,\\,p(y|\\theta,\\sigma^2)\\,\\,,$$\n",
    "\n",
    "where we recall that the marginal likelihood $p(y|\\theta,\\sigma^2)$ is given for any $\\theta\\in\\Theta_\\theta$ and $\\sigma^2\\in\\Theta_{\\sigma^2}$ by \n",
    "\n",
    "$$p(y|\\theta,\\sigma^2) = \\int_{\\mathbb{R}^{d}}p(y|\\tilde{x},\\sigma^2)p(\\tilde{x}|\\theta)d\\tilde{x}\\,\\,.$$\n",
    "\n",
    "***So, if we can calculate $\\theta^{*}$, we can then plug it into our imaging MAP algorithm!***\n",
    "\n",
    "**Suppose for now that $p(y|\\theta,\\sigma^2)$ was tractable** and that we had access to the gradients $\\nabla_{\\theta} \\log p(y|\\theta,\\sigma^2)$, and  $\\nabla_{\\sigma^2} \\log p(y|\\theta,\\sigma^2)$. Recalling that $\\Theta_\\theta$ and $\\Theta_{\\sigma^2}$ are a convex compact set, we could use the projected gradient algorithm. For $\\theta_0 \\in \\Theta_\\theta$ and $\\sigma^2_0 \\in \\Theta_{\\sigma^2}$, for all $n \\in \\mathbb{N}$\n",
    "\n",
    "$$\\theta_{n+1} = \\Pi_{\\Theta_\\theta}[\\theta_{n} - \\delta_{n}\\nabla_{\\theta}\\log p(y|\\theta_{n},\\sigma_n^2)]$$\n",
    "$$\\sigma^2_{n+1} = \\Pi_{\\Theta_{\\sigma^2}}[\\sigma^2_{n} - \\delta_{n}\\nabla_{\\sigma^2}\\log p(y|\\theta_{n},\\sigma_n^2)].$$\n",
    "Here $\\Pi_{\\Theta_\\theta}$ and $\\Pi_{\\Theta_{\\sigma^2}}$ are projection operators such that $\\Pi_{\\Theta_\\theta}(\\theta)$ is the closest point in $\\Theta$ to $\\theta$ ($\\Pi_{\\Theta_\\sigma^2}$ is defined similarly). \n",
    "**Since $\\nabla_{\\theta} \\log p(y|\\theta,\\sigma^2)$ and $\\nabla_{\\sigma^2} \\log p(y|\\theta,\\sigma^2)$ are not tractable** (due to the high dimensional integral in $p(y|\\theta,\\sigma^2)$), we cannot directly use the above algorithm and we choose a stochastic variant of the projected gradient algorithm. For this purpose, we use the Fisher Identity to rewrite the gradient as follows \n",
    "\n",
    "$$\\nabla_{\\theta}\\log p(y|\\theta,\\sigma^2) = -\\int_{\\mathbb{R}^{d}}g(\\tilde{x})p(\\tilde{x}|y,\\theta,\\sigma^2)d\\tilde{x} - \\nabla_{\\theta}\\log (Z(\\theta))\\,\\,,$$\n",
    "\n",
    "and by using Monte Carlo simulations\n",
    "\n",
    "$$\\nabla_{\\theta}\\log p(y|\\theta\\sigma^2) \\approx - \\dfrac{1}{m}\\sum_{i=1}^{m}g(X_{i}) - \\nabla_{\\theta}\\log (Z(\\theta)) = \\Delta_{m,\\theta}\\,\\,,$$\n",
    "\n",
    "where $(X_{k})_{k=1}^{m}$ is a sample of size $m\\in\\mathbb{N}^{*}$ generated by using a Markov Chain targeting $p(x|y,\\theta)$, or a regularised approximation of this density.\n",
    "\n",
    "**The last term that remains** is the $\\nabla_{\\theta}\\log (Z(\\theta))$. Assume that there exists an integer $\\alpha$ such that $g$ is $\\alpha$ positively homogeneous function, i.e. for any $x\\in\\mathbb{R}^{d}$ and $t>0$, $g(tx) = t^{\\alpha}g(x)$, and recalling that $ \\Theta_\\theta \\subseteq (0,+\\infty)$ we have for any $\\theta\\in\\Theta_\\theta$\n",
    "\n",
    "$$Z(\\theta) = \\int_{\\mathbb{R}^{d}}e^{-\\theta g(\\tilde{x})}d\\tilde{x} = \\int_{\\mathbb{R}^{d}}e^{- g(\\theta^{1/\\alpha}\\tilde{x})}d\\tilde{x} =\n",
    "\\theta^{-d/\\alpha}\\int_{\\mathbb{R}^{d}}e^{g(\\tilde{x})}d\\tilde{x}\\,\\,,$$\n",
    "\n",
    "and so,\n",
    "\n",
    "$$\\dfrac{d}{d\\theta}\\log Z(\\theta) = -\\dfrac{d}{\\alpha\\theta}.$$\n",
    "\n",
    "Now, we can calculate $\\Delta_{m,\\theta}= - \\dfrac{1}{m}\\sum_{i=1}^{m}g(X_{i}) + \\dfrac{d}{\\alpha\\theta}\\,\\,.$\n",
    "\n",
    "Similarly, we approximate the gradient $\\nabla_{\\sigma^2}\\log p(y|\\theta\\sigma^2)$ by\n",
    "\n",
    "$$\\nabla_{\\sigma^2}\\log p(y|\\theta,\\sigma^2) \\approx \\Delta_{m,\\sigma^2}= \\dfrac{1}{m}\\sum_{i=1}^{m}\\nabla_{\\sigma^2}f_y^{\\sigma^2}(X_i) - \\dfrac{d}{2\\sigma^2}$$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5IneXb5RSGU2",
   "metadata": {
    "id": "5IneXb5RSGU2"
   },
   "source": [
    "# Markov Chain Monte Carlo (MCMC) methodology in imaging inverse problems\n",
    "\n",
    "The question that arises here is how can we sample from the posterior? One common approach is to solve the overdamped Langevin stochastic differential equation (SDE)\n",
    "\n",
    "$$dX_{t} = \\nabla \\log p(X_{t}|y,\\theta,\\sigma^2) dt + \\sqrt{2}W_{t}=\n",
    "\\nabla \\log p(y|X_{t},\\sigma^2)dt + \\nabla \\log p(X_{t}|\\theta)dt+ \\sqrt{2}W_{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,\\theta,\\sigma^2)$ as unique stationary density.\n",
    "\n",
    "In high dimensions, it is typically infeasible to solve the above SDE, and therefore we need to use a numerical approximation. A natural choice of numerical scheme is the Unadjusted Langevin algorithm (ULA) Markov chain $(X_{k})_{k>0}$ obtained from an Euler-Maruyama discretisation of the above SDE\n",
    "\n",
    "$$X_{k+1} = X_{k} + \\gamma_{k}\\nabla \\log p(y|X_{k},\\sigma^2) + \\gamma_{k}\\nabla \\log p(X_{k}|\\theta)+ \\sqrt{2\\gamma_{k}}Z_{k+1}\\,\\,,$$\n",
    "\n",
    "where $\\{Z_{k} : k\\in \\mathbb{N}\\}$ is a family of i.i.d Gaussian random variables with zero mean and identity covariance matrix and $(\\gamma_{k})_{k\\in\\mathbb{N}}$ is a sequence of positive step-sizes.\n",
    "\n",
    "However, as we aforementioned, $\\log p(X_{k}|\\theta) = -\\theta \\cdot g(x)$ might not be smooth, and so the gradient is not available. Suppose $\\tilde{g}(x) = \\theta\\cdot g(x)$. To smooth $\\tilde{g}$, we will consider the Moreau-Yosida (MY) envelope of $\\tilde{g}$ defined as\n",
    "\n",
    "$$\\tilde{g}^{\\lambda}(x)=\\min_{u\\in\\mathbb{R}^{M}}\\{\\tilde{g}(u)+(2\\lambda)^{-1}||u-x||^{2}\\}$$\n",
    "\n",
    "The MY envelope is convex and always $L_{\\tilde{g}^{\\lambda}}-$ differentiable with $L_{\\tilde{g}^{\\lambda}}= \\dfrac{1}{\\lambda}$ and :\n",
    "\n",
    "$$\\nabla \\tilde{g}^{\\lambda}(x)=\\lambda^{-1}(x-\\mathrm{prox}_{\\tilde{g}}^{\\lambda}(x))$$\n",
    "\n",
    "where \n",
    "\n",
    "$$\\mathrm{prox}_{\\tilde{g}}^{\\lambda}(x)=\\mathrm{argmin}_{u\\in\\mathbb{R}^{M}}\\{g(u)+(2\\lambda\\theta)^{-1}||u-x||^{2}\\}$$\n",
    "\n",
    "Then, under a smoothed $\\log p^{\\lambda}(X_{k}|\\theta)$, we get the Moreau-Yosida ULA (MYULA) algorithm:\n",
    "\n",
    "$$X_{k+1} = X_{k} - \\gamma_{k}\\nabla f^{\\sigma^2}_{y}(x) - \\gamma_{k}\\nabla \\tilde{g}^{\\lambda}(x)+ \\sqrt{2\\gamma_{k}}Z_{k+1}\\,\\,,$$\n",
    "\n",
    "and a Markov Chain $(X_{k})_{k>0}$ with $p(x|y,\\theta,\\sigma^2)$ as the invariant measure can be sampled!\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cMqYxg3BGDxn",
   "metadata": {
    "id": "cMqYxg3BGDxn"
   },
   "source": [
    "## Getting Started\n",
    "\n",
    "To begin, we load the required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8089ea5d",
   "metadata": {
    "id": "8089ea5d"
   },
   "outputs": [],
   "source": [
    "# Load necessary packages\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from scipy.linalg import norm\n",
    "import math\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "from sampling_tools import * \n",
    "\n",
    "from tqdm.auto import tqdm\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": "MCueGfG2Gaq-",
   "metadata": {
    "id": "MCueGfG2Gaq-"
   },
   "source": [
    "## Import the ground truth image "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a33578fc",
   "metadata": {
    "id": "a33578fc"
   },
   "outputs": [],
   "source": [
    "# Load and show the image\n",
    "\n",
    "x = np.array(Image.open(\"cman.png\")).astype(np.float64)\n",
    "plot_im(x, \"Ground truth image x\")\n",
    "\n",
    "# Image dimension\n",
    "nx,ny = [x.shape[0],x.shape[1]]\n",
    "dimx = nx*ny\n",
    "\n",
    "# Convert to torch tensor\n",
    "x = torch.Tensor(x).to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4j3fceX-Gk5f",
   "metadata": {
    "id": "4j3fceX-Gk5f"
   },
   "source": [
    "## Define a forward operator.\n",
    "\n",
    "In this case, $A$ will be a linear operator applying a [box blur](https://en.wikipedia.org/wiki/Box_blur) of size 9 by 9 pixels."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7e803e6c",
   "metadata": {},
   "source": [
    "Note that the forward operator $A$ is a circulant matrix. Therefore, we express the convolution in the Fourier domain as follows \n",
    "$$ h*x = Ax = \\mathcal{F}^{-1}\\left(\\mathcal{F}(A)\\cdot\\mathcal{F}(x)\\right),$$\n",
    "where \n",
    "* $\\mathcal{F}$ and $\\mathcal{F}^{-1}$ are respectively the forward and inverse discrete Fourier transform operators;\n",
    "* $h$ is the blur kernel used to form the forward operator $A$;\n",
    "* The dot $\\cdot$ denotes pointwise multiplication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1673c932",
   "metadata": {
    "id": "1673c932"
   },
   "outputs": [],
   "source": [
    "kernel_len = [9,9]\n",
    "type_blur = \"uniform\"\n",
    "A, AT, AAT_norm = blur_operators(kernel_len, (nx, ny), type_blur, device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1-pURXytG0V6",
   "metadata": {
    "id": "1-pURXytG0V6"
   },
   "source": [
    "## Degraded measurement\n",
    "\n",
    "Apply the filter and add some noise to obtain the measurements $y = Ax + \\omega$ and achieve a blurred signal-to-noise ratio (BSNR) of $30$ dB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8f28246",
   "metadata": {
    "id": "c8f28246"
   },
   "outputs": [],
   "source": [
    "# generate the blurred and noisy observation 'y'\n",
    "Ax = A(x)\n",
    "\n",
    "BSNR = 30 # we will use this noise level\n",
    "\n",
    "min_BSNR = 15\n",
    "max_BSNR = 45\n",
    "\n",
    "# Define a convex set for sigma2.\n",
    "min_sigma2= (torch.linalg.matrix_norm(Ax-torch.mean(Ax), ord='fro')/math.sqrt(dimx*10**(max_BSNR/10)))**2\n",
    "max_sigma2= (torch.linalg.matrix_norm(Ax-torch.mean(Ax), ord='fro')/math.sqrt(dimx*10**(min_BSNR/10)))**2\n",
    "sigma2_init = (min_sigma2 + max_sigma2) / 2\n",
    "\n",
    "# sigma\n",
    "sigma = torch.linalg.matrix_norm(Ax-torch.mean(Ax), ord='fro')/math.sqrt(dimx*10**(BSNR/10))\n",
    "\n",
    "# observed data\n",
    "y = Ax + sigma*torch.randn_like(x)\n",
    "\n",
    "plot_im(y, \"noisy and blurry observation y\")\n",
    "\n",
    "# noise variance\n",
    "print(f\"\\t the noise variance is: {sigma**2}\\tmin = {min_sigma2} \\t max = {max_sigma2}\\n\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "gAbo3PMCG8U3",
   "metadata": {
    "id": "gAbo3PMCG8U3"
   },
   "source": [
    "## Define the likelihood\n",
    "The log likelihood also termed data fidelity is given by\n",
    "$$f^{\\sigma^2}_y(x) = \\dfrac{||y - Ax||^2_F}{2\\sigma^2}.$$\n",
    "Here $\\lVert \\cdot\\rVert_F$ denotes the Frobenius norm, note that although we introduced the image as a vector in $\\mathbb{R}^d$ it is more convienent here to consider the image as a $2$-dimensional array.\n",
    "\n",
    "The gradient of the log likelihood \n",
    "\n",
    "$$\\log p(y|x,\\sigma^2) = -f^{\\sigma^2}_y(x) - \\frac{d}{2}\\log(\\sigma^2) + C \\quad (\\text{$C$ is a constant})$$\n",
    "\n",
    "w.r.t. $\\sigma^2$ is defined by\n",
    "\n",
    "$$\\nabla_{\\sigma^2}\\log p(y|x,\\sigma^2) = \\dfrac{||y - Ax||^2_F}{2(\\sigma^2)^2} - \\dfrac{d}{2\\sigma^2}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2BMHlE6oHDBt",
   "metadata": {
    "id": "2BMHlE6oHDBt"
   },
   "outputs": [],
   "source": [
    "# Likelihood\n",
    "f = lambda x, sigma2: (torch.linalg.norm(y-A(x),'fro')**2)/(2*sigma2)   # Negative log-likelihood -logp(t|x,sigma^2)\n",
    "\n",
    "# Gradient the negative log-likelihood w.r.t. sigma^2\n",
    "gradf_wrt_sigma2 = lambda x, sigma2: (torch.linalg.norm(y - A(x),'fro')**2) / (2*sigma2**2) - dimx / (2*sigma2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "865e598c-28be-4d98-b7b0-e7edae19660d",
   "metadata": {},
   "source": [
    "The gradient $ x \\mapsto \\nabla f^{\\sigma^2}_y(x)$ is defined by\n",
    "$$\\nabla f^{\\sigma^2}_y(x) = \\dfrac{A^T(Ax - y)}{\\sigma^2}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a266b4cb",
   "metadata": {
    "id": "a266b4cb"
   },
   "outputs": [],
   "source": [
    "# gradient of the negative log-likelihood w.r.t. x\n",
    "ATy = AT(y)\n",
    "gradf_wrt_x = lambda x, sigma2: (AT(A(x)) - ATy) / sigma2        "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "Q0p6LRQWHItH",
   "metadata": {
    "id": "Q0p6LRQWHItH"
   },
   "source": [
    "## Define a prior\n",
    "\n",
    "In this assignment, we will consider a prior based on the total variation (TV) norm\n",
    "\n",
    "$$p(x) \\propto e^{- \\theta ||x|| _{TV}}\\,\\,,$$\n",
    "\n",
    "and so $g(x)=||x|| _{TV}$ where $||x|| _{TV}=\\sum_{i,j}^{}\\sqrt{|x_{i+1,j}-x_{i,j}|^2+|x_{i,j+1}-x_{i,j}|^2}$. The intuition behind TV norm is that it preserves the edges and smooths the flatter regions of the image. The function `g_fun()` calculates the $||\\cdot|| _{TV}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824a7961",
   "metadata": {
    "id": "824a7961"
   },
   "outputs": [],
   "source": [
    "#define the TV norm function for monitoring\n",
    "g_fun = lambda x: TVnorm(x)           "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "VeB-T2C4KzRm",
   "metadata": {
    "id": "VeB-T2C4KzRm"
   },
   "source": [
    "We should note that $g(x)=||\\cdot||_{TV}$ is not a smooth function. To smooth it, we will consider the Moreau-Yosida (MY) envelope as we discussed above.\n",
    "\n",
    "For a function $g$, the function `proxg()` below calculates the proximity operator \n",
    "\n",
    "$$\\mathrm{prox}_{g}^{\\lambda}(x)=\\mathrm{argmin}_{u\\in\\mathbb{R}^{M}}\\{g(u)+(2\\lambda)^{-1}||u-x||^{2}\\}\\,\\,.$$\n",
    "\n",
    "For a function $g$, the function `gradg()` calculates $$\\nabla g^{\\lambda}(x)=\\lambda^{-1}(x-\\mathrm{prox}_{g}^{\\lambda}(x))$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1IYnV51cLBbl",
   "metadata": {
    "id": "1IYnV51cLBbl"
   },
   "outputs": [],
   "source": [
    "proxg = lambda x, lam: chambolle_prox_TV(x, device, {'lambda' : lam, 'MaxIter' : 25}) # proximity operator\n",
    "gradg = lambda x, lam, lambda_prox: (x - proxg(x,lam)) / lambda_prox    # gradient of the prior"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "EGIk-O9Ma7dl",
   "metadata": {
    "id": "EGIk-O9Ma7dl"
   },
   "source": [
    "## Define the log-posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "_oX8t3rVa-xT",
   "metadata": {
    "id": "_oX8t3rVa-xT"
   },
   "outputs": [],
   "source": [
    "logPi = lambda x, sigma2, theta:   (- f(x,sigma2) - theta * g_fun(x))    # Log of posterior distribution"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "Lb03V7Fhaw6s",
   "metadata": {
    "id": "Lb03V7Fhaw6s"
   },
   "source": [
    "## Setting the algorithm parameters\n",
    "\n",
    "Since an MCMC algorithm needs to be implemented- as this is explained above - we will need to choose the step-size $\\gamma_{k}$. In order to converge, we need to take $\\gamma_{k} < 2/L$ where $L=L_{f^{\\sigma^2}_{y}} + L_{\\tilde{g}^{\\lambda}}$ with $L_{f^{\\sigma^2}_{y}}$ being the Lipschitz constant of $f^{\\sigma^2}_{y}$ and $L_{\\tilde{g}^{\\lambda}}$ the Lipschitz constant of $\\tilde{g}^{\\lambda}$.\n",
    "\n",
    "In this assignment, we will consider $\\gamma_{k} = 1.98/L$ and $\\lambda=1/L$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd70c52c",
   "metadata": {
    "id": "dd70c52c"
   },
   "outputs": [],
   "source": [
    "## Lipschitz Constants\n",
    "\n",
    "# Lipschitz constant of f\n",
    "Lp_sigma = lambda sigma2: AAT_norm**2 / sigma2  \n",
    "L_f =  min(Lp_sigma(min_sigma2), Lp_sigma(max_sigma2))  \n",
    "\n",
    "# regularization parameter of proximity operator (lambda).\n",
    "lambdaMax = 2\n",
    "lambda_prox = min((5/L_f), lambdaMax)   \n",
    "\n",
    "# Lipschitz constant of g.\n",
    "L_g =  1/lambda_prox \n",
    "\n",
    "# Lipschitz constant of g + f\n",
    "L =  L_f + L_g\n",
    "\n",
    "# Stepsize of MCMC algorithm.\n",
    "gamma = 1.98*1/L\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "O5cErOxTzxGQ",
   "metadata": {
    "id": "O5cErOxTzxGQ"
   },
   "source": [
    "Regarding the projected gradient algorithm parameters and without providing more details, we will consider  $\\delta_{n} = c_{0} \\frac{n^{-0.8}}{d}$ as a stepsize for the optimization algorithm where $c_{0}$ is a constant to scale the stepsize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d_YKn5ogvHTy",
   "metadata": {
    "id": "d_YKn5ogvHTy"
   },
   "outputs": [],
   "source": [
    "# Initialization of parameter theta\n",
    "th_init = 0.01\n",
    "\n",
    "# Admissible set for theta (min and max values).\n",
    "min_th = 0.001\n",
    "max_th = 1\n",
    "\n",
    "# define stepsize delta \n",
    "d_exp = 0.8\n",
    "delta = lambda i: (i**(-d_exp)) / dimx \n",
    "\n",
    "# constant for scaling the stepsize of each parameter\n",
    "c_eta = 10\n",
    "c_sigma2 = 10000"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "LcWhk8zX0niY",
   "metadata": {
    "id": "LcWhk8zX0niY"
   },
   "source": [
    "## Stochastic Approximation Proximal Gradient (SAPG) algorithm\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ZeHdbomM0zw-",
   "metadata": {
    "id": "ZeHdbomM0zw-"
   },
   "source": [
    "Initializations for the MCMC and optimization algorithms"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eIG5WlbO0274",
   "metadata": {
    "id": "eIG5WlbO0274"
   },
   "outputs": [],
   "source": [
    "\n",
    "# Warmup period for the MCMC sampling\n",
    "warmupSteps = 1000\n",
    "\n",
    "# total number of iterations for the optimization algorithm on theta\n",
    "total_iter = 2500\n",
    "\n",
    "# burn-in period for the optimization algorithm on theta\n",
    "burnIn = int(total_iter * 0.7)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "o2IY4u13GEw5",
   "metadata": {
    "id": "o2IY4u13GEw5"
   },
   "source": [
    "For a fixed value $\\theta_{0}\\in \\Theta_\\theta$ and $\\sigma^2_0\\in\\Theta_{\\sigma^2}$, we run the MCMC sampler as this is defined above to warm it up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ELfiAlVKFNqY",
   "metadata": {
    "id": "ELfiAlVKFNqY"
   },
   "outputs": [],
   "source": [
    "# Initialization of the warm-up chain\n",
    "X_wu = y.to(device).detach().clone()\n",
    "\n",
    "#Run MYULA sampler with fix theta and fix sigma^2 to warm up the markov chain\n",
    "\n",
    "fix_sigma2 = sigma2_init\n",
    "fix_theta = th_init\n",
    "\n",
    "print('Running Warm up     \\n')\n",
    "\n",
    "for k in tqdm(range(1,warmupSteps)):\n",
    "    # --- Gradients\n",
    "    gradf_X_wu = gradf_wrt_x(X_wu, fix_sigma2)\n",
    "    gradg_X_wu = gradg(X_wu, lambda_prox*fix_theta, lambda_prox)\n",
    "    # --- end (gradients)\n",
    "    \n",
    "    # --- MYULA warm-up\n",
    "    X_wu =  X_wu - gamma*gradg_X_wu - gamma*gradf_X_wu + math.sqrt(2*gamma)*torch.randn_like(X_wu)\n",
    "    # --- end (warm-up)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "zzCQwAi71Zcp",
   "metadata": {
    "id": "zzCQwAi71Zcp"
   },
   "source": [
    "For stability reasons, we will work on a logarithmic scale. So we define an auxiliary variable $\\eta$ such that $\\theta_{n} = e^{\\eta_{n}} \\iff\t\\eta_{n} = \\log\\theta_{n}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "h2XpE6XOO9X2",
   "metadata": {
    "id": "h2XpE6XOO9X2"
   },
   "outputs": [],
   "source": [
    "# Keeping track of the reg. parameter's trace\n",
    "theta_trace = torch.zeros(total_iter)\n",
    "theta_trace[0] = th_init\n",
    "\n",
    "sigma2_trace= torch.zeros(total_iter)\n",
    "sigma2_trace[0]= sigma2_init\n",
    "\n",
    "# We work on a logarithmic scale, so we define an axiliary variable \n",
    "#eta such that theta=exp{eta}. \n",
    "\n",
    "eta_init = math.log(th_init)\n",
    "min_eta = math.log(min_th)\n",
    "max_eta = math.log(max_th)\n",
    "\n",
    "eta_trace = torch.zeros(total_iter)\n",
    "eta_trace[0] = eta_init\n",
    "\n",
    "# Stop criteria (relative change tolerance) for the proximal gradient algorithm\n",
    "\n",
    "stopTol=1e-5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a8gjJmCqSCEU",
   "metadata": {
    "id": "a8gjJmCqSCEU"
   },
   "source": [
    "***Note:*** The intuition behind the SAPG algorithm is that for each update of $\\theta$, we need to \"update\"\n",
    "the MCMC sampler for the new value. Therefore, to compute\n",
    "$\\theta^{*}$, we can build a sequence $(\\theta_{n})_{n\\in\\mathbb{N}}$ associated with the following recursion for any $n\\in\\mathbb{N}$\n",
    "\n",
    "$$\\theta_{n+1} = \\Pi_{\\Theta_\\theta}[\\theta_{n} + \\delta_{n}\\nabla_\\theta \\log p(y|\\theta_n,\\sigma_n^2)]\\,\\,,$$\n",
    "\n",
    "where $$\\nabla_\\theta \\log p(y|\\theta_n,\\sigma_n^2) \\approx - \\dfrac{1}{m_{n}}\\sum_{i=1}^{m_{n}}g(X_{i}) - \\nabla_{\\theta}\\log (Z(\\theta_{n})).$$\n",
    "\n",
    "Sarting from some $\\theta_{0}\\in\\Theta_\\theta$, and where $(m_{n})_{n\\in\\mathbb{N}}$ is a sequence of non-decreasing sample sizes.\n",
    "\n",
    "Similarly, we generate a sequence the noise variance $(\\sigma^2_{n})_{n\\in\\mathbb{N}}$ as follows\n",
    "$$\\sigma^2_{n+1} = \\Pi_{\\Theta_{\\sigma^2}}[\\sigma^2_{n} + \\delta_{n}\\nabla_{\\sigma^2} \\log p(y|\\theta_n,\\sigma_n^2)]\\,\\,,$$\n",
    "\n",
    "where $$\\nabla_{\\sigma^2} \\log p(y|\\theta_n,\\sigma_n^2) \\approx  \\dfrac{1}{m_{n}}\\sum_{i=1}^{m_{n}}\\dfrac{||y - AX_i||^2_F}{2(\\sigma^2)^2)} - \\dfrac{d}{2\\sigma^2}.$$\n",
    "\n",
    "Empirically, we have seen that taking $m_{n} = 1$ $\\forall n\\in\\mathbb{N}$ is enough. This means that after each update of $\\theta$ and $\\sigma^2$, we need to run 1 iteration of the MCMC sampler."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aVXM6ktQQqFp",
   "metadata": {
    "id": "aVXM6ktQQqFp"
   },
   "source": [
    "To monitor the behaviour of the SAPG algorithm, we need to store two quantities: the log-posterior trace (without the normalizing constant) and the log-prior trace (without the normalizing constant)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "JgMWTq6rVGz0",
   "metadata": {
    "id": "JgMWTq6rVGz0"
   },
   "outputs": [],
   "source": [
    "print('\\nRunning SAPG algorithm     \\n')\n",
    "# We want to keep track of two traces: the log-likelihood and the TV function to monitor the behaviour of the algorithm.\n",
    "\n",
    "# to monitor convergence\n",
    "logPiTraceX = []      \n",
    "# to monitor how the regularisation function evolves\n",
    "gtrace = []          \n",
    "\n",
    "mean_theta =[]\n",
    "mean_sigma2 = []\n",
    "\n",
    "# start MYULA markov chain from last sample after warmup\n",
    "X = X_wu.clone()       \n",
    "\n",
    "# for demo purposes we want the algorithm to run for the full amount of iterations\n",
    "# in practice you'd like to use the stopping criteria - set this to True\n",
    "stopping_active = False  \n",
    "\n",
    "\n",
    "for k in tqdm(range(1,total_iter)): \n",
    "\n",
    "    ################################################################################\n",
    "    # MYULA SAMPLER\n",
    "    ################################################################################\n",
    "\n",
    "    # Number of samples\n",
    "    m = 1\n",
    "\n",
    "    # If we run the MCMC sampler for m times to get m samples X_m, therefore we need to average \n",
    "    # gradients w.r.t. \\theta and \\sigma^2 before the update\n",
    "    g_trace_m = torch.zeros(m).to(device) \n",
    "    grad_sigma2_trace = torch.zeros(m).to(device) \n",
    "    \n",
    "    #Sample from posterior with MYULA:\n",
    "    \n",
    "    for ii in range(m):\n",
    "\n",
    "        # Calculate the gradient related to g for the current theta\n",
    "        gradgX = gradg(X,lambda_prox*theta_trace[k-1],lambda_prox)  \n",
    "        \n",
    "        # --- Calculate the gradient related to f for the current theta\n",
    "        gradfX = gradf_wrt_x(X,sigma2_trace[k-1])\n",
    "        # --- end\n",
    "        \n",
    "        # --- MYULA update\n",
    "        X =  X - gamma*gradgX - gamma*gradfX + math.sqrt(2*gamma)*torch.randn_like(X)\n",
    "        # --- end\n",
    "        \n",
    "        # --- g(X) \n",
    "        g_trace_m[ii] = g_fun(X)\n",
    "        \n",
    "        # --- Gradient w.r.t. sigma2\n",
    "        grad_sigma2_trace[ii] = gradf_wrt_sigma2(X, sigma2_trace[k-1])\n",
    "        # --- end\n",
    "        \n",
    "    # --- Save current state to monitor convergence\n",
    "    logPiTraceX.append(logPi(X, sigma2_trace[k-1],theta_trace[k-1]).cpu().numpy())\n",
    "    gtrace.append(g_trace_m[-1].cpu().numpy())\n",
    "    # --- end (monitoring)\n",
    "    \n",
    "    # ################################################################################\n",
    "    #  PROJECTED GRADIENT ALGORITHM\n",
    "    # ################################################################################\n",
    "\n",
    "    # --- update \\eta and \\theta\n",
    "    # -- Gradient\n",
    "    grad_eta = (dimx / torch.exp(eta_trace[k-1]) - torch.mean(g_trace_m)) * torch.exp(eta_trace[k-1]) \n",
    "    \n",
    "    # -- Update\n",
    "    etak_temp = eta_trace[k-1] + c_eta * delta(k)  * grad_eta\n",
    "    # project \\eta onto the admissible set of value\n",
    "    eta_trace[k] = min(max(etak_temp, min_eta), max_eta)\n",
    "    \n",
    "    # Save the value of theta\n",
    "    theta_trace[k] = torch.exp(eta_trace[k])\n",
    "    # --- end (update)\n",
    "    \n",
    "    # --- Update sigma^2\n",
    "    sigma2_k_temp = sigma2_trace[k-1] + c_sigma2 * delta(k) * torch.mean(grad_sigma2_trace)\n",
    "    \n",
    "    # Save the value of sigma^2\n",
    "    sigma2_trace[k] = min(max(sigma2_k_temp, min_sigma2), max_sigma2)\n",
    "    # --- end (update)\n",
    "    \n",
    "    # --- Check stop criteria. If relative error is smaller than op.stopTol stop\n",
    "    if k>burnIn+1:\n",
    "        mean_theta_k = torch.mean(theta_trace[burnIn:k]).cpu().numpy()\n",
    "        mean_theta_k_next = torch.mean(theta_trace[burnIn:(k+1)]).cpu().numpy()\n",
    "        mean_sigma2_k = torch.mean(sigma2_trace[burnIn:k]).cpu().numpy()\n",
    "        mean_sigma2_k_next = torch.mean(sigma2_trace[burnIn:(k+1)]).cpu().numpy()\n",
    "        mean_theta.append(mean_theta_k_next)\n",
    "        mean_sigma2.append(mean_sigma2_k_next)\n",
    "        \n",
    "        relErrTh1 = np.abs(mean_theta_k_next - mean_theta_k) / mean_theta_k\n",
    "        \n",
    "        relErrSi1 = np.abs(mean_sigma2_k_next - mean_sigma2_k) / mean_sigma2_k\n",
    "        \n",
    "        if (relErrTh1<stopTol) and (relErrSi1<stopTol) and stopping_active :\n",
    "            \n",
    "            print(\"Toleration reached!\")\n",
    "            break\n",
    "     # --- end (stop criteria)       \n",
    "\n",
    "# --- Collecting data\n",
    "last_samp = k\n",
    "\n",
    "logPiTraceX = logPiTraceX[:last_samp+1]\n",
    "gXTrace = gtrace[:last_samp+1]\n",
    "\n",
    "theta_EB = torch.exp(torch.mean(eta_trace[burnIn:last_samp+1]))\n",
    "last_theta = theta_trace[last_samp]\n",
    "thetas = theta_trace[:last_samp+1]\n",
    "\n",
    "sigma2_EB = torch.mean(sigma2_trace[burnIn:last_samp+1])\n",
    "sigmas=sigma2_trace[:last_samp+1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "EqSnk3_Qvtip",
   "metadata": {
    "id": "EqSnk3_Qvtip"
   },
   "source": [
    "# Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4jjIG1U3kSe1",
   "metadata": {
    "id": "4jjIG1U3kSe1"
   },
   "outputs": [],
   "source": [
    "print(\"Estimated theta: \", theta_EB)\n",
    "print(\"Last theta: \", last_theta)\n",
    "print(\"Estimated value of σ2 \",sigma2_EB, sigma**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad23a38",
   "metadata": {
    "id": "0ad23a38"
   },
   "outputs": [],
   "source": [
    "# Plot the results\n",
    "\n",
    "plot1 = plt.figure()\n",
    "plt.plot(thetas[:].cpu().numpy(),linestyle=\"-\")\n",
    "plt.xlabel(\"$Iterations$\")\n",
    "plt.ylabel(\"$\\\\theta$\")\n",
    "\n",
    "plot1 = plt.figure()\n",
    "plt.plot( logPiTraceX,linestyle=\"-\")\n",
    "plt.xlabel(\"$Iterations$\")\n",
    "plt.ylabel(\"$log(p(x|y))$\")\n",
    "\n",
    "plot1, ax1 = plt.subplots()\n",
    "plt.plot( gXTrace[burnIn:],linestyle=\"-\", label=\"$g(x)$\")\n",
    "plt.plot( dimx/thetas[burnIn:].cpu().numpy(),linestyle=\"-\", label=\"d/$\\\\theta$\")\n",
    "plt.xlabel(\"$Iterations$\")\n",
    "plt.ylabel(\"$g(x) \\,\\,vs.\\,\\, d/\\\\theta$\")\n",
    "ax1.legend()\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(sigmas.cpu().numpy(),linestyle=\"-\",label=\"$\\\\sigma_{n}^{2}$\")\n",
    "plt.axhline(y=sigma.cpu().numpy()**2, color='r', linestyle='--',label=\"$\\\\sigma_{\\dagger}^{2}$\")\n",
    "ax.set_xlabel(\"$Iterations\\,\\,(n)$\")\n",
    "ax.set_ylabel(\"$\\\\sigma^{2}$\")\n",
    "ax.legend()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "empirical_bayesian_estimation.ipynb",
   "provenance": []
  },
  "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
}