{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torchvision.transforms.functional import rotate\n",
    "from pytomography.utils import rotate_detector_z\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTomography is built around PyTorch, which uses the data type `torch.tensor` (very analogous to the `numpy.array`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 4,  9, 16])"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([2,3,4])\n",
    "x**2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All objects and images are stored using `torch.tensor`s. The shape of the tensor depends on the imaging modality\n",
    "\n",
    "* For SPECT, objects have shape `[batch_size, Lx, Ly, Lz]` and images have shape `[batch_size, Ltheta, Lr, Lz]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128, 128])"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.ones((1,128,128,128))\n",
    "x.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create a simple cylinder using a meshgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([128, 128, 128])"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.linspace(-1,1,128)\n",
    "xv, yv, zv = torch.meshgrid(x,x,x, indexing='ij')\n",
    "obj = (xv**2 + 0.9*zv**2 < 0.5) * (torch.abs(yv)<0.8)\n",
    "obj = obj.to(torch.float)\n",
    "obj.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to ensure we have the `batch_size` dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128, 128])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obj = obj.unsqueeze(dim=0)\n",
    "obj.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot from a coronal, sagital, and axial perspective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23204 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, _ = plt.subplots(1,3,figsize=(10,3))\n",
    "plt.subplot(131)\n",
    "plt.pcolormesh(x,x,obj[0].sum(axis=0).T, cmap='Greys_r')\n",
    "plt.xlabel('y')\n",
    "plt.ylabel('z')\n",
    "plt.title('Coronal')\n",
    "plt.subplot(132)\n",
    "plt.pcolormesh(x,x,obj[0].sum(axis=1).T, cmap='Greys_r')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('z')\n",
    "plt.title('Sagittal')\n",
    "plt.subplot(133)\n",
    "plt.pcolormesh(obj[0].sum(axis=2).T, cmap='Greys_r')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Axial')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fundamental Operations"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a 3D object $f$, we may want to turn this into projections $g$ using\n",
    "\n",
    "$$g = Hf$$\n",
    "\n",
    "Assuming no attenuation, PSF modeling, scattering, we can model this by projecting the object at different angles (think of it like taking an X-Ray at a number of different angles)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Idea\n",
    "\n",
    "Rotate the object first, always project along the x-axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proj = obj.sum(axis=1)\n",
    "proj.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an projection"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Rotate Function"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The object is a tensor with 4 dimensions (including batch size). We can use the rotate function, which rotates the object in the $xy$-plane so long as $x$ and $y$ are the final two dimensions\n",
    "\n",
    "* Right now $z$ is the last dimension\n",
    "\n",
    "So we need to transpose the array then rotate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 45\n",
    "phi = 270 - beta\n",
    "obj_rotated = rotate(obj.permute(0,3,1,2), -phi).permute(0,2,3,1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We rotate by $-\\phi$ so that the detector angle is $\\phi$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proj_45 = obj_rotated.sum(axis=1)\n",
    "proj_45.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All this functionality is contained in the `rotate_detector_z` function [here](https://pytomography.readthedocs.io/en/latest/_modules/pytomography/utils/helper_functions.html#rotate_detector_z)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.QuadMesh at 0x7fb33e5c2e80>"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16836 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.pcolormesh(proj_45[0].T, cmap='Greys')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is equivalent to what we would call $g_{45^{\\circ}}$ in the SPECT example of the manual [here](https://pytomography.readthedocs.io/en/latest/notebooks/conventions.html#Mappings). We can the compute $g = \\sum_{\\theta} g_{\\theta} \\otimes \\hat{\\theta}$ by computing these projections for a number of angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "angles = np.arange(0,360.,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = torch.zeros((1,len(angles),128,128))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128])"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image[:,2].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i,angle in enumerate(angles):\n",
    "    object_i = rotate_detector_z(obj,angle)\n",
    "    image[:,i] = object_i.sum(axis=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can look at the image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'r')"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 30940 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.pcolormesh(image[0,:,:,64].T)\n",
    "plt.xlabel('Angle')\n",
    "plt.ylabel('r')"
   ]
  },
  {
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
