{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "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": [
    "I've really been pushing the notation\n",
    "\n",
    "$$H = \\sum_{\\theta} P(\\theta) A(\\theta) \\otimes \\hat{\\theta}$$\n",
    "\n",
    "but its often better to reframe it as\n",
    "\n",
    "$$H = \\sum_{\\theta} P(\\theta) A(\\theta) R(\\theta) \\otimes \\hat{\\theta}$$\n",
    "\n",
    "where $R$ is a rotation matrix that rotates the object before applying corrections (this is what we've been doing in the code).\n",
    "\n",
    "* $A$ applies some sort of correction (attenuation/PSF) to the object before $P$ projects it to image space. Lets call $\\mathbb{U}$ the vector space of objects, $\\mathbb{V}$ the vector space of images, and $\\mathbb{V}_{\\theta}$ as the vector space of projections. Note: an image consists of multiple projections, so $\\mathbb{V}_{\\theta}$ is a subspace of $\\mathbb{V}$, and $\\bigcup_{\\theta} \\mathbb{V}_{\\theta} = \\mathbb{V}$, so in mathematical terms, the $\\mathbb{V}_{\\theta}$ **partition** the image space.\n",
    "\n",
    "then we can express the mappings as \n",
    "* $R:\\mathbb{U}\\to\\mathbb{U}$\n",
    "* $A:\\mathbb{U}\\to\\mathbb{U}$\n",
    "* $P:\\mathbb{U} \\to \\mathbb{V}_{\\theta}$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, we can link these mappings $A$ together to perform multiple corrections. For example, if we want to perform CT **and** PSF correction, then\n",
    "\n",
    "$$H = \\sum_{\\theta} P(\\theta) A_2(\\theta) A_1(\\theta) R(\\theta) \\otimes \\hat{\\theta}$$\n",
    "\n",
    "where $A_2$ applies PSF blurring and $A_1$ applies attenuation correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "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": [
    "All the functions defined in the past few tutorials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = 0.3\n",
    "radius = 25\n",
    "Lx = 128\n",
    "# PSF Stuff\n",
    "collimator_slope = 0.07\n",
    "collimator_intercept = 0.1\n",
    "\n",
    "def rev_cumsum(x: torch.Tensor):\n",
    "    return torch.cumsum(x.flip(dims=(1,)), dim=1).flip(dims=(1,)) - x/2\n",
    "\n",
    "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\n",
    "\n",
    "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\n",
    "# Get Attenuation Stuff\n",
    "\n",
    "\n",
    "# Get PSF Stuff\n",
    "sigma = get_sigma(radius=radius, dx=dx, Lx=Lx, collimator_slope=collimator_slope, collimator_intercept=collimator_intercept)   \n",
    "N = len(sigma)\n",
    "kernel_size = 45\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 128])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obj.sum(axis=1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "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",
    "    # Apply operation R\n",
    "    object_i = rotate_detector_z(obj,angle)\n",
    "    # Apply operation A1 (CT Correction)\n",
    "    mu_i = rotate_detector_z(mu, angle)\n",
    "    p_i = torch.exp(-rev_cumsum(mu_i * dx))\n",
    "    object_i  = object_i * p_i\n",
    "    # Apply operation A2 (PSF)\n",
    "    object_i = layer(object_i)\n",
    "    # Apply operation P\n",
    "    projection = object_i.sum(axis=1)\n",
    "    # Add to image (effectively $\\otimes \\hat{\\theta}$)\n",
    "    image[:,i] = projection\n",
    "image = image.detach() #because was passed through network"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we've applied both attenuation and PSF correction to the image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83144 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].T) for (a, i) in zip(ax, iis)]\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the individual components can be compartmentalized:\n",
    "* Matrices $A$ are referred to as **transforms** and they're implementation is in `pytomography.transforms`\n",
    "* The system matrix $H$ has implementation in `pytomography.projections`. This system matrix takes in the corresponding transforms, and implements forward/back projection\n",
    "* Any metadata about the object, image, or projection, is stored in seperate objects with implementation in `pytomography.metadata`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "pytomography.device = device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "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(len(angles)))\n",
    "psf_meta = PSFMeta(collimator_slope = 0.07, collimator_intercept=0.1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_transform = SPECTAttenuationTransform(mu)\n",
    "psf_transform = SPECTPSFTransform(psf_meta)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Projections (system matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": [
    "Then everything can be ran quite simply:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = system_matrix.forward(obj)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plotted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 107856 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()"
   ]
  }
 ],
 "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
}
