{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.310463Z",
     "start_time": "2019-09-16T16:08:20.027641Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RayVectors\n",
    "The most fundamental object in batoid is the `RayVector`, which you can roughly think of as and array of photons.  RayVectors are defined primarily by the position vectors `r` and velocity vectors `v` of their constituent photons.  Positions in batoid are always specified in meters, and velocities are always in units of the speed of light in vacuum.  This makes the magnitude of the velocity vector equal to $1/n$, where $n$ is the refractive index of the medium in which each ray is currently propagating.  If you're just planning on geometrically tracing rays through reflective optics, then positions and velocities are the only required RayVector parameters.  The main `RayVector` constructor requires `x`, `y`, and `z` components of position and `vx`, `vy`, `vz` components of velocity.\n",
    "\n",
    "The next most important `RayVector` attribute after position and velocity is `wavelength`, which is (almost always) required for tracing through refractive optics.  (The exception is if you define a refractive medium that doesn't depend on wavelength.  Real media do have wavelength-dependent refractive indices of course.)  Wavelengths are always specified in meters, and always in vacuum -- even when the rays being created aren't currently in vacuum.  If you don't specify a value for `wavelength` explicitly when constructing a `RayVector`, a (somewhat nonsensical) value of 0.0 meters will be used.\n",
    "\n",
    "Next in the list of `RayVector` attributes is time `t`.  Batoid doesn't track time directly, but rather tracks the product of time and the speed of light in vacuum.  That product has dimensions of length, so time in batoid is measured in meters.  Note that this convention is complementary to the convention of tracking velocity in units of the speed of light in vacuum -- we still have the familiar relation $\\Delta \\vec{r} = \\vec{v} \\Delta t$.  The default value of time `t` for a `RayVector` is 0.0 meters.\n",
    "\n",
    "The final fundamental attributes of a `RayVector` are a `flux` for each ray in arbitrary units and status indicators for whether rays in `RayVector` have become `vignetted` during tracing, or if batoid `failed` to successfully trace some rays in a `RayVector` for some reason.  These and a number of additional derived attributes will be discussed more below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.326107Z",
     "start_time": "2019-09-16T16:08:20.311898Z"
    }
   },
   "outputs": [],
   "source": [
    "# Create a (length-1) RayVector from position and velocity 3-tuples (or 3-elements lists, arrays, ...)\n",
    "x, y, z = 0.1, 0.2, 0.3  # meters\n",
    "vx, vy, vz = 0.1, 0.2, np.sqrt(1 - 0.1**2 - 0.2**2)  # in units of c\n",
    "r = (x, y, z)\n",
    "v = (vx, vy, vz)\n",
    "rv = batoid.RayVector(x, y, z, vx, vy, vz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.341380Z",
     "start_time": "2019-09-16T16:08:20.327887Z"
    }
   },
   "outputs": [],
   "source": [
    "# Printing the ray reveals its attributes.\n",
    "print(rv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.358551Z",
     "start_time": "2019-09-16T16:08:20.343082Z"
    }
   },
   "outputs": [],
   "source": [
    "# You can access the position and velocity attributes either individually (x, y, z) or as\n",
    "# 3-vectors (r)\n",
    "print(rv.x)\n",
    "print(rv.vz)\n",
    "\n",
    "print(rv.r)\n",
    "print(rv.v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.373456Z",
     "start_time": "2019-09-16T16:08:20.359940Z"
    }
   },
   "outputs": [],
   "source": [
    "# Rays have a number of optional fields too\n",
    "print(rv.t)  # The time in meters (see explanation above)\n",
    "print(rv.wavelength)  # Vacuum wavelength in meters.\n",
    "print(rv.flux)  # Flux in arbitrary units.\n",
    "print(rv.vignetted)  # Whether or not the ray has been vignetted\n",
    "print(rv.failed)  # Whether or not the ray is failed.  \n",
    "                  # Usually this means an intersection between the Ray and a Surface could not be found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.388153Z",
     "start_time": "2019-09-16T16:08:20.374817Z"
    }
   },
   "outputs": [],
   "source": [
    "# You can create rays with specific values for these optional fields too \n",
    "rv = batoid.RayVector(*r, *v, t=0.3, wavelength=400e-9, flux=2.3, vignetted=True)\n",
    "print(rv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Derived Ray attributes\n",
    "Two additional attributes derivable from the fundamental attributes detailed above are the wavevector $\\vec{k}$ and the temporal frequency $\\omega$.  The wavevector is defined by\n",
    "\n",
    "$\\vec{k} = \\frac{2 \\pi \\vec{v}}{\\lambda |\\vec{v}|^2}$\n",
    "\n",
    "(with $\\lambda$ being the vacuum wavelength) and has units of radians per meter.\n",
    "\n",
    "The normal physics definition of angular temporal frequency is $\\omega = \\frac{2 \\pi c}{\\lambda}$, but in batoid, we divide out the speed of light in vacuum to get\n",
    "\n",
    "$\\omega = \\frac{2 \\pi}{\\lambda}$\n",
    "\n",
    "The units are again radians per meter (but in contrast to $\\vec{k}$ this time in a scalar variable).  This choice for the frequency means that $\\omega t$ is a phase angle in batoid (just as $\\vec{k} \\cdot \\vec{r}$ is)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.403797Z",
     "start_time": "2019-09-16T16:08:20.389306Z"
    }
   },
   "outputs": [],
   "source": [
    "print(rv.omega)  # angular frequency\n",
    "print(rv.k)  # wave vector\n",
    "print(rv.kx)  # individual components of the wave vector are also accessible directly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RayVector methods\n",
    "The first set of methods available to `RayVector` is for propagation of photons through time.  The `positionAtTime` method accepts a time (still in meters) and returns the position of the `RayVector` photons at that past, present, or future time.  The `propagate` method similarly propagates the rays to a given time.  Note that the return value of `propagate` is a reference to the original `RayVector` (not a copy).\n",
    "\n",
    "For the second set of `RayVector` methods, we change our conception of what a ray is: from a propagating infinitesimal point to a propagating plane wave.  I.e., we imagine that the velocity `v` attribute is that of a plane wave, and that the position `r` and time `t` together indicate a point in space-time where the wave amplitude is a maximum.  Note that these conditions will remain true as we propagate the wave in time.\n",
    "\n",
    "The `phase` method accepts a position and time and returns the phase $\\phi$ of these plane waves in radians.  The `amplitude` method returns the complex amplitudes $\\exp(i \\phi)$ ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.447139Z",
     "start_time": "2019-09-16T16:08:20.405582Z"
    }
   },
   "outputs": [],
   "source": [
    "rv = batoid.RayVector(0,0,0, 0,0,1)\n",
    "print(rv.r)\n",
    "print(rv.positionAtTime(1.0))  # propagated from (0,0,0) to (0,0,1)\n",
    "rv2 = rv.propagate(1.0)  # propagate ray forward to t=1\n",
    "print(rv.r)  # propagate works in-place\n",
    "print(rv2.r)  # the return value is a reference to the original\n",
    "print(rv == rv2)  # so they're equal\n",
    "print(rv is rv2)  # in fact, they're the same object\n",
    "# If you want to preserver the original ray when performing a propagation, first make a copy.\n",
    "rv3 = rv.copy().propagate(2.0)\n",
    "print(rv.r)\n",
    "print(rv3.r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.608757Z",
     "start_time": "2019-09-16T16:08:20.449119Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot one period of the plane wave with wavelength 800e-9 m\n",
    "wavelength = 800e-9\n",
    "rv = batoid.RayVector(0,0,0, 0,0,1, t=0, wavelength=wavelength)\n",
    "ts = np.linspace(0, wavelength)\n",
    "plt.plot(ts, [np.cos(rv.phase([0,0,0], t)) for t in ts])\n",
    "plt.axhline(0, c='k')\n",
    "plt.xlabel(\"c t (meters)\")\n",
    "plt.ylabel(\"amplitude\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:20.768786Z",
     "start_time": "2019-09-16T16:08:20.610269Z"
    }
   },
   "outputs": [],
   "source": [
    "# We can also look at a fixed time over a small region of space\n",
    "# Since the wave is traveling in the +z direction, its amplitude\n",
    "# varies with z, but not with x or y\n",
    "wavelength = 800e-9\n",
    "rv = batoid.RayVector(0,0,0, 0,0,1, t=0, wavelength=wavelength)\n",
    "xs = ys = zs = np.linspace(0, wavelength)\n",
    "plt.plot(xs, [np.cos(rv.phase([x,0,0], 0)) for x in xs], label='x')\n",
    "plt.plot(ys, [np.cos(rv.phase([0,y,0], 0))+0.1 for y in ys], label='y+0.1')\n",
    "plt.plot(zs, [np.cos(rv.phase([0,0,z], 0)) for z in zs], label='z')\n",
    "plt.xlabel(\"position (meters)\")\n",
    "plt.ylabel(\"amplitude\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RayVector factory functions\n",
    "It can be slow to construct a large grid of positions in python just to then assemble them into `RayVector`s.  To help speed up this construction, batoid includes a few `RayVector` factory functions that can be used to rapidly create commonly desired grids of rays.  These include `RayVector.asGrid`, `RayVector.asPolar`, and `RayVector.asSpokes`.  In each case, the goal is to sample either a plane wave or spherical wave incident upon some optical system by creating a grid of either parallel or co-radial co-phased `Ray`s.  For `RayVector.asGrid`, the `Ray`s are assembled onto a parallelogram (commonly just a square), for `RayVector.asPolar` the pattern is more circularly symmetric.  The `RayVector.asSpokes` function is a bit special, in that it does not uniformly sample the region of interest.  It does have a use, however, for integrating functions on a circle or annulus using Gaussian quadrature.  There are a number of different options for each of these factory functions to support different use cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:32:51.013189Z",
     "start_time": "2019-09-16T16:32:50.764667Z"
    }
   },
   "outputs": [],
   "source": [
    "# An example of a square grid of rays:\n",
    "rv1 = batoid.RayVector.asGrid(\n",
    "    backDist=10,  # roughly how far back from the origin the rays are created\n",
    "    nx=15,  # number of rays on a side\n",
    "    lx=10.0,  # length of one side of the grid in meters\n",
    "    theta_x=np.deg2rad(0.0),  # Field angle of rays in x-direction\n",
    "    theta_y=np.deg2rad(0.0),  # Field angle of rays in y-direction\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "# Polar RayVector with similar extent\n",
    "rv2 = batoid.RayVector.asPolar(\n",
    "    backDist=10,\n",
    "    nrad=20,  # number of radii\n",
    "    naz=50,  # number of azimuths on outermost ring.\n",
    "    outer=5.0,  # outer radius.  Default inner radius is 0.0\n",
    "    theta_x=np.deg2rad(0.0),\n",
    "    theta_y=np.deg2rad(0.0),\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "# Spokes RayVector with similar extent\n",
    "rv3 = batoid.RayVector.asSpokes(\n",
    "    backDist=10,\n",
    "    rings=10,  # number of rings\n",
    "    spokes=10,  # number of spokes\n",
    "    outer=5.0,  # outer radius\n",
    "    inner=2.5,  # inner radius\n",
    "    theta_x=np.deg2rad(0.0),\n",
    "    theta_y=np.deg2rad(0.0),\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "plt.scatter(rv1.x, rv1.y, c='b', s=25)\n",
    "plt.scatter(rv2.x, rv2.y, c='r', s=25)\n",
    "plt.scatter(rv3.x, rv3.y, c='c', s=25)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.xlim(-6, 6)\n",
    "plt.ylim(-6, 6)\n",
    "plt.gca().set_aspect(\"equal\")\n",
    "plt.show()\n",
    "\n",
    "plt.scatter(rv1.x, rv1.z, c='b', s=25)\n",
    "plt.scatter(rv2.x, rv2.z, c='r', s=10)\n",
    "plt.scatter(rv3.x, rv3.z, c='c', s=5)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"z\")\n",
    "plt.xlim(-6, 6)\n",
    "plt.ylim(-1, 11)\n",
    "plt.gca().set_aspect(\"equal\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:38:22.550725Z",
     "start_time": "2019-09-16T16:38:22.302373Z"
    }
   },
   "outputs": [],
   "source": [
    "# Repeat with a non-zero field angle.\n",
    "\n",
    "# An example of a square grid of rays:\n",
    "rv1 = batoid.RayVector.asGrid(\n",
    "    backDist=10,  # roughly how far back from the origin the rays are created\n",
    "    nx=15,  # number of rays on a side\n",
    "    lx=10.0,  # length of one side of the grid in meters\n",
    "    theta_x=np.deg2rad(5),  # Field angle of rays in x-direction\n",
    "    theta_y=np.deg2rad(10),  # Field angle of rays in y-direction\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "# Polar RayVector with similar extent\n",
    "rv2 = batoid.RayVector.asPolar(\n",
    "    backDist=10,\n",
    "    nrad=20,  # number of radii\n",
    "    naz=50,  # number of azimuths on outermost ring.\n",
    "    outer=5.0,  # outer radius.  Default inner radius is 0.0\n",
    "    theta_x=np.deg2rad(5),\n",
    "    theta_y=np.deg2rad(10),\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "# Spokes RayVector with similar extent\n",
    "rv3 = batoid.RayVector.asSpokes(\n",
    "    backDist=10,\n",
    "    rings=10,  # number of rings\n",
    "    spokes=10,  # number of spokes\n",
    "    outer=5.0,  # outer radius\n",
    "    inner=2.5,  # inner radius\n",
    "    theta_x=np.deg2rad(5),\n",
    "    theta_y=np.deg2rad(10),\n",
    "    wavelength=500e-9\n",
    ")\n",
    "\n",
    "plt.scatter(rv1.x, rv1.y, c='b', s=25)\n",
    "plt.scatter(rv2.x, rv2.y, c='r', s=25)\n",
    "plt.scatter(rv3.x, rv3.y, c='c', s=25)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.xlim(-6, 6)\n",
    "plt.ylim(-6, 6)\n",
    "plt.gca().set_aspect(\"equal\")\n",
    "plt.show()\n",
    "\n",
    "plt.scatter(rv1.x, rv1.z, c='b', s=25)\n",
    "plt.scatter(rv2.x, rv2.z, c='r', s=10)\n",
    "plt.scatter(rv3.x, rv3.z, c='c', s=5)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"z\")\n",
    "plt.xlim(-6, 6)\n",
    "plt.ylim(-1, 11)\n",
    "plt.gca().set_aspect(\"equal\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:08:21.588727Z",
     "start_time": "2019-09-16T16:08:21.291406Z"
    }
   },
   "outputs": [],
   "source": [
    "# If you have ipyvolume installed, you can visualize the ray grids in 3D.\n",
    "has_ipv = True\n",
    "try:\n",
    "    import ipyvolume as ipv\n",
    "except ImportError:\n",
    "    has_ipv = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-16T16:42:18.951978Z",
     "start_time": "2019-09-16T16:42:18.866891Z"
    }
   },
   "outputs": [],
   "source": [
    "if has_ipv:\n",
    "    ipv.figure()\n",
    "    ipv.quiver(rv1.x, rv1.y, rv1.z, rv1.vx, rv1.vy, rv1.vz, size=5, color='red')\n",
    "    ipv.quiver(rv2.x, rv2.y, rv2.z, rv2.vx, rv2.vy, rv2.vz, size=5, color='blue')\n",
    "    ipv.quiver(rv3.x, rv3.y, rv3.z, rv3.vx, rv3.vy, rv3.vz, size=5, color='cyan')\n",
    "    ipv.xyzlim(10)\n",
    "    ipv.view(0,90)\n",
    "    ipv.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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
