{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7991a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deltamic import Gaussian_psf, compute_spatial_frequency_grid,normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,compute_box_size,generate_gaussian_psf\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",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "#assert torch.cuda.is_available()\n",
    "\n",
    "device = 'cpu'\n",
    "box_shape = np.array([20]*3)\n",
    "\n",
    "filename = \"../data/spot.obj\"\n",
    "\n",
    "Mesh_gt = trimesh.load(filename)\n",
    "faces = np.array(Mesh_gt.faces)\n",
    "verts = np.array(Mesh_gt.vertices)\n",
    "\n",
    "box_size = compute_box_size(verts, offset = 0.2)\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",
    "sigma_matrix = (1e-3*torch.eye(3,device=device))\n",
    "OTF = generate_gaussian_psf(sigma_matrix,meshFT.xi0,meshFT.xi1,meshFT.xi2).to(device)\n",
    "\n",
    "\"\"\"\n",
    "Image creation:\n",
    "\"\"\"\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(\"Spot.tif\",img_as_ubyte(image_ft.detach().cpu().numpy()))\n",
    "print(\"forward pass computed successfully in\", t2-t1,\"seconds\", \"Image saved as Spot.tif\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d99a095",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deltamic import normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,generate_gaussian_psf, Gaussian_psf, center_verts_in_box\n",
    "from tqdm import tqdm\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "import skimage.io as io \n",
    "from largesteps.optimize import AdamUniform\n",
    "from largesteps.geometry import compute_matrix\n",
    "from largesteps.solvers import CholeskySolver\n",
    "\n",
    "\n",
    "folder_result = \"Results/\"\n",
    "def create_dir(folder_name):\n",
    "    try: os.mkdir(folder_name)\n",
    "    except: pass\n",
    "  \n",
    "create_dir(folder_result)\n",
    "create_dir(folder_result+\"Images\")\n",
    "create_dir(folder_result+\"Meshes\")\n",
    "\n",
    "\"\"\"\n",
    "Parameters:\n",
    "\"\"\"\n",
    "\n",
    "device = 'cpu'\n",
    "name_micim = \"Spot.tif\"\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "micim = normalize_tensor(torch.tensor(io.imread(name_micim).astype(np.float32),dtype=torch.float)).to(device)\n",
    "box_size = np.array([[-0.9058,  1.2208],\n",
    "                     [-0.9058,  1.2208],\n",
    "                     [-0.9058,  1.2208]])\n",
    "mesh = trimesh.primitives.Sphere(subdivisions = 4)\n",
    "verts, faces = np.array(mesh.vertices)*2,np.array(mesh.faces)\n",
    "verts = center_verts_in_box(verts, box_size, offset=.4)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device)\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",
    "box_shape= np.array([20]*3)\n",
    "\n",
    "meshFT = Fourier3dMesh(box_size,box_shape,device=device, dtype = torch.float32,narrowband_thresh=1e-2)\n",
    "\n",
    "\n",
    "#Gaussian PSF\n",
    "model = Gaussian_psf(box_shape,box_size[:,1],sigma = 1e3,device = device)\n",
    "\n",
    "print(\"box_size: \",box_size, \"Box_ft_shape\",box_shape)\n",
    "\n",
    "\"\"\"\n",
    "Large \n",
    "Steps \n",
    "Optimization\n",
    "\"\"\"\n",
    "\n",
    "lr_base = 0.01\n",
    "lambda_=50.0\n",
    "alpha = 1.0\n",
    "Verts.requires_grad = True \n",
    "optimizer_geometry = AdamUniform([{'params': Verts}],lr=lr_base)\n",
    "optimizer_psf = torch.optim.Adam(model.parameters(), lr = lr_base)\n",
    "\n",
    "losses = []\n",
    "\n",
    "#M = compute_matrix(Verts, Faces, lambda_)\n",
    "#solver = CholeskySolver(M@M)\n",
    "    \n",
    "for k in (pbar:=tqdm(range(10000))):\n",
    "    \n",
    "    \n",
    "    #Loop of the optimization of the position of the vertices once we got a good approximation of the PSF\n",
    "    ###############################\n",
    "    optimizer_geometry.zero_grad()\n",
    "    ###############################\n",
    "    ftmesh = meshFT(Verts, Faces,Faces_coeff)\n",
    "    PSF = model.forward()\n",
    "    OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "    meshFT.OTF = OTF\n",
    "    \n",
    "    image_ft=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "    \n",
    "    loss_mse = torch.mean(((image_ft-micim)**2)*(micim**alpha))\n",
    "    loss = loss_mse\n",
    "    loss.backward()\n",
    "    \n",
    "    pbar.set_description(\"Current_Loss: \"+str([loss.item()]))\n",
    "    \n",
    "    \n",
    "    #with torch.no_grad():\n",
    "    #    Verts.grad = solver.solve(Verts.grad)\n",
    "    \n",
    "    np.save(folder_result+'Meshes/'+str(k)+\".npy\",(Verts.detach().cpu().numpy(),Faces.detach().cpu().numpy()))\n",
    "    ###############################\n",
    "    optimizer_geometry.step()\n",
    "    ###############################\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        loss_mse = torch.mean((image_ft-micim)**2)\n",
    "    \n",
    "    losses.append(loss_mse.item())\n",
    "    np.save(folder_result+'Meshes/loss.npy',losses)\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b260777",
   "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
}
