create_test_case.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cupy as cp\n",
    "import xara\n",
    "import xaosim\n",
    "from tqdm import  tqdm\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "from time import sleep\n",
    "from time import time\n",
    "import scipy.sparse as sparse\n",
    "\n",
    "from xara import fft, ifft, shift\n",
    "import astropy.units as u\n",
    "from scipy.ndimage import median_filter\n",
    "from scipy.signal import correlate2d\n",
    "#from sneks import FigSaver\n",
    "#fsave = FigSaver(\"figsave.json\")\n",
    "import saro"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run /home/rlaugier/Documents/kernel/adk/detection_maps_cleaned.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shifter(im0,vect, buildmask = True, sg_rad=40.0, verbose=False, nbit=10):\n",
    "\n",
    "\n",
    "    szh = im0.shape[1] # horiz\n",
    "    szv = im0.shape[0] # vertic\n",
    "\n",
    "    temp = np.max(im0.shape) # max dimension of image\n",
    "\n",
    "    for sz in [64, 128, 256, 512, 1024, 2048]:\n",
    "        if sz >= temp: break\n",
    "\n",
    "    dz = sz//2.           # image half-size\n",
    "    if buildmask:\n",
    "        #print(\"We have to make a new mask here\")\n",
    "        imcenter = xara.find_psf_center(im0, verbose=verbose)\n",
    "        sgmask = xara.super_gauss(sz, sz, imcenter[1], imcenter[0], sg_rad)\n",
    "    else:\n",
    "        #print(\"Mask already exists\")\n",
    "        print(\"ERROR: HERE you should build the relevant mask\")\n",
    "        return\n",
    "    x,y = np.meshgrid(np.arange(sz)-dz, np.arange(sz)-dz)\n",
    "    wedge_x, wedge_y = x*np.pi/dz, y*np.pi/dz\n",
    "    offset = np.zeros((sz, sz), dtype=complex) # to Fourier-center array\n",
    "\n",
    "    # insert image in zero-padded array (dim. power of two)\n",
    "    im = np.zeros((sz, sz))\n",
    "    orih, oriv = (sz-szh)//2, (sz-szv)//2\n",
    "    im[oriv:oriv+szv,orih:orih+szh] = im0\n",
    "    \n",
    "    #print(vect[1],vect[0])\n",
    "\n",
    "    (x0, y0) = (vect[1], vect[0])\n",
    "    \n",
    "    im -= np.median(im)\n",
    "\n",
    "    dx, dy = x0, y0\n",
    "    im = np.roll(np.roll(im, -int(dx), axis=1), -int(dy), axis=0)\n",
    "\n",
    "    #print(\"recenter: dx=%.2f, dy=%.2f\" % (dx, dy))\n",
    "    dx -= np.int(dx)\n",
    "    dy -= np.int(dy)\n",
    "\n",
    "    temp = im * sgmask\n",
    "    mynorm = temp.sum()\n",
    "\n",
    "    # array for Fourier-translation\n",
    "    dummy = shift(dx * wedge_x + dy * wedge_y)\n",
    "    offset.real, offset.imag = np.cos(dummy), np.sin(dummy)\n",
    "    dummy = np.abs(shift(ifft(offset * fft(shift(temp)))))\n",
    "\n",
    "    #dummy = im\n",
    "    # image masking, and set integral to right value\n",
    "    dummy *= sgmask\n",
    "\n",
    "    return (dummy * mynorm / dummy.sum())\n",
    "\n",
    "def intro_companion(image, params, pscale):\n",
    "    rho = params[0]\n",
    "    theta = params[1]\n",
    "    c = params[2]\n",
    "    xshift = (- rho * np.sin(np.deg2rad(theta))) / pscale\n",
    "    yshift = ( rho * np.cos(np.deg2rad(theta)) ) / pscale\n",
    "    \n",
    "    compagim = shifter(image, np.array([-yshift, -xshift]))\n",
    "    return image + compagim / c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    myinstrument.start()\n",
    "    myinstrument.stop()\n",
    "    myinstrument.close()\n",
    "except:\n",
    "    print(\"no Instrument was found\")\n",
    "myinstrument = xaosim.instrument()\n",
    "\n",
    "thepupil = myinstrument.cam.pupil[:,:]\n",
    "#thepupil = myinstrument.cam.pupil[:,:]\n",
    "\n",
    "wl = 1.6e-6\n",
    "ron = 0.05\n",
    "centerpup = thepupil.shape[0]//2\n",
    "resolpup = 206 #Careful: this relies on the instrument defaults\n",
    "pscale = myinstrument.cam.pscale\n",
    "radius = resolpup // 2\n",
    "padpup = 5\n",
    "resolker = 36\n",
    "teldiam = 7.92\n",
    "binary = True\n",
    "ppscale = teldiam / resolpup\n",
    "croppedpup = thepupil[centerpup - radius-padpup:centerpup + radius+padpup,\n",
    "                      centerpup - radius-padpup:centerpup +radius+padpup]\n",
    "\n",
    "imsize = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = wl / (2.*0.23) * u.rad.to(u.mas) / pscale\n",
    "print(\"A good value for r0 is %.1f\"%(p))\n",
    "mask = xara.super_gauss(imsize, imsize, imsize/2, imsize/2, p)\n",
    "wferror = 50.  #RMS wavefront error in nm\n",
    "Nphot = 1e4  #Total number of photons per frame\n",
    "Nreal = 10000\n",
    "ron = 0.\n",
    "apod = False\n",
    "\n",
    "myinstrument.cam.pupil = thepupil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Definition of the parameters for simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Creation of phase screens\n",
    "sz = myinstrument.atmo.rndarr.shape[0]\n",
    "myinstrument.atmo.update_rms(wferror)\n",
    "phs = []\n",
    "myinstrument.atmo.update_screen()\n",
    "for i in tqdm(range(Nreal)):\n",
    "    myinstrument.atmo.rndarr = np.random.rand(sz,sz)\n",
    "    myinstrument.atmo.update_screen()\n",
    "    phs.append(myinstrument.atmo.shm_phs.get_data().astype(np.float32))\n",
    "#    plt.figure()\n",
    "#    plt.imshow(phs[i])\n",
    "#    plt.colorbar()\n",
    "#    print(\"RMS = %.1f\"%(np.std(phs[i])))\n",
    "phs = np.array(phs, dtype=np.float32)\n",
    "np.save(\"test_data/phs_%d_%.0f_2.npy\"%(Nreal,wferror), phs)\n",
    "\n",
    "#loading the phase screen\n",
    "\n",
    "#phs2 = np.load(\"phs_100_50.npy\")\n",
    "#phs = np.load(\"phs_100_10.npy\") + phs2[3]\n",
    "#plotim(phs2[0], name=\"Phase screen1 (nm)\")\n",
    "#plotim(phs2[20], name=\"Phase screen2 (nm)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"test_data/test_case_phs\", phs[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = pupil_tool(pupil_array, binary=True, step=0.3, thr=0.5, blim=0.9, name=\"Scexao_normal_pupil\")\n",
    "model = xara.create_discrete_model(apert=myinstrument.cam.pupil,step=0.35, ppscale=ppscale, binary=False, tmin=0.2)\n",
    "\n",
    "print(\"Building original reduced pupil\")\n",
    "thekpo = xara.KPO(array=model)#,bmax=7.5\n",
    "a = plt.figure()\n",
    "plt.imshow(myinstrument.cam.pupil*1)\n",
    "plt.show()\n",
    "a = thekpo.kpi.plot_pupil_and_uv(cmap=\"viridis\")\n",
    "m2pix = xara.mas2rad( myinstrument.cam.pscale* 128 / 1.6e-6)#pscale supposed 16.7\n",
    "thekpo.CWAVEL = wl\n",
    "\n",
    "np.savetxt(\"test_data/pupil_model.txt\", model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phs = np.load(\"phs_%d_%.0f_2.npy\"%(Nreal,wferror))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading an apodizing phase plate\n",
    "The apodizing phase plates are built using a modified Gerchberg-Saxton algorithm (https://www.researchgate.net/publication/241486004_Pupil_phase_apodization_for_achromatic_imaging_of_extra-_solar_planets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_noise(image, Nphot=1e6, ron=1.):\n",
    "    ronframe = np.random.normal(loc=0, scale=ron,size=image.shape)\n",
    "    rescaled = image * Nphot / np.max(image)\n",
    "    noisy = ronframe + np.random.poisson(lam=rescaled)\n",
    "    noisy = np.clip(noisy, 0, np.inf)\n",
    "    return noisy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "cleanscreen = np.zeros_like(phs[0])\n",
    "x0, y0, hw = 320//2, 256//2, 128//2\n",
    "cleanimages = []\n",
    "parimages = []\n",
    "cvissclean = []\n",
    "cvisspar = []\n",
    "cvissnosig = []\n",
    "for i in tqdm(range(Nreal)):\n",
    "    #myinstrument.atmo.shm_phs.set_data((cleanscreen + cropapp).astype(np.float64))\n",
    "    #therawclean = myinstrument.snap()[:,x0-128:x0+128]\n",
    "    #compagim = intro_companion(therawclean, params=params[i],pscale=pscale)\n",
    "    #cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    #cleanimages.append(cropped)\n",
    "    #cvissclean.append(thekpo.extract_cvis_from_img(cropped,m2pix,method=\"LDFT1\"))\n",
    "    \n",
    "    myinstrument.atmo.shm_phs.set_data((phs[i] + cropapp).astype(np.float64))\n",
    "    therawpar = myinstrument.snap()[:,x0-128:x0+128]\n",
    "    cropped = therawpar[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    aberrated = add_noise(cropped, Nphot=Nphot, ron=ron)\n",
    "    cvissnosig.append(thekpo.extract_cvis_from_img(aberrated,m2pix,method=\"LDFT1\"))\n",
    "    #compagim = intro_companion(therawpar, params=params[i],pscale=pscale)\n",
    "    #cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    #parimages.append(cropped)\n",
    "    #cvisspar.append(thekpo.extract_cvis_from_img(cropped, m2pix, method=\"LDFT1\"))\n",
    "parimages = np.array(parimages)\n",
    "cleanimages = np.array(cleanimages)\n",
    "cvissclean = np.array(cvissclean)\n",
    "cvisspar = np.array(cvisspar)\n",
    "cvissnosig = np.array(cvissnosig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kerns = thekpo.kpi.KPM.dot(np.angle(cvissnosig).T).T\n",
    "covk = np.cov(kerns.T)\n",
    "a = plt.figure()\n",
    "plt.imshow(covk)\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "np.save(\"test_data/covk_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,apod),covk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = scipy.linalg.sqrtm(np.linalg.inv(covk))\n",
    "thekpo.Mp = Mp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating the test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ndata=2\n",
    "cleanscreen = np.zeros_like(phs[0])\n",
    "x0, y0, hw = 320//2, 256//2, 128//2\n",
    "\n",
    "params_single = np.array([75., 127., 10])\n",
    "\n",
    "detpas = np.array([0., 20., 32., 40.])\n",
    "params_series = np.array([np.array([75., 127.+detpa, 10]) for detpa in detpas])\n",
    "\n",
    "myinstrument.atmo.shm_phs.set_data((phs[0] + cropapp).astype(np.float64))\n",
    "therawpar = myinstrument.snap()[:,x0-128:x0+128]\n",
    "cropped = therawpar[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "calibrator = add_noise(cropped, Nphot=Nphot, ron=ron)\n",
    "myinstrument.atmo.shm_phs.set_data((phs[1] + cropapp).astype(np.float64))\n",
    "therawpar = myinstrument.snap()[:,x0-128:x0+128]\n",
    "compagim = intro_companion(therawpar, params=params_single,pscale=pscale)\n",
    "cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "tarimage_single = add_noise(cropped, Nphot=Nphot, ron=ron)\n",
    "\n",
    "tarimage_series =[]\n",
    "for i in range(detpas.shape[0]):\n",
    "    myinstrument.atmo.shm_phs.set_data((phs[2+i] + cropapp).astype(np.float64))\n",
    "    therawpar = myinstrument.snap()[:,x0-128:x0+128]\n",
    "    compagim = intro_companion(therawpar, params=params_series[i],pscale=pscale)\n",
    "    cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    tarimage_series.append(add_noise(cropped, Nphot=Nphot, ron=ron))\n",
    "tarimage_series = np.array(tarimage_series)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"test_data/calibrator\",calibrator)\n",
    "np.save(\"test_data/tarimage_single\",tarimage_single)\n",
    "np.save(\"test_data/tarimage_series\",tarimage_series)\n",
    "np.save(\"test_data/params_single\", params_single)\n",
    "np.save(\"test_data/params_series\", params_series)\n",
    "np.save(\"test_data/detpas\", detpas)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cleanscreen = np.zeros_like(phs[0])\n",
    "x0, y0, hw = 320//2, 256//2, 128//2\n",
    "cleanimages = []\n",
    "parimages = []\n",
    "cvissclean = []\n",
    "cvisspar = []\n",
    "cvissnosig = []\n",
    "for i in tqdm(range(Nreal)):\n",
    "    #myinstrument.atmo.shm_phs.set_data((cleanscreen + cropapp).astype(np.float64))\n",
    "    #therawclean = myinstrument.snap()[:,x0-128:x0+128]\n",
    "    #compagim = intro_companion(therawclean, params=params[i],pscale=pscale)\n",
    "    #cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    #cleanimages.append(cropped)\n",
    "    #cvissclean.append(thekpo.extract_cvis_from_img(cropped,m2pix,method=\"LDFT1\"))\n",
    "    \n",
    "    myinstrument.atmo.shm_phs.set_data((phs[i] + cropapp).astype(np.float64))\n",
    "    therawpar = myinstrument.snap()[:,x0-128:x0+128]\n",
    "    cropped = therawpar[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    aberrated = add_noise(cropped, Nphot=Nphot, ron=ron)\n",
    "    cvissnosig.append(thekpo.extract_cvis_from_img(aberrated,m2pix,method=\"LDFT1\"))\n",
    "    #compagim = intro_companion(therawpar, params=params[i],pscale=pscale)\n",
    "    #cropped = compagim[y0-hw:y0+hw,y0-hw:y0+hw]\n",
    "    #parimages.append(cropped)\n",
    "    #cvisspar.append(thekpo.extract_cvis_from_img(cropped, m2pix, method=\"LDFT1\"))\n",
    "parimages = np.array(parimages)\n",
    "cleanimages = np.array(cleanimages)\n",
    "cvissclean = np.array(cvissclean)\n",
    "cvisspar = np.array(cvisspar)\n",
    "cvissnosig = np.array(cvissnosig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resol = 80\n",
    "step = 10.\n",
    "#totalrot.append(dtheta)\n",
    "pgrid = build_grid(resol, step)\n",
    "conts = thekpo.get_sadk_contrast_cpu(pgrid,0.01, 0.95, wl, W=None, dthetas=[0]).reshape((resol, resol))\n",
    "np.save(\"test_data/detection_map_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,apod),conts)\n",
    "logcontrast = np.log10(conts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"~/Documents/kernel/tools/colormaps/kernel_cm.py\"\n",
    "logcontrast_apod = np.log10(np.load(\"detection_map_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,True)))\n",
    "logcontrast_classical = np.log10(np.load(\"detection_map_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,False)))\n",
    "amin = 1.\n",
    "amax = np.nanmax([logcontrast_apod, logcontrast_classical])\n",
    "a = plt.figure(figsize=(12,6))\n",
    "plt.subplot(121)\n",
    "plt.imshow(logcontrast_apod, cmap=\"inferno\", vmin=amin, vmax=amax)\n",
    "plt.colorbar()\n",
    "plt.title(\"Apodized logcontrast\")\n",
    "plt.subplot(122)\n",
    "plt.imshow(logcontrast_classical, cmap=\"inferno\", vmin=amin, vmax=amax)\n",
    "plt.colorbar()\n",
    "plt.title(\"Normal logcontrast\")\n",
    "plt.show()\n",
    "fsave.save(a,\"Both_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,apod))\n",
    "a = plt.figure()\n",
    "plt.imshow(logcontrast_apod - logcontrast_classical, cmap=bbr, vmin=-1., vmax=1.)\n",
    "plt.colorbar()\n",
    "plt.title(\"Compared log-contrast normalized by peak count\")\n",
    "plt.show()\n",
    "fsave.save(a,\"Compared_%d_%.0f_2_apod_%r.npy\"%(Nreal,wferror,apod))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "np.save(\"phasapod_cvissclean_50\", cvissclean)\n",
    "np.save(\"phasapod_cvisspar_50\", cvisspar)\n",
    "np.save(\"phasapod_cvissnosig\", cvissnosig)\n",
    "np.save(\"phasapod_params\", params)\n",
    "#cvissclean = np.load(\"phasapod_cvissclean_25.npy\")\n",
    "#cvisspar = np.load(\"phasapod_cvisspar_25.npy\")\n",
    "#cvissnosig = np.load(\"phasapod_cvissnosig.npy\")\n",
    "#params = np.load(\"phasapod_params.npy\")"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}