{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pytomography\n",
    "from pytomography.metadata import ObjectMeta, ImageMeta, PSFMeta\n",
    "from pytomography.transforms import SPECTAttenuationTransform, SPECTPSFTransform\n",
    "from pytomography.projections import SystemMatrix\n",
    "# Set device so its faster\n",
    "import torch\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "pytomography.device = device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Last Time"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We created an object and attenuation map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = 0.3 #cm\n",
    "x = torch.linspace(-1,1,128)\n",
    "xv, yv, zv = torch.meshgrid(x,x,x, indexing='ij')\n",
    "# SPECT Object\n",
    "obj = (xv**2 + 0.9*zv**2 < 0.5) * (torch.abs(yv)<0.8)\n",
    "obj = obj.to(torch.float).unsqueeze(dim=0)\n",
    "# CT Object\n",
    "mu = (xv**2 + 0.9*zv**2 < 0.3) * (torch.abs(yv)<0.6)\n",
    "mu = mu.to(torch.float).unsqueeze(dim=0) * 0.05 #cm^-1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specified the appropriate object/image metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_meta = ObjectMeta(dr=(0.3,0.3,0.3), shape=(128,128,128))\n",
    "image_meta = ImageMeta(object_meta, angles=np.arange(0,360.,3), radii=25*np.ones(120))\n",
    "psf_meta = PSFMeta(collimator_slope = 0.03, collimator_intercept=0.1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Specified the transforms used in the system matrix, and created the system matrix $H$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_transform = SPECTAttenuationTransform(mu)\n",
    "psf_transform = SPECTPSFTransform(psf_meta)\n",
    "system_matrix = SystemMatrix(\n",
    "    obj2obj_transforms = [att_transform,psf_transform],\n",
    "    im2im_transforms = [],\n",
    "    object_meta = object_meta,\n",
    "    image_meta = image_meta)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then simulated the image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = system_matrix.forward(obj)\n",
    "image = torch.poisson(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 205764 characters]",
      "text/plain": [
       "<Figure size 1000x200 with 7 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "iis = [0,10,20,30,40,50,60]\n",
    "fig, ax = plt.subplots(1,7,figsize=(10,2))\n",
    "[a.pcolormesh(image[0,i].cpu().T) for (a, i) in zip(ax, iis)]\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This Time"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown previously, back projection $H^T g$ alone does not give a good estimate of the initial object $f$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_bp = system_matrix.backward(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32580 characters]",
      "text/plain": [
       "<Figure size 600x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,2,figsize=(6,3))\n",
    "ax[0].pcolormesh(obj[0][:,64].T)\n",
    "ax[1].pcolormesh(obj_bp.cpu()[0][:,64].T)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MLEM algorithm can be written as\n",
    "\n",
    "$$\\boxed{f^{(n+1)} = \\left[\\frac{1}{H^T \\vec{1}} H^T \\left( \\frac{g}{Hf^{(n)}}\\right) \\right]f^{(n)}}$$\n",
    "\n",
    "This is an iterative algorithm which requires an initial object guess $f^{(0)}$. The more times you run the algorithm, the better (though the algorithm has a tendency to amplify noise in later iterations, so usually only a couple iterations are used)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial prediction\n",
    "object_prediction = torch.ones((image.shape[0], *system_matrix.object_meta.shape)).to(device)\n",
    "# number of iterations\n",
    "n_iters = 30\n",
    "# parameter delta to avoid division by 0\n",
    "delta = 1e-11\n",
    "# Run basic MLEM\n",
    "for j in range(n_iters):\n",
    "    ratio = (image+delta) / (system_matrix.forward(object_prediction) + delta)\n",
    "    object_prediction = object_prediction * system_matrix.backward(ratio, normalize=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can plot a coronal slice of the reconstructed object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33572 characters]",
      "text/plain": [
       "<Figure size 600x300 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,2,figsize=(6,3))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(obj[0][:,64].T, vmax=1.2)\n",
    "plt.axis('off')\n",
    "plt.colorbar()\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(object_prediction.cpu()[0][:,64].T, vmax=1.2)\n",
    "plt.axis('off')\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a 1 dimensional profile through the center:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f60ac4fcca0>]"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 27232 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(obj[0][:,64,64])\n",
    "plt.plot(object_prediction.cpu()[0][:,64,64])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytomographytest",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
