test

Dockerfile

@@ -0,0 +1,40 @@
+FROM mambaorg/micromamba:1.5.8
+SHELL ["/bin/bash", "-lc"]
+
+USER root
+RUN mkdir -p /var/lib/apt/lists/partial && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends git cmake make g++ && \
+    rm -rf /var/lib/apt/lists/*
+
+WORKDIR /work
+
+# Your conda env file is in binder/environment.yml inside the repo
+COPY binder/environment.yml /tmp/environment.yml
+RUN micromamba env create -f /tmp/environment.yml && micromamba clean -a -y
+
+# Build StochasticCIL
+RUN git clone https://github.com/epapoutsellis/StochasticCIL.git && \
+    cd StochasticCIL && \
+    git fetch --all --tags && \
+    git checkout svrg && \
+    git config user.email "hf@local" && \
+    git config user.name "HF Build" && \
+    (git tag -a v1.0 -m "Version 1.0" || true) && \
+    mkdir -p build && cd build && \
+    PREFIX="$(micromamba run -n ssp python -c 'import sys, os; print(os.path.dirname(os.path.dirname(sys.executable)))')" && \
+    cmake .. -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \
+      -DCONDA_BUILD=OFF \
+      -DCMAKE_BUILD_TYPE=Release \
+      -DLIBRARY_LIB="${PREFIX}/lib" \
+      -DLIBRARY_INC="${PREFIX}" \
+      -DCMAKE_INSTALL_PREFIX="${PREFIX}" \
+      -DPython_EXECUTABLE="${PREFIX}/bin/python" && \
+    make -j"$(nproc)" && \
+    make install
+
+COPY start.sh /start.sh
+RUN chmod +x /start.sh
+
+EXPOSE 7860
+CMD ["/start.sh"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vaggelis Papoutsellis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

PnP_BM3D_VR_Skip.ipynb

@@ -0,0 +1,687 @@
+{
+ "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
+}

ProxSkip.py

@@ -0,0 +1,124 @@
+from cil.optimisation.algorithms import Algorithm
+import numpy as np
+import logging
+
+
+class ProxSkip(Algorithm):
+    
+
+    r"""Proximal Skip  (ProxSkip) algorithm, see "ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!†"
+    
+        Parameters
+        ----------
+
+        initial : DataContainer
+                  Initial point for the ProxSkip algorithm. 
+        f : Function
+            A smooth function with Lipschitz continuous gradient.
+        g : Function
+            A convex function with a "simple" proximal.
+        prob : positive :obj:`float`
+             Probability to skip the proximal step. If :code:`prob=1`, proximal step is used in every iteration.             
+        step_size : positive :obj:`float`
+            Step size for the ProxSkip algorithm and is equal to 1./L where L is the  Lipschitz constant for the gradient of f.          
+            
+     """
+
+
+    def __init__(self, initial, f, g, step_size, prob, seed=None, **kwargs):
+        """ Set up of the algorithm
+        """        
+
+        super(ProxSkip, self).__init__(**kwargs)
+
+        self.f = f # smooth function
+        self.g = g # proximable
+        self.step_size = step_size
+        self.prob = prob
+        self.rng = np.random.default_rng(seed)
+        self.set_up(initial, f, g, step_size, prob, **kwargs)
+        self.thetas = []
+        self.prox_iterates = []
+ 
+                  
+    def set_up(self, initial, f, g, step_size, prob, **kwargs):
+        
+        logging.info("{} setting up".format(self.__class__.__name__, ))        
+        
+        self.initial = initial[0]
+
+        self.x = initial[0].copy()   
+        self.xhat_new = initial[0].copy()
+        self.x_new = initial[0].copy()
+        self.ht = initial[1].copy() #self.f.gradient(initial)      
+        # self.ht = self.f.gradient(initial) 
+
+        self.configured = True
+        
+        # count proximal and non proximal steps
+        self.use_prox = 0
+        # self.no_use_prox = 0
+
+        logging.info("{} configured".format(self.__class__.__name__, ))
+                                  
+
+    def update(self):
+        r""" Performs a single iteration of the ProxSkip algorithm        
+        """
+        
+        self.f.gradient(self.x, out=self.xhat_new)
+        self.xhat_new -= self.ht
+        self.x.sapyb(1., self.xhat_new, -self.step_size, out=self.xhat_new)
+
+        theta = self.rng.random() < self.prob 
+        if self.iteration==0:
+            theta = 1
+        self.thetas.append(theta)
+
+        if theta==1:
+            # print("here")
+            # Proximal step is used
+            self.g.proximal(self.xhat_new - (self.step_size/self.prob)*self.ht, self.step_size/self.prob, out=self.x_new)
+            self.ht.sapyb(1., (self.x_new - self.xhat_new), (self.prob/self.step_size), out=self.ht)            
+            self.use_prox+=1
+            # self.prox_iterates.append(self.x.copy())
+        else:
+            # Proximal step is skipped
+            # print("here1")
+            self.x_new.fill(self.xhat_new)
+
+    def _update_previous_solution(self):
+        """ Swaps the references to current and previous solution based on the :func:`~Algorithm.update_previous_solution` of the base class :class:`Algorithm`.
+        """
+        tmp = self.x_new
+        self.x = self.x_new
+        self.x = tmp
+
+    def get_output(self):
+        " Returns the current solution. "
+        return self.x        
+      
+                  
+    def update_objective(self):
+
+        """ Updates the objective
+
+        .. math:: f(x) + g(x)
+
+        """        
+
+        fun_g = self.g(self.x)
+        fun_f = self.f(self.x)
+        p1 = fun_f + fun_g
+        self.loss.append( p1 )
+
+    def proximal_evaluations(self):
+
+        prox_evals = []
+
+        for i in self.iterations:
+            if self.rng.random() < self.prob:
+                prox_evals.append(1)
+            else:
+                prox_evals.append(0)
+        return prox_evals

README.md

@@ -1,10 +1,74 @@
 ---
-title: Split Skip And Play
-emoji: 🌍
-colorFrom: pink
-colorTo: purple
+title: Split-Skip-and-Play
 sdk: docker
-pinned: false
+app_port: 7860
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Split-Skip-and-Play
+
+Code to reproduce results for paper ["SPLIT, SKIP AND PLAY: VARIANCE-REDUCED PROXSKIP FOR TOMOGRAPHY
+RECONSTRUCTION IS EXTREMELY FAST"](https://arxiv.org/abs/2602.09527) by Evangelos Papoutsellis, Zeljko Kereta, Kostas Papafitsoros. This is an extension of the paper ["Why do we regularise in every iteration for imaging inverse problems?"](https://arxiv.org/abs/2411.00688).
+
+![](PnP_BM3D_1min_reconstruction.gif)
+
+[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/epapoutsellis/Split-Skip-and-Play/HEAD)  
+
+### Abstract
+Many modern iterative solvers for large-scale tomographic reconstruction incur two major computational costs per iteration: expensive forward/adjoint projections to update the data fidelity term and costly proximal computations for the regulariser, often done via inner iterations. This paper studies for the first time the application of methods that couple randomised skipping of the proximal with variance-reduced subset-based optimisation of data-fit term, to simultaneously reduce both costs in challenging tomographic reconstruction tasks. We provide a series of experiments using both synthetic and real data, demonstrating striking speed-ups of the order 5x--20x compared to the non-skipped counterparts which have been so far the standard approach for efficiently solving these problems. Our work lays the groundwork for broader adoption of these methods in inverse problems.
+
+
+### Installation
+
+We use the [Core Imaging Library (CIL)](https://github.com/TomographicImaging/CIL) with some additional [new functionalities](#Appendix). Code and installation tested on macOS (Apple M2 Pro), Linux, and Windows 10.
+
+```
+conda create --name ssp -c conda-forge python=3.12 "numpy<2.0" cmake scipy six cython numba pillow jupyterlab scikit-learn dask "zarr<3" pywavelets astra-toolbox tqdm nb_conda_kernels
+
+conda activate ssp
+pip install bm3d xdesign scikit-image
+pip install "setuptools<82" --upgrade
+
+git clone https://github.com/epapoutsellis/StochasticCIL.git
+cd StochasticCIL
+git checkout svrg
+git tag -a v1.0 -m "Version 1.0"
+mkdir build
+cd build
+cmake ../ -DCONDA_BUILD=OFF -DCMAKE_BUILD_TYPE="Release" -DLIBRARY_LIB=$CONDA_PREFIX/lib -DLIBRARY_INC=$CONDA_PREFIX -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
+make install
+```
+For windows: `cmake ../ -DCONDA_BUILD=OFF`, `cmake --build . --target install`
+
+### Citation
+
+```
+@article{2602.09527,
+Author = {Evangelos Papoutsellis and Zeljko Kereta and Kostas Papafitsoros},
+Title = {Split, Skip and Play: Variance-Reduced ProxSkip for Tomography Reconstruction is Extremely Fast},
+Year = {2026},
+Eprint = {arXiv:2602.09527},
+}
+```
+
+### New CIL Functionalities
+
+1) Splitting methods for Acquisition Data compatible with CIL and [SIRF](https://github.com/SyneRBI/SIRF).
+2) Sampling methods for CIL Stochastic Functions (used also by SPDHG).
+3) ApproximateGradientSumFunction (Base class for CIL Stochastic Functions)
+4) SGFunction
+5) SAGFunction
+6) SAGAFunction
+7) SVRGFunction
+8) LSVRGFunction
+9) PGA (Proximal Gradient Algorithm), base class for GD, ISTA, FISTA.  Designed for fixed and/or adaptive Preconditioners and step sizes.
+10) Preconditioner (base class). An instance of Preconditioner can be passed to GD, ISTA, FISTA. 
+11) StepSizeMethod (base class for step size search), including standard Armijo/Backtracking.
+12) Callback utilities, including standard metrics (compared to a reference) and statistics of iterates. Any function (CIL), or callables from other libraries can be used.
+13) PD3O (Primal Dual Three Operator Splitting Algorithm) which can be combined with a Stochastic CIL function.
+14) SIRF Priors classes can be used in the CIL Optimisation Framework for free. In addition, SIRF ObjectiveFunction Classes can be used. This allows more flexibility for SIRF users to have control and monitor a CIL algorithm.
+
+### Acknowledgements    
+
+E. Papoutsellis acknowledges funding through the Innovate UK Analysis for Innovators (A4i) program: "Denoising of chemical imaging and tomography data" under which the experiments were initially conducted, CCPi (EPSRC grant EP/T026677/1), CCP SyneRBI (EPSRC grant EP/T026693/1).
+

TV_Tomography_Reconstruction_VR_Skip.ipynb

@@ -0,0 +1,556 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e82fbdab-f56e-471c-a314-63c97d08d197",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from cil.framework import AcquisitionGeometry\n",
+    "from cil.optimisation.algorithms import ISTA, FISTA\n",
+    "from cil.optimisation.functions import SVRGFunction, LSVRGFunction, LeastSquares\n",
+    "from cil.optimisation.utilities import RandomSampling, MetricsDiagnostics\n",
+    "from cil.plugins.astra import FBP, ProjectionOperator\n",
+    "from TotalVariation import TotalVariationNew\n",
+    "from ProxSkip import ProxSkip\n",
+    "from utils import create_circular_mask, StoppingCriterion\n",
+    "\n",
+    "import numpy as np\n",
+    "import zarr\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.rcParams['lines.linewidth'] = 5\n",
+    "plt.rcParams['lines.markersize'] = 20\n",
+    "plt.rcParams['font.size'] = 40\n",
+    "plt.rcParams[\"image.cmap\"]  = \"inferno\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e5e68083-d6ff-4cc8-9f4e-6e7fbcd38d43",
+   "metadata": {},
+   "source": [
+    "### Total Variation Reconstruction on a real dataset\n",
+    "\n",
+    "In this notebook, we play the following game. We start from a high-accuracy reference solution of\n",
+    "\n",
+    "$$ \n",
+    "\\begin{equation}\n",
+    "\\min_{x\\in\\mathbb{R}^n} \\; \\frac{1}{2}\\|Ax - b\\|_{2}^{2}\n",
+    "+ \\alpha \\mathrm{TV}(x)\n",
+    "+ \\mathbb{I}_{\\{x \\ge 0\\}}(x),\n",
+    "\\end{equation}\n",
+    "$$\n",
+    "\n",
+    "and we fix a target accuracy threshold $\\varepsilon>0$. The goal is to compare algorithms by **time-to-accuracy**: namely, determine which algorithm reaches $\\frac{\\|x_k - x^*\\|_2}{\\|x^*\\|_2}$ below $\\varepsilon$ in the **shortest runtime**.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1d21a834-36d9-4a3c-b036-c9fdf5c3dfbe",
+   "metadata": {},
+   "source": [
+    "### Load real dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee55e26e-95ac-48d0-91fa-8ffbe38825fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sino = zarr.load(\"data/NiPd_spent_01_microct_rings_removed_2D.zarr\")\n",
+    "\n",
+    "print(f\"sinogram shape is {sino.shape}\")\n",
+    "\n",
+    "_, horizontal = sino.shape\n",
+    "\n",
+    "angles_list = np.linspace(0, np.pi, 800)[::2]\n",
+    "ag2D = AcquisitionGeometry.create_Parallel2D().\\\n",
+    "        set_panel(horizontal).\\\n",
+    "        set_angles(angles_list, angle_unit=\"radian\").\\\n",
+    "        set_labels(['angle','horizontal'])\n",
+    "ig2D = ag2D.get_ImageGeometry()\n",
+    "\n",
+    "data2D = ag2D.allocate()\n",
+    "data2D.fill(sino)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "41b24ec4-f012-4fa1-a2de-95b162228de1",
+   "metadata": {},
+   "source": [
+    "### FBP reconstruction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f350f391-d445-432e-a772-24810e3a6266",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fbp = FBP(ig2D, ag2D, device=\"cpu\")(data2D)\n",
+    "fbp.array[fbp.array<0] = 0\n",
+    "plt.imshow(fbp.array)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b48aee65-6bae-4520-ae24-35d14555adbf",
+   "metadata": {},
+   "source": [
+    "### Load high accuracy solution "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79240538-ca62-4f0b-b1e6-df38b8d48f4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "alpha = 0.1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f251eeb-3298-46ab-a8f3-bb8782d9ffd9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pdhg_optimal_info = zarr.open_group(\"data/pdhg_optimal_finden_tv_alpha_{}_explicit_precond_maxiterations_200000.zarr\".format(alpha))\n",
+    "\n",
+    "pdhg_optimal_np = pdhg_optimal_info[\"solution\"][:]\n",
+    "pdhg_optimal_cil = ig2D.allocate()\n",
+    "pdhg_optimal_cil.fill(pdhg_optimal_np)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44e0cb98-199d-46b1-9250-d876446a36f1",
+   "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": "6079afd7-9f2b-43bf-9a9a-59a589c52b8e",
+   "metadata": {},
+   "source": [
+    "### Define Stopping Criteria and Metrics (PSNR/SSIM)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "baa0f235-7dab-430a-87c6-4b6d5a72f31a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "h, w = pdhg_optimal_cil.shape\n",
+    "mask = create_circular_mask(h, w, center=(170,165), radius=150)\n",
+    "\n",
+    "def NRSE(x, y, **kwargs):\n",
+    "    return np.sqrt(np.sum(np.abs(x - y*mask.ravel())**2))/np.sqrt(np.sum(x**2))\n",
+    "\n",
+    "cb_metrics = MetricsDiagnostics(reference_image=pdhg_optimal_cil*mask, \n",
+    "                         metrics_dict={\"rse\":NRSE}) \n",
+    "\n",
+    "epsilon = 0.99e-5\n",
+    "inner_its = 10\n",
+    "epochs = 100\n",
+    "max_iteration = 5000\n",
+    "seed_skip = 42\n",
+    "seed_sampling = 42"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4146831-c967-407d-a1f0-aa0ff7dca979",
+   "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": "743c77b3-49a0-46c7-bf43-b322f78b1753",
+   "metadata": {},
+   "source": [
+    "### Run deterministic algorithms: No data splitting, no skipping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86570298-273c-4ab5-a04b-96e70df402ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "K = ProjectionOperator(ig2D, ag2D, device=\"cpu\")\n",
+    "F = LeastSquares(A=K, b=data2D, c=0.5)\n",
+    "G = alpha * TotalVariationNew(max_iteration = inner_its, tolerance=None, correlation='Space',\n",
+    "                                     backend='c', lower=0, upper=np.infty, isotropic=True, \n",
+    "                                     split=False, info=False, strong_convexity_constant=0,\n",
+    "                                     warm_start=True)\n",
+    "initial = ig2D.allocate()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "71eca541-0265-4af8-bbbf-c7b791bbe362",
+   "metadata": {},
+   "source": [
+    "### FISTA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f5a3e483-7bdc-45da-b35b-c8147f1c6ce2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "step_size = 1./F.L\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon) \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_error])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "47e90863-2b50-4d61-a9b3-e0918f3bf316",
+   "metadata": {},
+   "source": [
+    "### ISTA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f11fc38d-cca0-4fad-bd5c-86f3e3402e41",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "step_size = 1.99/F.L\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon) \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_error])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9dd6d3f8-7418-44b4-9a6a-f831f116cced",
+   "metadata": {},
+   "source": [
+    "### Run ProxSkip: Skip the regulariser"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "72f458c9-6f9c-4a2d-8877-19a9dc5a638f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prob = 0.1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "805fd84c-9f20-45ba-8188-957f044cee89",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "step_size = 1.99/F.L\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon) \n",
+    "proxskip = ProxSkip(initial = [initial, initial], f = F, step_size = step_size, g=G, \n",
+    "        update_objective_interval = 1, prob=prob, seed = seed_skip,\n",
+    "        max_iteration = max_iteration) \n",
+    "proxskip.run(verbose=0, callback=[cb_metrics, cb_error])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f6d2c3a3-ec14-4b9c-ad87-76e62f1161fe",
+   "metadata": {},
+   "source": [
+    "### Run stochastic algorithms: Data splitting, no skipping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8fbc2e91-2a5d-4f50-9f60-c2ced3b41038",
+   "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": "2364fecd-aab6-4653-8ba2-cf855a9ebda9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nsub = 200\n",
+    "data_split, method = data2D.split_to_subsets(nsub, method= \"ordered\", info=True)\n",
+    "list_func = list_of_functions(data_split) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a25d5716-cba8-4c78-9327-043d351ed2b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selection = RandomSampling(len(list_func), nsub, seed=seed_sampling)\n",
+    "Fsvrg = SVRGFunction(list_func, selection = selection, update_frequency=len(list_func))\n",
+    "Fsvrg.initial = initial\n",
+    "step_size = 1./(Fsvrg.L)\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon, epochs = 100) \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_error])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1dce5d0d-4873-40f9-8339-d35fd5eac23f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selection = RandomSampling(len(list_func), nsub, seed=seed_sampling)\n",
+    "Flsvrg = LSVRGFunction(list_func, selection = selection, update_prob=1./len(list_func))\n",
+    "Flsvrg.initial = initial\n",
+    "step_size = 1./(Fsvrg.L)\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon, epochs = 100) \n",
+    "prox_lsvrg = ISTA(initial = initial, f = Flsvrg, step_size = step_size, g=G, \n",
+    "            update_objective_interval = 1, \n",
+    "            max_iteration = max_iteration) \n",
+    "prox_lsvrg.run(verbose=0, callback=[cb_metrics, cb_error])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "59988152-5f09-4950-a770-a8b4d435ee44",
+   "metadata": {},
+   "source": [
+    "### Run stochastic algorithms: Data splitting and skipping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4d74a1ae-4f52-4e04-a4a4-03b4488f5fef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selection = RandomSampling(len(list_func), nsub, seed=seed_sampling)\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",
+    "cb_error = StoppingCriterion(epsilon=epsilon, epochs = 100) \n",
+    "prox_svrg_skip = ProxSkip(initial = [initial, initial], f = Fsvrg_skip, step_size = step_size, g=G, \n",
+    "            update_objective_interval = 1, prob = prob,\n",
+    "            max_iteration = max_iteration) \n",
+    "prox_svrg_skip.run(verbose=0, callback=[cb_metrics, cb_error])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2244d1db-1338-46ca-8355-bfcc03ce7e56",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "selection = RandomSampling(len(list_func), nsub, seed=seed_sampling)\n",
+    "Flsvrg_skip = LSVRGFunction(list_func, selection = selection, update_prob=1./len(list_func))\n",
+    "Flsvrg_skip.initial = initial\n",
+    "step_size = 1./(Flsvrg_skip.L)\n",
+    "cb_error = StoppingCriterion(epsilon=epsilon, epochs = 100) \n",
+    "prox_lsvrg_skip = ProxSkip(initial = [initial, initial], f = Flsvrg_skip, step_size = step_size, g=G, \n",
+    "            update_objective_interval = 1, prob = prob,\n",
+    "            max_iteration = max_iteration) \n",
+    "prox_lsvrg_skip.run(verbose=0, callback=[cb_metrics, cb_error])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f0d1a611-d16c-432e-a6e5-56cae3b7e640",
+   "metadata": {},
+   "source": [
+    "### Plot PSNR/SSIM progress\n",
+    "- with respect to iteration\n",
+    "- with respect to time"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee4475aa-9018-427e-82c4-63d0f487e899",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def find_iter_to_error(algo, error):\n",
+    "    rse = np.asarray(algo.rse)\n",
+    "    idx = np.where(rse < error)[0]\n",
+    "    return np.nan if idx.size == 0 else int(idx[0])\n",
+    "\n",
+    "def time_to_iter(algo, it):\n",
+    "    \"\"\"Time up to and including iteration it.\"\"\"\n",
+    "    if it is None or (isinstance(it, float) and np.isnan(it)):\n",
+    "        return np.nan\n",
+    "    it = int(it)\n",
+    "    return float(np.sum(np.asarray(algo.timing)[:it+1]))\n",
+    "\n",
+    "def time_to_error(algo, error):\n",
+    "    it = find_iter_to_error(algo, error)\n",
+    "    return time_to_iter(algo, it)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a74d852a-744f-46ee-927b-2152c130bd6d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "t_ista = time_to_error(ista, epsilon)\n",
+    "t_proxskip = time_to_error(proxskip, epsilon)\n",
+    "t_prox_svrg =time_to_error(prox_svrg, epsilon)\n",
+    "t_prox_lsvrg = time_to_error(prox_lsvrg, epsilon)\n",
+    "t_prox_svrg_skip = time_to_error(prox_svrg_skip, epsilon)\n",
+    "t_prox_lsvrg_skip = time_to_error(prox_lsvrg_skip, epsilon)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f28efd88-80de-494b-a991-207df0444ad8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(2,1,figsize=(30, 25))\n",
+    "\n",
+    "fig.subplots_adjust(top=0.80)\n",
+    "\n",
+    "ax[0].semilogy(prox_svrg.rse[:-1],label=f\"ProxSVRG (N=200)\")\n",
+    "ax[0].semilogy(prox_lsvrg.rse[:-1],label=f\"ProxLSVRG (N=200)\")\n",
+    "ax[0].semilogy(proxskip.rse[:-1],label=f\"ProxSkip (p=0.1)\")\n",
+    "ax[0].semilogy(prox_svrg_skip.rse[:-1],label=f\"ProxSVRGSkip (N=200, p=0.05)\")\n",
+    "ax[0].semilogy(prox_lsvrg_skip.rse[:-1],label=f\"ProxLSVRGSkip (N=200, p=0.05)\")\n",
+    "ax[0].semilogy(ista.rse[:-1],label=f\"ISTA\")\n",
+    "ax[0].semilogy(fista.rse[:-1],label=f\"FISTA\")\n",
+    "ax[0].set_xlabel(\"Iteration\")\n",
+    "ax[0].set_ylabel(r\"$\\frac{\\|x_{k} - x^{*}\\|_{2}}{\\|x^{*}\\|_{2}}$\")\n",
+    "ax[0].grid(True, which=\"major\")\n",
+    "\n",
+    "ax[0].text(\n",
+    "    0.25, 0.9, \"Proximal-TV with 10 iterations\",\n",
+    "    transform=ax[0].transAxes,\n",
+    "    ha=\"left\", va=\"top\",\n",
+    "    fontsize=45,\n",
+    "    bbox=dict(boxstyle=\"round,pad=0.35\", facecolor=\"white\", edgecolor=\"none\")\n",
+    ")\n",
+    "ax[0].legend(ncols=2, loc=\"lower center\", bbox_to_anchor=(0.5, 1.00), frameon=True)\n",
+    "\n",
+    "ax[1].semilogy(np.cumsum(prox_svrg.timing), prox_svrg.rse[:-1],label=f\"ProxSVRG (N=200)\")\n",
+    "ax[1].semilogy(np.cumsum(prox_lsvrg.timing), prox_lsvrg.rse[:-1],label=f\"ProxLSVRG (N=200)\")\n",
+    "ax[1].semilogy(np.cumsum(proxskip.timing), proxskip.rse[:-1],label=f\"ProxSkip (p=0.1)\")\n",
+    "ax[1].semilogy(np.cumsum(prox_svrg_skip.timing), prox_svrg_skip.rse[:-1],label=f\"ProxSVRGSkip (N=200, p=0.05)\")\n",
+    "ax[1].semilogy(np.cumsum(prox_lsvrg_skip.timing), prox_lsvrg_skip.rse[:-1],label=f\"ProxLSVRGSkip (N=200, p=0.05)\")\n",
+    "ax[1].semilogy(np.cumsum(ista.timing), ista.rse[:-1],label=f\"ISTA\")\n",
+    "ax[1].semilogy(np.cumsum(fista.timing), fista.rse[:-1],label=f\"FISTA\")\n",
+    "ax[1].set_xlabel(\"Time (sec)\")\n",
+    "ax[1].set_ylabel(r\"$\\frac{\\|x_{k} - x^{*}\\|_{2}}{\\|x^{*}\\|_{2}}$\")\n",
+    "ax[1].grid(True, which=\"major\")\n",
+    "\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "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
+}

TotalVariation.py

@@ -0,0 +1,399 @@
+# -*- coding: utf-8 -*-
+#  Copyright 2020 United Kingdom Research and Innovation
+#  Copyright 2020 The University of Manchester
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+# Authors:
+# CIL Developers, listed at: https://github.com/TomographicImaging/CIL/blob/master/NOTICE.txt
+# Claire Delplancke (University of Bath)
+
+
+from cil.optimisation.functions import Function, IndicatorBox, MixedL21Norm, MixedL11Norm
+from cil.optimisation.operators import GradientOperator
+import numpy as np
+from numbers import Number
+import warnings
+import logging
+
+
+class TotalVariationNew(Function):
+
+    r""" Total variation Function
+
+    .. math:: \mathrm{TV}(u) := \|\nabla u\|_{2,1} = \sum \|\nabla u\|_{2},\, (\mbox{isotropic})
+
+    .. math:: \mathrm{TV}(u) := \|\nabla u\|_{1,1} = \sum \|\nabla u\|_{1}\, (\mbox{anisotropic})
+
+    Notes
+    -----
+
+    The :code:`TotalVariation` (TV) :code:`Function` acts as a composite function, i.e.,
+    the composition of the :class:`.MixedL21Norm` function and the :class:`.GradientOperator` operator,
+
+    .. math:: f(u) = \|u\|_{2,1}, \Rightarrow (f\circ\nabla)(u) = f(\nabla x) = \mathrm{TV}(u)
+
+    In that case, the proximal operator of TV does not have an exact solution and we use an iterative 
+    algorithm to solve:
+
+    .. math:: \mathrm{prox}_{\tau \mathrm{TV}}(b) := \underset{u}{\mathrm{argmin}} \frac{1}{2\tau}\|u - b\|^{2} + \mathrm{TV}(u)
+
+    The algorithm used for the proximal operator of TV is the Fast Gradient Projection algorithm (or FISTA)
+    applied to the _dual problem_ of the above problem, see :cite:`BeckTeboulle_b`, :cite:`BeckTeboulle_a`, :cite:`Zhu2010`.
+
+    See also "Multicontrast MRI Reconstruction with Structure-Guided Total Variation", Ehrhardt, Betcke, 2016.
+
+
+    Parameters
+    ----------
+
+    max_iteration : :obj:`int`, default = 5
+        Maximum number of iterations for the FGP algorithm to solve to solve the dual problem 
+        of the Total Variation Denoising problem (ROF). If warm_start=False, this should be around 100,
+        or larger, with a set tolerance. 
+    tolerance : :obj:`float`, default = None
+        Stopping criterion for the FGP algorithm used to to solve the dual problem 
+        of the Total Variation Denoising problem (ROF). If the difference between iterates in the FGP algorithm is less than the tolerance
+        the iterations end before the max_iteration number. 
+
+        .. math:: \|x^{k+1} - x^{k}\|_{2} < \mathrm{tolerance}
+
+    correlation : :obj:`str`, default = `Space`
+        Correlation between `Space` and/or `SpaceChannels` for the :class:`.GradientOperator`.
+    backend :  :obj:`str`, default = `c`      
+        Backend to compute the :class:`.GradientOperator`
+    lower : :obj:`'float`, default = None
+        A constraint is enforced using the :class:`.IndicatorBox` function, e.g., :code:`IndicatorBox(lower, upper)`.
+    upper : :obj:`'float`, default = None
+        A constraint is enforced using the :class:`.IndicatorBox` function, e.g., :code:`IndicatorBox(lower, upper)`.  
+    isotropic : :obj:`boolean`, default = True
+        Use either isotropic or anisotropic definition of TV.
+
+        .. math:: |x|_{2} = \sqrt{x_{1}^{2} + x_{2}^{2}},\, (\mbox{isotropic})
+
+        .. math:: |x|_{1} = |x_{1}| + |x_{2}|\, (\mbox{anisotropic})
+
+    split : :obj:`boolean`, default = False
+        Splits the Gradient into spatial gradient and spectral or temporal gradient for multichannel data.
+
+    info : :obj:`boolean`, default = False
+        Information is printed for the stopping criterion of the FGP algorithm used to solve the dual problem
+        of the Total Variation Denoising problem (ROF).
+
+    strong_convexity_constant : :obj:`float`, default = 0
+        A strongly convex term weighted by the :code:`strong_convexity_constant` (:math:`\gamma`) parameter is added to the Total variation. 
+        Now the :code:`TotalVariation` function is :math:`\gamma` - strongly convex and the proximal operator is
+
+        .. math:: \underset{u}{\mathrm{argmin}} \frac{1}{2\tau}\|u - b\|^{2} + \mathrm{TV}(u) + \frac{\gamma}{2}\|u\|^{2} \Leftrightarrow
+
+        .. math:: \underset{u}{\mathrm{argmin}} \frac{1}{2\frac{\tau}{1+\gamma\tau}}\|u - \frac{b}{1+\gamma\tau}\|^{2} + \mathrm{TV}(u) 
+
+    warm_start : :obj:`boolean`, default = True
+        If set to true, the FGP algorithm used to solve the dual problem of the Total Variation Denoising problem (ROF) is initiated by the final value from the previous iteration and not at zero. 
+        This allows the max_iteration value to be reduced to 5-10 iterations. 
+
+
+    Note
+    ----
+
+    With warm_start set to the default, True, the TV function will keep in memory the range of the gradient of the image to be denoised, i.e. N times the dimensionality of the image. This increases the memory requirements. 
+    However, during the evaluation of `proximal` the memory requirements will be unchanged as the same amount of memory will need to be allocated and deallocated. 
+
+    Note
+    ----
+
+    In the case where the Total variation becomes a :math:`\gamma` - strongly convex function, i.e.,
+
+    .. math:: \mathrm{TV}(u) + \frac{\gamma}{2}\|u\|^{2}
+
+    :math:`\gamma` should be relatively small, so as the second term above will not act as an additional regulariser.
+    For more information, see :cite:`Rasch2020`, :cite:`CP2011`.
+
+
+
+
+    Examples
+    --------
+
+    .. math:: \underset{u}{\mathrm{argmin}} \frac{1}{2}\|u - b\|^{2} + \alpha\|\nabla u\|_{2,1}
+
+    >>> alpha = 2.0
+    >>> TV = TotalVariation()
+    >>> sol = TV.proximal(b, tau = alpha)
+
+    Examples
+    --------
+
+    .. math:: \underset{u}{\mathrm{argmin}} \frac{1}{2}\|u - b\|^{2} + \alpha\|\nabla u\|_{1,1} + \mathbb{I}_{C}(u)
+
+    where :math:`C = \{1.0\leq u\leq 2.0\}`.
+
+    >>> alpha = 2.0
+    >>> TV = TotalVariation(isotropic=False, lower=1.0, upper=2.0)
+    >>> sol = TV.proximal(b, tau = alpha)    
+
+
+    Examples
+    --------
+
+    .. math:: \underset{u}{\mathrm{argmin}} \frac{1}{2}\|u - b\|^{2} + (\alpha\|\nabla u\|_{2,1} + \frac{\gamma}{2}\|u\|^{2})
+
+    >>> alpha = 2.0
+    >>> gamma = 1e-3
+    >>> TV = alpha * TotalVariation(isotropic=False, strong_convexity_constant=gamma)
+    >>> sol = TV.proximal(b, tau = 1.0)    
+
+    """
+
+    def __init__(self,
+                 max_iteration=10,
+                 tolerance=None,
+                 correlation="Space",
+                 backend="c",
+                 lower=None,
+                 upper=None,
+                 isotropic=True,
+                 split=False,
+                 info=False,
+                 strong_convexity_constant=0,
+                 warm_start=True):
+
+        super(TotalVariationNew, self).__init__(L=None)
+
+        # Regularising parameter = alpha
+        self.regularisation_parameter = 1.
+
+        self.iterations = max_iteration
+
+        self.tolerance = tolerance
+
+        # Total variation correlation (isotropic=Default)
+        self.isotropic = isotropic
+
+        # correlation space or spacechannels
+        self.correlation = correlation
+        self.backend = backend
+
+        # Define orthogonal projection onto the convex set C
+        if lower is None:
+            lower = -np.inf
+        if upper is None:
+            upper = np.inf
+        self.lower = lower
+        self.upper = upper
+        self.projection_C = IndicatorBox(lower, upper).proximal
+
+        # Setup GradientOperator as None. This is to avoid domain argument in the __init__
+        self._gradient = None
+        self._domain = None
+
+        self.info = info
+        if self.info:
+            warnings.warn(" `info` is deprecate. Please use logging instead.")
+
+        # splitting Gradient
+        self.split = split
+
+        # For the warm_start functionality
+        self.warm_start = warm_start
+        self._p2 = None
+
+        # Strong convexity for TV
+        self.strong_convexity_constant = strong_convexity_constant
+
+        # Define Total variation norm
+        if self.isotropic:
+            self.func = MixedL21Norm()
+        else:
+            self.func = MixedL11Norm()
+
+    def _get_p2(self):
+        r"""The initial value for the dual in the proximal calculation - allocated to zero in the case of warm_start=False
+          or initialised as the last iterate seen in the proximal calculation in the case warm_start=True ."""
+
+        if self._p2 is None:
+            return self.gradient.range_geometry().allocate(0)
+        else:
+            return self._p2
+
+    @property
+    def regularisation_parameter(self):
+        return self._regularisation_parameter
+
+    @regularisation_parameter.setter
+    def regularisation_parameter(self, value):
+        if not isinstance(value, Number):
+            raise TypeError(
+                "regularisation_parameter: expected a number, got {}".format(type(value)))
+        self._regularisation_parameter = value
+
+    def __call__(self, x):
+        r""" Returns the value of the TotalVariation function at :code:`x` ."""
+
+        try:
+            self._domain = x.geometry
+        except:
+            self._domain = x
+
+        # Compute Lipschitz constant provided that domain is not None.
+        # Lipschitz constant dependes on the GradientOperator, which is configured only if domain is not None
+        if self._L is None:
+            self.calculate_Lipschitz()
+
+        if self.strong_convexity_constant > 0:
+            strongly_convex_term = (
+                self.strong_convexity_constant/2)*x.squared_norm()
+        else:
+            strongly_convex_term = 0
+
+        return self.regularisation_parameter * self.func(self.gradient.direct(x)) + strongly_convex_term
+
+    def proximal(self, x, tau, out=None):
+        r""" Returns the proximal operator of the TotalVariation function at :code:`x` ."""
+
+        if self.strong_convexity_constant > 0:
+
+            strongly_convex_factor = (1 + tau * self.strong_convexity_constant)
+            x /= strongly_convex_factor
+            tau /= strongly_convex_factor
+
+        if out is None:
+            solution = self._fista_on_dual_rof(x, tau)
+        else:
+            self._fista_on_dual_rof(x, tau, out=out)
+
+        if self.strong_convexity_constant > 0:
+            x *= strongly_convex_factor
+            tau *= strongly_convex_factor
+
+        if out is None:
+            return solution
+
+    def _fista_on_dual_rof(self, x, tau, out=None):
+        r""" Runs the Fast Gradient Projection (FGP) algorithm to solve the dual problem 
+        of the Total Variation Denoising problem (ROF).
+
+        .. math: \max_{\|y\|_{\infty}<=1.} \frac{1}{2}\|\nabla^{*} y + x \|^{2} - \frac{1}{2}\|x\|^{2}
+
+        """
+        try:
+            self._domain = x.geometry
+        except:
+            self._domain = x
+
+        # Compute Lipschitz constant provided that domain is not None.
+        # Lipschitz constant depends on the GradientOperator, which is configured only if domain is not None
+        if self._L is None:
+            self.calculate_Lipschitz()
+
+        # initialise
+        t = 1
+
+        # dual variable - its content is overwritten during iterations
+        p1 = self.gradient.range_geometry().allocate(None)
+        p2 = self._get_p2()
+        tmp_q = p2.copy()
+
+        # multiply tau by -1 * regularisation_parameter here so it's not recomputed every iteration
+        # when tau is an array this is done inplace so reverted at the end
+        if isinstance(tau, Number):
+            tau_reg_neg = -self.regularisation_parameter * tau
+        else:
+            tau_reg_neg = tau
+            tau.multiply(-self.regularisation_parameter, out=tau_reg_neg)
+
+        should_return = False
+        if out is None:
+            should_return = True
+            out = self.gradient.domain_geometry().allocate(0)
+
+        for k in range(self.iterations):
+
+            t0 = t
+            self.gradient.adjoint(tmp_q, out=out)
+            out.sapyb(tau_reg_neg, x, 1.0, out=out)
+            self.projection_C(out, tau=None, out=out)
+
+            self.gradient.direct(out, out=p1)
+
+            multip = (-self.L)/tau_reg_neg
+
+            tmp_q.sapyb(1., p1, multip, out=tmp_q)
+
+            if self.tolerance is not None and k % 5 == 0:
+                p1 *= multip
+                error = p1.norm()
+                error /= tmp_q.norm()
+                if error < self.tolerance:
+                    break
+
+            self.func.proximal_conjugate(tmp_q, 1.0, out=p1)
+
+            t = (1 + np.sqrt(1 + 4 * t0 ** 2)) / 2
+            p1.subtract(p2, out=tmp_q)
+            tmp_q *= (t0-1)/t
+            tmp_q += p1
+
+            # switch p1 and p2 references
+            tmp = p1
+            p1 = p2
+            p2 = tmp
+        if self.warm_start:
+            self._p2 = p2
+
+        if self.info:
+            if self.tolerance is not None:
+                logging.info(
+                    "Stop at {} iterations with tolerance {} .".format(k, error))
+            else:
+                logging.info("Stop at {} iterations.".format(k))
+
+        # return tau to its original state if it was modified
+        if id(tau_reg_neg) == id(tau):
+            tau_reg_neg.divide(-self.regularisation_parameter, out=tau)
+
+        if should_return:
+            return out
+
+    def convex_conjugate(self, x):
+        r""" Returns the value of convex conjugate of the TotalVariation function at :code:`x` ."""
+        return 0.0
+
+    def calculate_Lipschitz(self):
+        r""" Default value for the Lipschitz constant."""
+
+        # Compute the Lipschitz parameter from the operator if possible
+        # Leave it initialised to None otherwise
+        self._L = (1./self.gradient.norm())**2
+
+    @property
+    def gradient(self):
+        r""" GradientOperator is created if it is not instantiated yet. The domain of the `_gradient`,
+        is created in the `__call__` and `proximal` methods. 
+
+        """
+        if self._domain is not None:
+            self._gradient = GradientOperator(
+                self._domain, correlation=self.correlation, backend=self.backend)
+        else:
+            raise ValueError(
+                " The domain of the TotalVariation is {}. Please use the __call__ or proximal methods first before calling gradient.".format(self._domain))
+
+        return self._gradient
+
+    def __rmul__(self, scalar):
+        if not isinstance(scalar, Number):
+            raise TypeError(
+                "scalar: Expected a number, got {}".format(type(scalar)))
+        self.regularisation_parameter *= scalar
+        return self

binder/.ipynb_checkpoints/environment-checkpoint.yml

@@ -0,0 +1,27 @@
+name: ssp
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  - "numpy<2.0"
+  - "zarr<3"
+
+  - cmake
+  - scipy
+  - six
+  - cython
+  - numba
+  - pillow
+  - pywavelets
+  - astra-toolbox
+  - tqdm
+  - "setuptools<82"
+  - git
+
+  # pip deps
+  - pip
+  - wheel
+  - pip:
+      - bm3d
+      - xdesign
+      - scikit-image

binder/.ipynb_checkpoints/postBuild-checkpoint

@@ -0,0 +1,46 @@
+#!/bin/bash
+set -euxo pipefail
+
+python -m ipykernel install --user --name ssp --display-name "Python (ssp)"
+
+# Pick the active environment prefix in Binder
+PREFIX="${CONDA_PREFIX:-${NB_PYTHON_PREFIX:-}}"
+if [ -z "${PREFIX}" ]; then
+  PREFIX="$(python -c 'import sys, os; print(os.path.dirname(os.path.dirname(sys.executable)))')"
+fi
+echo "Using PREFIX=${PREFIX}"
+
+if [ ! -d "StochasticCIL" ]; then
+  git clone https://github.com/epapoutsellis/StochasticCIL.git
+fi
+
+cd StochasticCIL
+git fetch --all --tags
+git checkout svrg
+
+# Ensure annotated tag for `git describe`
+git config user.email "binder@local"
+git config user.name "Binder Build"
+
+if git rev-parse -q --verify refs/tags/v1.0 >/dev/null; then
+  if [ "$(git cat-file -t v1.0)" != "tag" ]; then
+    git tag -d v1.0
+    git tag -a v1.0 -m "Version 1.0"
+  fi
+else
+  git tag -a v1.0 -m "Version 1.0"
+fi
+
+mkdir -p build
+cd build
+
+cmake ../ -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \
+  -DCONDA_BUILD=OFF \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DLIBRARY_LIB="${PREFIX}/lib" \
+  -DLIBRARY_INC="${PREFIX}" \
+  -DCMAKE_INSTALL_PREFIX="${PREFIX}" \
+  -DPython_EXECUTABLE="${PREFIX}/bin/python"
+
+make -j"$(nproc)"
+make install

binder/environment.yml

@@ -0,0 +1,28 @@
+name: ssp
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  - "numpy<2.0"
+  - "zarr<3"
+  - jupyterlab
+  - ipykernel
+  - cmake
+  - scipy
+  - six
+  - cython
+  - numba
+  - pillow
+  - pywavelets
+  - astra-toolbox
+  - tqdm
+  - "setuptools<82"
+  - git
+
+  # pip deps
+  - pip
+  - wheel
+  - pip:
+      - bm3d
+      - xdesign
+      - scikit-image

binder/postBuild

@@ -0,0 +1,46 @@
+#!/bin/bash
+set -euxo pipefail
+
+python -m ipykernel install --user --name ssp --display-name "Python (ssp)"
+
+# Pick the active environment prefix in Binder
+PREFIX="${CONDA_PREFIX:-${NB_PYTHON_PREFIX:-}}"
+if [ -z "${PREFIX}" ]; then
+  PREFIX="$(python -c 'import sys, os; print(os.path.dirname(os.path.dirname(sys.executable)))')"
+fi
+echo "Using PREFIX=${PREFIX}"
+
+if [ ! -d "StochasticCIL" ]; then
+  git clone https://github.com/epapoutsellis/StochasticCIL.git
+fi
+
+cd StochasticCIL
+git fetch --all --tags
+git checkout svrg
+
+# Ensure annotated tag for `git describe`
+git config user.email "binder@local"
+git config user.name "Binder Build"
+
+if git rev-parse -q --verify refs/tags/v1.0 >/dev/null; then
+  if [ "$(git cat-file -t v1.0)" != "tag" ]; then
+    git tag -d v1.0
+    git tag -a v1.0 -m "Version 1.0"
+  fi
+else
+  git tag -a v1.0 -m "Version 1.0"
+fi
+
+mkdir -p build
+cd build
+
+cmake ../ -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \
+  -DCONDA_BUILD=OFF \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DLIBRARY_LIB="${PREFIX}/lib" \
+  -DLIBRARY_INC="${PREFIX}" \
+  -DCMAKE_INSTALL_PREFIX="${PREFIX}" \
+  -DPython_EXECUTABLE="${PREFIX}/bin/python"
+
+make -j"$(nproc)"
+make install

environment.yml

@@ -0,0 +1,27 @@
+name: ssp
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  - "numpy<2.0"
+  - "zarr<3"
+
+  - cmake
+  - scipy
+  - six
+  - cython
+  - numba
+  - pillow
+  - pywavelets
+  - astra-toolbox
+  - tqdm
+  - "setuptools<82"
+  - git
+
+  # pip deps
+  - pip
+  - wheel
+  - pip:
+      - bm3d
+      - xdesign
+      - scikit-image

start.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+TOKEN="${JUPYTER_TOKEN:-}"
+
+exec micromamba run -n ssp jupyter lab \
+  --allow-root \
+  --ip=0.0.0.0 --port=7860 \
+  --no-browser \
+  --ServerApp.allow_remote_access=True \
+  --ServerApp.root_dir=/work \
+  --IdentityProvider.token="${TOKEN}"

utils.py

@@ -0,0 +1,128 @@
+from cil.optimisation.utilities import AlgorithmDiagnostics
+import numpy as np
+from bm3d import bm3d, BM3DStages
+from cil.optimisation.functions import Function
+
+class StoppingCriterionTime(AlgorithmDiagnostics):
+
+    def __init__(self, time):
+
+        self.time = time
+        super(StoppingCriterionTime, self).__init__(verbose=0)
+
+        self.should_stop = False
+        
+    def _should_stop(self):
+
+        return self.should_stop       
+            
+    def __call__(self, algo):
+
+        if algo.iteration==0:
+            algo.should_stop = self._should_stop  
+
+        stop_crit = np.sum(algo.timing)>self.time
+        if stop_crit:
+            self.should_stop = True                 
+            print("Stop at {} time {}".format(algo.iteration, np.sum(algo.timing)))
+
+
+class BM3DFunction(Function):
+    """
+    PnP 'regulariser' whose proximal applies BM3D denoising.
+
+    In PnP-ISTA/FISTA we typically use a FIXED BM3D sigma (regularization strength),
+    independent of the gradient step-size tau.
+    Optionally apply damping: (1-gamma) z + gamma * BM3D(z).
+    """
+
+    def __init__(self, sigma, gamma=1.0, profile="np",
+                 stage_arg=BM3DStages.ALL_STAGES, positivity=True):
+        self.sigma = float(sigma)      # BM3D noise parameter
+        self.gamma = float(gamma)      # damping in (0,1]
+        if not (0.0 < self.gamma <= 1.0):
+            raise ValueError("gamma must be in (0,1].")
+        self.profile = profile
+        self.stage_arg = stage_arg
+        self.positivity = positivity
+        super().__init__()
+
+    def __call__(self, x):
+        return 0.0
+
+    def convex_conjugate(self, x):
+        return 0.0
+
+    def _denoise(self, znp: np.ndarray) -> np.ndarray:
+        z = np.asarray(znp, dtype=np.float32)
+        # BM3D expects sigma as noise std (same units as the image)
+        return bm3d(z, sigma_psd=self.sigma,
+                    profile=self.profile,
+                    stage_arg=self.stage_arg).astype(np.float32)
+
+    def proximal(self, x, tau, out=None):
+        z = x.array.astype(np.float32, copy=False)
+        d = self._denoise(z)
+
+        # damping/relaxation (recommended if you see oscillations)
+        u = (1.0 - self.gamma) * z + self.gamma * d
+
+        if self.positivity:
+            u = np.maximum(u, 0.0)
+
+        if out is None:
+            out = x * 0.0
+        out.fill(u)
+        return out
+
+
+def create_circular_mask(h, w, center=None, radius=None):
+
+    if center is None: 
+        center = (int(w/2), int(h/2))
+    if radius is None: 
+        radius = min(center[0], center[1], w-center[0], h-center[1])
+
+    Y, X = np.ogrid[:h, :w]
+    dist_from_center = np.sqrt((X - center[0])**2 + (Y-center[1])**2)
+
+    mask = dist_from_center <= radius
+    return mask
+
+
+class StoppingCriterion(AlgorithmDiagnostics):
+    def __init__(self, epsilon, epochs=None):
+        self.epsilon = epsilon
+        self.epochs = epochs
+        super().__init__(verbose=0)
+        self.should_stop = False
+        self.rse_reached = False
+
+    def _should_stop(self):
+        return self.should_stop
+
+    def __call__(self, algo):
+
+        if algo.iteration == 0:
+            algo.should_stop = self._should_stop
+
+        stop_rse = (algo.rse[-1] <= self.epsilon)
+
+        stop_epochs = False
+        if self.epochs is not None:
+            try:
+                dp = algo.f.data_passes
+                dp_last = dp[-1] if hasattr(dp, "__len__") else dp
+                stop_epochs = (dp_last > self.epochs)
+            except AttributeError:
+                stop_epochs = False  
+                
+        stop = stop_rse or stop_epochs
+
+        if algo.iteration < algo.max_iteration:
+            if stop:
+                self.rse_reached = stop_rse
+                self.should_stop = True
+                print(f"Accuracy reached at {algo.iteration}, time = {np.sum(algo.timing):.4f}, NRSE = {algo.rse[-1]:.4e}")
+        else:
+            print(f"Failed to reach accuracy. Stop at {algo.iteration}, time = {np.sum(algo.timing):.4f}, NRSE = {algo.rse[-1]:.4e}")