{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# SPECTrecon OS-EM\n",
    "\n",
    "This page illustrates\n",
    "Ordered-subset expectation-maximization (OS-EM)\n",
    "image reconstruction with the Julia package\n",
    "[`SPECTrecon`](https://github.com/JuliaImageRecon/SPECTrecon.jl)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Setup"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Packages needed here."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using SPECTrecon: psf_gauss, SPECTplan, project!, backproject!, Ablock\n",
    "using SPECTrecon: osem, osem!, mlem!\n",
    "using LinearMapsAA: LinearMapAA, LinearMapAO\n",
    "using LinearAlgebra: mul!\n",
    "using Distributions: Poisson\n",
    "using MIRTjim: jim, prompt\n",
    "using Plots: scatter, plot!, default; default(markerstrokecolor=:auto)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following line is helpful when running this file as a script;\n",
    "this way it will prompt user to hit a key after each figure is displayed."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "isinteractive() ? jim(:prompt, true) : prompt(:draw);"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Overview\n",
    "\n",
    "Ordered-subset expectation-maximization (OS-EM)\n",
    "is a commonly used algorithm for performing SPECT image reconstruction\n",
    "because of its favorable combination of image quality and speed.\n",
    "See\n",
    "[Hudson and Larkin, 1994](https://doi.org/10.1109/42.363108)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Simulation data"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "nx,ny,nz = 64,64,50\n",
    "T = Float32\n",
    "xtrue = zeros(T, nx,ny,nz) # simple phantom\n",
    "xtrue[(1nx÷4):(2nx÷3), 1ny÷5:(3ny÷5), 2nz÷6:(3nz÷6)] .= 1\n",
    "xtrue[(2nx÷5):(3nx÷5), 1ny÷5:(2ny÷5), 4nz÷6:(5nz÷6)] .= 2\n",
    "\n",
    "average(x) = sum(x) / length(x)\n",
    "function mid3(x::AbstractArray{<:Number,3}) # 3 central planes\n",
    "    (nx,ny,nz) = size(x)\n",
    "    xy = x[:,:,ceil(Int, nz÷2)]\n",
    "    xz = x[:,ceil(Int,end/2),:]\n",
    "    zy = x[ceil(Int, nx÷2),:,:]'\n",
    "    return [xy xz; zy fill(average(xy), nz, nz)]\n",
    "end\n",
    "jim(mid3(xtrue), \"Middle slices of xtrue\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## PSF"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Create a synthetic depth-dependent PSF for a single view"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "px = 11\n",
    "psf1 = psf_gauss( ; ny, px)\n",
    "jim(psf1, \"PSF for each of $ny planes\"; ratio=1)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "In general the PSF can vary from view to view\n",
    "due to non-circular detector orbits.\n",
    "For simplicity, here we illustrate the case\n",
    "where the PSF is the same for every view."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "nview = 60\n",
    "psfs = repeat(psf1, 1, 1, 1, nview)\n",
    "size(psfs)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## SPECT system model using `LinearMapAA`"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "dy = 8 # transaxial pixel size in mm\n",
    "mumap = zeros(T, size(xtrue)) # zero μ-map just for illustration here\n",
    "plan = SPECTplan(mumap, psfs, dy; T)\n",
    "\n",
    "forw! = (y,x) -> project!(y, x, plan)\n",
    "back! = (x,y) -> backproject!(x, y, plan)\n",
    "idim = (nx,ny,nz)\n",
    "odim = (nx,nz,nview)\n",
    "A = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Noisy data"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "if !@isdefined(ynoisy) # generate (scaled) Poisson data\n",
    "    ytrue = A * xtrue\n",
    "    target_mean = 20 # aim for mean of 20 counts per ray\n",
    "    scale = target_mean / average(ytrue)\n",
    "    scatter_fraction = 0.1 # 10% uniform scatter for illustration\n",
    "    scatter_mean = scatter_fraction * average(ytrue) # uniform for simplicity\n",
    "    background = scatter_mean * ones(T,nx,nz,nview)\n",
    "    ynoisy = rand.(Poisson.(scale * (ytrue + background))) / scale\n",
    "end\n",
    "jim(ynoisy, \"$nview noisy projection views\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## OS-EM algorithm - basic version"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "x0 = ones(T, nx, ny, nz) # initial uniform image\n",
    "\n",
    "niter = 8\n",
    "nblocks = 4\n",
    "Ab = Ablock(plan, nblocks) # create a linear map for each block\n",
    "\n",
    "if !@isdefined(xhat1)\n",
    "    xhat1 = osem(x0, ynoisy, background, Ab; niter)\n",
    "end;"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "This preferable OS-EM version preallocates the output `xhat2`:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "if !@isdefined(xhat2)\n",
    "    xhat2 = copy(x0)\n",
    "    osem!(xhat2, x0, ynoisy, background, Ab; niter)\n",
    "end\n",
    "\n",
    "@assert xhat1 ≈ xhat2"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Compare with ML-EM"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Run 30 iterations of ML-EM algorithm."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "niter_mlem = 30\n",
    "if !@isdefined(xhat3)\n",
    "    xhat3 = copy(x0)\n",
    "    mlem!(xhat3, x0, ynoisy, background, A; niter=niter_mlem)\n",
    "end\n",
    "\n",
    "jim(jim(mid3(xhat2), \"OS-EM at $niter iterations\"),\n",
    "    jim(mid3(xhat3), \"ML-EM at $niter_mlem iterations\"))"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.11.1"
  },
  "kernelspec": {
   "name": "julia-1.11",
   "display_name": "Julia 1.11.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}