{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:44:06.807887Z",
     "start_time": "2019-09-22T21:44:06.490976Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:44:06.937321Z",
     "start_time": "2019-09-22T21:44:06.823159Z"
    }
   },
   "outputs": [],
   "source": [
    "telescope = batoid.Optic.fromYaml(\"HSC.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:45:37.057223Z",
     "start_time": "2019-09-22T21:45:37.052363Z"
    }
   },
   "outputs": [],
   "source": [
    "def pupil(thx, thy, nside=512):\n",
    "    rays = batoid.RayVector.asGrid(\n",
    "        optic=telescope, wavelength=600e-9,\n",
    "        theta_x=thx, theta_y=thy,\n",
    "        nx=nside, ny=nside\n",
    "    )\n",
    "    rays2 = rays.copy()\n",
    "    telescope.stopSurface.interact(rays2)\n",
    "    telescope.trace(rays)\n",
    "    w = ~rays.vignetted\n",
    "    return rays2.x[w], rays2.y[w]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:45:40.444732Z",
     "start_time": "2019-09-22T21:45:38.811487Z"
    }
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 12))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(*pupil(np.deg2rad(0.75),0), s=0.1)\n",
    "ax.set_xlim(-4.2, 4.2)\n",
    "ax.set_ylim(-4.2, 4.2)\n",
    "ax.set_aspect(1)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:45:43.983992Z",
     "start_time": "2019-09-22T21:45:43.980595Z"
    }
   },
   "outputs": [],
   "source": [
    "def spanRange(x, nside=512):\n",
    "    xmin, xmax = np.min(x), np.max(x)\n",
    "    xspan = xmax - xmin\n",
    "    xmin = xmin - 0.8*xspan\n",
    "    xmax = xmax + 0.8*xspan\n",
    "    return np.linspace(xmin, xmax, nside)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:12.340579Z",
     "start_time": "2019-09-22T21:54:12.330951Z"
    }
   },
   "outputs": [],
   "source": [
    "def pinhole(thx, thy, nside=256):\n",
    "    # reset skips\n",
    "    for item in telescope.itemDict:\n",
    "        telescope[item].skip = False\n",
    "\n",
    "    # First, need to determine where on the filter to constrain rays.  We'll use the average position of the \n",
    "    # pupil beam that would have intersected the filter.\n",
    "    rays = batoid.RayVector.asGrid(\n",
    "        optic=telescope, lx=10, theta_x=thx, theta_y=thy,\n",
    "        nx=nside, wavelength=600e-9\n",
    "    )\n",
    "    tf = telescope.traceFull(rays)\n",
    "    surface = tf['F_entrance']\n",
    "    w = ~surface['out'].vignetted\n",
    "    rs = surface['out'][w]\n",
    "    xmean, ymean = np.mean(rs.x), np.mean(rs.y)\n",
    "    # Now we need to generate a bunch of rays that all pass through the above part of the filter, but over \n",
    "    # a range of angles.\n",
    "    # What is the range of angles for the pupil beam?  \n",
    "    vx = spanRange(rs.vx, nside=nside)\n",
    "    vy = spanRange(rs.vy, nside=nside)\n",
    "    vx, vy = np.meshgrid(vx, vy)\n",
    "    vz = np.sqrt(1-vx*vx+vy*vy)\n",
    "    # Now need to make a RayVector with appropriate x,y,vx,vy,...\n",
    "#     rv = batoid.RayVector([\n",
    "#         batoid.Ray([xmean, ymean, 0], [vx_, vy_, vz_], 0, 600e-9)\n",
    "#         for vx_, vy_, vz_ in zip(vx.ravel(), vy.ravel(), vz.ravel())])\n",
    "    rv = batoid.RayVector(\n",
    "        xmean*np.ones(nside*nside, dtype=float),\n",
    "        ymean*np.ones(nside*nside, dtype=float),\n",
    "        np.zeros(nside*nside, dtype=float),\n",
    "        vx.ravel(), vy.ravel(), vz.ravel(),\n",
    "        np.zeros(nside*nside, dtype=float),\n",
    "        600e-9*np.ones(nside*nside, dtype=float),\n",
    "        coordSys = surface['out'].coordSys\n",
    "    )\n",
    "    # trace forward from filter.  So temporarily skip everything before the filter.\n",
    "    before_items = ['SubaruHSC.POPT2', \n",
    "                    'SubaruHSC.FEU',\n",
    "                    'SubaruHSC.TopRing',\n",
    "                    'SubaruHSC.BottomRing',\n",
    "                    'SubaruHSC.TertiarySpiderFirstPass',\n",
    "                    'SubaruHSC.PM',\n",
    "                    'SubaruHSC.TertiarySpiderSecondPass',\n",
    "                    'SubaruHSC.HSC.WFC.G1',\n",
    "                    'SubaruHSC.HSC.WFC.G2',\n",
    "                    'SubaruHSC.HSC.WFC.ADC',\n",
    "                    'SubaruHSC.HSC.WFC.G3',\n",
    "                    'SubaruHSC.HSC.WFC.G4',\n",
    "                    'SubaruHSC.HSC.WFC.G5',\n",
    "                   ]\n",
    "    for item in before_items:\n",
    "        telescope[item].skip = True\n",
    "    forward_rays = telescope.trace(rv.copy())\n",
    "    # reset skips\n",
    "    for item in telescope.itemDict:\n",
    "        telescope[item].skip = False\n",
    "    # Now skip everything that happens *after* and including the filter\n",
    "    after_items = ['SubaruHSC.HSC.CAM.F',\n",
    "                   'SubaruHSC.HSC.CAM.W',\n",
    "                   'SubaruHSC.HSC.CAM.D',\n",
    "                  ]\n",
    "    for item in after_items:\n",
    "        telescope[item].skip = True\n",
    "    rv = batoid.RayVector(\n",
    "        rv.x, rv.y, rv.z,\n",
    "        -rv.vx, -rv.vy, -rv.vz,\n",
    "        rv.t, rv.wavelength,\n",
    "        coordSys = rv.coordSys\n",
    "    )\n",
    "    reverse_rays = telescope.trace(rv.copy(), reverse=True)\n",
    "\n",
    "    # reset skips\n",
    "    for item in telescope.itemDict:\n",
    "        telescope[item].skip = False\n",
    "        \n",
    "    w = ~reverse_rays.vignetted\n",
    "    return forward_rays.x[w], forward_rays.y[w]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:13.774290Z",
     "start_time": "2019-09-22T21:54:13.001905Z"
    }
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 12))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(*pinhole(0,0), s=1)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:21.391367Z",
     "start_time": "2019-09-22T21:54:21.385603Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot(thx, thy):\n",
    "    fig = plt.figure(figsize=(6, 6))\n",
    "    ax = fig.add_subplot(111)\n",
    "\n",
    "    pux, puy = pupil(thx, thy)\n",
    "    xspan = np.max(pux) - np.min(pux)\n",
    "    yspan = np.max(puy) - np.min(puy)\n",
    "    span = max(xspan, yspan)\n",
    "    pux = (pux - np.mean(pux))/span\n",
    "    puy = (puy - np.mean(puy))/span\n",
    "    \n",
    "    phx, phy = pinhole(thx, thy)\n",
    "    xspan = np.max(phx) - np.min(phx)\n",
    "    yspan = np.max(phy) - np.min(phy)\n",
    "    span = max(xspan, yspan)\n",
    "\n",
    "    phx = -(phx - np.mean(phx))/span\n",
    "    phy = -(phy - np.mean(phy))/span\n",
    "\n",
    "    ax.scatter(pux, puy, s=2, alpha=0.1, c='r', label='pupil')\n",
    "    ax.scatter(phx, phy, s=2, alpha=0.2, c='b', label='pinhole')\n",
    "    ax.legend()\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:26.486667Z",
     "start_time": "2019-09-22T21:54:23.689643Z"
    }
   },
   "outputs": [],
   "source": [
    "plot(0, np.deg2rad(0.75))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:29.710612Z",
     "start_time": "2019-09-22T21:54:29.705415Z"
    }
   },
   "outputs": [],
   "source": [
    "def both(thx, thy):\n",
    "    pux, puy = pupil(thx, thy)\n",
    "    xspan = np.max(pux) - np.min(pux)\n",
    "    yspan = np.max(puy) - np.min(puy)\n",
    "    span = max(xspan, yspan)\n",
    "    pux = (pux - np.mean(pux))/span\n",
    "    puy = (puy - np.mean(puy))/span\n",
    "\n",
    "    phx, phy = pinhole(thx, thy)\n",
    "    xspan = np.max(phx) - np.min(phx)\n",
    "    yspan = np.max(phy) - np.min(phy)\n",
    "    span = max(xspan, yspan)\n",
    "    phx = -(phx - np.mean(phx))/span\n",
    "    phy = -(phy - np.mean(phy))/span\n",
    "\n",
    "    return pux, puy, phx, phy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-22T21:54:40.232328Z",
     "start_time": "2019-09-22T21:54:31.673884Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8,8))\n",
    "for ax, thy in zip(axes.ravel(), [0.0, 0.25, 0.5, 0.75]):\n",
    "    pux, puy, phx, phy = both(0.0, np.deg2rad(thy))\n",
    "\n",
    "    ax.scatter(pux, puy, s=2, alpha=0.1, c='r', label='pupil')\n",
    "    ax.scatter(phx, phy, s=2, alpha=0.2, c='b', label='pinhole')\n",
    "    ax.set_title(r\"$\\theta_y$ = {:5.2f}\".format(thy))\n",
    "    ax.legend(loc=\"upper right\")\n",
    "fig.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.7.8-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}