{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3745/3448874680.py:17: MatplotlibDeprecationWarning: The seaborn styles shipped by Matplotlib are deprecated since 3.6, as they no longer correspond to the styles shipped by seaborn. However, they will remain available as 'seaborn-v0_8-<style>'. Alternatively, directly use the seaborn API instead.\n",
      "  plt.style.use('seaborn-white')\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from __future__ import annotations\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "sys.path.append('/home/gpuvmadm/PyTomography/src')\n",
    "import pytomography\n",
    "from pytomography.metadata import PSFMeta\n",
    "from pytomography.mappings import MapNet\n",
    "from pytomography.utils import pad_object, pad_image, unpad_object, unpad_image, rotate_detector_z\n",
    "from pytomography.metadata import ObjectMeta, ImageMeta\n",
    "from pytomography.projections import ForwardProjectionNet, BackProjectionNet\n",
    "from pytomography.algorithms import OSEMOSL\n",
    "from pytomography.priors import RelativeDifferencePrior\n",
    "from pytomography.mappings import PETAttenuationNet, PETPSFNet\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('seaborn-white')\n",
    "from matplotlib import colors\n",
    "import torch\n",
    "from scipy.ndimage import map_coordinates\n",
    "from scipy.interpolate import interpn\n",
    "from scipy.ndimage import zoom\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "pytomography.device = device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/gpuvmadm/PyTomography/test_files' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "GT = np.fromfile(os.path.join(path,'body2_F18.bin'), dtype=np.float32)\n",
    "CT = np.fromfile(os.path.join(path,'body_atn_511keV.bin'), dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "GT = GT.reshape((751,512,512))\n",
    "CT = CT.reshape((751,512,512))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "GT = zoom(GT, (384/751, 256/512, 256/512), order=1)\n",
    "CT = zoom(CT, (384/751, 256/512, 256/512), order=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "GT = np.transpose(GT, (2,1,0))\n",
    "CT = np.transpose(CT, (2,1,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "GT = torch.tensor(GT).unsqueeze(dim=0)\n",
    "CT = torch.tensor(CT).unsqueeze(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_meta = ObjectMeta(dr=(0.15,0.15,0.767), shape=[256,256,384])\n",
    "image_meta = ImageMeta(object_meta, angles=np.linspace(0,180,574, endpoint=False), radii=44.3*np.ones(574))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx, dy, dz = object_meta.dr\n",
    "scale_factor = dx*dy*dz * 1.76 #1.76 cps / Bq\n",
    "GT = GT * scale_factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kernel_noncol(x,r,R, delta=1e-8):\n",
    "    if r**2<R**2:\n",
    "        sigma = torch.sqrt(R**2 - r**2)/4 * np.pi / 180\n",
    "    else:\n",
    "        sigma = torch.zeros(r.shape) + delta\n",
    "    result = torch.exp(-x**2/sigma**2 / 2)\n",
    "    return result / (torch.sum(result)+delta)\n",
    "\n",
    "def kernel_penetration(x,r,R,mu=0.87, delta=1e-8):\n",
    "    result = torch.exp(-torch.abs(mu*x / ((r/R)*torch.sqrt(1-(r/R)**2) + delta)))\n",
    "    if r>=0:\n",
    "        result*= x <= 0\n",
    "    else:\n",
    "        result*= x >= 0\n",
    "    return result / (torch.sum(result)+delta)\n",
    "\n",
    "def kernel_scattering(x,r,R,scatter_fact=0.327, delta=1e-8):\n",
    "    sigma = scatter_fact * torch.sqrt(1-(r/R)**2) / (2 * np.sqrt(2*np.log(2))) # fwhm -> sigma\n",
    "    result = torch.exp(-x**2/sigma**2 / 2)\n",
    "    return result / (torch.sum(result)+delta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernels = [kernel_noncol, kernel_penetration, lambda x, r, R: kernel_penetration(2*x, r, R), kernel_scattering]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/anaconda/envs/torch/lib/python3.9/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2894.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "ctnet = PETAttenuationNet(CT = CT, device=device)\n",
    "psfnet = PETPSFNet(kernels, device=device)\n",
    "fp_net = ForwardProjectionNet(obj2obj_nets=[],\n",
    "                              im2im_nets=[ctnet,psfnet],\n",
    "                              object_meta=object_meta,\n",
    "                              image_meta=image_meta)\n",
    "bp_net = BackProjectionNet(obj2obj_nets=[],\n",
    "                           im2im_nets=[ctnet,psfnet],\n",
    "                           object_meta=object_meta,\n",
    "                           image_meta=image_meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13076 characters]",
      "text/plain": [
       "<Figure size 1000x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(10,2))\n",
    "plt.plot(psfnet.PSF_matrix[0][0][0][300].cpu())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "projections_analytic = fp_net(torch.poisson(GT))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f2140597fd0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 112588 characters]",
      "text/plain": [
       "<Figure size 300x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(3,6))\n",
    "plt.pcolormesh(projections_analytic[0,0].cpu().T, cmap='nipy_spectral')\n",
    "plt.colorbar()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a reconstruction network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdp = RelativeDifferencePrior(beta=20, gamma=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "osem_net = OSEMOSL(image = projections_analytic,\n",
    "                   forward_projection_net=fp_net,\n",
    "                   back_projection_net=bp_net)\n",
    "reconstructed_object = osem_net(n_iters=2, n_subsets=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "osem_net = OSEMOSL(image = projections_analytic,\n",
    "                   forward_projection_net=fp_net,\n",
    "                   back_projection_net=bp_net,\n",
    "                   prior=rdp)\n",
    "reconstructed_object_rdp = osem_net(n_iters=10, n_subsets=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f214045d790>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 97800 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplots(1,2,figsize=(10,10))\n",
    "plt.subplot(121)\n",
    "plt.title('2it_7sub_noprior')\n",
    "plt.pcolormesh(reconstructed_object[0][:,128].cpu().T, cmap='nipy_spectral', vmax=4.5)\n",
    "plt.colorbar()\n",
    "plt.subplot(122)\n",
    "plt.title('10it_7sub_rdpbeta20gamma2')\n",
    "plt.pcolormesh(reconstructed_object_rdp[0][:,128].cpu().T, cmap='nipy_spectral', vmax=4.5)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn.functional import pad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(125).reshape(1,5,5,5)\n",
    "x = pad(x, [1,1,1,1,1,1])\n",
    "x = torch.roll(x, shifts=(-1,1,1), dims=(1,2,3))\n",
    "x = x[:,1:-1,1:-1,1:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[  0,   0,   0,   0,   0],\n",
       "          [  0,  25,  26,  27,  28],\n",
       "          [  0,  30,  31,  32,  33],\n",
       "          [  0,  35,  36,  37,  38],\n",
       "          [  0,  40,  41,  42,  43]],\n",
       "\n",
       "         [[  0,   0,   0,   0,   0],\n",
       "          [  0,  50,  51,  52,  53],\n",
       "          [  0,  55,  56,  57,  58],\n",
       "          [  0,  60,  61,  62,  63],\n",
       "          [  0,  65,  66,  67,  68]],\n",
       "\n",
       "         [[  0,   0,   0,   0,   0],\n",
       "          [  0,  75,  76,  77,  78],\n",
       "          [  0,  80,  81,  82,  83],\n",
       "          [  0,  85,  86,  87,  88],\n",
       "          [  0,  90,  91,  92,  93]],\n",
       "\n",
       "         [[  0,   0,   0,   0,   0],\n",
       "          [  0, 100, 101, 102, 103],\n",
       "          [  0, 105, 106, 107, 108],\n",
       "          [  0, 110, 111, 112, 113],\n",
       "          [  0, 115, 116, 117, 118]],\n",
       "\n",
       "         [[  0,   0,   0,   0,   0],\n",
       "          [  0,   0,   0,   0,   0],\n",
       "          [  0,   0,   0,   0,   0],\n",
       "          [  0,   0,   0,   0,   0],\n",
       "          [  0,   0,   0,   0,   0]]]])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
