{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac1d7dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python3 setup.py install"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6d80fc9",
   "metadata": {},
   "source": [
    "# Gaussian PSF "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87258f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deltamic import Gaussian_psf, compute_spatial_frequency_grid,normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,compute_box_size\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import torch\n",
    "import skimage.io as io \n",
    "from skimage import img_as_ubyte\n",
    "from time import time\n",
    "\n",
    "\"\"\"\n",
    "Parameters:\n",
    "\"\"\"\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"Computations done in CUDA!\")\n",
    "    device = 'cuda:0'\n",
    "    n_subdivisions = 2\n",
    "    box_shape = np.array([100]*3)\n",
    "else : \n",
    "    print(\"No CUDA devices available, Computations done in CPU\")\n",
    "    device = 'cpu'\n",
    "    n_subdivisions = 1\n",
    "    box_shape = np.array([60]*3).tolist()\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "mesh = trimesh.primitives.Sphere(subdivisions = n_subdivisions)\n",
    "verts, faces = np.array(mesh.vertices),np.array(mesh.faces)\n",
    "verts-=verts.min(axis=0)\n",
    "verts*=0.5\n",
    "box_size = compute_box_size(verts, offset = 0.35)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device,requires_grad = True)\n",
    "Faces = torch.tensor(faces, dtype = torch.long, device = device)\n",
    "Faces_coeff = torch.ones(len(Faces),dtype = torch.float, device = device)\n",
    "box_size = torch.tensor(box_size) \n",
    "\n",
    "print(\"Vertices shape: \",verts.shape,\"Faces shape: \", faces.shape)\n",
    "print(\"Min/Max x/y/z position of vertices: \",verts.min(axis=0),verts.max(axis=0))\n",
    "print(\"box_size: \",box_size, \"box_shape\",box_shape)\n",
    "\n",
    "narrowband_thresh = torch.tensor(0,dtype = torch.float, device = device)\n",
    "meshFT = Fourier3dMesh(box_size,box_shape,device=device, dtype = torch.float32)\n",
    "\n",
    "model = Gaussian_psf(box_shape,box_size[:,1],sigma = 1e3,device = device)\n",
    "PSF = model.forward()\n",
    "OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "\n",
    "\"\"\"\n",
    "Image creation:\n",
    "\"\"\"\n",
    "from deltamic.optimize import VectorAdam\n",
    "\n",
    "t1 = time()\n",
    "ftmesh = meshFT(Verts,Faces, Faces_coeff)\n",
    "\n",
    "image_ft=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "t2 = time()\n",
    "\n",
    "io.imsave(\"Test_sphere_Gaussian.tif\",img_as_ubyte(image_ft.detach().cpu().numpy()))\n",
    "print(\"forward pass computed successfully in\", t2-t1,\"seconds\", \"Image saved as Test_sphere_Gaussian.tif\")\n",
    "\n",
    "loss = torch.sum(torch.abs(ftmesh))\n",
    "t1 = time()\n",
    "loss.backward()\n",
    "t2 = time()\n",
    "print(\"backward pass computed successfully in\", t2-t1,\"seconds\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22b3b8c2",
   "metadata": {},
   "source": [
    "# Pupil Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a38854",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deltamic import Hanser_psf, compute_spatial_frequency_grid,normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,compute_box_size\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import torch\n",
    "import skimage.io as io \n",
    "from skimage import img_as_ubyte\n",
    "from time import time\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "Parameters:\n",
    "\"\"\"\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"Computations done in CUDA!\")\n",
    "    device = 'cuda:0'\n",
    "    n_subdivisions = 2\n",
    "    box_shape = np.array([100]*3)\n",
    "else : \n",
    "    print(\"No CUDA devices available, Computations done in CPU\")\n",
    "    device = 'cpu'\n",
    "    n_subdivisions = 1\n",
    "    box_shape = np.array([60]*3).tolist()\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "mesh = trimesh.primitives.Sphere(subdivisions = n_subdivisions)\n",
    "verts, faces = np.array(mesh.vertices),np.array(mesh.faces)\n",
    "verts-=verts.min(axis=0)\n",
    "verts*=0.5\n",
    "box_size = compute_box_size(verts, offset = 0.35)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device,requires_grad = True)\n",
    "Faces = torch.tensor(faces, dtype = torch.long, device = device)\n",
    "Faces_coeff = torch.ones(len(Faces),dtype = torch.float, device = device)\n",
    "box_size = torch.tensor(box_size) \n",
    "\n",
    "print(\"Vertices shape: \",verts.shape,\"Faces shape: \", faces.shape)\n",
    "print(\"Min/Max x/y/z position of vertices: \",verts.min(axis=0),verts.max(axis=0))\n",
    "print(\"box_size: \",box_size, \"box_shape\",box_shape)\n",
    "\n",
    "narrowband_thresh = torch.tensor(0,dtype = torch.float, device = device)\n",
    "\n",
    "\n",
    "psf_params = dict(na=1.27, ni=1.33, wl=0.561, vec_corr=\"none\")\n",
    "\n",
    "# Set the Nyquist sampling rate\n",
    "nyquist_sampling = psf_params[\"wl\"] / psf_params[\"na\"] /4\n",
    "\n",
    "# our oversampling factor, must be odd for easy integration (for peaked symmetrical funtions the kernel must be odd)\n",
    "oversample_factor = 1\n",
    "\n",
    "psf_params[\"res\"] = nyquist_sampling / oversample_factor\n",
    "\n",
    "meshFT = Fourier3dMesh(box_size,box_shape,device=device, dtype = torch.float32)\n",
    "model = Hanser_psf(box_shape,meshFT.xi0,meshFT.xi1,meshFT.xi2,NA=psf_params['na'], wavelength=psf_params['wl'], ni=psf_params['ni'], res_lateral=psf_params['res'],res_axial=psf_params['res'], num_zernike = 21, device = device)\n",
    "PSF = model.forward()\n",
    "OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "\n",
    "\"\"\"\n",
    "Image creation:\n",
    "\"\"\"\n",
    "from deltamic.optimize import VectorAdam\n",
    "\n",
    "t1 = time()\n",
    "ftmesh = meshFT(Verts,Faces, Faces_coeff)\n",
    "\n",
    "image_ft=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "t2 = time()\n",
    "\n",
    "io.imsave(\"Test_sphere_Hanser.tif\",img_as_ubyte(image_ft.detach().cpu().numpy()))\n",
    "print(\"forward pass computed successfully in\", t2-t1,\"seconds\", \"Image saved as Test_sphere_Hanser.tif\")\n",
    "\n",
    "loss = torch.sum(torch.abs(ftmesh))\n",
    "t1 = time()\n",
    "loss.backward()\n",
    "t2 = time()\n",
    "print(\"backward pass computed successfully in\", t2-t1,\"seconds\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b676d89",
   "metadata": {},
   "source": [
    "# Gibson-Lanni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100952c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deltamic import Gibson_Lanni_psf,compute_spatial_frequency_grid,normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,compute_box_size\n",
    "import trimesh\n",
    "import numpy as np\n",
    "from skimage import img_as_ubyte\n",
    "import torch\n",
    "import skimage.io as io \n",
    "from skimage import img_as_ubyte\n",
    "from time import time\n",
    "\n",
    "\"\"\"\n",
    "Parameters:\n",
    "\"\"\"\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"Computations done in CUDA!\")\n",
    "    device = 'cuda:0'\n",
    "    n_subdivisions = 2\n",
    "    box_shape = np.array([100]*3)\n",
    "else : \n",
    "    print(\"No CUDA devices available, Computations done in CPU\")\n",
    "    device = 'cpu'\n",
    "    n_subdivisions = 1\n",
    "    box_shape = np.array([60]*3).tolist()\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "mesh = trimesh.primitives.Sphere(subdivisions = n_subdivisions)\n",
    "verts, faces = np.array(mesh.vertices),np.array(mesh.faces)\n",
    "verts-=verts.min(axis=0)\n",
    "verts*=0.5\n",
    "box_size = compute_box_size(verts, offset = 0.35)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device,requires_grad = True)\n",
    "Faces = torch.tensor(faces, dtype = torch.long, device = device)\n",
    "Faces_coeff = torch.ones(len(Faces),dtype = torch.float, device = device)\n",
    "box_size = torch.tensor(box_size) \n",
    "\n",
    "print(\"Vertices shape: \",verts.shape,\"Faces shape: \", faces.shape)\n",
    "print(\"Min/Max x/y/z position of vertices: \",verts.min(axis=0),verts.max(axis=0))\n",
    "print(\"box_size: \",box_size, \"box_shape\",box_shape)\n",
    "\n",
    "narrowband_thresh = torch.tensor(0,dtype = torch.float, device = device)\n",
    "\n",
    "\n",
    "psf_params = dict(na=1.27, ni=1.33, wl=0.561, vec_corr=\"none\")\n",
    "\n",
    "# Set the Nyquist sampling rate\n",
    "nyquist_sampling = psf_params[\"wl\"] / psf_params[\"na\"] /4\n",
    "\n",
    "# our oversampling factor, must be odd for easy integration (for peaked symmetrical funtions the kernel must be odd)\n",
    "oversample_factor = 1\n",
    "\n",
    "psf_params[\"res\"] = nyquist_sampling / oversample_factor\n",
    "\n",
    "meshFT = Fourier3dMesh(box_size,box_shape,device=device, dtype = torch.float32)\n",
    "model = Gibson_Lanni_psf(box_shape,pZ = 0, wavelength =psf_params[\"wl\"], NA = psf_params[\"na\"],ni = psf_params['ni'], \n",
    "                         ni0=psf_params['ni'],res_axial=psf_params['res'],res_lateral=psf_params['res'] ,device = device)\n",
    "PSF = model.forward()\n",
    "OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "\n",
    "\"\"\"\n",
    "Image creation:\n",
    "\"\"\"\n",
    "from deltamic.optimize import VectorAdam\n",
    "\n",
    "t1 = time()\n",
    "ftmesh = meshFT(Verts,Faces, Faces_coeff)\n",
    "\n",
    "image_ft=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "t2 = time()\n",
    "\n",
    "io.imsave(\"Test_sphere.tif\",img_as_ubyte(image_ft.detach().cpu().numpy()))\n",
    "print(\"forward pass computed successfully in\", t2-t1,\"seconds\", \"Image saved as Test_sphere.tif\")\n",
    "\n",
    "loss = torch.sum(torch.abs(ftmesh))\n",
    "t1 = time()\n",
    "loss.backward()\n",
    "t2 = time()\n",
    "print(\"backward pass computed successfully in\", t2-t1,\"seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391b257a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
