{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6026a74b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "\n",
    "isdark = False\n",
    "\n",
    "from rayoptics.environment import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c7d740ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "opm = OpticalModel()\n",
    "sm  = opm['seq_model']\n",
    "osp = opm['optical_spec']\n",
    "pm = opm['parax_model']\n",
    "\n",
    "osp['pupil'] = PupilSpec(osp, key=['object', 'pupil'], value=16)\n",
    "osp['fov'] = FieldSpec(osp, key=['object', 'angle'], value=1.5, flds=[0., 0.707, 1.], is_relative=True)\n",
    "osp['wvls'] = WvlSpec([('F', 0.5), (587.5618, 1.0), ('C', 0.5)], ref_wl=1)\n",
    "\n",
    "opm.radius_mode = True\n",
    "\n",
    "sm.gaps[0].thi=1e10\n",
    "\n",
    "use_schyrle = True\n",
    "if use_schyrle:\n",
    "    sm.add_surface([246, 3, 'N-LASF9', 'Schott', 42/2])\n",
    "    sm.add_surface([1e9, 290+36])\n",
    "    sm.add_surface([1e9, 3, 'N-LASF9', 'Schott', 16/2])\n",
    "    sm.add_surface([-30.6, 36+36])\n",
    "    sm.add_surface([30.6, 3, 'N-LASF9', 'Schott', 16/2])\n",
    "    sm.add_surface([1e9, 36+36])\n",
    "    sm.add_surface([1e9, 3, 'N-LASF9', 'Schott', 16/2])\n",
    "    sm.add_surface([-30.6, 36])\n",
    "else:\n",
    "    sm.add_surface([246, 2, 'N-LASF9', 'Schott', 42/2])\n",
    "    sm.add_surface([1e9, 290+2*36])\n",
    "    sm.add_surface([1e9, 2, 'N-LASF9', 'Schott', 16])\n",
    "    sm.add_surface([-30.6, 36+36+36])\n",
    "    sm.add_surface([30.6, 2, 'N-LASF9', 'Schott', 16])\n",
    "    sm.add_surface([1e9, 96])\n",
    "# add a focusing mirror\n",
    "#sm.add_surface([30.6, 2, 'N-LASF9', 'Schott', 16])\n",
    "#sm.add_surface([1e9, 36])\n",
    "#sm.add_surface([-24.456, .975, 'N-SF5,Schott'])\n",
    "#sm.set_stop()\n",
    "#sm.add_surface([21.896, 4.822])\n",
    "#sm.add_surface([86.759, 5.127, 'N-LAK9', 'Schott'])\n",
    "#sm.add_surface([-30.4942, 41.2365])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a26f0ff4",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c905e0317df438b9be9f0e4f16961b1",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 14176 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 14176 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "opm.update_model()\n",
    "\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "85fd9885",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "efl                1448\n",
      "ffl          -1.193e+04\n",
      "pp1          -1.048e+04\n",
      "bfl               216.1\n",
      "ppk                1232\n",
      "f/#               90.51\n",
      "m            -1.448e-07\n",
      "red          -6.905e+06\n",
      "obj_dist          1e+10\n",
      "obj_ang             1.5\n",
      "enp_dist             -0\n",
      "enp_radius            8\n",
      "na obj            8e-10\n",
      "n obj                 1\n",
      "img_dist          216.1\n",
      "img_ht            37.92\n",
      "exp_dist          220.5\n",
      "exp_radius        0.971\n",
      "na img        -0.005524\n",
      "n img                 1\n",
      "optical invariant       0.2095\n"
     ]
    }
   ],
   "source": [
    "pm.first_order_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8f39b24f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "296797f5b1b14d0d9d67a63d24e00540",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 13592 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 13592 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# add a focusing mirror\n",
    "sm.insert_surface_and_gap()\n",
    "cur = sm.cur_surface\n",
    "sm.ifcs[cur] = ThinLens(power=1/30.)\n",
    "sm.gaps[cur].thi = 30.\n",
    "\n",
    "opm.update_model()\n",
    "\n",
    "opm.seq_model.gaps[-1].thi = opm.optical_spec.parax_data.fod.bfl\n",
    "opm.update_model()\n",
    "\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "fbb733ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ddd9627c930414b87d0c7e0ed0c5169",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 53136 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 53136 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, scale_type=Fit.User_Scale,\n",
    "#                      user_scale_value=0.1, is_dark=isdark).plot()\n",
    "abr_plt = plt.figure(FigureClass=RayFanFigure, opt_model=opm, data_type='Ray',\n",
    "                     scale_type=Fit.All_Same).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "95deb34e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f354af5f6e9544778b41419b1ac19d63",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 48048 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 48048 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "opm.seq_model.gaps[-1].thi = opm.optical_spec.parax_data.fod.bfl\n",
    "opm.update_model()\n",
    "\n",
    "#spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, scale_type=Fit.User_Scale,\n",
    "#                      user_scale_value=0.1, is_dark=isdark).plot()T\n",
    "spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "a6ed5037",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "23.655706218096086"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hey it's me from the future, I made these functions\n",
    "\n",
    "import rayoptics.optical.model_constants as mc\n",
    "\n",
    "def dump_dist(p, wi, ray_pkg, fld, wvl, foc):\n",
    "    if ray_pkg is not None:\n",
    "        image_pt = fld.ref_sphere[0]\n",
    "        ray = ray_pkg[mc.ray]\n",
    "        dist = foc / ray[-1][mc.d][2]\n",
    "        defocused_pt = ray[-1][mc.p] + dist*ray[-1][mc.d]\n",
    "        t_abr = defocused_pt - image_pt\n",
    "        return np.sqrt(np.sum(t_abr*t_abr))\n",
    "    \n",
    "def spot_rms(sm, fld_idx=0):\n",
    "    return np.sqrt(np.mean(np.square(sm.trace_grid(dump_dist, fld_idx, form='list', append_if_none=False)[0]), axis=1))\n",
    "\n",
    "def dump_rays(p, wi, ray_pkg, fld, wvl, foc):\n",
    "    if ray_pkg is not None:\n",
    "        image_pt = fld.ref_sphere[0]\n",
    "        ray = ray_pkg[mc.ray]\n",
    "        v = ray[-1][mc.d][0:2] / ray[-1][mc.d][2]\n",
    "        return [ray[-1][mc.p][0:2] - image_pt[0:2], v]\n",
    "    \n",
    "def get_focus(sm, offset=False, weights=np.array([0.25, 0.5, 0.25])):\n",
    "    offsets = np.linspace(-40, 40, 2000)\n",
    "    vals, colors = sm.trace_grid(dump_rays, 0, form='list', append_if_none=False)\n",
    "\n",
    "    rms = np.zeros((offsets.size, len(vals)), dtype=float)\n",
    "\n",
    "    for i in range(offsets.size):\n",
    "        for j in range(len(vals)):\n",
    "            rms[i][j] = np.sqrt(np.mean(np.square(vals[j][:,0]+vals[j][:,1]*offsets[i])))\n",
    "\n",
    "    if offset:\n",
    "        return offsets[np.argmin(np.sum(rms*weights, axis=1))]\n",
    "    else:\n",
    "        return sm.gaps[-1].thi + offsets[np.argmin(np.sum(rms*weights, axis=1))]\n",
    "\n",
    "get_focus(sm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d6099313",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.00418596, 0.01665802, 0.00529431]),\n",
       " array([0.25632774, 0.0865608 , 0.10988704]),\n",
       " array([0.47439537, 0.14863163, 0.25140339]))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99b74a4f3edb4672b76acf563827ae2d",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 34736 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 34736 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "opm.seq_model.gaps[-1].thi = get_focus(sm)\n",
    "opm.update_model()\n",
    "\n",
    "#spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, scale_type=Fit.User_Scale,\n",
    "#                      user_scale_value=0.1, is_dark=isdark).plot()T\n",
    "spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, is_dark=isdark).plot()\n",
    "spot_rms(sm, 0), spot_rms(sm, 1), spot_rms(sm, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b96dd72",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
