RandomNoise_V8.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from random import uniform, randint, shuffle\n",
    "from NPSmethods2 import readInImages, pupilFunc, getFreq\n",
    "from numpy.fft import ifftshift, irfft2\n",
    "import pickle as pkl\n",
    "import matplotlib.cm as cm\n",
    "from komm import AWGNChannel\n",
    "from sklearn.preprocessing import QuantileTransformer, StandardScaler, RobustScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Geranate random assortment of parameters\n",
    "def make_Paras(seedParasMin, seedParasMax, numToGen):\n",
    "    outPara = np.empty(shape =(numToGen, len(seedParasMin)))\n",
    "    paraTemp = [None] * len(seedParasMin)\n",
    "    for i in range(numToGen):\n",
    "        for j in range(len(seedParasMin)):\n",
    "            if j == 5:\n",
    "                if seedParasMin[j] < 0:\n",
    "                    low = 1.3\n",
    "                else:\n",
    "                    low = 0.77\n",
    "                if seedParasMax[j] < 0:\n",
    "                    high = 0.77\n",
    "                else:\n",
    "                    high = 1.3\n",
    "            else:\n",
    "                low = 1\n",
    "                high = 1\n",
    "            para = uniform(seedParasMin[j] * low, seedParasMax[j] * high)\n",
    "            paraTemp[j] = para\n",
    "        outPara[i] = paraTemp\n",
    "    return outPara\n",
    "        \n",
    "\n",
    "#Generate M2k_Fit\n",
    "def make_M2k_Fit(paras, imgSysData):\n",
    "    \n",
    "    A, tau, S0, alpha, phi, beta, delta_s = paras\n",
    "    _, _, K_X, K_Y = getFreq(imgSysData[\"CCDPixelSize\"],\n",
    "                                      imgSysData[\"magnification\"], (100,100))\n",
    "    d = imgSysData[\"wavelen\"] / (2*np.pi*imgSysData[\"NA\"])\n",
    "    R_p, Theta_p = np.abs(K_X + 1j*K_Y) * d, np.angle(K_X + 1j*K_Y)\n",
    "    p1 = pupilFunc(R_p, Theta_p + np.pi, tau, S0, alpha, phi, beta)\n",
    "    p2 = np.conj(pupilFunc(R_p, Theta_p, tau, S0, alpha, phi, beta)) * \\\n",
    "            np.exp(-2*1j*delta_s)\n",
    "    PSF = (p1 + p2) / (2 * np.cos(delta_s))\n",
    "    M2k = np.abs(PSF)**2\n",
    "    M2k_Fit = A * M2k\n",
    "    M2k_Fit[M2k_Fit.shape[0]//2, M2k_Fit.shape[1]//2] = 0\n",
    "    \n",
    "    return M2k_Fit\n",
    "\n",
    "#Function that gets largest values in an array\n",
    "def Get_Max(array):\n",
    "    max1 = np.amax(array) * 3\n",
    "    max2 = np.amax(array[array < max1]) * 3\n",
    "    \n",
    "    return array, max1, max2\n",
    "\n",
    "#Function that mirrors an array with dimensions of odd values\n",
    "def Mirror(array):\n",
    "    array2 = np.flip(np.flip(array, 0), 1)\n",
    "    array2 = np.delete(array2, (0), axis=1)\n",
    "    newarray = np.concatenate((array, array2), 1)\n",
    "    \n",
    "    return newarray\n",
    "\n",
    "#Turn one NPS to one OD\n",
    "def NPSavg_To_Single(NPSavgs, numNPSs):\n",
    "    NPSs = []\n",
    "    \n",
    "    #Reduce dimensions down to 49x49 image\n",
    "    NPSavgs = np.delete(NPSavgs, np.s_[75:100], axis = 1)\n",
    "    NPSavgs = np.delete(NPSavgs, np.s_[0:26], axis = 1)\n",
    "    NPSavgs = np.delete(NPSavgs, np.s_[75:100], axis = 2)\n",
    "    NPSavgs = np.delete(NPSavgs, np.s_[0:26], axis = 2)\n",
    "\n",
    "    \n",
    "    for count, NPSavg in enumerate(NPSavgs):\n",
    "        print(\"Parameters Set: %d\" % (count + 1))\n",
    "        for i in range(numNPSs):\n",
    "            NPStemp = NPSavg.copy()\n",
    "            NPStemp[NPStemp < 0] = 0 #Truncate at 0\n",
    "            NSR = np.random.exponential(scale=1.1, size=(49,25))**1.1 - 1 #Noise to signal ratio\n",
    "            NSR[NSR < -1] = -1\n",
    "            NSR, max1, max2 = Get_Max(NSR) #Get the maximum values for noise to signal ratio\n",
    "            NSR = Mirror(NSR) #Mirror the noise to signal ratio\n",
    "            Scale = np.random.uniform(bkgMin, bkgMax)\n",
    "            Bkg = np.random.exponential(scale=Scale, size=(49,25)) #Generate background noise\n",
    "            Bkg[Bkg < 0] = 0\n",
    "            Bkg = Mirror(Bkg) #Mirror the background noise\n",
    "            \n",
    "            randc = randint(1,3) // 3 #Randomly select brightness of pixels that make a plus sign around center (noise feature)\n",
    "            randx = randint(1,3) // 2 #Randomly determine whether the image has a horizontal line noise feature\n",
    "            randy = randint(1,3) // 2 #Randomly determine whether the image has a vertical line noise feature\n",
    "            \n",
    "            mean = np.mean(NPStemp) #Mean value for the array used in creating noise features\n",
    "            amax = np.amax(NPStemp) #Maximum value for the array used in creating noise features\n",
    "            mean_bkg = np.mean(Bkg) #Mean value of the photon shot noise used in creating noise features\n",
    "            \n",
    "            for j in range(len(NPStemp)):\n",
    "                for k in range(len(NPStemp.T)):\n",
    "                    if (NPSavg[j][k] == 0):#Add background noise\n",
    "                        NPStemp[j][k] += Bkg[j][k]\n",
    "                    else:\n",
    "                        NPStemp[j][k] += mean_bkg\n",
    "#                     NPStemp[j][k] += mean_bkg\n",
    "                    if ((j == 23 or j == 25) and k == 24): #Add center bright points\n",
    "                        if randc == 1:\n",
    "                            NPStemp[j][k] = mean + mean * (max2) \n",
    "                        else :\n",
    "                            NPStemp[j][k] = amax + amax * (max1)\n",
    "                    elif (j == 24 and (k == 23 or k == 25)): #Add center bright points\n",
    "                        if randc == 1:\n",
    "                            NPStemp[j][k] = amax + amax * (max1)\n",
    "                        else:\n",
    "                            NPStemp[j][k] = mean + mean * (max2)\n",
    "                    elif (NPSavg[j][k] != 0): #Add noise scaled by noise to signal ratio\n",
    "                        NPStemp[j][k] += NPStemp[j][k] * NSR[j][k]\n",
    "                    if j == 24: #Add horizontal and vertial lines\n",
    "                        NPStemp[j][k] += .5 * mean * randx \n",
    "                    if k == 24:\n",
    "                        NPStemp[j][k] += .5 * mean * randy\n",
    "            NPStemp[24][24] = 0 #Zero the center dot\n",
    "            NPStemp[NPStemp < 0] = 0 #Floor negative values to 0\n",
    "            NPStemp[NPSavg == 0] = 0\n",
    "            NPSs.append(NPStemp)\n",
    "            \n",
    "    return np.array(NPSs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "savepath = './results/' #Savepath for images and parameters\n",
    "\n",
    "numNPSs = 2 #Number of NPS images to make for each parameter set\n",
    "numParas = 25000 #Number of parameter sets to generate\n",
    "\n",
    "#Import minimum and maximum seed parameters\n",
    "seedParasMax = pkl.load(open(savepath + \"seedParasMax.pkl\", \"rb\"))\n",
    "seedParasMin = pkl.load(open(savepath + \"seedParasMin.pkl\", \"rb\"))\n",
    "bkgMax = pkl.load(open(savepath + \"bkgMax.pkl\", \"rb\"))\n",
    "bkgMin = pkl.load(open(savepath + \"bkgMin.pkl\", \"rb\"))\n",
    "\n",
    "NPSavgs = np.empty((numParas,100,100))\n",
    "fakeNPSs = np.empty((numNPSs*numParas,49,49))\n",
    "savedParas = []\n",
    "\n",
    "imgSysData = { \n",
    "    \"CCDPixelSize\": 13,      # pixel size of the CCD, in micron \n",
    "    \"magnification\": 27,      # 799.943 / 29.9099, # magnification of the imaging system \n",
    "    \"wavelen\"     : 0.852,     # wavelength of the imaging beam, in micron \n",
    "    \"NA\"          : 0.37,      # numerical aperture of the objective \n",
    "    \"ODtoAtom\"    : 13\n",
    "} \n",
    "\n",
    "#Create image and parameter arrays\n",
    "for i, paras in enumerate(make_Paras(seedParasMin, seedParasMax, numParas)):\n",
    "    if (i + 1) % 1000 == 0:\n",
    "        print(i + 1)\n",
    "    NPSavgs[i] = make_M2k_Fit(paras, imgSysData)\n",
    "    for j in range(numNPSs):\n",
    "        savedParas.append(paras)\n",
    "fakeNPSs = NPSavg_To_Single(NPSavgs, numNPSs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(np.array(fakeNPSs).flatten(), density=1, range=(0,0.25))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Save fake average images\n",
    "K_x, K_y, _, _ = getFreq(imgSysData[\"CCDPixelSize\"], imgSysData[\"magnification\"], (100,100))\n",
    "for imnum in range(100):\n",
    "    fig = plt.figure(figsize=(5.5, 4.5))\n",
    "    ax = fig.add_subplot(111)\n",
    "    im = ax.pcolor(K_x, K_y, NPSavgs[imnum], cmap=cm.jet, vmin=0, vmax=.5)\n",
    "    plt.colorbar(im)\n",
    "    fig.savefig(\"C:\\\\Users\\\\eric0\\\\OneDrive\\\\Desktop\\\\Cold Physics\\\\fake_NPSavg_ims\\\\fake_NPSavg_\" + str(imnum + 1))\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save fake single-OD images\n",
    "K_x, K_y, _, _ = getFreq(imgSysData[\"CCDPixelSize\"], imgSysData[\"magnification\"], (49,49))\n",
    "for imnum in range(100):\n",
    "    fig = plt.figure(figsize=(5.5, 4.5))\n",
    "    ax = fig.add_subplot(111)\n",
    "    im = ax.pcolor(K_x, K_y, fakeNPSs[imnum], cmap=cm.jet, vmin=0, vmax=.5)\n",
    "    plt.colorbar(im)\n",
    "    fig.savefig(\"C:\\\\Users\\\\eric0\\\\OneDrive\\\\Desktop\\\\Cold Physics\\\\fake_NPS_ims\\\\fake_NPS_\" + str(imnum + 1))\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save images and parameters\n",
    "pkl.dump(NPSavgs, open(savepath + \"fakeNPSavgs\", \"wb\"))\n",
    "pkl.dump(fakeNPSs, open(savepath + \"fakeNPSs\", \"wb\"))\n",
    "pkl.dump(savedParas, open(savepath + \"savedParas\", \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Delete extreme outlier images\n",
    "fakeNPSs = np.array(fakeNPSs).reshape((len(fakeNPSs),-1))\n",
    "savedParas = np.array(savedParas)\n",
    "savedParas = np.delete(savedParas, np.where(np.amax(fakeNPSs, axis=1) > 20), 0)\n",
    "fakeNPSs = np.delete(fakeNPSs, np.where(np.amax(fakeNPSs, axis=1) > 20), 0)\n",
    "savedParas = np.delete(savedParas, np.where(np.amin(fakeNPSs, axis=1) < 0), 0)\n",
    "fakeNPSs = np.delete(fakeNPSs, np.where(np.amin(fakeNPSs, axis=1) < 0), 0)\n",
    "\n",
    "#Normalize images and split into training and validation data\n",
    "# fakeNPSs /= 20\n",
    "defocus = savedParas[:,[5]]\n",
    "otherparas = np.delete(savedParas, 5, axis=1)\n",
    "trainin, testin, trainout, testout, otherparasin, otherparasout = train_test_split(fakeNPSs, defocus, otherparas, test_size=0.25, shuffle=True)\n",
    "for i in range(len(trainin)):\n",
    "    trainin[i] -= np.mean(trainin[i])\n",
    "    trainin[i] /= np.std(trainin[i])\n",
    "for i in range(len(testin)):\n",
    "    testin[i] -= np.mean(testin[i])\n",
    "    testin[i] /= np.std(testin[i])\n",
    "# trainin[trainin < 0.001] = 0.001\n",
    "# testin[testin < 0.001] = 0.001\n",
    "# trainin = np.log10(trainin)\n",
    "# testin = np.log10(testin)\n",
    "# trainin /= 3\n",
    "# testin /= 3\n",
    "scaler = RobustScaler(unit_variance=True)\n",
    "yscaler = scaler.fit(trainout)\n",
    "pkl.dump(yscaler, open(savepath + \"yscaler\", \"wb\"))\n",
    "trainout = yscaler.transform(trainout)\n",
    "testout = yscaler.transform(testout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image in trainin[0:20]:\n",
    "    plt.hist(image.flatten(), density=1)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check number of samples. Should be close to numNPSs*numParas\n",
    "print(np.shape(fakeNPSs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Convert training and validation data to tensors\n",
    "trainin = tf.convert_to_tensor(trainin)\n",
    "trainout = tf.convert_to_tensor(trainout)\n",
    "testin = tf.convert_to_tensor(testin)\n",
    "testout = tf.convert_to_tensor(testout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Store training and testing data\n",
    "pkl.dump(trainin, open(savepath + \"trainin\", \"wb\"))\n",
    "pkl.dump(trainout, open(savepath + \"trainout\", \"wb\"))\n",
    "pkl.dump(testin, open(savepath + \"testin\", \"wb\"))\n",
    "pkl.dump(testout, open(savepath + \"testout\", \"wb\"))\n",
    "pkl.dump(otherparasin, open(savepath + \"otherparasin\", \"wb\"))\n",
    "pkl.dump(otherparasout, open(savepath + \"otherparasout\", \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create images with same parameters as a real image set\n",
    "NPSavgs = []\n",
    "paras = []\n",
    "for i in range(100):\n",
    "    paras.append([0.3307113143177933, 1.0000000000000002, 0.009999999999999986, 2.4599999999999995, -1.926865182387927, 0.3640009284887366, -0.03788728477733443])\n",
    "for para in paras:\n",
    "    NPSavgs.append(make_M2k_Fit(para, imgSysData))\n",
    "fakeNPSs = NPSavg_To_Single(NPSavgs, numNPSs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save fake avg nps images with same parameters as real image set\n",
    "K_x, K_y, _, _ = getFreq(imgSysData[\"CCDPixelSize\"], imgSysData[\"magnification\"], (100,100))\n",
    "for imnum in range(100):\n",
    "    fig = plt.figure(figsize=(5.5, 4.5))\n",
    "    ax = fig.add_subplot(111)\n",
    "    im = ax.pcolor(K_x, K_y, NPSavgs[imnum], cmap=cm.jet, vmin=0, vmax=.5)\n",
    "    plt.colorbar(im)\n",
    "    fig.savefig(\"C:\\\\Users\\\\eric0\\\\OneDrive\\\\Desktop\\\\Cold Physics\\\\fake_NPSavg_ims_SamePara\\\\fake_NPSavg_\" + str(imnum + 1))\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save fake single-OD NPS images with same parameters as real image set\n",
    "K_x, K_y, _, _ = getFreq(imgSysData[\"CCDPixelSize\"], imgSysData[\"magnification\"], (49,49))\n",
    "for imnum in range(100):\n",
    "    fig = plt.figure(figsize=(5.5, 4.5))\n",
    "    ax = fig.add_subplot(111)\n",
    "    im = ax.pcolor(K_x, K_y, fakeNPSs[imnum], cmap=cm.jet, vmin = 0, vmax=.5)\n",
    "    plt.colorbar(im)\n",
    "    fig.savefig(\"C:\\\\Users\\\\eric0\\\\OneDrive\\\\Desktop\\\\Cold Physics\\\\fake_NPS_ims_SamePara\\\\fake_NPS_\" + str(imnum + 1))\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for fakeNPS in fakeNPSs:\n",
    "    plt.hist(np.array(fakeNPS).flatten(), density=1, range=(0,0.25))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}