{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:01.646086Z",
     "start_time": "2019-09-15T19:50:01.146856Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\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\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:02.189469Z",
     "start_time": "2019-09-15T19:50:02.154006Z"
    }
   },
   "outputs": [],
   "source": [
    "telescope = batoid.Optic.fromYaml(\"LSST_r.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:04.080809Z",
     "start_time": "2019-09-15T19:50:04.074120Z"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "rays = batoid.RayVector.asPolar(\n",
    "    optic=telescope, wavelength=625e-9,\n",
    "    theta_x=np.deg2rad(1.5), theta_y=0.0,\n",
    "    nrandom=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:05.224812Z",
     "start_time": "2019-09-15T19:50:05.217338Z"
    }
   },
   "outputs": [],
   "source": [
    "outRays = telescope.trace(rays.copy())\n",
    "np.sum(~outRays.vignetted), len(outRays)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:05.965610Z",
     "start_time": "2019-09-15T19:50:05.945701Z"
    }
   },
   "outputs": [],
   "source": [
    "def randomRaySumXYZ(outRays, nray, xsize=5*10e-6, ysize=5*10e-6, zsize=5*10e-6, nx=100, plotk=False, alpha=0.1):\n",
    "    w = np.nonzero(~outRays.vignetted)[0]\n",
    "    xmean = np.mean(outRays.x[w])\n",
    "    ymean = np.mean(outRays.y[w])\n",
    "    zmean = np.mean(outRays.z[w])\n",
    "    tmean = np.mean(outRays.t[w])\n",
    "\n",
    "    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(10, 10))\n",
    "    for ax in axes.ravel():\n",
    "        ax.tick_params(axis='both', which='major', labelsize=18)\n",
    "        ax.tick_params(axis='both', which='minor', labelsize=18)\n",
    "    \n",
    "    # XY plots first\n",
    "    xrange = xmean + xsize*np.array([-0.5, 0.5])\n",
    "    yrange = ymean + ysize*np.array([-0.5, 0.5])\n",
    "\n",
    "    xs = np.linspace(xrange[0], xrange[1], nx)\n",
    "    ys = np.linspace(yrange[0], yrange[1], nx)\n",
    "    xs, ys = np.meshgrid(xs, ys)\n",
    "    zs = np.ones_like(xs)*zmean\n",
    "    \n",
    "    points = np.concatenate([aux[..., None] for aux in (xs, ys, zs)], axis=-1)\n",
    "    rng = np.random.default_rng(123)\n",
    "    indices = rng.choice(w, nray)\n",
    "    tmpRays = outRays[indices]\n",
    "    \n",
    "    amplitudes = np.empty_like(xs, dtype=np.complex128)\n",
    "    for idx in np.ndindex(amplitudes.shape):\n",
    "        amplitudes[idx] = tmpRays.sumAmplitude(points[idx], tmean)\n",
    "    \n",
    "    extent = np.array([-xsize, xsize, -ysize, ysize])/10e-6\n",
    "    i0 = axes[0,0].imshow(np.real(amplitudes), extent=extent)\n",
    "    axes[0,0].set_title(r\"$Re[E]$\", fontsize=18)\n",
    "    axes[0,0].set_xlabel(\"X\", fontsize=18)\n",
    "    axes[0,0].set_ylabel(\"Y\", fontsize=18)\n",
    "\n",
    "    i1 = axes[0,1].imshow(np.abs(amplitudes)**2, extent=extent)\n",
    "    axes[0,1].set_title(r\"$|E|^2$\", fontsize=18)\n",
    "    axes[0,1].set_xlabel(\"X\", fontsize=18)\n",
    "    axes[0,1].set_ylabel(\"Y\", fontsize=18)\n",
    "\n",
    "    if plotk:\n",
    "        for ray in tmpRays:\n",
    "            axes[0,0].plot(\n",
    "                (np.array([0, ray.vx])*10e-6 + ray.x - xmean)/10e-6,\n",
    "                (np.array([0, ray.vy])*10e-6 + ray.y - ymean)/10e-6,\n",
    "                c='r', alpha=alpha\n",
    "            )            \n",
    "\n",
    "    # Now XZ plots    \n",
    "    xrange = xmean+xsize*np.array([-0.5, 0.5])\n",
    "    zrange = zmean+zsize*np.array([-0.5, 0.5])\n",
    "\n",
    "    xs = np.linspace(xrange[0], xrange[1], nx)\n",
    "    zs = np.linspace(zrange[0], zrange[1], nx)\n",
    "    xs, zs = np.meshgrid(xs, zs)\n",
    "    ys = np.ones_like(xs)*ymean\n",
    "    \n",
    "    points = np.concatenate([aux[..., None] for aux in (xs, ys, zs)], axis=-1)\n",
    "\n",
    "    amplitudes = np.empty_like(xs, dtype=np.complex128)\n",
    "    for idx in np.ndindex(amplitudes.shape):\n",
    "        amplitudes[idx] = tmpRays.sumAmplitude(points[idx], tmean)\n",
    "    \n",
    "    extent = np.array([-xsize, xsize, -zsize, zsize])/10e-6\n",
    "    \n",
    "    i0 = axes[1,0].imshow(np.real(amplitudes), extent=extent)\n",
    "    axes[1,0].set_title(r\"$Re[E]$\", fontsize=18)\n",
    "    axes[1,0].axhline(0, c='w', lw=4)\n",
    "    axes[1,0].set_xlabel(\"X\", fontsize=18)\n",
    "    axes[1,0].set_ylabel(\"Z\", fontsize=18)\n",
    "    \n",
    "    i1 = axes[1,1].imshow(np.abs(amplitudes)**2, extent=extent)\n",
    "    axes[1,1].set_title(r\"$|E|^2$\", fontsize=18)\n",
    "    axes[1,1].axhline(0, c='w', lw=4)\n",
    "    axes[1,1].set_xlabel(\"X\", fontsize=18)\n",
    "    axes[1,1].set_ylabel(\"Z\", fontsize=18)\n",
    "\n",
    "    if plotk:\n",
    "        for ray in tmpRays:\n",
    "            axes[1,0].plot(\n",
    "                (np.array([0, ray.vx])*10e-6 + ray.x - xmean)/10e-6,\n",
    "                (np.array([0, ray.vz])*10e-6 + ray.z - zmean)/10e-6,\n",
    "                c='r', alpha=alpha\n",
    "            )            \n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-15T19:50:10.599671Z",
     "start_time": "2019-09-15T19:50:08.329515Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "randomRaySumXYZ(outRays, len(outRays), xsize=2*10e-6, ysize=2*10e-6, zsize=2*10e-6, plotk=True, alpha=0.01)"
   ]
  },
  {
   "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
}