{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# SPECTrecon PSF\n",
    "\n",
    "This page explains the PSF portion of the Julia package\n",
    "[`SPECTrecon.jl`](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, plan_psf, fft_conv!, fft_conv_adj!\n",
    "using MIRTjim: jim, prompt\n",
    "using Plots: scatter, scatter!, plot!, default\n",
    "default(markerstrokecolor=:auto, markersize=3)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following line is helpful when running this example.jl 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",
    "After rotating the image and the attenuation map,\n",
    "second step in SPECT image forward projection\n",
    "is to apply depth-dependent point spread function (PSF).\n",
    "Each (rotated) image plane\n",
    "is a certain distance from the SPECT detector\n",
    "and must be convolved with the 2D PSF appropriate\n",
    "for that plane.\n",
    "\n",
    "Because SPECT has relatively poor spatial resolution,\n",
    "the PSF is usually fairly wide,\n",
    "so convolution using FFT operations\n",
    "is typically more efficient\n",
    "than direct spatial convolution.\n",
    "\n",
    "Following other libraries like\n",
    "[FFTW.jl](https://github.com/JuliaMath/FFTW.jl),\n",
    "the PSF operations herein start with a `plan`\n",
    "where work arrays are preallocated\n",
    "for subsequent use.\n",
    "The `plan` is a `Vector` of `PlanPSF` objects:\n",
    "one for each thread.\n",
    "(Parallelism is across planes for a 3D image volume.)\n",
    "The number of threads defaults to `Threads.nthreads()`.\n",
    "\n",
    "\n",
    "## Example\n",
    "\n",
    "Start with a 3D image volume."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "T = Float32 # work with single precision to save memory\n",
    "nx = 32\n",
    "nz = 30\n",
    "image = zeros(T, nx, nx, nz) # ny = nx required\n",
    "image[1nx÷4, 1nx÷4, 3nz÷4] = 1\n",
    "image[2nx÷4, 2nx÷4, 2nz÷4] = 1\n",
    "image[3nx÷4, 3nx÷4, 1nz÷4] = 1\n",
    "jim(image, \"Original image\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Create a synthetic gaussian depth-dependent PSF for a single view"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "px = 11\n",
    "nview = 1 # for simplicity in this illustration\n",
    "psf = repeat(psf_gauss( ; ny=nx, px), 1, 1, 1, nview)\n",
    "jim(psf, \"PSF for each of $nx planes\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now plan the PSF modeling\n",
    "by specifying\n",
    "* the image size (must be square)\n",
    "* the PSF size: must be `px × pz × ny × nview`\n",
    "* the `Type` used for the work arrays."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "plan = plan_psf( ; nx, nz, px, T)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Here are the internals for the plan for the first thread:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "plan[1]"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "With this `plan` pre-allocated, now we can apply the depth-dependent PSF\n",
    "to the image volume (assumed already rotated here)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "result = similar(image) # allocate memory for the result\n",
    "fft_conv!(result, image, psf[:,:,:,1], plan) # mutates the first argument\n",
    "jim(result, \"After applying PSF\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Adjoint\n",
    "\n",
    "To ensure adjoint consistency between SPECT forward- and back-projection,\n",
    "there is also an adjoint routine:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "adj = similar(result)\n",
    "fft_conv_adj!(adj, result, psf[:,:,:,1], plan)\n",
    "jim(adj, \"Adjoint of PSF modeling\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The adjoint is *not* the same as the inverse\n",
    "so one does not expect the output here to match the original image!"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## LinearMap\n",
    "\n",
    "One can form a linear map corresponding to PSF modeling using `LinearMapAA`.\n",
    "Perhaps the main purpose is simply for verifying adjoint correctness."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using LinearMapsAA: LinearMapAA\n",
    "\n",
    "nx, nz, px = 10, 7, 5 # small size for illustration\n",
    "psf3 = psf_gauss( ; ny=nx, px)\n",
    "plan = plan_psf( ; nx, nz, px, T)\n",
    "idim = (nx,nx,nz)\n",
    "odim = (nx,nx,nz)\n",
    "forw! = (y,x) -> fft_conv!(y, x, psf3, plan)\n",
    "back! = (x,y) -> fft_conv_adj!(x, y, psf3, plan)\n",
    "A = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)\n",
    "\n",
    "Afull = Matrix(A)\n",
    "Aadj = Matrix(A')\n",
    "jim(cat(dims=3, Afull, Aadj'), \"Linear map for PSF modeling and its adjoint\")"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following check verifies adjoint consistency:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@assert Afull ≈ Aadj'"
   ],
   "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
}