{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "423b35e8-db0b-4164-825e-f47704b68092",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c1682e3-945c-4174-93fb-9d32c7b910cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def colorbar(mappable, **kwargs):\n",
    "    from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "    import matplotlib.pyplot as plt\n",
    "    last_axes = plt.gca()\n",
    "    ax = mappable.axes\n",
    "    fig = ax.figure\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    cbar = fig.colorbar(mappable, cax=cax, **kwargs)\n",
    "    plt.sca(last_axes)\n",
    "    return cbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ed339c6-9ac7-4f64-adef-d18236beb483",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "telescope = batoid.Optic.fromYaml(\"LSST_g.yaml\")\n",
    "scale = 0.2/10e-6  # arcsec / meter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbcfb76a-2443-42a3-8df1-51cf8f8ae311",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Make refractive interfaces partially reflective\n",
    "reflect = 0.02\n",
    "transmit = 0.98\n",
    "for surface in telescope.itemDict.values():\n",
    "    if isinstance(surface, batoid.RefractiveInterface):\n",
    "        surface.forwardCoating = batoid.SimpleCoating(reflect, transmit)\n",
    "        surface.reverseCoating = batoid.SimpleCoating(reflect, transmit)\n",
    "    if isinstance(surface, batoid.Detector):\n",
    "        surface.forwardCoating = batoid.SimpleCoating(reflect, transmit)\n",
    "# Allow filter entrance surface to be trickier\n",
    "telescope['Filter_entrance'].forwardCoating = batoid.SimpleCoating(0.5, 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e4637c1-e218-48cb-8583-551d2bece5aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "angle = 1.0\n",
    "dirCos = batoid.utils.gnomonicToDirCos(0.0, np.deg2rad(angle))\n",
    "\n",
    "rays = batoid.RayVector.asPolar(\n",
    "    optic=telescope, wavelength=500e-9,\n",
    "    theta_x=0.0, theta_y=np.deg2rad(angle),\n",
    "    # nrad=300, naz=900\n",
    "    nrandom=int(1e5)\n",
    ")\n",
    "\n",
    "rForward, rReverse = telescope.traceSplit(rays, minFlux=1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f38dc8a-49ba-47a7-97ce-d5a594f70ac8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"# input rays          = {}\".format(len(rays)))\n",
    "print(\"# forward output rays = {}\".format(sum(len(rr) for rr in rForward)))\n",
    "print(\"# reverse output rays = {}\".format(sum(len(rr) for rr in rReverse)))\n",
    "print(\"input flux          = {}\".format(np.sum(rays.flux)))\n",
    "forwardFlux = np.sum([np.sum(rr.flux) for rr in rForward])\n",
    "reverseFlux = np.sum([np.sum(rr.flux) for rr in rReverse])\n",
    "print(\"forward output flux = {}\".format(forwardFlux))\n",
    "print(\"reverse output flux = {}\".format(reverseFlux))\n",
    "print(\"destroyed flux      = {}\".format(\n",
    "    np.sum(rays.flux) - forwardFlux - reverseFlux\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924f46eb-4baf-4260-8ba0-a49aef4680f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Separate out the direct path\n",
    "direct_path_idx = np.argmin([len(r.path) for r in rForward])\n",
    "direct_path = rForward[direct_path_idx]\n",
    "del rForward[direct_path_idx]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6deeb1dc-ef25-48bd-9781-3b585c845f07",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "direct_flux = np.sum(direct_path.flux)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1ec9c3c-eb0a-4def-a903-e4575efe5572",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = np.concatenate([rr.x for rr in rForward])\n",
    "y = np.concatenate([rr.y for rr in rForward])\n",
    "flux = np.concatenate([rr.flux for rr in rForward])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48eac48d-9ccc-4787-bc99-f0dabcdc892e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "H, xedges, yedges = np.histogram2d(\n",
    "    x, y, \n",
    "    weights=flux, \n",
    "    bins=200, \n",
    "    range=[[-0.38, 0.38], [-0.38, 0.38]]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6042fa-1352-4bd6-aa3f-10fcc426c974",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# What is the area of each pixel?\n",
    "dx = np.mean(np.diff(xedges))\n",
    "dy = np.mean(np.diff(yedges))\n",
    "area = dx*dy # in m^2\n",
    "area *= (1/10e-6)**2 # pix^2\n",
    "area *= (1/0.2)**2 # arcsec^2\n",
    "print(area)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f49787b6-a957-4ca2-b09f-739c6fd46d46",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(5, 4))\n",
    "colorbar(ax.imshow(\n",
    "    -2.5*np.log10((H.T+1e-300)/area/direct_flux), \n",
    "    vmin=25, vmax=33,\n",
    "    cmap='viridis_r',\n",
    "    extent=[-0.38, 0.38, -0.38, 0.38]\n",
    "), label='mag/sq arcsec')\n",
    "ax.set_aspect(1)\n",
    "ax.set_title(\"0th magnitude ghosts\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9f0efb2-2061-46e5-9e82-dc2d60933f60",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
