Grating search for 2D SIM

SIM_lattice_search.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from itertools import combinations\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import random\n",
    "from joblib import Parallel, delayed\n",
    "import multiprocessing\n",
    "import json\n",
    "import pickle\n",
    "import datetime\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class grating():\n",
    "    def __init__(self, h, x, y):\n",
    "        self.h = h\n",
    "        self.theta_x, self.theta_y = x, y\n",
    "        self.a_h = np.array([self.h, 0]) \n",
    "        self.a_theta = np.array([self.theta_x, self.theta_y])\n",
    "        self.theta_p = self._theta_p()\n",
    "        self.g = self.a_h[0]*np.sin(self.theta_p)\n",
    "    \n",
    "    def _theta_p(self):\n",
    "        return np.arctan2(self.theta_y,self.theta_x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_min, h_max, h_step = 0, 36*2, 6 \n",
    "theta_x_min, theta_x_max, theta_x_step = -50, 50, 1\n",
    "theta_y_min, theta_y_max, theta_y_step = -50, 50, 1\n",
    "h_range = np.arange(h_min, h_max, h_step)\n",
    "theta_x_range = np.arange(theta_x_min, theta_x_max, theta_x_step)\n",
    "theta_y_range = np.arange(theta_y_min, theta_y_max, theta_y_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gratings = []\n",
    "for x in theta_x_range:\n",
    "    for y in theta_y_range:\n",
    "        for h in h_range:\n",
    "            gratings.append(grating(h, x, y)) \n",
    "# gratings.append(grating(7, 7, 25))\n",
    "# gratings.append(grating(10, 10, -9))\n",
    "# gratings.append(grating(25, 25, 7))\n",
    "random.seed(2)\n",
    "random.shuffle(gratings)\n",
    "print(len(gratings))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(list(combinations(gratings[::200], 3))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def angle_test(gratings, delta_theta, plot):\n",
    "    comb = combinations(gratings, 3) \n",
    "    accepted = []\n",
    "    for i in tqdm(list(comb)): \n",
    "        alpha = np.zeros(3)\n",
    "        for h, p in enumerate(list(combinations(i,2))):\n",
    "            alpha[h] = np.abs(p[0].theta_p-p[1].theta_p)\n",
    "        if sum((np.abs(np.rad2deg(alpha)-60))<delta_theta)==2 and sum((np.abs(np.rad2deg(alpha)-120))<delta_theta)==1:\n",
    "            accepted.append(i)\n",
    "    return accepted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gratings_ang = angle_test(gratings[::100], 2, False)\n",
    "print(len(gratings_ang))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in gratings_ang[1:2]:\n",
    "    plt.polar(p[0].__dict__['theta_p'],2,'o') \n",
    "    plt.polar(p[1].__dict__['theta_p'],2,'o') \n",
    "    plt.polar(p[2].__dict__['theta_p'],2,'o') \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def constant_test(gratings, delta_g, p, dp):\n",
    "    accepted = []\n",
    "    for i in tqdm(gratings):\n",
    "        g = np.array([i[0].g, i[1].g, i[2].g])\n",
    "        g_bar = (np.max(g)+np.min(g))/2\n",
    "        delta = np.abs(np.max(g_bar-g)/g_bar)\n",
    "        if delta<delta_g and sum(abs((abs(g)-p))<dp)==3:\n",
    "            accepted.append(i)\n",
    "    return accepted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gratings_ang = pickle.load( open( f'C:/Data/SLM_Patterns/27_04_2021/20210427-0856-49_gratings_ang_list.bin', \"rb\" ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gratings_const = constant_test(gratings_ang[:], delta_g=2, p=24, dp=1.5)\n",
    "print(len(gratings_const))\n",
    "# for i in range(30):\n",
    "#     plt.plot(i, len(constant_test(gratings_ang[:], delta_g=5, p=i, dp=3)),'ro')\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# timestamp = datetime.datetime.now()\n",
    "# t = timestamp.strftime(\"%Y%m%d-%H%M-%S\")\n",
    "# with open(f'C:/Data/SLM_Patterns/{t}_gratings_ang_list.bin',\"wb\") as f:\n",
    "#     pickle.dump(gratings_ang, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for p in gratings_const:\n",
    "    print('...')\n",
    "    print(p[0].__dict__)\n",
    "    print(p[1].__dict__)\n",
    "    print(p[2].__dict__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "def adjust_gamma(image, gamma=1.0):\n",
    "    # build a lookup table mapping the pixel values [0, 255] to\n",
    "    # their adjusted gamma values\n",
    "    invGamma = 1.0 / gamma\n",
    "    table = np.array([((i / 255.0) ** invGamma) * 255\n",
    "        for i in np.arange(0, 256)]).astype(\"uint8\")\n",
    "    # apply gamma correction using the lookup table\n",
    "    return cv2.LUT(image, table) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Matlab code converted\n",
    "xx = 1920\n",
    "yy = 1920\n",
    "N = 18\n",
    "rgb_raw = np.zeros((N,xx,yy,3), dtype=np.uint8)\n",
    "[50,51,23,77,91,121]\n",
    "for n in range(1):\n",
    "    n = 1\n",
    "    timestamp = datetime.datetime.now()\n",
    "    t = timestamp.strftime(\"%Y%m%d-%H%M-%S\")\n",
    "    d = timestamp.strftime(\"%d_%m_%Y\")\n",
    "    p = 6\n",
    "    for i in range(3):\n",
    "        for phasestep in range(p):\n",
    "            grat = gratings_const[n][i]\n",
    "    #         print(grat.__dict__)\n",
    "            gratdir = grat.theta_p\n",
    "            gratper = grat.g\n",
    "\n",
    "            k = np.array([np.sin(gratdir), np.cos(gratdir)]) * 2*np.pi/gratper\n",
    "            phase = 2*np.pi/p\n",
    "            patt = np.sin(np.tile(np.arange(xx), (xx,1))*k[0] + np.tile(np.arange(yy), (yy,1)).T*k[1] + phase*phasestep)\n",
    "            patt = ((patt>0)*128).astype(np.uint8)\n",
    "            img = Image.fromarray(patt[:1080,:])\n",
    "#             img.show()\n",
    "            Path(f'C:/Data/SLM_Patterns/{d}/{t}/').mkdir(exist_ok=True, parents=True)\n",
    "            img.save(f'C:/Data/SLM_Patterns/{d}/{t}/{t}_grat_dir_{i}_{phasestep}_h_{grat.h}_xy_{grat.a_theta}_g_{grat.g:.3f}.png')\n",
    "            img_fft = np.fft.fft2(patt)\n",
    "            fft = np.log(np.fft.fftshift(np.abs(img_fft)))\n",
    "    #         plt.figure(figsize=(8,8))        \n",
    "    #         plt.imshow(fft,cmap='inferno')#, vmax=0.25*1e8)#, vmin)\n",
    "    #         plt.colorbar()\n",
    "    #         plt.show()\n",
    "            fft = cv2.normalize(fft, None, alpha = 0, beta = 255, norm_type = cv2.NORM_MINMAX, dtype = cv2.CV_8U) \n",
    "            rgb_raw[n,:,:,i] = fft#.astype(np.uint8)\n",
    "    #     plt.figure(figsize=(8,8))\n",
    "    #     plt.figtext(0.5, 0.05, f'{n}', wrap=True, horizontalalignment='center', fontsize=12)\n",
    "    #     plt.imshow(adjust_gamma(rgb_raw[n,400:700,400:700], 0.3), cmap='inferno')\n",
    "# N1 = 0\n",
    "# rgb_raw = rgb_raw[N1]\n",
    "# rgb = Image.fromarray(rgb_raw)\n",
    "# # rgb.save(f'grat_fft_rgb.tiff')\n",
    "# for i in range(3):\n",
    "#     grat = gratings_const[N1][i]\n",
    "#     print(grat.__dict__)"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}