{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": [
    "# Last Time"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implemented\n",
    "\n",
    "$$H = \\sum_{\\theta} P(\\theta) A(\\theta) \\otimes \\hat{\\theta}$$\n",
    "\n",
    "and\n",
    "\n",
    "$$H^T = \\sum_{\\theta} A^T(\\theta)P^T(\\theta)  \\otimes \\hat{\\theta}^T$$\n",
    "\n",
    "where $A(\\theta)$ applies attenuation correction"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Create objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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.1 #cm^-1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Forward and back project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rev_cumsum(x: torch.Tensor):\n",
    "    return torch.cumsum(x.flip(dims=(1,)), dim=1).flip(dims=(1,)) - x/2\n",
    "\n",
    "# Forward Projection\n",
    "angles = np.arange(0,360.,3)\n",
    "image = torch.zeros((1,len(angles),128,128))\n",
    "for i,angle in enumerate(angles):\n",
    "    mu_i = rotate_detector_z(mu, angle)\n",
    "    p_i = torch.exp(-rev_cumsum(mu_i * dx))\n",
    "    object_i = rotate_detector_z(obj,angle)\n",
    "    image[:,i] = (object_i*p_i).sum(axis=1)\n",
    "    \n",
    "# Back Projection\n",
    "obj_bp = torch.zeros([1, 128, 128, 128])\n",
    "for i,angle in enumerate(angles):\n",
    "    obj_bp_i = torch.ones([1, 128, 128, 128]) * image[:,i].unsqueeze(dim=1)\n",
    "    mu_i = rotate_detector_z(mu, angle)\n",
    "    p_i = torch.exp(-rev_cumsum(mu_i * dx))\n",
    "    obj_bp_i = obj_bp_i * p_i\n",
    "    obj_bp_i = rotate_detector_z(obj_bp_i, angle, negative=True)\n",
    "    obj_bp += obj_bp_i"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This Time"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to add PSF blurring to $H$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"../images/psf.png\" alt= “” width=\"600\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The red distributions to the right represent probability distributions of where the emission is detected on the detector. The closer the emission is to the detector, the narrower the response function.\n",
    "\n",
    "* For many realistic cases, the curve is Gaussian, and the relationship between $\\sigma$ and the distance to the detector $z$ follows a linear relationship:\n",
    "\n",
    "$$\\sigma(r) = ar+b$$\n",
    "\n",
    "So how do we simulate this in $H$?\n",
    "\n",
    "1. Use different blurring kernels for each location, dependent on location from the detector with width given by $\\sigma=ar+b$\n",
    "2. Forward project"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"../images/psf2.png\" alt= “” width=\"600\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this we'll need a few functions:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Get radial distances from projection to scanner, given $R$ (detector radius) and $dx$ (spacing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_distance(Lx: int, R: float, dx: float):\n",
    "    if Lx%2==0:\n",
    "        r = R + (Lx/2 - np.arange(Lx) - 1/2) * dx\n",
    "    else:\n",
    "        r = R + ((Lx-1)/2 - np.arange(Lx) ) * dx\n",
    "    # Correction for if radius of scanner is inside the the bounds\n",
    "    r[r<0] = 0\n",
    "    return r"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* This function uses `get_distance` to get $\\sigma$ at each point in the array. `collimator_slope` is $a$, and `collimator_intercept` is $b$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sigma(\n",
    "        radius: float,\n",
    "        dx: float,\n",
    "        Lx: tuple,\n",
    "        collimator_slope: float,\n",
    "        collimator_intercept: float\n",
    "    ) -> np.array:\n",
    "        distances = get_distance(Lx, radius, dx)\n",
    "        sigma = collimator_slope * distances + collimator_intercept\n",
    "        return sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.1835, 3.1625, 3.1415, 3.1205, 3.0995, 3.0785, 3.0575, 3.0365,\n",
       "       3.0155, 2.9945, 2.9735, 2.9525, 2.9315, 2.9105, 2.8895, 2.8685,\n",
       "       2.8475, 2.8265, 2.8055, 2.7845, 2.7635, 2.7425, 2.7215, 2.7005,\n",
       "       2.6795, 2.6585, 2.6375, 2.6165, 2.5955, 2.5745, 2.5535, 2.5325,\n",
       "       2.5115, 2.4905, 2.4695, 2.4485, 2.4275, 2.4065, 2.3855, 2.3645,\n",
       "       2.3435, 2.3225, 2.3015, 2.2805, 2.2595, 2.2385, 2.2175, 2.1965,\n",
       "       2.1755, 2.1545, 2.1335, 2.1125, 2.0915, 2.0705, 2.0495, 2.0285,\n",
       "       2.0075, 1.9865, 1.9655, 1.9445, 1.9235, 1.9025, 1.8815, 1.8605,\n",
       "       1.8395, 1.8185, 1.7975, 1.7765, 1.7555, 1.7345, 1.7135, 1.6925,\n",
       "       1.6715, 1.6505, 1.6295, 1.6085, 1.5875, 1.5665, 1.5455, 1.5245,\n",
       "       1.5035, 1.4825, 1.4615, 1.4405, 1.4195, 1.3985, 1.3775, 1.3565,\n",
       "       1.3355, 1.3145, 1.2935, 1.2725, 1.2515, 1.2305, 1.2095, 1.1885,\n",
       "       1.1675, 1.1465, 1.1255, 1.1045, 1.0835, 1.0625, 1.0415, 1.0205,\n",
       "       0.9995, 0.9785, 0.9575, 0.9365, 0.9155, 0.8945, 0.8735, 0.8525,\n",
       "       0.8315, 0.8105, 0.7895, 0.7685, 0.7475, 0.7265, 0.7055, 0.6845,\n",
       "       0.6635, 0.6425, 0.6215, 0.6005, 0.5795, 0.5585, 0.5375, 0.5165])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigma = get_sigma(radius=25, dx=0.3, Lx=128, collimator_slope=0.07, collimator_intercept=0.1)\n",
    "sigma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember the object shape is (1,128,128,128) and we want to blur the first axis. We can build a convolutional neural network to apply a seperate Gaussian blurring kernel for each slice on axis 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = len(sigma)\n",
    "kernel_size = 17\n",
    "delta = 1e-9\n",
    "layer = torch.nn.Conv2d(N, N, kernel_size, groups=N, padding='same',\n",
    "                        padding_mode='zeros', bias=0)\n",
    "x_grid, y_grid = torch.meshgrid(2*[torch.arange(-int(kernel_size//2), int(kernel_size//2)+1)], indexing='ij')\n",
    "x_grid = x_grid.unsqueeze(dim=0).repeat((N,1,1))\n",
    "y_grid = y_grid.unsqueeze(dim=0).repeat((N,1,1))\n",
    "sigma = torch.tensor(sigma, dtype=torch.float32).reshape((N,1,1))\n",
    "kernel = torch.exp(-(x_grid**2 + y_grid**2) / (2*sigma**2 + delta))\n",
    "kernel = kernel / kernel.sum(axis=(1,2)).reshape(N,1,1)\n",
    "layer.weight.data = kernel.unsqueeze(dim=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the blurring is independent of the detector angle (unless the detector follows a non-circular path, and in the source code, we adjust the code to allow for this generality)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Forward Projection\n",
    "angles = np.arange(0,360.,3)\n",
    "image = torch.zeros((1,len(angles),128,128))\n",
    "for i,angle in enumerate(angles):\n",
    "    object_i = layer(rotate_detector_z(obj,angle))\n",
    "    image[:,i] = object_i.sum(axis=1)\n",
    "image = image.detach() #because was passed through network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets look at some resulting images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88524 characters]",
      "text/plain": [
       "<Figure size 1500x300 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=(15,3))\n",
    "[a.pcolormesh(image[0,i].T) for (a, i) in zip(ax, iis)]\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_bp = torch.zeros([1, 128, 128, 128])\n",
    "for i,angle in enumerate(angles):\n",
    "    obj_bp_i = torch.ones([1, 128, 128, 128]) * image[:,i].unsqueeze(dim=1)\n",
    "    obj_bp_i = layer(obj_bp_i)\n",
    "    obj_bp_i = rotate_detector_z(obj_bp_i, angle, negative=True)\n",
    "    obj_bp += obj_bp_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
