{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "884c47b0-1a10-4c48-989f-ea7996ed94e4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ipyvolume as ipv\n",
    "import batoid\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01cfd4df-5060-4bff-9c98-445a184493af",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "az = -120  # degrees\n",
    "alt = 42  # degrees\n",
    "rotTelPos = 0  # degrees\n",
    "\n",
    "# By default, batoid will load the LSST optics in the Optical Coordinate System described here:\n",
    "# https://sitcomtn-003.lsst.io/#the-optical-coordinate-system.  I.e., the origin is at the \n",
    "# theoretical vertex of M1, +z points towards the sky, +x is parallel to the elevation bearing.\n",
    "\n",
    "# For LSST + CBP, it's more convenient to make the global coordinate system fixed to the ground,\n",
    "# something like the Terrestrial Coordinate System described in \n",
    "# https://docushare.lsst.org/docushare/dsweb/Get/Document-18499.\n",
    "\n",
    "# We shift around the telescope to realize this change.  We also \"point\" the telescope according \n",
    "# to the above azimuth and altitude.\n",
    "\n",
    "elevBearingHeight = 5.425  # Height of elevation bearing above azimuth ring\n",
    "\n",
    "lsst = batoid.Optic.fromYaml(\"LSST_r.yaml\")\n",
    "lsst = lsst.withGlobalShift([0, 0, 3.53])  # Height of M1 vertex above azimuth ring\n",
    "lsst = lsst.withLocallyRotatedOptic(\"LSSTCamera\", batoid.RotZ(np.deg2rad(rotTelPos)))  # Camera rotator\n",
    "lsst = lsst.withLocalRotation(batoid.RotZ(np.deg2rad(90-az)))  # Point in Az\n",
    "lsst = lsst.withLocalRotation(batoid.RotX(np.deg2rad(90-alt)), rotCenter=[0, 0, elevBearingHeight], coordSys=batoid.globalCoordSys)  # Point in Alt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3844ebf-0a83-49b3-9b3a-29884e559c23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Collimated Beam Projector and place it in the right position in the dome.\n",
    "# The rotations below orient the CBP with respect to its default description such that\n",
    "# it's aiming roughly at the LSST M1 mirror.  These were found through trial-and-error.\n",
    "# The GlobalShift though should be pretty accurate.\n",
    "cbp = batoid.Optic.fromYaml(\"CBP.yaml\")\n",
    "cbp = cbp.withLocalRotation(batoid.RotX(np.deg2rad(180)))\n",
    "cbp = cbp.withLocalRotation(batoid.RotZ(np.deg2rad(150)))\n",
    "cbp = cbp.withGlobalShift([-np.sin(np.deg2rad(30))*12.4, np.cos(np.deg2rad(30))*12.4, 12.135])\n",
    "# cbp = cbp.withGlobalRotation(batoid.RotZ(np.deg2rad(daz)))  # optional 10 cm dome error\n",
    "cbp = cbp.withLocalRotation(batoid.RotX(np.deg2rad(48)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a0caa47-a5d2-4674-bc88-ca93141d0115",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Launch some rays, starting at the positions of some holes in the CBP focal plane.\n",
    "\n",
    "nx = 20  # number of rays in ray-fan\n",
    "\n",
    "# CBP focal plane hole positions in meters.\n",
    "holes = [\n",
    "    (0.0, 0.0),\n",
    "    (-1.e-2, 0.0),\n",
    "    (1.e-2, 0.0),\n",
    "    (0.0, -1.e-2),\n",
    "    (0.0, 1.e-2)\n",
    "]\n",
    "rvs = []\n",
    "for hole in holes:\n",
    "    # Set ray positions\n",
    "    x = np.full(2*nx, hole[0])\n",
    "    y = np.full(2*nx, hole[1])\n",
    "    z = np.zeros_like(x)\n",
    "    # Set ray angles.  We just sample uniformly between -0.3, 0.3 radians\n",
    "    thx = np.concatenate([np.linspace(-0.3, 0.3, nx), np.zeros(nx)])\n",
    "    thy = np.concatenate([np.zeros(nx), np.linspace(-0.3, 0.3, nx)])\n",
    "    # Turn angles into velocities\n",
    "    vx, vy, vz = batoid.utils.fieldToDirCos(thx, thy)\n",
    "    # wavelength\n",
    "    w = np.full_like(x, 500e-9)\n",
    "    rvs.append(\n",
    "        batoid.RayVector(x, y, z, vx, vy, vz, wavelength=w, coordSys=cbp['Detector'].coordSys)\n",
    "    )\n",
    "rvs = batoid.concatenateRayVectors(rvs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30ecd684-9409-48a4-a887-fda76b00de52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trace!\n",
    "# Start by tracing CBP optics in reverse; from focal plane to the Schmidt plate \"entrance\".\n",
    "tf_cbp = cbp.traceFull(rvs, reverse=True)\n",
    "# Then take the results of that and push them through the LSST optics.\n",
    "tf_lsst = lsst.traceFull(tf_cbp['Schmidt_plate_entrance']['out'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83073113-093e-432b-95a9-3e8028b7799d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Visualize\n",
    "fig = ipv.figure(height=600, width=800)\n",
    "ipv.style.set_style_dark()\n",
    "ipv.style.box_off()\n",
    "# ipv.style.axes_off()\n",
    "lsst.draw3d(ipv, color='cyan')\n",
    "cbp.draw3d(ipv, color='red')\n",
    "batoid.drawTrace3d(ipv, tf_cbp, color='yellow')\n",
    "batoid.drawTrace3d(ipv, tf_lsst, color='yellow')\n",
    "\n",
    "ipv.xlim(-10, 10)\n",
    "ipv.ylim(-10, 10)\n",
    "ipv.zlim(0, 20)\n",
    "ipv.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8887dbd5-c08e-428c-8926-ff56cb430259",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84f308c2-38e2-42c4-99cb-5ac041ad3375",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Where did the rays impact the LSST Focal Plane?\n",
    "\n",
    "# Just plot rays for one hole at a time\n",
    "ihole = 0\n",
    "\n",
    "optic = 'Detector'\n",
    "\n",
    "w = ~tf_lsst[optic]['out'][ihole*40:(ihole+1)*40].vignetted\n",
    "x = tf_lsst[optic]['out'][ihole*40:(ihole+1)*40].x\n",
    "y = tf_lsst[optic]['out'][ihole*40:(ihole+1)*40].y\n",
    "\n",
    "mx = np.mean(x)\n",
    "my = np.mean(y)\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(x[w], y[w])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da622e9d-8deb-45f6-8bbd-541c17004083",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show all spots.  Zoom in to see individual rays\n",
    "optic = 'Detector'\n",
    "\n",
    "w = ~tf_lsst[optic]['out'].vignetted\n",
    "x = tf_lsst[optic]['out'].x\n",
    "y = tf_lsst[optic]['out'].y\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(x[w], y[w])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25fec1d3-d0c9-4916-8e1d-53c89ee23113",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
