{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reconstructing SIMIND Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pytomography\n",
    "from pytomography.io.SPECT import simind\n",
    "from pytomography.metadata import PSFMeta\n",
    "from pytomography.priors import RelativeDifferencePrior\n",
    "from pytomography.projections import SPECTSystemMatrix\n",
    "from pytomography.metadata import PSFMeta\n",
    "from pytomography.transforms import SPECTAttenuationTransform, SPECTPSFTransform\n",
    "from pytomography.algorithms import OSEMBSR\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the default device for all reconstruction for pytomography:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytomography.device = device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the folder location for downloaded files (you will need to modify this to be the directory where you saved the files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/disk1/pytomography_tutorial_data/simind_tutorial/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case 1: Single Projection File\n",
    "\n",
    "When running SIMIND simulations, it is commonly the case that organs, background, and lesions are simulated seperately and then added together in projection space after the fact. For now, lets consider the case of a single SIMIND file containing all regions of the body, we'll move to seperate files later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the object and image space metadata, along with the projection data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_meta, image_meta, projections = simind.get_projections(os.path.join(path, 'single_projections', 'body1t2ew6_tot_w2.hdr'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the estimate for the scatter using the `get_scatter_from_TEW` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "scatter = simind.get_scatter_from_TEW(\n",
    "    headerfile_peak = os.path.join(path, 'single_projections', 'body1t2ew6_tot_w2.hdr'),\n",
    "    headerfile_lower = os.path.join(path, 'single_projections', 'body1t2ew6_tot_w1.hdr'),\n",
    "    headerfile_upper = os.path.join(path, 'single_projections', 'body1t2ew6_tot_w3.hdr')\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The CT data is loaded in a similar fashion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "CT = simind.get_atteuation_map(os.path.join(path, 'single_projections', 'body1.hct'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is good practice to ensure the CT and projections are aligned\n",
    "* Note: the sagittal view of the CT is the mirror of the projection at 90 degrees; this has to do with the coordinate system defined in the manual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "CT_slice_coronal = CT[0,:,64,:].T.cpu()\n",
    "projection_coronal = projections[0][0].T.cpu()\n",
    "CT_slice_sagittal = CT[0,64,:,:].T.cpu()\n",
    "projection_sagittal = projections[0][30].T.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 190256 characters]",
      "text/plain": [
       "<Figure size 1400x700 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1, 4, figsize=(14,7))\n",
    "plt.subplot(141)\n",
    "plt.pcolormesh(CT_slice_coronal, cmap='Greys_r')\n",
    "plt.title('CT; Coronal View')\n",
    "plt.subplot(142)\n",
    "plt.pcolormesh(projection_coronal, cmap='nipy_spectral')\n",
    "plt.title(r'Spect; $\\beta=0^{\\circ}$')\n",
    "plt.subplot(143)\n",
    "plt.pcolormesh(CT_slice_sagittal, cmap='Greys_r')\n",
    "plt.title('CT; Sagittal View')\n",
    "plt.subplot(144)\n",
    "plt.pcolormesh(projection_sagittal, cmap='nipy_spectral')\n",
    "plt.title(r'Spect; $\\beta=90^{\\circ}$')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to define all the required mappings to model $g = Hf$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Attenuation Modeling\n",
    "\n",
    "* This transform corresponds to a mapping $A_i(\\theta)$ as specified in the user manual and takes in an object/corresponding projection angle and returns a modified object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_transform = SPECTAttenuationTransform(CT)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PSF Modeling\n",
    "\n",
    "* This transform corresponds to a mapping $A_i(\\theta)$ as specified in the users manual. The PSF collimator information is obtained from the SIMIND header file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gpuvmadm/PyTomography/src/pytomography/io/io_utils.py:43: RuntimeWarning: overflow encountered in exp\n",
      "  return c1*np.exp(-d1*np.sqrt(energy)) + c2*np.exp(-d2*np.sqrt(energy))\n"
     ]
    }
   ],
   "source": [
    "psf_meta = simind.get_psfmeta_from_header(os.path.join(path, 'single_projections', 'body1t2ew6_tot_w2.hdr'))\n",
    "psf_transform = SPECTPSFTransform(psf_meta)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can build our system matrix $H$:\n",
    "\n",
    "* The forward method `system_matrix.forward(f)` corresponds to $Hf$ where $H = \\sum_{\\theta} P(\\theta)A_1(\\theta)A_2(\\theta) \\otimes \\hat{\\theta}$\n",
    "* The backward method `system_matrix.backward(g)` corresponds to $H^T g$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_matrix = SPECTSystemMatrix(\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": [
    "We can also define the Bayesian Prior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior = RelativeDifferencePrior(beta=1, gamma=5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The system matrix is then used to create the reconstruction algorithm, which in this case is a variant of OSEM using the block sequential regularizer technique to encorporate the prior defined above. Note that if the prior is not given, this reconstruction algorithm corresponds to regular OSEM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstruction_algorithm = OSEMBSR(\n",
    "    image = projections,\n",
    "    system_matrix = system_matrix,\n",
    "    scatter = scatter,\n",
    "    prior = prior)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the reconstructed object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_object = reconstruction_algorithm(n_iters=10, n_subsets=8)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the reconstructed object next to the CT.\n",
    "* The reconstructed object has units of counts, and would need to be adjusted by a proportionality factor if one wants to obtain units of MBq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 71488 characters]",
      "text/plain": [
       "<Figure size 800x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1,2,figsize=(8,6))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(reconstructed_object[0].cpu()[:,70].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.title('Reconstructed Object')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(CT.cpu()[0][:,70].T, cmap='Greys_r')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.title('$\\mu$ (CT scan)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case 2: Multiple Projections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When running simulations with SIMIND, it is standard to run different regions of the body seperately and then combine them together in projection space after the fact. This enables one to create arbtrary organ activities using the same set of simulated data. A few noteworthy things\n",
    "\n",
    "* SIMIND projection data is typically given in units of CPS/MBq where CPS=counts/second. Real life SPECT data consists of units of counts. In order to simulate a realistic scenario, therefor, we need to multiply the projections by some value in MBq, and some value in s. The value in MBq corresponds to the total activity of the organ/region you've simulated. The value in s corresponds to the total projection time.\n",
    "\n",
    "We'll start by opening up SIMIND projections corresponding to 8 distinct regions of the body, each with activities representative of a true clinical case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "organs = ['bkg', 'liver', 'l_lung', 'r_lung', 'l_kidney', 'r_kidney','salivary', 'bladder']\n",
    "activities = [2500, 450, 7, 7, 100, 100, 20, 90] # MBq\n",
    "headerfiles = [os.path.join(path, 'multi_projections', organ, 'photopeak.h00') for organ in organs]\n",
    "headerfiles_lower = [os.path.join(path, 'multi_projections', organ, 'lowerscatter.h00') for organ in organs]\n",
    "headerfiles_upper = [os.path.join(path, 'multi_projections', organ, 'upperscatter.h00') for organ in organs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the projection/scatter data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_meta, image_meta, projections = simind.combine_projection_data(headerfiles, activities)\n",
    "scatter = simind.combine_scatter_data_TEW(headerfiles, headerfiles_lower, headerfiles_upper, activities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that our `weights` were in units of MBq, so our projection data currently has units of counts/second."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f23707097c0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 232828 characters]",
      "text/plain": [
       "<Figure size 600x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1, 2, figsize=(6,6))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(projections[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts/Second')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(scatter[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts/Second')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to decide how long the scan was taken for. We'll assume each projection was taken for 15 seconds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "dT = 15 #s\n",
    "projections *= dT\n",
    "scatter *= dT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f237406db50>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 231024 characters]",
      "text/plain": [
       "<Figure size 600x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1, 2, figsize=(6,6))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(projections[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(scatter[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In realtity, SPECT projection data is noisy. Provided enough events are ran during the SIMIND simulation, the SIMIND data is essentially noise-free. We need to add Poisson noise if we are to simulate a realistic clinical imaging scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "projections = torch.poisson(projections)\n",
    "scatter = torch.poisson(scatter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f23743d2730>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 178148 characters]",
      "text/plain": [
       "<Figure size 600x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1, 2, figsize=(6,6))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(projections[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(scatter[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.colorbar(label='Counts')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the projection and scatter data is representative of a realistic clinical scenario. Now we reconstruct as we did earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "CT = simind.get_atteuation_map(os.path.join(path, 'multi_projections', 'mu208.hct'))\n",
    "att_transform = SPECTAttenuationTransform(CT)\n",
    "psf_meta = simind.get_psfmeta_from_header(headerfiles[0])\n",
    "psf_transform = SPECTPSFTransform(psf_meta)\n",
    "system_matrix = SPECTSystemMatrix(\n",
    "    obj2obj_transforms = [att_transform,psf_transform],\n",
    "    im2im_transforms = [],\n",
    "    object_meta = object_meta,\n",
    "    image_meta = image_meta)\n",
    "prior = RelativeDifferencePrior(beta=1, gamma=5)\n",
    "reconstruction_algorithm = OSEMBSR(\n",
    "    image = projections,\n",
    "    system_matrix = system_matrix,\n",
    "    scatter = scatter,\n",
    "    prior = prior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_object = reconstruction_algorithm(n_iters=10, n_subsets=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets see the reconstructed object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 76404 characters]",
      "text/plain": [
       "<Figure size 800x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1,2,figsize=(8,6))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(reconstructed_object[0].cpu()[:,70].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.title('Reconstructed Object')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(CT.cpu()[0][:,70].T, cmap='Greys_r')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.title('$\\mu$ (CT scan)')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "324f5aa1485c6e81c0125c86ad1340c603dab08bc11babd97bd730cfafa2ec4a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
