{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reconstructing Simind Data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The purpose of this tutorial is to show how to open SPECT/CT datafiles obtained from the SIMIND Monte Carlo program, configure proper alignment of images, and reconstruct using no attenuation correction, and attenuation correction. Note that in this tutorial, no PSF modeling is used. For a complete example using attenuation correction, PSF correction, scatter correction, and prior modeling, see the Quantitative tutorial. \n",
    "\n",
    "We'll build the OSEM network from scratch (unlike in what was done in the `Introduction` tutorial); the purpose of this is to show the architecture of the network. For your own cases where you want to do quick reconstruction, I'd recommend using the `get_osem_net` function to quickly create a network."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we'll import all required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'get_osem_net' from 'pytomography.algorithms.osem' (/home/gpuvmadm/PyTomography/src/pytomography/algorithms/osem.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[1;32m/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb Cell 4\u001b[0m in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m sys\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mappend(\u001b[39m'\u001b[39m\u001b[39m/home/gpuvmadm/PyTomography/src\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mos\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=3'>4</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytomography\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39malgorithms\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mosem\u001b[39;00m \u001b[39mimport\u001b[39;00m OSEMOSL, get_osem_net\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytomography\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mprojections\u001b[39;00m \u001b[39mimport\u001b[39;00m ForwardProjectionNet, BackProjectionNet\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#W4sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpytomography\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mmappings\u001b[39;00m \u001b[39mimport\u001b[39;00m SPECTAttenuationNet\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'get_osem_net' from 'pytomography.algorithms.osem' (/home/gpuvmadm/PyTomography/src/pytomography/algorithms/osem.py)"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('/home/gpuvmadm/PyTomography/src')\n",
    "import os\n",
    "from pytomography.algorithms.osem import OSEMOSL, get_osem_net\n",
    "from pytomography.projections import ForwardProjectionNet, BackProjectionNet\n",
    "from pytomography.mappings import SPECTAttenuationNet\n",
    "from pytomography.io import simind_projections_to_data, simind_CT_to_data\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like in regular PyTorch, we choose our computation device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'torch' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb Cell 6\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bgpuvm00004jhubvm01.canadacentral.cloudapp.azure.com/home/gpuvmadm/PyTomography/docs/source/notebooks/t_siminddata.ipynb#X10sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m device \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mdevice(\u001b[39m\"\u001b[39m\u001b[39mcuda\u001b[39m\u001b[39m\"\u001b[39m \u001b[39mif\u001b[39;00m torch\u001b[39m.\u001b[39mcuda\u001b[39m.\u001b[39mis_available() \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39mcpu\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'torch' is not defined"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can load in some projection data produced by SIMIND. It's important to check the header file to determine if the radial distances are specified in units of mm or cm (usually the detector is at least 15cm away from the patient)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 128, 512])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filepath = os.path.join(path, 'xcatwb_tot_w2.h00')\n",
    "object_meta, image_meta, projections = simind_projections_to_data(filepath, distance='mm')\n",
    "projections.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tells us we have 64 projections of matrix size 128x512. We can look at the image meta data to get more information. Information about the corresponding angle and radius for each projections is stored in `image_meta.radii` and `image_meta.angles` respectively. We can plot the projections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 266844 characters]",
      "text/plain": [
       "<Figure size 1500x600 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(1,5,figsize=(15,6))\n",
    "for i, proj in enumerate([0,10,20,30,40]):\n",
    "    axes[i].pcolormesh(projections[0][proj].T, cmap='nipy_spectral')\n",
    "    axes[i].set_title(f'r={image_meta.radii[proj]}; Angle={image_meta.angles[proj]}')\n",
    "    axes[i].axis('off')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `object_meta` gives corresponding metadata about the shape and voxel spacing of the object which is to be constructed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(128, 128, 512)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "object_meta.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.47950000762939454, 0.47950000762939454, 0.47950000762939454)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "object_meta.dr"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this information to create an `OSEMNet` object, which will be used for ordered subset expectation maximum reconstruction.\n",
    "\n",
    "$$f_i^{(n+1)} = \\frac{f_i^{(n)}}{\\sum_j c_{ij} + \\beta \\frac{\\partial V}{\\partial f_r}|_{f_i=f_i^{(n)}}} \\sum_j c_{ij}\\frac{g_j}{\\sum_i c_{ij}f_i^{(n)}} $$\n",
    "\n",
    "First we need to give all required information for the forward projection operator $\\sum_i c_{ij} a_i$ and the back projection operator $\\sum_j c_{ij} b_j$ work.\n",
    "\n",
    "* This amounts to defining the different image correction techniques we want to use. In some cases we'll want to adjust for phenomenon such as attenuation as PSF blurring."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## Basic Reconstruction\n",
    "\n",
    "To start, we won't focus on any image correction techniques. We can define our forward and back projection nets as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_net = ForwardProjectionNet(obj2obj_nets=[],\n",
    "                              im2im_nets=[],\n",
    "                              object_meta=object_meta,\n",
    "                              image_meta=image_meta,\n",
    "                              device=device)\n",
    "bp_net = BackProjectionNet(obj2obj_nets=[],\n",
    "                           im2im_nets=[],\n",
    "                           object_meta=object_meta,\n",
    "                           image_meta=image_meta,\n",
    "                           device=device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Projection networks work as follows: `fp_net` converts an object in to an image and `bp_net` converts an image into an object. We'll cover the specifics of this in later tutorials. All you need to understand for now is that they're used as mathematical tools inside the `OSEM` algorithm."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since osem is an iterative algorithm, we need an initial object guess. A common initial guess is just an array of all ones. We can create this by using the `shape` attribute of the `object_meta` (which tells us the exact dimensions are objects will be)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([128, 128, 512])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "initial_guess = torch.ones(object_meta.shape)\n",
    "initial_guess.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But as you may have noticed, the arrays constructed are typically 4 dimensional; there is a batch dimension out front. This dimension is included so that multiple images can be reconstructed at the same time, or multiple energy windows of the same image can be reconstructed at the same time. So even though we're only reconstructing one image, we need to add that extra dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128, 512])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "initial_guess = initial_guess.unsqueeze(dim=0)\n",
    "initial_guess.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can combine everything we need in the `OSEMNet`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "osem_net = OSEMOSL(image = projections,\n",
    "                   object_initial=initial_guess,\n",
    "                   forward_projection_net=fp_net,\n",
    "                   back_projection_net=bp_net)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use this algorithm to reconstruct our object with the desired number of iterations and subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_object = osem_net(n_iters=2, n_subsets=8)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned object is a pytorch tensor; we need to convert it to a numpy array before we can do any plotting with it. We'll also index using [0] to get the first element of the batch (in this case there's only one object in the batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_object = reconstructed_object[0].cpu().numpy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot some slices of the object. Note that there is no correction techniques used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 92456 characters]",
      "text/plain": [
       "<Figure size 800x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1,2,figsize=(8,10))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(reconstructed_object[64].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(reconstructed_object[:,64].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Attenuation Correction"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To encorporate CT Correction, you must ensure your CT object is aligned with the SPECT data. In particular:\n",
    "1. If the dimension of the projections are `[batch_size, Ltheta, Lr, LZ]` then the dimension of the CT should be `[1, Lr, Lr, Lz]`. \n",
    "2. The projections should be aligned with the CT. This can be checked by comparing (i) the projection at 0 degrees with a coronal view of the CT and (ii) the projection at 90 degrees with a sagital view of the CT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CT = simind_CT_to_data(os.path.join(path, 'xcat_wb.hct'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This returns a 3D tensor, but remember we need all tensors to be 4 dimensional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CT = CT.unsqueeze(dim=0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 154280 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1, 4, figsize=(10,6))\n",
    "plt.subplot(141)\n",
    "plt.pcolormesh(CT[0].sum(axis=1).T, cmap='Greys_r')\n",
    "plt.axis('off')\n",
    "plt.subplot(142)\n",
    "plt.pcolormesh(projections[0][0].T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.subplot(143)\n",
    "plt.pcolormesh(CT[0].sum(axis=0).T, cmap='Greys_r')\n",
    "plt.axis('off')\n",
    "plt.subplot(144)\n",
    "plt.pcolormesh(projections[0][16].T, cmap='nipy_spectral')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the sagital views are opposite to eachother: but this is what we expect. The horizontal axis of the projection (SPECT image) is the $r$ axis, which is aligned opposite to the $y$ axis at a projection angle of 90 degrees. For more information, see the user's manual on coordinate system conventions."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this CT image, we can create a correction network to model attenuation correction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ct_net = SPECTAttenuationNet(CT, device=device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we simply need to add this correction to the forward and back projection networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp_net = ForwardProjectionNet(obj2obj_nets=[ct_net],\n",
    "                              im2im_nets=[],\n",
    "                              object_meta=object_meta,\n",
    "                              image_meta=image_meta,\n",
    "                              device=device)\n",
    "bp_net = BackProjectionNet(obj2obj_nets=[ct_net],\n",
    "                           im2im_nets=[],\n",
    "                           object_meta=object_meta,\n",
    "                           image_meta=image_meta,\n",
    "                           device=device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then do everything the same as above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_guess = torch.ones(object_meta.shape).unsqueeze(dim=0)\n",
    "osem_net = OSEMOSL(image = projections,\n",
    "                   object_initial=initial_guess,\n",
    "                   forward_projection_net=fp_net,\n",
    "                   back_projection_net=bp_net)\n",
    "reconstructed_object = osem_net(n_iters=4, n_subsets=8)\n",
    "reconstructed_object = reconstructed_object[0].cpu().numpy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can plot our new, attenuation corrected, reconstructions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 96824 characters]",
      "text/plain": [
       "<Figure size 800x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1,2,figsize=(8,10))\n",
    "plt.subplot(121)\n",
    "plt.pcolormesh(reconstructed_object[64].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.subplot(122)\n",
    "plt.pcolormesh(reconstructed_object[:,64].T, cmap='nipy_spectral')\n",
    "plt.colorbar()\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared to the images reconstructed without attenuation correction, the number of predicted counts is much higher, and the concentration is relatively higher in bones."
   ]
  }
 ],
 "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "324f5aa1485c6e81c0125c86ad1340c603dab08bc11babd97bd730cfafa2ec4a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
