{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T05:08:18.589482Z",
     "start_time": "2019-09-16T05:08:18.328879Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import os\n",
    "import yaml\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact\n",
    "import ipywidgets as widgets\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T05:08:18.616504Z",
     "start_time": "2019-09-16T05:08:18.591020Z"
    }
   },
   "outputs": [],
   "source": [
    "D = 0.1\n",
    "m0 = batoid.ConstMedium(1.000277)  # air\n",
    "w, n = np.genfromtxt(os.path.join(batoid.datadir, \"media\", \"silica_dispersion.txt\")).T\n",
    "w *= 1e-6\n",
    "m1 = batoid.TableMedium(w, n)\n",
    "wavelength = 500e-9\n",
    "n0 = m0.getN(wavelength)\n",
    "n1 = m1.getN(wavelength)\n",
    "R1 = 0.5  # convex entrance surface\n",
    "R2 = -0.5  # convex exit surface\n",
    "d = 0.01\n",
    "# Lens maker equation:\n",
    "f_inv = (n1-n0)*(1./R1 - 1./R2 + (n1-n0)*d/R1/R2)\n",
    "f = 1/f_inv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.000277 1.4623265011100002\n"
     ]
    }
   ],
   "source": [
    "print(n0, n1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T05:11:37.703648Z",
     "start_time": "2019-09-16T05:11:37.700031Z"
    }
   },
   "outputs": [],
   "source": [
    "def makeSystem(focus):\n",
    "    systemStr = \"\"\"\n",
    "    opticalSystem:\n",
    "      type: CompoundOptic\n",
    "      inMedium: 1.000277\n",
    "      items:\n",
    "        -\n",
    "          type: RefractiveInterface\n",
    "          surface:\n",
    "            type: Sphere\n",
    "            R: {R1}\n",
    "          obscuration:\n",
    "            type: ClearCircle\n",
    "            radius: {Do2}\n",
    "          inMedium: 1.000277\n",
    "          outMedium: &silica\n",
    "            type: SellmeierMedium\n",
    "            B1: 0.6961663\n",
    "            B2: 0.4079426\n",
    "            B3: 0.8974794\n",
    "            C1: 0.00467914825849\n",
    "            C2: 0.013512063073959999\n",
    "            C3: 97.93400253792099\n",
    "          name: L1\n",
    "        -\n",
    "          type: RefractiveInterface\n",
    "          surface:\n",
    "            type: Sphere\n",
    "            R: {R2}\n",
    "          obscuration:\n",
    "            type: ClearCircle\n",
    "            radius: {Do2}\n",
    "          coordSys:\n",
    "            z: {d}\n",
    "          inMedium: *silica\n",
    "          outMedium: 1.000277\n",
    "          name: L2\n",
    "        -\n",
    "          type: Detector\n",
    "          surface:\n",
    "            type: Plane\n",
    "          obscuration:\n",
    "            type: ClearCircle\n",
    "            radius: {Do2}\n",
    "          coordSys:\n",
    "            z: {f}\n",
    "          inMedium: 1.000277          \n",
    "          name: D\n",
    "    \"\"\".format(**dict(R1=R1, R2=R2, Do2=D/2, d=d, f=f+focus))\n",
    "    config = yaml.safe_load(systemStr)\n",
    "    system = batoid.parse.parse_optic(config['opticalSystem'])\n",
    "    return system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T05:14:02.405398Z",
     "start_time": "2019-09-16T05:14:02.250439Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d00fef05d25544f3a382620d7e01ec59",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=1.1, continuous_update=False, description='theta_x', max=5.0, min=-5.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@interact(theta_x=widgets.FloatSlider(min=-5,max=5,step=1,value=1.1, continuous_update=False),\n",
    "          theta_y=widgets.FloatSlider(min=-5,max=5,step=1,value=2.2, continuous_update=False),\n",
    "          wavelength=widgets.FloatSlider(min=0.4,max=0.7,step=0.05,value=0.5, continuous_update=False),\n",
    "          focus=widgets.FloatSlider(min=-20, max=20, step=1,value=0.0, continuous_update=False))\n",
    "def spot(theta_x, theta_y, wavelength, focus):\n",
    "    \"\"\"Display a spot diagram for a Newtonian telescope.\n",
    "\n",
    "    @param theta_x  Field angle in degrees\n",
    "    @param theta_y  Field angle in degrees\n",
    "    @param wavelength in microns\n",
    "    @param focus    Defocus distance in mm\n",
    "    \"\"\"\n",
    "    system = makeSystem(focus*1e-3)\n",
    "    dirCos = batoid.utils.fieldToDirCos(np.deg2rad(theta_x), np.deg2rad(theta_y))\n",
    "    # Flip the z-cosine since we constructed our telescope to point down\n",
    "    # instead of up, and hence need the initial rays to point up instead\n",
    "    # of down.\n",
    "    dirCos = dirCos[0:2]+(-dirCos[2],)\n",
    "    rays = batoid.RayVector.asPolar(\n",
    "        backDist=1.0, nrad=16, naz=64, wavelength=wavelength*1e-6,\n",
    "        dirCos=dirCos,\n",
    "        outer=D/2*0.999, medium=batoid.ConstMedium(1.000277)\n",
    "    )\n",
    "    \n",
    "    system.trace(rays)\n",
    "    w = ~rays.vignetted\n",
    "    fig = plt.figure()\n",
    "    ax = plt.axes()\n",
    "    ax.scatter(1e3*(rays.x[w]-np.mean(rays.x[w])), 1e3*(rays.y[w]-np.mean(rays.y[w])), s=1)\n",
    "    ax.set_xlim(-1, 1)\n",
    "    ax.set_ylim(-1, 1)\n",
    "    ax.set_xlabel(\"mm\")\n",
    "    ax.set_ylabel(\"mm\")\n",
    "    ax.set_title(r\"$\\theta_x = {:4.2f}\\,,\\theta_y = {:4.2f}\\,,\\lambda = {:4g}\\,, f={:4.2f}$\".format(theta_x, theta_y, wavelength, focus))\n",
    "    ax.set_aspect('equal')    \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipyvolume as ipv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f248346044374bdbbbac65dbe8623855",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Figure(camera=PerspectiveCamera(fov=45.0, position=(0.0, 0.0, 2.0), projectionMatrix=(1.0, 0.0,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ipv.figure(width=1000)\n",
    "system = makeSystem(0)\n",
    "theta_x = theta_y = 0\n",
    "dirCos = batoid.utils.fieldToDirCos(np.deg2rad(theta_x), np.deg2rad(theta_y))\n",
    "dirCos = dirCos[0:2]+(-dirCos[2],)\n",
    "rays = batoid.RayVector.asPolar(\n",
    "    backDist=1.0, nrad=2, naz=12, wavelength=500e-9,\n",
    "    dirCos=dirCos,\n",
    "    outer=D/2*0.99, medium=batoid.ConstMedium(1.000277)\n",
    ")\n",
    "tf = system.traceFull(rays)\n",
    "system.draw3d(ipv)\n",
    "batoid.drawTrace3d(ipv, tf, color='blue')\n",
    "ipv.xlim(-0.5, 0.5)\n",
    "ipv.ylim(-0.5, 0.5)\n",
    "ipv.zlim(0.0, 1.0)\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.26676440384479083"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# should need a phase delay of\n",
    "# (n1-n0) / (2 sqrt(3) R)\n",
    "(n1-n0) / (2 * np.sqrt(3) * R1) # meters of Z4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def makePhaseScreen(amplitude):\n",
    "    systemStr = \"\"\"\n",
    "    opticalSystem:\n",
    "      type: CompoundOptic\n",
    "      inMedium: 1.000277\n",
    "      items:\n",
    "        -\n",
    "          type: OPDScreen\n",
    "          surface:\n",
    "            type: Plane\n",
    "          screen:\n",
    "            type: Zernike\n",
    "            coef: [0.0, 0.0, 0.0, 0.0, {amplitude}]\n",
    "            # R_outer: 0.05\n",
    "          obscuration:\n",
    "            type: ClearCircle\n",
    "            radius: 0.05\n",
    "          inMedium: 1.000277\n",
    "          outMedium: 1.000277\n",
    "          name: screen\n",
    "        -\n",
    "          type: Detector\n",
    "          surface:\n",
    "            type: Plane\n",
    "          obscuration:\n",
    "            type: ClearCircle\n",
    "            radius: 0.2\n",
    "          coordSys:\n",
    "            z: 0.5\n",
    "          inMedium: 1.000277          \n",
    "          name: D\n",
    "    \"\"\".format(amplitude=amplitude)\n",
    "    config = yaml.safe_load(systemStr)\n",
    "    system = batoid.parse.parse_optic(config['opticalSystem'])\n",
    "    return system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6687e7dd5c42404ca756759659dbe756",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Figure(camera=PerspectiveCamera(fov=45.0, position=(0.0, 0.0, 2.0), projectionMatrix=(1.0, 0.0,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 12808 characters]\n",
      "text/plain": [
       "<Figure size 216x216 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ipv.figure(width=1000)\n",
    "# system = makePhaseScreen(-0.7e-3)\n",
    "system = makePhaseScreen(-0.26676440384479083)\n",
    "theta_x = theta_y = 0\n",
    "dirCos = batoid.utils.fieldToDirCos(np.deg2rad(theta_x), np.deg2rad(theta_y))\n",
    "dirCos = dirCos[0:2]+(-dirCos[2],)\n",
    "rays = batoid.RayVector.asPolar(\n",
    "    backDist=1.0, nrad=4, naz=12, wavelength=500e-9,\n",
    "    dirCos=dirCos,\n",
    "    outer=D/2*0.99, medium=batoid.ConstMedium(1.000277)\n",
    ")\n",
    "tf = system.traceFull(rays)\n",
    "system.draw3d(ipv)\n",
    "batoid.drawTrace3d(ipv, tf, color='blue')\n",
    "ipv.xlim(-0.5, 0.5)\n",
    "ipv.ylim(-0.5, 0.5)\n",
    "ipv.zlim(0.0, 1.0)\n",
    "ipv.show()\n",
    "\n",
    "fig = plt.figure(figsize=(3, 3))\n",
    "ax = fig.add_subplot()\n",
    "s = ax.scatter(tf['screen']['out'].x, tf['screen']['out'].y, c=tf['screen']['out'].t-tf['screen']['in'].t)\n",
    "plt.colorbar(s)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 70244 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = np.linspace(-0.05, 0.05, 20)\n",
    "x, y = np.meshgrid(x, x)\n",
    "r = np.hypot(x, y)\n",
    "w = r < 0.05\n",
    "x, y = x[w], y[w]\n",
    "c = system['.screen'].screen.sag(x, y)\n",
    "plt.scatter(x, y, c=c)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "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": {
    "801e0613f9f44fc5b66b8534af8dd6e4": {
     "views": [
      {
       "cell_index": 3
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
