{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:51:58.660252Z",
     "start_time": "2019-09-16T04:51:58.428690Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import galsim\n",
    "import numpy as np\n",
    "from IPython.display import display\n",
    "from ipywidgets import interact, interactive_output, interact_manual\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:51:59.542479Z",
     "start_time": "2019-09-16T04:51:59.507864Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fiducial_telescope = batoid.Optic.fromYaml(\"AuxTel.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:51:59.905273Z",
     "start_time": "2019-09-16T04:51:59.900103Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def spotPlot(telescope, wavelength, theta_x, theta_y, logscale, ax):\n",
    "    rays = batoid.RayVector.asPolar(\n",
    "        optic=telescope,\n",
    "        inner=telescope.pupilSize/2*telescope.pupilObscuration,\n",
    "        theta_x=np.deg2rad(theta_x/60), theta_y=np.deg2rad(theta_y/60),\n",
    "        nrad=48, naz=192, wavelength=wavelength*1e-9\n",
    "    )\n",
    "\n",
    "    telescope.trace(rays)\n",
    "    w = ~rays.vignetted\n",
    "    spots = np.vstack([rays.x[w], rays.y[w]])\n",
    "    spots -= np.mean(spots, axis=1)[:,None]\n",
    "    spots *= 1e6 # meters -> microns\n",
    "\n",
    "    ax.scatter(spots[0], spots[1], s=1, alpha=0.5)\n",
    "    ax.set_xlim(-0.2*10**logscale, 0.2*10**logscale)\n",
    "    ax.set_ylim(-0.2*10**logscale, 0.2*10**logscale)\n",
    "    ax.set_title(r\"$\\theta_x = {:4.2f}\\,,\\theta_y = {:4.2f}$\".format(theta_x, theta_y))\n",
    "    ax.set_xlabel(\"microns\")\n",
    "    ax.set_ylabel(\"microns\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:00.263620Z",
     "start_time": "2019-09-16T04:52:00.259915Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def wavefrontPlot(telescope, wavelength, theta_x, theta_y, ax):\n",
    "    wf = batoid.wavefront(\n",
    "        telescope, \n",
    "        np.deg2rad(theta_x/60), np.deg2rad(theta_y/60), \n",
    "        wavelength*1e-9, nx=128\n",
    "    )    \n",
    "    wfplot = ax.imshow(\n",
    "        wf.array,\n",
    "        extent=np.r_[-1,1,-1,1]*telescope.pupilSize/2\n",
    "    )\n",
    "    ax.set_xlabel(\"meters\")\n",
    "    ax.set_ylabel(\"meters\")\n",
    "    plt.colorbar(wfplot, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:00.657493Z",
     "start_time": "2019-09-16T04:52:00.652420Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fftPSFPlot(telescope, wavelength, theta_x, theta_y, ax):\n",
    "    fft = batoid.fftPSF(\n",
    "        telescope, \n",
    "        np.deg2rad(theta_x/60), np.deg2rad(theta_y/60), \n",
    "        wavelength*1e-9, nx=32\n",
    "    )\n",
    "    # We should be very close to primitive vectors that are a multiple of\n",
    "    # [1,0] and [0,1].  If the multiplier is negative though, then this will\n",
    "    # make it look like our PSF is upside-down.  So we check for this here and \n",
    "    # invert if necessary.  This will make it easier to compare with the spot \n",
    "    # diagram, for instance\n",
    "    if fft.primitiveVectors[0,0] < 0:\n",
    "        fft.array = fft.array[::-1,::-1]\n",
    "\n",
    "    scale = np.sqrt(np.abs(np.linalg.det(fft.primitiveVectors)))\n",
    "    nxout = fft.array.shape[0]\n",
    "    fft.array /= np.sum(fft.array)\n",
    "    fftplot = ax.imshow(\n",
    "        fft.array,\n",
    "        extent=np.r_[-1,1,-1,1]*scale*nxout/2*1e6\n",
    "    )\n",
    "    ax.set_title(\"FFT PSF\")\n",
    "    ax.set_xlabel(\"micron\")\n",
    "    ax.set_ylabel(\"micron\")    \n",
    "    plt.colorbar(fftplot, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:01.120718Z",
     "start_time": "2019-09-16T04:52:01.116117Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def huygensPSFPlot(telescope, wavelength, theta_x, theta_y, ax):\n",
    "    huygensPSF = batoid.huygensPSF(\n",
    "        telescope, \n",
    "        np.deg2rad(theta_x/60), np.deg2rad(theta_y/60),\n",
    "        wavelength*1e-9, nx=32\n",
    "    )\n",
    "    # We should be very close to primitive vectors that are a multiple of\n",
    "    # [1,0] and [0,1].  If the multiplier is negative though, then this will\n",
    "    # make it look like our PSF is upside-down.  So we check for this here and \n",
    "    # invert if necessary.  This will make it easier to compare with the spot \n",
    "    # diagram, for instance\n",
    "    if huygensPSF.primitiveVectors[0,0] < 0:\n",
    "        huygensPSF.array = huygensPSF.array[::-1,::-1]\n",
    "\n",
    "    huygensPSF.array /= np.sum(huygensPSF.array)    \n",
    "    scale = np.sqrt(np.abs(np.linalg.det(huygensPSF.primitiveVectors)))\n",
    "    nxout = huygensPSF.array.shape[0]\n",
    "    \n",
    "    huygensplot = plt.imshow(\n",
    "        huygensPSF.array,\n",
    "        extent=np.r_[-1,1,-1,1]*scale*nxout/2*1e6\n",
    "    )\n",
    "    ax.set_title(\"Huygens PSF\")\n",
    "    ax.set_xlabel(\"micron\")\n",
    "    ax.set_ylabel(\"micron\")    \n",
    "    plt.colorbar(huygensplot, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:02.687372Z",
     "start_time": "2019-09-16T04:52:01.806203Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "what = dict(\n",
    "    do_spot = widgets.Checkbox(value=True, description='Spot'),\n",
    "    do_wavefront = widgets.Checkbox(value=True, description='Wavefront'),\n",
    "    do_fftPSF = widgets.Checkbox(value=True, description='FFT PSF'),\n",
    "    do_huygensPSF = widgets.Checkbox(value=True, description='Huygens PSF')\n",
    ")\n",
    "where = dict(\n",
    "    wavelength=widgets.FloatSlider(min=300.0,max=1100.0,step=25.0,value=625.0, description=\"$\\lambda$ (nm)\"),\n",
    "    theta_x=widgets.FloatSlider(min=-4.0,max=4.0,step=0.1,value=0.0, description=\"$\\\\theta_x (arcmin)$\"),\n",
    "    theta_y=widgets.FloatSlider(min=-4.0,max=4.0,step=0.1,value=0.0, description=\"$\\\\theta_y (arcmin)$\"),\n",
    "    logscale=widgets.FloatSlider(min=1.5, max=3, step=0.1, value=1, description=\"scale\")\n",
    ")\n",
    "perturb = dict(\n",
    "    optic=widgets.Dropdown(\n",
    "        options=fiducial_telescope.itemDict.keys(), \n",
    "        value='AuxTel.M2'\n",
    "    ),\n",
    "    dx=widgets.FloatSlider(min=-0.2, max=0.2, step=0.05, value=0.0, description=\"dx ($mm$)\"),\n",
    "    dy=widgets.FloatSlider(min=-0.2, max=0.2, step=0.05, value=0.0, description=\"dy ($mm$)\"),\n",
    "    dz=widgets.FloatSlider(min=-100, max=100, step=1, value=0.0, description=\"dz ($\\mu m$)\"),\n",
    "    dthx=widgets.FloatSlider(min=-1, max=1, step=0.1, value=0.0, description=\"d$\\phi_x$ (arcmin)\"),\n",
    "    dthy=widgets.FloatSlider(min=-1, max=1, step=0.1, value=0.0, description=\"d$\\phi_y$ (arcmin)\"),\n",
    ")\n",
    "\n",
    "def f(do_spot, do_wavefront, do_fftPSF, do_huygensPSF,\n",
    "    wavelength, theta_x, theta_y, optic, dx, dy, dz, dthx, dthy, logscale, **kwargs):\n",
    "    fiducial_telescope = batoid.Optic.fromYaml(\"AuxTel.yaml\")\n",
    "    telescope = (fiducial_telescope\n",
    "            .withGloballyShiftedOptic(optic, [dx*1e-3, dy*1e-3, dz*1e-6])\n",
    "            .withLocallyRotatedOptic(optic, batoid.RotX(np.deg2rad(dthx/60)).dot(batoid.RotY(np.deg2rad(dthy/60))))\n",
    "    )    \n",
    "    nplot = sum([do_spot, do_wavefront, do_fftPSF, do_huygensPSF])\n",
    "    \n",
    "    if nplot > 0:\n",
    "        fig, axes = plt.subplots(ncols=nplot, figsize=(4*nplot, 4), squeeze=False)\n",
    "\n",
    "        iax = 0\n",
    "        if do_spot:\n",
    "            ax = axes.ravel()[iax]\n",
    "            spotPlot(telescope, wavelength, theta_x, theta_y, logscale, ax)\n",
    "            iax += 1\n",
    "\n",
    "        if do_wavefront:\n",
    "            ax = axes.ravel()[iax]\n",
    "            wavefrontPlot(telescope, wavelength, theta_x, theta_y, ax)\n",
    "            iax += 1\n",
    "\n",
    "        if do_fftPSF:\n",
    "            ax = axes.ravel()[iax]\n",
    "            fftPSFPlot(telescope, wavelength, theta_x, theta_y, ax)\n",
    "            iax += 1\n",
    "\n",
    "        if do_huygensPSF:\n",
    "            ax = axes.ravel()[iax]\n",
    "            huygensPSFPlot(telescope, wavelength, theta_x, theta_y, ax)\n",
    "\n",
    "        fig.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "all_widgets = {}\n",
    "for d in [what, where, perturb]:\n",
    "    for k in d:\n",
    "        all_widgets[k] = d[k]\n",
    "\n",
    "output = interactive_output(f, all_widgets)\n",
    "display(\n",
    "    widgets.VBox([\n",
    "        widgets.HBox([\n",
    "            widgets.VBox([v for v in what.values()]), \n",
    "            widgets.VBox([v for v in where.values()]), \n",
    "            widgets.VBox([v for v in perturb.values()])\n",
    "        ]),\n",
    "        output\n",
    "    ]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:13.478276Z",
     "start_time": "2019-09-16T04:52:13.342499Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@interact(wavelen=widgets.FloatSlider(min=300.0,max=1100.0,step=25.0,value=625.0,\n",
    "                                      description=\"$\\lambda$ (nm)\"),\n",
    "          theta_x=widgets.FloatSlider(min=-4.0,max=4.0,step=0.1,value=1.0, \n",
    "                                      description=\"$\\\\theta_x (arcmin)$\"),\n",
    "          theta_y=widgets.FloatSlider(min=-4.0,max=4.0,step=0.1,value=0.0, \n",
    "                                      description=\"$\\\\theta_y (arcmin)$\"),\n",
    "          optic=widgets.Dropdown(\n",
    "              options=fiducial_telescope.itemDict.keys(), \n",
    "              value='AuxTel.M2'\n",
    "          ),\n",
    "          dx=widgets.FloatSlider(min=-5, max=5, step=0.05, value=0.0,\n",
    "                                 description=\"dx ($mm$)\"),\n",
    "          dy=widgets.FloatSlider(min=-5, max=5, step=0.05, value=0.0,\n",
    "                                 description=\"dy ($mm$)\"),\n",
    "          dz=widgets.FloatSlider(min=-2000, max=2000, step=1, value=0.0,\n",
    "                                 description=\"dz ($\\mu m$)\"),\n",
    "          dthx=widgets.FloatSlider(min=-10, max=10, step=0.1, value=0.0,\n",
    "                                   description=\"d$\\phi_x$ (arcmin)\"),\n",
    "          dthy=widgets.FloatSlider(min=-10, max=10, step=0.1, value=0.0,\n",
    "                                   description=\"d$\\phi_y$ (arcmin)\"))\n",
    "def zernike(wavelen, theta_x, theta_y, optic, dx, dy, dz, dthx, dthy):\n",
    "    telescope = (fiducial_telescope\n",
    "            .withGloballyShiftedOptic(optic, [dx*1e-3, dy*1e-3, dz*1e-6])\n",
    "            .withLocallyRotatedOptic(optic, batoid.RotX(np.deg2rad(dthx/60)).dot(batoid.RotY(np.deg2rad(dthy/60))))\n",
    "    )\n",
    "\n",
    "    z = batoid.zernikeGQ(\n",
    "        telescope, np.deg2rad(theta_x/60), np.deg2rad(theta_y/60), wavelen*1e-9,\n",
    "        jmax=22, eps=telescope.pupilObscuration, reference='chief'\n",
    "    )*wavelen\n",
    "    for i in range(1, len(z)//2+1):\n",
    "        print(\"{:6d}   {:10.3f}      {:6d}  {:10.3f}\".format(i, z[i], i+11, z[i+11]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T04:52:33.647764Z",
     "start_time": "2019-09-16T04:52:33.574470Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@interact_manual(\n",
    "    wavelen=widgets.FloatSlider(min=300.0,max=1100.0,step=25.0,value=625.0,\n",
    "                                description=\"$\\lambda$ (nm)\"),\n",
    "    optic=widgets.Dropdown(\n",
    "        options=fiducial_telescope.itemDict.keys(), \n",
    "        value='AuxTel.LATISS'\n",
    "    ),\n",
    "    z_coef=widgets.Dropdown(\n",
    "        options=list(range(1, 56)), value=1,\n",
    "        description=\"Zernike coefficient\"\n",
    "    ),\n",
    "    z_amp=widgets.FloatSlider(min=-0.1, max=0.1, step=0.01, value=0.0,\n",
    "                              description=\"Zernike amplitude\"),\n",
    "    dx=widgets.FloatSlider(min=-0.2, max=0.2, step=0.05, value=0.0,\n",
    "                           description=\"dx ($mm$)\"),\n",
    "    dy=widgets.FloatSlider(min=-0.2, max=0.2, step=0.05, value=0.0,\n",
    "                           description=\"dy ($mm$)\"),\n",
    "    dz=widgets.FloatSlider(min=-500, max=500, step=10, value=0.0,\n",
    "                           description=\"dz ($\\mu m$)\"),\n",
    "    dthx=widgets.FloatSlider(min=-1, max=1, step=0.1, value=0.0,\n",
    "                             description=\"d$\\phi_x$ (arcmin)\"),\n",
    "    dthy=widgets.FloatSlider(min=-1, max=1, step=0.1, value=0.0,\n",
    "                             description=\"d$\\phi_y$ (arcmin)\"),\n",
    "    do_resid=widgets.Checkbox(value=False, description=\"residual?\"))\n",
    "def zFoV(wavelen, optic, z_coef, z_amp, dx, dy, dz, dthx, dthy, do_resid):\n",
    "    telescope = (fiducial_telescope\n",
    "            .withGloballyShiftedOptic(optic, [dx*1e-3, dy*1e-3, dz*1e-6])\n",
    "            .withLocallyRotatedOptic(\n",
    "                    optic,\n",
    "                    batoid.RotX(dthx*np.pi/180/60).dot(batoid.RotY(dthy*np.pi/180/60))\n",
    "            )\n",
    "    )\n",
    "    if z_amp != 0:\n",
    "        try:\n",
    "            interface = telescope[optic]\n",
    "            s0 = interface.surface\n",
    "        except:\n",
    "            pass\n",
    "        else:\n",
    "            s1 = batoid.Sum([\n",
    "                s0,\n",
    "                batoid.Zernike(\n",
    "                    [0]*z_coef+[z_amp*wavelen*1e-9], \n",
    "                    R_outer=interface.outRadius,\n",
    "                    R_inner=interface.inRadius,\n",
    "                )\n",
    "            ])\n",
    "            telescope = telescope.withSurface(optic, s1)\n",
    "\n",
    "    thmax = 3.5  # arcmin\n",
    "    thxs = np.linspace(-thmax, thmax, 15)\n",
    "    thys = np.linspace(-thmax, thmax, 15)\n",
    "\n",
    "    img = np.zeros((15, 15), dtype=float)\n",
    "    vmin = -0.3\n",
    "    vmax = 0.3\n",
    "    zs = []\n",
    "    thxplot = []\n",
    "    thyplot = []\n",
    "    for ix, thx in enumerate(thxs):\n",
    "        for iy, thy in enumerate(thys):\n",
    "            if np.hypot(thx, thy) > thmax: \n",
    "                continue\n",
    "            z = batoid.zernike(\n",
    "                telescope, np.deg2rad(thx/60), np.deg2rad(thy/60), wavelen*1e-9,\n",
    "                jmax=21, eps=telescope.pupilObscuration, nx=16\n",
    "            )\n",
    "            thxplot.append(thx)\n",
    "            thyplot.append(thy)\n",
    "            if do_resid:\n",
    "                vmin = -0.05\n",
    "                vmax = 0.05\n",
    "                z -= batoid.zernike(\n",
    "                    fiducial_telescope, np.deg2rad(thx/60), np.deg2rad(thy/60), 625e-9,\n",
    "                    jmax=21, eps=telescope.pupilObscuration, nx=16\n",
    "                )\n",
    "            zs.append(z)\n",
    "    zs = np.array(zs).T\n",
    "    thxplot = np.array(thxplot)\n",
    "    thyplot = np.array(thyplot)\n",
    "    fig = plt.figure(figsize=(13, 8))\n",
    "    batoid.plotUtils.zernikePyramid(thxplot, thyplot, zs[4:], vmin=vmin, vmax=vmax, fig=fig)\n",
    "    plt.show()\n",
    "\n",
    "    # Compute double Zernike \n",
    "    fBasis = galsim.zernike.zernikeBasis(22, thxplot, thyplot, thmax)\n",
    "    dzs, _, _, _ = np.linalg.lstsq(fBasis.T, zs.T, rcond=None)\n",
    "    dzs = dzs[:,4:]\n",
    "    asort = np.argsort(np.abs(dzs).ravel())[::-1]\n",
    "    focal_idx, pupil_idx = np.unravel_index(asort[:10], dzs.shape)\n",
    "    cumsum = 0.0\n",
    "    for fid, pid in zip(focal_idx, pupil_idx):\n",
    "        val = dzs[fid, pid]\n",
    "        cumsum += val**2\n",
    "        print(\"{:3d} {:3d} {:8.4f} {:8.4f}\".format(fid, pid+4, val, np.sqrt(cumsum)))\n",
    "    print(\"sum sqr dz {:8.4f}\".format(np.sqrt(np.sum(dzs**2))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  },
  "widgets": {
   "state": {
    "bf032f0236b24ec8b012732196da0aeb": {
     "views": [
      {
       "cell_index": 7
      }
     ]
    },
    "e58d73c5a90d4ac2b12b81970dba2174": {
     "views": [
      {
       "cell_index": 5
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
