{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e570205c-8d15-409a-a80a-2deabe8093db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cil.optimisation.functions import LeastSquares, SVRGFunction, SGFunction\n",
    "from cil.framework import AcquisitionGeometry\n",
    "from cil.plugins.astra import ProjectionOperator\n",
    "from cil.plugins.astra import FBP\n",
    "from cil.optimisation.utilities import MetricsDiagnostics, RandomSampling\n",
    "from cil.optimisation.algorithms import FISTA, ISTA\n",
    "\n",
    "from ProxSkip import ProxSkip\n",
    "from utils import StoppingCriterionTime, BM3DFunction\n",
    "from skimage.metrics import peak_signal_noise_ratio as PSNR\n",
    "from skimage.metrics import structural_similarity as SSIM\n",
    "from xdesign import Foam, discrete_phantom\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"image.cmap\"]  = \"inferno\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95eac1e1-ad82-4977-9576-1b8cc99f8d57",
   "metadata": {},
   "source": [
    "### PnP Variance Reduced ProxSkip on a simulated dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1dd614-7f15-4894-b13d-a1b32c82c969",
   "metadata": {},
   "source": [
    "In this notebook, we play the following game. We compare algorithms by **time-to-quality**: that is, we measure how quickly each method achieves high reconstruction quality—quantified by **PSNR** and **SSIM**—with respect to the ground truth. Here, the ground truth is available, and we impose a fixed computational budget. The winner is the algorithm that attains the highest PSNR/SSIM within this time limit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d25fd695-7af3-4992-bf6e-0ce6fe299870",
   "metadata": {},
   "source": [
    "### Generate Foam Phantom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7154574f-40e1-4b60-828b-ddc4b1be8d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 347  # image size\n",
    "density = 0.025  # attenuation density of the image\n",
    "np.random.seed(1234)\n",
    "gt = discrete_phantom(Foam(size_range=[0.075, 0.005], gap=2e-3, porosity=1.0), size=N - 10)\n",
    "gt = gt / np.max(gt) * density\n",
    "gt = np.pad(gt, 5)\n",
    "gt[gt < 0] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7963e21e-24c7-46e0-9cda-3c9c6aaf9678",
   "metadata": {},
   "source": [
    "### Generate Acquisition Data with simulated noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "226ded3c-a540-49d0-b0a3-3c560de87c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_angles = 400 \n",
    "\n",
    "angles = np.linspace(0, 180, num_angles, endpoint=False, dtype=np.float32)\n",
    "ag2D = AcquisitionGeometry.create_Parallel2D()\\\n",
    "        .set_panel(num_pixels=N)\\\n",
    "        .set_angles(angles=angles)\n",
    "ig2D = ag2D.get_ImageGeometry()\n",
    "x_cil = ig2D.allocate()\n",
    "x_cil.fill(gt)\n",
    "A = ProjectionOperator(ig2D, ag2D, device=\"cpu\")\n",
    "\n",
    "# noiseless sinogram\n",
    "noiseless = A.direct(x_cil)\n",
    "\n",
    "## simulated noise\n",
    "max_intensity = 2000\n",
    "expected_counts = max_intensity * np.exp(-noiseless.array)\n",
    "noisy_counts = np.random.poisson(expected_counts).astype(np.float32)\n",
    "noisy_counts[noisy_counts == 0] = 1  # deal with 0s\n",
    "y_np = -np.log(noisy_counts / max_intensity)\n",
    "noisy = ag2D.allocate()\n",
    "noisy.fill(y_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faa9bb7a-9454-4c6e-802e-80c991cc0da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(12, 10), constrained_layout=True)\n",
    "\n",
    "axs[0].imshow(gt)\n",
    "axs[0].set_title(\"Ground Truth\")\n",
    "\n",
    "axs[1].imshow(noiseless.array)\n",
    "axs[1].set_title(\"Noiseless Sinogram\")\n",
    "\n",
    "axs[2].imshow(noisy.array)\n",
    "axs[2].set_title(\"Noisy Sinogram\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b138b305-26db-4c0e-a40f-687cff75d883",
   "metadata": {},
   "source": [
    "### FBP reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78e70056-aea9-4780-89e2-52351f51b9ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "fbp = FBP(ig2D, ag2D, device=\"cpu\")(noisy)\n",
    "fbp.array[fbp.array<0] = 0\n",
    "plt.imshow(fbp.array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "498ad7aa-540b-426b-8810-991a72312e5c",
   "metadata": {},
   "source": [
    "### Family of considered algorithms (ProxSkip template)\n",
    "\n",
    "**Parameters:** $\\gamma>0$, probability $p\\in(0,1]$, data subsets $N$\n",
    "**Initialize:** $x_0, h_0 \\in \\mathbb{R}^n$\n",
    "\n",
    "For $k=0,1,\\dots,K-1$:\n",
    "\n",
    "1. Compute $G_k$ (an unbiased estimator of $\\nabla f(x_k)$).\n",
    "2. $$\\hat x_{k+1} = x_k - \\gamma\\big(G_k(x_k) - h_k\\big).$$\n",
    "3. Sample $\\theta_k\\sim\\mathrm{Bernoulli}(p)$, $\\theta_k\\in{0,1}$.\n",
    "4. If $\\theta_k=1$:\n",
    "   $$x_{k+1}=\\mathrm{prox}_{\\frac{\\gamma}{p}g}\\left(\\hat x_{k+1}-\\frac{\\gamma}{p}h_k\\right),$$\n",
    "   else:\n",
    "   $$x_{k+1}=\\hat x_{k+1}.$$\n",
    "5. Update:\n",
    "   $$h_{k+1}=h_k+\\frac{p}{\\gamma}\\big(x_{k+1}-\\hat x_{k+1}\\big).$$\n",
    "\n",
    "---\n",
    "\n",
    "| $p=1$     | $0<p<1$       | $G_k$                                                                |\n",
    "| --------- | ------------- | -------------------------------------------------------------------- |\n",
    "| ISTA      | ProxSkip      | $\\nabla f(x_k)$                                                      |\n",
    "| ProxSGD   | ProxSGDSkip   | $N\\nabla f_{i_k}(x_k)$                                               |\n",
    "| ProxSAGA  | ProxSAGASkip  | $N(\\nabla f_{i_k}(x_k)-v_k^{,i_k})+\\bar v_k$                         |\n",
    "| ProxSVRG  | ProxSVRGSkip  | $N(\\nabla f_{i_k}(x_k)-\\nabla f_{i_k}(\\tilde x))+\\nabla f(\\tilde x)$ |\n",
    "| ProxLSVRG | ProxLSVRGSkip | as above, updated with $p=1/N$                                       |\n",
    "\n",
    "*Note:* FISTA is ISTA with an acceleration step (Beck & Teboulle).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfd1a9e7-f591-4391-8c4b-71228b3c49e1",
   "metadata": {},
   "source": [
    "### Define Stopping Criteria and Metrics (PSNR/SSIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b4287b-06a1-4d96-baa5-db52f3b8e4b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_cil = ig2D.allocate()\n",
    "gt_cil.fill(gt)\n",
    "\n",
    "SSIM_ = lambda x, y: SSIM(x, y, data_range=x.max()-x.min())\n",
    "PSNR_ = lambda x, y: PSNR(x, y, data_range=x.max()-x.min())\n",
    "\n",
    "time_stop = 3*60 # 3 mins\n",
    "max_iteration = 100_000\n",
    "\n",
    "cb_metrics = MetricsDiagnostics(reference_image=gt_cil, \n",
    "                         metrics_dict={\"psnr\":PSNR_,\"ssim\":SSIM_}) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e03af7b-468a-43c1-9344-e66dcea5ce58",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Avoid computing objectives, not needed for this demo\n",
    "def update_objective(self):\n",
    "    return 0.\n",
    "\n",
    "ISTA.update_objective = update_objective\n",
    "ProxSkip.update_objective = update_objective\n",
    "FISTA.update_objective = update_objective"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f22e499-849f-47b5-b959-b2b8254af0e6",
   "metadata": {},
   "source": [
    "### Run deterministic algorithms: No data splitting, no skipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e67a015c-1e41-40fe-84c2-383e0558d4f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_no_skip = 0.0002 \n",
    "F = LeastSquares(A=A, b=noisy, c=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070838d2-505b-4265-846b-cee8a2a96450",
   "metadata": {},
   "source": [
    "### FISTA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02a9d83c-3a9a-4935-8dd0-c6c6ff3dea2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial = ig2D.allocate()\n",
    "step_size = 1./F.L\n",
    "G = BM3DFunction(sigma_no_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "fista = FISTA(initial = initial, f = F, step_size = step_size, g=G, \n",
    "        update_objective_interval = 1,\n",
    "        max_iteration = max_iteration) \n",
    "fista.run(verbose=0, callback=[cb_metrics, cb_time])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a7833e5-095e-4524-b7f6-d3f4d5be473f",
   "metadata": {},
   "source": [
    "### ISTA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b02973-d57c-4f69-b6d1-4af18fa6b160",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = BM3DFunction(sigma_no_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "ista = ISTA(initial = initial, f = F, step_size = step_size, g=G, \n",
    "        update_objective_interval = 1,\n",
    "        max_iteration = max_iteration) \n",
    "ista.run(verbose=0, callback=[cb_metrics, cb_time])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57bd7320-c2cb-4e4c-a3f3-f382a0769d6a",
   "metadata": {},
   "source": [
    "### Run ProxSkip: Skip the regulariser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a9b36d3-dd8a-43df-999c-acb414c3ba51",
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = 0.05\n",
    "sigma_skip = sigma_no_skip/np.sqrt(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd6d1c7b-77c8-44e4-a71c-8169a097567e",
   "metadata": {},
   "outputs": [],
   "source": [
    "G =  BM3DFunction(sigma_skip)\n",
    "\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "proxskip = ProxSkip(initial = [initial, initial], f = F, step_size = step_size, g=G, \n",
    "        update_objective_interval = 1, prob=prob, seed = 42, \n",
    "        max_iteration = max_iteration) \n",
    "proxskip.run(verbose=0, callback=[cb_metrics, cb_time])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84f96bea-e56d-4fc7-a8b9-7b288430f079",
   "metadata": {},
   "source": [
    "### Run stochastic algorithms: Data splitting, no skipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750978d9-af9e-4366-aeaf-810155d22f1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_of_functions(data):\n",
    "    \n",
    "    list_funcs = []\n",
    "    ig = data[0].geometry.get_ImageGeometry()\n",
    "    \n",
    "    for d in data:\n",
    "        ageom_subset = d.geometry        \n",
    "        Ai = ProjectionOperator(ig, ageom_subset, device = 'cpu')    \n",
    "        fi = LeastSquares(Ai, b = d, c = 0.5)\n",
    "        list_funcs.append(fi)   \n",
    "        \n",
    "    return list_funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce7998a1-9c98-49fe-b35a-89be65498fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of subsets\n",
    "nsub = 50\n",
    "data_split, method = noisy.split_to_subsets(nsub, method= \"ordered\", info=True)\n",
    "\n",
    "# list of fis for finite sum \n",
    "list_func = list_of_functions(data_split) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd9e598-bafd-4bdf-b9b0-33f1646ce191",
   "metadata": {},
   "source": [
    "### ProxSVRG and ProxSGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb2e2d8-6054-4759-bd84-d5cc7d2bbc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = RandomSampling(len(list_func), nsub, seed=42)\n",
    "Fsvrg = SVRGFunction(list_func, selection = selection, update_frequency=len(list_func))\n",
    "Fsvrg.initial = initial\n",
    "step_size = 1./(Fsvrg.L)\n",
    "\n",
    "G =  BM3DFunction(sigma_no_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "prox_svrg = ISTA(initial = initial, f = Fsvrg, step_size = step_size, g=G, \n",
    "            update_objective_interval = 1, \n",
    "            max_iteration = max_iteration) \n",
    "prox_svrg.run(verbose=0, callback=[cb_metrics, cb_time]) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1565f4c-162c-41bd-919a-bfd61a4a57cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = RandomSampling(len(list_func), nsub, seed=42)\n",
    "Fsgd = SGFunction(list_func, selection = selection) \n",
    "step_size = 1./(Fsgd.L)\n",
    "\n",
    "G =  BM3DFunction(sigma_no_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "prox_sgd = ISTA(initial = initial, f = Fsgd, step_size = step_size, g=G, \n",
    "            update_objective_interval = 1, \n",
    "            max_iteration = max_iteration) \n",
    "prox_sgd.run(verbose=0, callback=[cb_metrics, cb_time]) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e35518e7-70d8-4c7d-955b-f062c133371e",
   "metadata": {},
   "source": [
    "### Run stochastic algorithms: Data splitting and skipping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94d2a7e7-b0d3-4418-8d9d-fa35f506c5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = RandomSampling(len(list_func), nsub, seed=42)\n",
    "Fsvrg_skip = SVRGFunction(list_func, selection = selection, update_frequency=len(list_func))\n",
    "Fsvrg_skip.initial = initial\n",
    "step_size = 1./(Fsvrg_skip.L)\n",
    "\n",
    "G =  BM3DFunction(sigma_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "prox_svrg_skip = ProxSkip(initial = [initial, initial], f = Fsvrg_skip, step_size = step_size, g=G, \n",
    "            update_objective_interval = 1, prob=prob, seed=42,\n",
    "            max_iteration = max_iteration) \n",
    "prox_svrg_skip.run(verbose=0, callback=[cb_metrics, cb_time]) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbfc107d-a6f7-45b8-8974-c1e766148bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "selection = RandomSampling(len(list_func), nsub, seed=42)\n",
    "Fsgd_skip = SGFunction(list_func, selection = selection)\n",
    "step_size = 1./(Fsgd_skip.L)\n",
    "\n",
    "G =  BM3DFunction(sigma_skip)\n",
    "cb_time = StoppingCriterionTime(time_stop)\n",
    "prox_sgd_skip = ProxSkip(initial = [initial, initial], f = Fsgd_skip, step_size = step_size, g=G, \n",
    "            update_objective_interval = 1, prob=prob, seed=42,\n",
    "            max_iteration = max_iteration) \n",
    "prox_sgd_skip.run(verbose=0, callback=[cb_metrics, cb_time]) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48181800-6930-454c-bc6c-e4fb9d4e2d42",
   "metadata": {},
   "source": [
    "### Plot PSNR/SSIM progress\n",
    "- with respect to iteration\n",
    "- with respect to time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a76c9c9-f023-419b-9f7c-fdc840c92301",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_ista  = np.cumsum(ista.timing)\n",
    "t_fista = np.cumsum(fista.timing)\n",
    "t_proxskip    = np.cumsum(proxskip.timing)\n",
    "t_prox_svrg  = np.cumsum(prox_svrg.timing)\n",
    "t_prox_svrg_skip = np.cumsum(prox_svrg_skip.timing)\n",
    "t_prox_sgd_skip   = np.cumsum(prox_sgd_skip.timing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca6e8583-5487-414e-ba5d-3573ae290595",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams['lines.linewidth'] = 5\n",
    "plt.rcParams['lines.markersize'] = 20\n",
    "plt.rcParams['font.size'] = 40\n",
    "\n",
    "fig, ax = plt.subplots(2, 2, figsize=(35, 25), constrained_layout=True)\n",
    "\n",
    "ax[0,0].plot(ista.psnr[:-1], label=f\"ISTA, K={ista.iteration}, #BM3D={ista.iteration}\")\n",
    "ax[0,0].plot(fista.psnr[:-1], label=f\"FISTA, K={fista.iteration}, #BM3D={fista.iteration}\")\n",
    "ax[0,0].plot(proxskip.psnr[:-1], label=f\"ProxSkip, K={proxskip.iteration},  #BM3D={np.sum(proxskip.thetas)}\")\n",
    "ax[0,0].plot(prox_svrg.psnr[:-1], label=f\"ProxSVRG, K={prox_svrg.iteration}, #BM3D={prox_svrg.iteration}\")\n",
    "ax[0,0].plot(prox_sgd_skip.psnr[:-1],label=f\"ProxSGDSkip, K={prox_svrg_skip.iteration}, #BM3D={np.sum(prox_svrg_skip.thetas)}\", alpha=0.65)\n",
    "ax[0,0].plot(prox_svrg_skip.psnr[:-1],label=f\"ProxSVRGSkip, K={prox_svrg_skip.iteration}, #BM3D={np.sum(prox_svrg_skip.thetas)}\")\n",
    "ax[0,0].grid(which=\"major\")\n",
    "ax[0,0].set_xlabel(\"Iteration\")\n",
    "ax[0,0].set_ylabel(\"PSNR\")\n",
    "\n",
    "ax[0,1].plot(ista.ssim[:-1], label=f\"ISTA, K={ista.iteration}, #BM3D={ista.iteration}\")\n",
    "ax[0,1].plot(fista.ssim[:-1], label=f\"FISTA, K={fista.iteration}, #BM3D={fista.iteration}\")\n",
    "ax[0,1].plot(proxskip.ssim[:-1], label=f\"ProxSkip, p={prob}, #prox={np.sum(proxskip.thetas)}\")\n",
    "ax[0,1].plot(prox_svrg.ssim[:-1], label=\"ProxSVRG\")\n",
    "ax[0,1].semilogy(prox_sgd_skip.ssim[:-1],label=f\"ProxSGDSkip, K={prox_svrg_skip.iteration}, N={nsub}, p={prob}, #BM3D={np.sum(prox_svrg_skip.thetas)}\", alpha=0.65)\n",
    "ax[0,1].plot(prox_svrg_skip.ssim[:-1],label=f\"ProxSVRGSkip, p={prob}, #prox={np.sum(prox_svrg_skip.thetas)}\")\n",
    "ax[0,1].grid(which=\"major\")\n",
    "ax[0,1].set_xlabel(\"Iteration\")\n",
    "ax[0,1].set_ylabel(\"SSIM\")\n",
    "\n",
    "ax[1,0].plot(t_ista,  ista.psnr[:-1], label=f\"ISTA, K={ista.iteration}, #BM3D={ista.iteration}\")\n",
    "ax[1,0].plot(t_fista, fista.psnr[:-1], label=f\"FISTA, K={fista.iteration}, #BM3D={fista.iteration}\")\n",
    "ax[1,0].plot(t_proxskip,    proxskip.psnr[:-1], label=f\"ProxSkip, K={proxskip.iteration},  #BM3D={np.sum(proxskip.thetas)}\")\n",
    "ax[1,0].plot(t_prox_svrg,  prox_svrg.psnr[:-1], label=f\"ProxSVRG, K={prox_svrg.iteration}, #BM3D={prox_svrg.iteration}\")\n",
    "ax[1,0].plot(t_prox_sgd_skip, prox_sgd_skip.psnr[:-1],\n",
    "               label=f\"ProxSGDSkip, K={prox_svrg_skip.iteration}, #BM3D={np.sum(prox_svrg_skip.thetas)}\", alpha=0.65)\n",
    "ax[1,0].plot(t_prox_svrg_skip, prox_svrg_skip.psnr[:-1],\n",
    "               label=f\"ProxSVRGSkip, K={prox_svrg_skip.iteration}, #BM3D={np.sum(prox_svrg_skip.thetas)}\")\n",
    "ax[1,0].grid(which=\"major\")\n",
    "ax[1,0].set_xlabel(\"Time (sec)\")\n",
    "ax[1,0].set_ylabel(\"PSNR\")\n",
    "\n",
    "ax[1,1].plot(t_ista,  ista.ssim[:-1], label=f\"ISTA, K={ista.iteration}, #BM3D={ista.iteration}\")\n",
    "ax[1,1].plot(t_fista, fista.ssim[:-1], label=f\"FISTA, K={fista.iteration}, #BM3D={fista.iteration}\")\n",
    "ax[1,1].plot(t_proxskip,    proxskip.ssim[:-1], label=f\"ProxSkip, p={prob}, #prox={np.sum(proxskip.thetas)}\")\n",
    "ax[1,1].plot(t_prox_svrg,  prox_svrg.ssim[:-1], label=\"ProxSVRG\")\n",
    "ax[1,1].semilogy(t_prox_sgd_skip, prox_sgd_skip.ssim[:-1],\n",
    "               label=f\"ProxSGDSkip, K={prox_svrg_skip.iteration}, N={nsub}, p={prob}, #BM3D={np.sum(prox_svrg_skip.thetas)}\", alpha=0.65)\n",
    "ax[1,1].plot(t_prox_svrg_skip, prox_svrg_skip.ssim[:-1],\n",
    "           label=f\"ProxSVRGSkip, p={prob}, #prox={np.sum(prox_svrg_skip.thetas)}\")\n",
    "ax[1,1].grid(which=\"major\")\n",
    "ax[1,1].set_xlabel(\"Time (sec)\")\n",
    "ax[1,1].set_ylabel(\"SSIM\")\n",
    "\n",
    "handles, labels = ax[0,1].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, loc=\"upper center\", bbox_to_anchor=(0.53, 1.14),\n",
    "           ncols=2, frameon=True)\n",
    "\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af5e1f06-ec65-4944-a84a-1681462e321e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "from matplotlib.patches import Rectangle\n",
    "\n",
    "imgs = [\n",
    "    x_cil.array,\n",
    "    fbp.array, \n",
    "    ista.solution.array,\n",
    "    fista.solution.array,\n",
    "    proxskip.solution.array,\n",
    "    prox_svrg.solution.array,\n",
    "    prox_sgd_skip.solution.array,\n",
    "    prox_svrg_skip.solution.array,\n",
    "]\n",
    "\n",
    "labels = [\n",
    "    \"Ground-Truth\",\n",
    "    f\"FBP\",\n",
    "    f\"ISTA\\nPSNR/SSIM = {ista.psnr[-1]:.2f}/{ista.ssim[-1]:.3f}\",\n",
    "    f\"FISTA\\nPSNR/SSIM = {fista.psnr[-1]:.2f}/{fista.ssim[-1]:.3f}\",\n",
    "    f\"ProxSkip\\nPSNR/SSIM = {proxskip.psnr[-1]:.2f}/{proxskip.ssim[-1]:.3f}\",\n",
    "    f\"ProxSVRG\\nPSNR/SSIM = {prox_svrg.psnr[-1]:.2f}/{prox_svrg.ssim[-1]:.3f}\",\n",
    "    f\"ProxSGDSkip\\nPSNR/SSIM = {prox_sgd_skip.psnr[-1]:.2f}/{prox_sgd_skip.ssim[-1]:.3f}\",\n",
    "    f\"ProxSVRGSkip\\nPSNR/SSIM = {prox_svrg_skip.psnr[-1]:.2f}/{prox_svrg_skip.ssim[-1]:.3f}\",\n",
    "]\n",
    "\n",
    "vmax = max(np.max(im) for im in imgs)\n",
    "vmin = 0.01\n",
    "vmax = 0.03\n",
    "\n",
    "h, w = imgs[0].shape[:2]\n",
    "img_aspect = w / h\n",
    "nrows, ncols = 2, 4\n",
    "\n",
    "fig_h = 15\n",
    "fig_w = fig_h * (ncols / nrows) * img_aspect\n",
    "\n",
    "fig = plt.figure(figsize=(fig_w, fig_h))\n",
    "gs = fig.add_gridspec(nrows, ncols, wspace=0, hspace=0)\n",
    "axes = np.array([fig.add_subplot(gs[i, j]) for i in range(nrows) for j in range(ncols)])\n",
    "fig.subplots_adjust(left=0, right=1, bottom=0, top=1)\n",
    "\n",
    "H, W = imgs[0].shape[:2]\n",
    "roi_size = 80\n",
    "cx, cy = 75, 120  \n",
    "\n",
    "x0 = cx - roi_size // 2\n",
    "x1 = x0 + roi_size\n",
    "y0 = cy - roi_size // 2\n",
    "y1 = y0 + roi_size\n",
    "\n",
    "x0 = max(0, min(x0, W - roi_size)); x1 = x0 + roi_size\n",
    "y0 = max(0, min(y0, H - roi_size)); y1 = y0 + roi_size\n",
    "\n",
    "inset_size = \"32%\"\n",
    "inset_borderpad = 0.6\n",
    "\n",
    "for ax, im, txt in zip(axes, imgs, labels):\n",
    "    ax.imshow(im, cmap=\"inferno\", vmin=vmin, vmax=vmax)\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "    ax.text(\n",
    "        0.02, 0.98, txt,\n",
    "        transform=ax.transAxes,\n",
    "        ha=\"left\", va=\"top\",\n",
    "        color=\"white\", fontsize=32,\n",
    "        bbox=dict(boxstyle=\"round,pad=0.3\", facecolor=\"black\", alpha=0.25, edgecolor=\"none\")\n",
    "    )\n",
    "\n",
    "    rect = Rectangle((x0, y0), roi_size, roi_size,\n",
    "                     linewidth=2.5, edgecolor=\"yellow\", facecolor=\"none\")\n",
    "    ax.add_patch(rect)\n",
    "\n",
    "    axins = inset_axes(ax, width=inset_size, height=inset_size,\n",
    "                       loc=\"lower right\", borderpad=inset_borderpad)\n",
    "    axins.imshow(im, cmap=\"inferno\", vmin=vmin, vmax=vmax)\n",
    "    axins.set_xlim(x0, x1)\n",
    "    axins.set_ylim(y1, y0) \n",
    "    axins.set_xticks([])\n",
    "    axins.set_yticks([])\n",
    "    for spine in axins.spines.values():\n",
    "        spine.set_linewidth(2)\n",
    "        spine.set_edgecolor(\"yellow\")\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa9861c4-aeaa-46db-92b7-3198ca2eaab3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:ssp]",
   "language": "python",
   "name": "conda-env-ssp-py"
  },
  "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}