{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "6026a74b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "\n",
    "isdark = False\n",
    "\n",
    "from rayoptics.environment import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "c7d740ee",
   "metadata": {
    "scrolled": true
   },
   "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=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",
    "def calc_curvature(n, fl):\n",
    "    return (n-1)*fl\n",
    "\n",
    "n_bk7 = 1.5168\n",
    "n_lasf9 = 1.85025\n",
    "n_f2 =  1.62005\n",
    "\n",
    "v_bk7 = 64.17\n",
    "v_lasf9 = 32.16\n",
    "v_f2 = 36.43\n",
    "# try for chaining a 3x telescope setup with a second 3x telescope setup\n",
    "f0 = 150\n",
    "f0_0 = (v_bk7-v_f2)*f0/v_bk7\n",
    "f0_1 = -f0_0*v_bk7/v_f2\n",
    "f0_actual = 1/(1/f0_0+1/f0_1)\n",
    "f1 = 60\n",
    "# 1/f = 1/f0 + 1/f1 = 1/f0 - v2/(f0*v1) = (v1-v2)/(v1*f0)\n",
    "# f0 = (v1-v2)*f/v1\n",
    "f1_0 = (v_bk7-v_f2)*f1/v_bk7\n",
    "f1_1 = -f1_0*v_bk7/v_f2\n",
    "\n",
    "f2 = 150\n",
    "f3 = 50\n",
    "\n",
    "f1_stacked = 2*f1\n",
    "\n",
    "r0 = calc_curvature(n_lasf9, f0)\n",
    "r0_0 = calc_curvature(n_bk7, f0_0)\n",
    "r0_1 = calc_curvature(n_f2, f0_1)\n",
    "r1 = calc_curvature(n_lasf9, f1)\n",
    "r1_0 = calc_curvature(n_bk7, f1_0)\n",
    "r1_1 = calc_curvature(n_f2, f1_1)\n",
    "r2 = calc_curvature(n_lasf9, f2)\n",
    "r3 = calc_curvature(n_lasf9, f3)\n",
    "\n",
    "#sm.add_surface([r0_0, 2, 'N-BK7', 'Schott', 42/2])\n",
    "#sm.add_surface([1e9, 2, 'N-F2', 'Schott', 42/2])\n",
    "#sm.add_surface([-r0_1, 180+2*f1])\n",
    "#sm.add_surface([-r0_1, 180+2*36.])\n",
    "\n",
    "sm.add_surface([r1_0, 4, 'N-BK7', 'Schott', 16])\n",
    "sm.add_surface([1e9, 2, 'N-F2', 'Schott', 16])\n",
    "sm.add_surface([-r1_1, 30])\n",
    "#sm.add_surface([1e9, 2, 'N-LASF9', 'Schott', 16])\n",
    "#sm.add_surface([-r1, f2+f3])\n",
    "\n",
    "\n",
    "#sm.add_surface([r2, 2, 'N-LASF9', 'Schott', 16])\n",
    "#sm.add_surface([1e9, f2+f3])\n",
    "#sm.add_surface([1e9, 2, 'N-LASF9', 'Schott', 16])\n",
    "#sm.add_surface([-r3, f3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "a26f0ff4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc33ba3f73f646adb605241decb1dcb5",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 53324 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, 53324 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",
    "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": 55,
   "id": "85fd9885",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "efl               50.17\n",
      "ffl              -54.42\n",
      "pp1              -4.251\n",
      "bfl               42.68\n",
      "ppk               7.488\n",
      "f/#               3.135\n",
      "m            -5.017e-09\n",
      "red          -1.993e+08\n",
      "obj_dist          1e+10\n",
      "obj_ang               5\n",
      "enp_dist             -0\n",
      "enp_radius            8\n",
      "na obj            8e-10\n",
      "n obj                 1\n",
      "img_dist          42.68\n",
      "img_ht            4.389\n",
      "exp_dist         -3.569\n",
      "exp_radius        7.375\n",
      "na img          -0.1575\n",
      "n img                 1\n",
      "optical invariant       0.6999\n"
     ]
    }
   ],
   "source": [
    "pm.first_order_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "fbb733ef",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "326d863c43bc4d3ba96074464b27f77a",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 47208 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, 47208 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": [
    "abr_plt = plt.figure(FigureClass=RayFanFigure, opt_model=opm, data_type='Ray',\n",
    "                     scale_type=Fit.All_Same).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "95deb34e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a88863352ccb4248815eecadf8289319",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 61160 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, 61160 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, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "fe1bd365",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca417a2888fa4fcc8eeebf1a55a1a65f",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 53324 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, 53324 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": [
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "0e1c606f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.71394997, 0.69072105, 0.67781734])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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[0]\n",
    "        dist = foc / ray[-1][1][2]\n",
    "        defocused_pt = ray[-1][0] + dist*ray[-1][1]\n",
    "        t_abr = defocused_pt - image_pt\n",
    "        return np.sqrt(np.sum(t_abr*t_abr))\n",
    "    \n",
    "def spot_rms(sm):\n",
    "    return np.sqrt(np.mean(np.square(sm.trace_grid(dump_dist, 0, form='list', append_if_none=False)[0]), axis=1))\n",
    "\n",
    "spot_rms(sm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "b0c10933",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a925ca79c19e4c789e2be0d62e8b8829",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 37624 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, 37624 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"
    },
    {
     "data": {
      "text/plain": [
       "-5.022511255627812"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import rayoptics.optical.model_constants as mc\n",
    "\n",
    "old_gap = opm.seq_model.gaps[-1].thi\n",
    "offsets = np.linspace(-40, 40, 2000)\n",
    "\n",
    "opm.seq_model.gaps[-1].thi = opm.optical_spec.parax_data.fod.bfl\n",
    "opm.update_model()\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",
    "vals, colors = sm.trace_grid(dump_rays, 0, form='list', append_if_none=False)\n",
    "v1, v2, v3 = vals\n",
    "v1_p = v1[:,0]\n",
    "v1_v = v1[:,1]\n",
    "v2_p = v2[:,0]\n",
    "v2_v = v2[:,1]\n",
    "v3_p = v3[:,0]\n",
    "v3_v = v3[:,1]\n",
    "\n",
    "v1_rms = np.zeros(offsets.size)\n",
    "v2_rms = np.zeros(offsets.size)\n",
    "v3_rms = np.zeros(offsets.size)\n",
    "\n",
    "for i in range(offsets.size):\n",
    "    v1_rms[i] = np.sqrt(np.mean(np.square(v1_p+v1_v*offsets[i])))\n",
    "    v2_rms[i] = np.sqrt(np.mean(np.square(v2_p+v2_v*offsets[i])))\n",
    "    v3_rms[i] = np.sqrt(np.mean(np.square(v3_p+v3_v*offsets[i])))\n",
    "\n",
    "opm.seq_model.gaps[-1].thi = old_gap\n",
    "opm.update_model()\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(offsets, v1_rms, color=colors[0])\n",
    "plt.plot(offsets, v2_rms, color=colors[1])\n",
    "plt.plot(offsets, v3_rms, color=colors[2])\n",
    "plt.show()\n",
    "\n",
    "min_o = offsets[np.argmin(0.25*v1_rms+0.50*v2_rms+0.25*v3_rms)]\n",
    "min_o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "5c091777",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.2159279  0.21609438 0.21714886]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9a64379f8f945ef9abbb9a773e53114",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 57296 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, 57296 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 + min_o\n",
    "opm.update_model()\n",
    "\n",
    "spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm, is_dark=isdark).plot()\n",
    "print(spot_rms(sm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "4f7eb3b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f9a7ddaee43945699b115b5232c0eb36",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 49128 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, 49128 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": [
    "abr_plt = plt.figure(FigureClass=RayFanFigure, opt_model=opm, data_type='Ray',\n",
    "                     scale_type=Fit.All_Same).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3f0c6344",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "37.675382106540546"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# let's try chaining the two lenses together and seeing how that looks\n",
    "\n",
    "# first let's wrap that focus calculation code into a function\n",
    "\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": 150,
   "id": "65db2cdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b3e9e9ed84b405f8429c5a9e70093ee",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 23428 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, 23428 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": [
    "opm2 = OpticalModel()\n",
    "sm2  = opm2['seq_model']\n",
    "osp2 = opm2['optical_spec']\n",
    "pm2 = opm2['parax_model']\n",
    "\n",
    "osp2['pupil'] = PupilSpec(osp2, key=['object', 'pupil'], value=16)\n",
    "osp2['fov'] = FieldSpec(osp2, key=['object', 'angle'], value=1.5, flds=[0., 0.707, 1.], is_relative=True)\n",
    "osp2['wvls'] = WvlSpec([('F', 0.5), (587.5618, 1.0), ('C', 0.5)], ref_wl=1)\n",
    "\n",
    "opm2.radius_mode = True\n",
    "\n",
    "sm2.gaps[0].thi=1e10\n",
    "sm2.add_surface([r0_0, 2, 'N-BK7', 'Schott', 42/2])\n",
    "sm2.add_surface([1e9, 2, 'N-F2', 'Schott', 42/2])\n",
    "sm2.add_surface([-r0_1, f0])\n",
    "first_gap_idx = sm2.cur_surface\n",
    "\n",
    "opm2.update_model()\n",
    "sm2.gaps[-1].thi = get_focus(sm2)\n",
    "\n",
    "opm2.update_model()\n",
    "\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm2, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "6b732394",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62f44eb8d991479d833cb6b7f5a7603e",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 22700 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, 22700 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": [
    "#use the gap from sm\n",
    "sm2.gaps[-1].thi += sm.gaps[-1].thi\n",
    "\n",
    "sm2.add_surface([r1_1, 1, 'N-F2', 'Schott', 16])\n",
    "sm2.add_surface([1e9, 2, 'N-BK7', 'Schott', 16])\n",
    "sm2.add_surface([-r1_0, sm.gaps[-1].thi])\n",
    "second_gap_idx = sm2.cur_surface\n",
    "\n",
    "opm2.update_model()\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm2, is_dark=isdark).plot()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "cb05ecc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "efl               505.1\n",
      "ffl                1173\n",
      "pp1                1678\n",
      "bfl              -136.6\n",
      "ppk               641.7\n",
      "f/#               31.57\n",
      "m            -5.051e-08\n",
      "red           -1.98e+07\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         -136.6\n",
      "img_ht            13.23\n",
      "exp_dist         -93.21\n",
      "exp_radius        3.446\n",
      "na img         -0.01584\n",
      "n img                 1\n",
      "optical invariant       0.2095\n"
     ]
    }
   ],
   "source": [
    "pm2.first_order_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "8d42a780",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07e104adb6e54cfe9748a449f8c28dfa",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 21456 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, 21456 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": [
    "sm2.insert_surface_and_gap()\n",
    "cur = sm2.cur_surface\n",
    "sm2.ifcs[cur] = ThinLens(power=1/30.)\n",
    "sm2.gaps[cur].thi = 30.\n",
    "\n",
    "opm2.update_model()\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm2, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "dc0aa747",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.51212369 0.14586092 0.10640832]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81d0ba94f2b94c4481d911a46bd55a7a",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 44452 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, 44452 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=opm2, is_dark=isdark).plot()\n",
    "print(spot_rms(sm2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "baba5f64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32.26113056528264, array([0.39920801, 0.02495399, 0.26426542]))"
      ]
     },
     "execution_count": 155,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e627f1fd290b4549bd0562ebe3b0db00",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 58208 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, 58208 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": [
    "adjust = get_focus(sm2)\n",
    "sm2.gaps[-1].thi = adjust\n",
    "\n",
    "opm2.update_model()\n",
    "spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm2, is_dark=isdark).plot()\n",
    "adjust, spot_rms(sm2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "ad4398b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "294c25f1b8e741bda22eb559c8329a52",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 61216 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, 61216 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": [
    "abr_plt = plt.figure(FigureClass=RayFanFigure, opt_model=opm2, data_type='Ray',\n",
    "                     scale_type=Fit.All_Same).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "981b6196",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7e0bf10>]"
      ]
     },
     "execution_count": 157,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e32152b6e79d4a0b9f74ce8cde921ec1",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 28368 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, 28368 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": [
    "def test_gap(opm, sm, gap_idx, start, end, npoints=10, weights=np.array([0.25, 0.5, 0.25])):\n",
    "    last_gap = sm.gaps[gap_idx].thi\n",
    "    focus_gap = sm.gaps[-1].thi\n",
    "\n",
    "    offsets = np.linspace(start, end, npoints)\n",
    "    test_rms = np.zeros(offsets.size)\n",
    "    for i in range(offsets.size):\n",
    "        sm.gaps[gap_idx].thi = last_gap + offsets[i]\n",
    "        opm.update_model()\n",
    "        sm.gaps[-1].thi = get_focus(sm)\n",
    "        opm.update_model()\n",
    "        test_rms[i] = np.sum(spot_rms(sm)*weights)\n",
    "    sm.gaps[gap_idx].thi = last_gap\n",
    "    sm.gaps[-1].thi = focus_gap\n",
    "    opm.update_model()\n",
    "    return offsets, test_rms\n",
    "    \n",
    "offsets, test_rms = test_gap(opm2, sm2, first_gap_idx, -10, 30)\n",
    "plt.figure()\n",
    "plt.plot(offsets, test_rms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "5585beb4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7e8b970>]"
      ]
     },
     "execution_count": 158,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40f76e5166924220a0dc4d51635bed85",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 34580 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, 34580 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": [
    "offsets, test_rms = test_gap(opm2, sm2, first_gap_idx, 15, 24, 30)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(offsets, test_rms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "3fe08147",
   "metadata": {},
   "outputs": [],
   "source": [
    "# turned off to see how optimizing the second gap affects the quality by itself\n",
    "\n",
    "#sm2.gaps[first_gap_idx].thi += offsets[np.argmin(test_rms)]\n",
    "#opm2.update_model()\n",
    "\n",
    "#offsets[np.argmin(test_rms)], np.min(test_rms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "7f469072",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7d14550>]"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c718250a551043a3a6546f6dfb2e94f7",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 29080 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, 29080 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": [
    "offsets2, test_rms2 = test_gap(opm2, sm2, second_gap_idx, -10, 10)\n",
    "plt.figure()\n",
    "plt.plot(offsets2, test_rms2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "d211a479",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7d7a020>]"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "22b4cedd881e44b3a073711c92bd81cc",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 28092 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, 28092 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": [
    "offsets2, test_rms2 = test_gap(opm2, sm2, second_gap_idx, -30, -10)\n",
    "plt.figure()\n",
    "plt.plot(offsets2, test_rms2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "2a527dbe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7de7700>]"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03cd52712e444af699f126d406548d7c",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 26496 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, 26496 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": [
    "offsets2, test_rms2 = test_gap(opm2, sm2, second_gap_idx, -40, -30)\n",
    "plt.figure()\n",
    "plt.plot(offsets2, test_rms2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "2d37f8aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "37.65537210153805"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# fuck how negative can it go?\n",
    "\n",
    "sm2.gaps[second_gap_idx].thi\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "36b1d249",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.38273603, 0.06067967, 0.16915498])"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a3bc8b1c06b4d9e812c10c72df7598b",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 38376 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, 38376 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": [
    "# let's place it at 2mm and see how that looks\n",
    "sm2.gaps[second_gap_idx].thi = 2\n",
    "\n",
    "sm2.gaps[-1].thi = get_focus(sm2)\n",
    "opm2.update_model()\n",
    "\n",
    "spot_plt = plt.figure(FigureClass=SpotDiagramFigure, opt_model=opm2, is_dark=isdark).plot()\n",
    "spot_rms(sm2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "d53a1d74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "badaf6fa5c7a46cfa2337a6550f9a6f1",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 25044 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, 25044 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": [
    "opm2.rebuild_from_seq()\n",
    "layout_plt = plt.figure(FigureClass=InteractiveLayout, opt_model=opm2, is_dark=isdark).plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "f8fc988d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fcfd7a0a860>]"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d61d5b22620c48a6a613806d3f72ea58",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 31064 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, 31064 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": [
    "# let's try messing with the first gap again\n",
    "\n",
    "offsets, test_rms = test_gap(opm2, sm2, first_gap_idx, -5, 5)\n",
    "plt.figure()\n",
    "plt.plot(offsets, test_rms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "9eed4ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# looks like it's pretty close to optimal, guess we'll have to mess with the first lens\n",
    "# to optimize for this second one\n",
    "\n",
    "# to be continued in part 2, 4D multivariable optimization\n",
    "# probably gonna switch optics libraries too,\n",
    "# this ray-optics is good but kind of slow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cddee23c",
   "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
}
