{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aef7b0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# here's an attempt to speed up raytracing so it takes a reasonable amount of\n",
    "# time to optimize systems\n",
    "# mainly by leveraging numpy to run all the calculations in a tighter for\n",
    "# loop than what Python can do\n",
    "\n",
    "# first let's copy the system from before to test against\n",
    "\n",
    "from rayoptics.environment import *\n",
    "\n",
    "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",
    "def achromatic(f, v0, v1):\n",
    "    return (v0-v1)*f/v0, -v1*f/(v0-v1)\n",
    "\n",
    "def calc_curvatures(ns, fls):\n",
    "    return tuple((n-1)*f for n, f in zip(ns, fls))\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, f0_1 = achromatic(f0, v_bk7, v_f2)\n",
    "\n",
    "f1 = 60\n",
    "f1_0, f1_1 = achromatic(f1, v_bk7, v_f2)\n",
    "\n",
    "r0_0, r0_1, r1_0, r1_1 = calc_curvatures((n_bk7, n_f2, n_bk7, n_f2), (f0_0, f0_1, f1_0, f1_1))\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",
    "\n",
    "opm.update_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "5c6bebd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we trace all the functions needed for raytracing until we get to the algorithm,\n",
    "# having them transfer arrays of values to evaluate instead of a single point\n",
    "\n",
    "import rayoptics.raytr.trace as raytr_trace\n",
    "from rayoptics.optical.model_constants import Intfc, Gap, Tfrm, Indx, Zdir\n",
    "from rayoptics.elem.profiles import Spherical\n",
    "from rayoptics.elem.surface import Surface\n",
    "\n",
    "from rayoptics.raytr.traceerror import TraceError, TraceMissedSurfaceError, TraceTIRError, TraceEvanescentRayError\n",
    "\n",
    "def super_trace_grid(sm, fi, wl=None, num_rays=21,\n",
    "                   append_if_none=True, **kwargs):\n",
    "    \"\"\" fct is applied to the raw grid and returned as a grid  \"\"\"\n",
    "    osp = sm.opt_model.optical_spec\n",
    "    wvls = osp.spectral_region\n",
    "    wvl = sm.central_wavelength()\n",
    "    wv_list = wvls.wavelengths if wl is None else [wvl]\n",
    "    fld = osp.field_of_view.fields[fi]\n",
    "    foc = osp.defocus.get_focus()\n",
    "\n",
    "    # make sure this is imported\n",
    "    rs_pkg, cr_pkg = raytr_trace.setup_pupil_coords(sm.opt_model,\n",
    "                                              fld, wvl, foc)\n",
    "    fld.chief_ray = cr_pkg\n",
    "    fld.ref_sphere = rs_pkg\n",
    "\n",
    "    grids = []\n",
    "    grid_start = np.array([-1., -1.])\n",
    "    grid_stop = np.array([1., 1.])\n",
    "    grid_def = [grid_start, grid_stop, num_rays]\n",
    "    results = None\n",
    "    for wi, wvl in enumerate(wv_list):\n",
    "        result = np.expand_dims(rly_trace_grid(sm.opt_model, grid_def, fld, wvl, foc,\n",
    "                                **kwargs), axis=0)\n",
    "        if results is None:\n",
    "            results = result\n",
    "        else:\n",
    "            results = np.concatenate((results, result), axis=0)\n",
    "    rc = wvls.render_colors\n",
    "    return results, rc\n",
    "\n",
    "# this generates an array of valid points from the grid\n",
    "def rly_trace_grid(opt_model, grid_rng, fld, wvl, foc, **kwargs): # from trace_grid\n",
    "    start = np.array(grid_rng[0])\n",
    "    stop = grid_rng[1]\n",
    "    num = grid_rng[2]\n",
    "    step = np.array((stop - start)/(num - 1))\n",
    "    grid = []\n",
    "    \n",
    "    valid_points = None\n",
    "    \n",
    "    for x in np.linspace(start[0], stop[0], num, dtype=np.float64):\n",
    "        ys = np.linspace(start[1], stop[1], num, dtype=np.float64)\n",
    "        valid_ys = ys[x**2 + ys**2 <= 1.0]\n",
    "        points = np.zeros((valid_ys.shape[0], 2))\n",
    "        points[:, 0] = x\n",
    "        points[:, 1] = valid_ys\n",
    "        if valid_points is None:\n",
    "            valid_points = points\n",
    "        else:\n",
    "            valid_points = np.concatenate((valid_points, points))\n",
    "    return rly_rly_trace(opt_model, valid_points, fld, wvl, **kwargs)\n",
    "\n",
    "def fld_apply_vignetting(fld, pupils):\n",
    "    ret = np.copy(pupils)\n",
    "    if fld.vlx != 0.0:\n",
    "        ret[pupils[:, 0] < 0.0, 0] *= (1. - fld.vlx)\n",
    "    if fld.vux != 0.0:\n",
    "        ret[pupils[:, 0] > 0.0, 0] *= (1. - fld.vux)\n",
    "    if fld.vly != 0.0:\n",
    "        ret[pupils[:, 1] < 0.0, 1] *= (1. - fld.vly)\n",
    "    if fld.vuy != 0.0:\n",
    "        ret[pupils[:, 1] > 0.0, 1] *= (1. - fld.vuy)\n",
    "    return ret\n",
    "\n",
    "def norm(v):\n",
    "    return np.sqrt(np.sum(v*v, axis=1))\n",
    "    \n",
    "# some final data conditioning before the good part\n",
    "def rly_rly_trace(opt_model, pupils, fld, wvl, **kwargs): # from trace_base\n",
    "    vig_pupils = fld_apply_vignetting(fld, pupils)\n",
    "    osp = opt_model.optical_spec\n",
    "    fod = osp.parax_data.fod\n",
    "    eprad = fod.enp_radius\n",
    "    aim_pt = np.array([0., 0.])\n",
    "    if hasattr(fld, 'aim_pt') and fld.aim_pt is not None:\n",
    "        aim_pt = np.array(fld.aim_pt)\n",
    "    pt1 = np.zeros((pupils.shape[0], 3))\n",
    "    pt1[:,0:2] = eprad*vig_pupils+aim_pt\n",
    "    pt1[:, 2] = fod.obj_dist+fod.enp_dist\n",
    "    pt0 = osp.obj_coords(fld)\n",
    "    dir0 = pt1 - pt0\n",
    "    length = norm(dir0)\n",
    "    dir0 = dir0/np.expand_dims(length, axis=1)\n",
    "    return rly_rly_rly_trace(opt_model.seq_model, pt0, dir0, wvl, **kwargs)\n",
    "\n",
    "def rly_rly_rly_trace(seq_model, pt0, dir0, wvl, eps=1.0e-12, **kwargs): # from raytr.trace and raytr.trace_raw\n",
    "    path = [v for v in seq_model.path(wvl)]\n",
    "    \n",
    "    #kwargs['first_surf'] = kwargs.get('first_surf', 1)\n",
    "    #kwargs['last_surf'] = kwargs.get('last_surf',\n",
    "    #                                 seq_model.get_num_surfaces()-2)\n",
    "    \n",
    "    rays = np.zeros((dir0.shape[0], len(path), 10))\n",
    "    # rays[nray][path][px, py, pz, dx, dy, dz, nx, ny, nz, dist]\n",
    "\n",
    "    #first_surf = kwargs.get('first_surf', 0)\n",
    "    #last_surf = kwargs.get('last_surf', None)\n",
    "    first_surf = kwargs.get('first_surf', 1)\n",
    "    last_surf = kwargs.get('last_surf', seq_model.get_num_surfaces()-2)\n",
    "    \n",
    "    # trace object surface\n",
    "    obj = path[0]\n",
    "    pt0 = np.expand_dims(pt0, axis=0)\n",
    "\n",
    "    srf_obj = obj[Intfc]\n",
    "    dst_b4, pt_obj = itfc_intersect(srf_obj, pt0, dir0, z_dir=obj[Zdir])\n",
    "\n",
    "    before = obj\n",
    "    before_pt = pt_obj\n",
    "    before_dir = dir0\n",
    "    before_normal = itfc_normal(srf_obj, before_pt)\n",
    "    tfrm_from_before = before[Tfrm]\n",
    "    z_dir_before = before[Zdir]\n",
    "\n",
    "    # loop remaining surfaces in path\n",
    "    for surf, after in enumerate(path[1:]):\n",
    "        try:\n",
    "            #np.tensordot(t, pts, (0, 1)).T seems to work for mass matrix multiplication\n",
    "            rt, t = tfrm_from_before\n",
    "            #b4_pt, b4_dir = rt.dot(before_pt - t), rt.dot(before_dir)\n",
    "            b4_pt = np.tensordot(rt, before_pt - t, (0, 1)).T\n",
    "            b4_dir = np.tensordot(rt, before_dir, (0, 1)).T\n",
    "\n",
    "            #pp_dst = -b4_pt.dot(b4_dir)\n",
    "            pp_dst = -np.expand_dims(np.sum(b4_pt*b4_dir, axis=1), axis=1)\n",
    "            pp_pt_before = b4_pt + pp_dst*b4_dir\n",
    "\n",
    "            ifc = after[Intfc]\n",
    "            z_dir_after = after[Zdir]\n",
    "\n",
    "            # intersect ray with profile\n",
    "            pp_dst_intrsct, inc_pt = itfc_intersect(ifc, pp_pt_before, b4_dir,\n",
    "                                                   eps=eps, z_dir=z_dir_before)\n",
    "            dst_b4 = pp_dst[:, 0] + pp_dst_intrsct\n",
    "            rays[:, surf, 0:3] = before_pt\n",
    "            rays[:, surf, 3:6] = before_dir\n",
    "            rays[:, surf, 6:9] = before_normal\n",
    "            rays[:, surf, 9] = dst_b4\n",
    "\n",
    "            normal = itfc_normal(ifc, inc_pt)\n",
    "\n",
    "\n",
    "            '''\n",
    "            # if the interface has a phase element, process that first\n",
    "            if hasattr(ifc, 'phase_element'):\n",
    "                doe_dir, phs = phase(ifc, inc_pt, b4_dir, normal, z_dir_before,\n",
    "                                     wvl, before[Indx], after[Indx])\n",
    "                # the output of the phase element becomes the input for the\n",
    "                #  refraction/reflection calculation\n",
    "                b4_dir = doe_dir\n",
    "                op_delta += phs\n",
    "            '''\n",
    "\n",
    "            # refract or reflect ray at interface\n",
    "            if ifc.interact_mode == 'reflect':\n",
    "                after_dir = reflect(b4_dir, normal)\n",
    "            elif ifc.interact_mode == 'transmit':\n",
    "                after_dir = bend(b4_dir, normal, before[Indx], after[Indx])\n",
    "            elif ifc.interact_mode == 'dummy':\n",
    "                after_dir = b4_dir\n",
    "            else:  # no action, input becomes output\n",
    "                after_dir = b4_dir\n",
    "\n",
    "            # Per `Hopkins, 1981 <https://dx.doi.org/10.1080/713820605>`_, the\n",
    "            #  propagation direction is given by the direction cosines of the\n",
    "            #  ray and therefore doesn't require the use of a negated\n",
    "            #  refractive index following a reflection. Thus we use the\n",
    "            #  (positive) refractive indices from the seq_model.rndx array.\n",
    "\n",
    "            before_pt = inc_pt\n",
    "            before_normal = normal\n",
    "            before_dir = after_dir\n",
    "            z_dir_before = z_dir_after\n",
    "            before = after\n",
    "            tfrm_from_before = before[Tfrm]\n",
    "\n",
    "        except TraceMissedSurfaceError as ray_miss:\n",
    "            ray_miss.surf = surf+1\n",
    "            ray_miss.ifc = ifc\n",
    "            ray_miss.prev_tfrm = before[Tfrm]\n",
    "            ray_miss.ray_pkg = rays, wvl\n",
    "            raise ray_miss\n",
    "\n",
    "        except TraceTIRError as ray_tir:\n",
    "            ray_tir.surf = surf+1\n",
    "            ray_tir.ifc = ifc\n",
    "            ray_tir.int_pt = inc_pt\n",
    "            ray_tir.ray_pkg = rays, wvl\n",
    "            raise ray_tir\n",
    "\n",
    "        except TraceEvanescentRayError as ray_evn:\n",
    "            ray_evn.surf = surf+1\n",
    "            ray_evn.ifc = ifc\n",
    "            ray_evn.int_pt = inc_pt\n",
    "            ray_evn.ray_pkg = rays, wvl\n",
    "            raise ray_evn\n",
    "\n",
    "    # lifted from the loop since it'll no longer hit the\n",
    "    # StopIteration exception\n",
    "\n",
    "    if len(path) > 1:\n",
    "        rays[:, surf + 1, 0:3] = inc_pt\n",
    "        rays[:, surf + 1, 3:6] = after_dir\n",
    "        rays[:, surf + 1, 6:9] = normal\n",
    "        rays[:, surf + 1, 9] = 0\n",
    "    return rays\n",
    "    \n",
    "def itfc_intersect(ifc, p, d, eps=1e-12, z_dir=1):\n",
    "    if isinstance(ifc, Surface):\n",
    "        if isinstance(ifc.profile, Spherical):\n",
    "            # copied from elem.profiles\n",
    "            ax2 = ifc.profile.cv\n",
    "            cx2 = ifc.profile.cv * np.sum(p*p, axis=1) - 2*p[:,2]\n",
    "            b = ifc.profile.cv * np.sum(d*p, axis=1) - d[:,2]\n",
    "            \n",
    "            discr = b*b-ax2*cx2\n",
    "            # Use z_dir to pick correct root\n",
    "            if np.any(discr < 0):\n",
    "                raise TraceMissedSurfaceError\n",
    "\n",
    "            s = cx2/(z_dir*np.sqrt(discr) - b)\n",
    "\n",
    "            p1 = p + np.expand_dims(s, axis=1)*d\n",
    "            return s, p1\n",
    "        else:\n",
    "            raise RuntimeError(\"intersection not implemented for profile {}\".format(ifc.profile))\n",
    "    else:\n",
    "        raise RuntimeError(\"intersection not implemented for {}\".format(ifc))\n",
    "\n",
    "def itfc_normal(ifc, pts):\n",
    "    if isinstance(ifc, Surface):\n",
    "        if isinstance(ifc.profile, Spherical):\n",
    "            # copied from elem.profiles\n",
    "            return np.stack([-ifc.profile.cv*pts[:,0], -ifc.profile.cv*pts[:,1], 1.0-ifc.profile.cv*pts[:,2]],\n",
    "                           axis=-1)\n",
    "        else:\n",
    "            raise RuntimeError(\"intersection not implemented for profile {}\".format(ifc.profile))\n",
    "    else:\n",
    "        raise RuntimeError(\"intersection not implemented for {}\".format(ifc))\n",
    "        \n",
    "#copied from raytrace.reflect\n",
    "def reflect(d_in, normals):\n",
    "    normal_len = norm(normals)\n",
    "    cosI = np.sum(d_in * normals, axis=1)/normal_len\n",
    "    d_out = d_in - 2.0*cosI*normals\n",
    "    return d_out\n",
    "\n",
    "#copied from raytrace.bend\n",
    "def bend(d_in, normal, n_in, n_out):\n",
    "    try:\n",
    "        normal_len = norm(normal)\n",
    "        cosI = np.sum(d_in * normal, axis=1)/normal_len\n",
    "        sinI_sqr = 1.0 - cosI*cosI\n",
    "        n_cosIp = np.copysign(np.sqrt(n_out*n_out - n_in*n_in*sinI_sqr), cosI)\n",
    "        alpha = np.expand_dims(n_cosIp - n_in*cosI, axis=1)\n",
    "        d_out = (n_in*d_in + alpha*normal)/n_out\n",
    "        return d_out\n",
    "    except ValueError:\n",
    "        raise TraceTIRError(d_in, normal, n_in, n_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "887c1b0d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.44231141, 0.4014741 , 0.38252552]),\n",
       " array([0.44778724, 0.40678566, 0.38775669]))"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "super_rays, _ = super_trace_grid(sm, 0)\n",
    "#print(super_rays[:,:,-1])\n",
    "\n",
    "def ray_err(sm, rays, fi):\n",
    "    osp = sm.opt_model.optical_spec\n",
    "    fld = osp.field_of_view.fields[fi]\n",
    "    foc = osp.defocus.get_focus()\n",
    "    \n",
    "    image_pt = fld.ref_sphere[0]\n",
    "    dist = np.expand_dims(foc/rays[:,:,-1,5], axis=-1)\n",
    "    defocused_pts = rays[:,:,-1,0:3] + dist*rays[:,:,-1,3:6]\n",
    "    t_abr = defocused_pts - image_pt\n",
    "    return np.sqrt(np.sum(t_abr*t_abr, axis=2))\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, num_rays=21):\n",
    "    return np.sqrt(np.mean(np.square(sm.trace_grid(dump_dist, fld_idx, form='list', num_rays=21, append_if_none=False)[0]), axis=1))\n",
    "\n",
    "def spot_rms2(sm, rays=None, fld_idx=0, num_rays=21):\n",
    "    if rays is None:\n",
    "        rays, _ = super_trace_grid(sm, fld_idx, num_rays=21)\n",
    "    return np.sqrt(np.mean(np.square(ray_err(sm, rays, fld_idx)), axis=1))\n",
    "\n",
    "spot_rms(sm), spot_rms2(sm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "b98b3263",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[[ 8.43971916e-01,  3.75098629e-01,  0.00000000e+00],\n",
       "         [ 2.36775522e-01,  1.05233565e-01,  9.65848461e-01]],\n",
       " \n",
       "        [[ 7.60036422e-01,  2.53345474e-01,  0.00000000e+00],\n",
       "         [ 2.33696527e-01,  7.78988423e-02,  9.69184040e-01]],\n",
       " \n",
       "        [[ 7.02253533e-01,  1.56056341e-01,  0.00000000e+00],\n",
       "         [ 2.31559445e-01,  5.14576544e-02,  9.71458869e-01]],\n",
       " \n",
       "        [[ 6.68405774e-01,  7.42673082e-02,  0.00000000e+00],\n",
       "         [ 2.30301018e-01,  2.55890020e-02,  9.72782938e-01]],\n",
       " \n",
       "        [[ 6.57255825e-01,  1.01347299e-16,  0.00000000e+00],\n",
       "         [ 2.29885412e-01,  3.54477885e-17,  9.73217703e-01]],\n",
       " \n",
       "        [[ 6.68405774e-01, -7.42673082e-02,  0.00000000e+00],\n",
       "         [ 2.30301018e-01, -2.55890020e-02,  9.72782938e-01]],\n",
       " \n",
       "        [[ 7.02253533e-01, -1.56056341e-01,  0.00000000e+00],\n",
       "         [ 2.31559445e-01, -5.14576544e-02,  9.71458869e-01]],\n",
       " \n",
       "        [[ 7.60036422e-01, -2.53345474e-01,  0.00000000e+00],\n",
       "         [ 2.33696527e-01, -7.78988423e-02,  9.69184040e-01]],\n",
       " \n",
       "        [[ 8.43971916e-01, -3.75098629e-01,  0.00000000e+00],\n",
       "         [ 2.36775522e-01, -1.05233565e-01,  9.65848461e-01]],\n",
       " \n",
       "        [[ 6.65190535e-01,  4.15744084e-01,  0.00000000e+00],\n",
       "         [ 2.07346720e-01,  1.29591700e-01,  9.69645981e-01]]]),\n",
       " array([[ 0.97899133,  0.        ,  0.        ,  0.26458622,  0.        ,\n",
       "          0.96436203, -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.84397192,  0.37509863,  0.        ,  0.23677552,  0.10523357,\n",
       "          0.96584846, -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.76003642,  0.25334547,  0.        ,  0.23369653,  0.07789884,\n",
       "          0.96918404, -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.70225353,  0.15605634,  0.        ,  0.23155944,  0.05145765,\n",
       "          0.97145887, -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.66840577,  0.07426731,  0.        ,  0.23030102,  0.025589  ,\n",
       "          0.97278294, -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.65725583,  0.        ,  0.        ,  0.22988541,  0.        ,\n",
       "          0.9732177 , -0.        , -0.        ,  1.        ,  0.        ],\n",
       "        [ 0.66840577, -0.07426731,  0.        ,  0.23030102, -0.025589  ,\n",
       "          0.97278294, -0.        ,  0.        ,  1.        ,  0.        ],\n",
       "        [ 0.70225353, -0.15605634,  0.        ,  0.23155944, -0.05145765,\n",
       "          0.97145887, -0.        ,  0.        ,  1.        ,  0.        ],\n",
       "        [ 0.76003642, -0.25334547,  0.        ,  0.23369653, -0.07789884,\n",
       "          0.96918404, -0.        ,  0.        ,  1.        ,  0.        ],\n",
       "        [ 0.84397192, -0.37509863,  0.        ,  0.23677552, -0.10523357,\n",
       "          0.96584846, -0.        ,  0.        ,  1.        ,  0.        ]]))"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# uhh why are they different?\n",
    "# let's check the rays themselves to see what they look like\n",
    "\n",
    "def dump_rays(p, wi, ray_pkg, fld, wvl, foc):\n",
    "    if ray_pkg is not None:\n",
    "        ray = ray_pkg[mc.ray]\n",
    "        return [ray[-1][mc.p], ray[-1][mc.d]]\n",
    "\n",
    "sm.trace_grid(dump_rays, 0, form='list', append_if_none=False)[0][0][0:10], super_rays[0,0:10,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "d2723737",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(311, (313, 10))"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# they look really similar...\n",
    "# let's make sure they're the same size\n",
    "ref_rays, _ = sm.trace_grid(dump_rays, 0, form='list', append_if_none=False)\n",
    "len(ref_rays[0]), super_rays[0,:,-1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "7ba8bcbc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-0.00000000e+00, -0.00000000e+00,  0.00000000e+00],\n",
       "        [-7.20000000e-10, -3.20000000e-10,  1.00000000e+00]],\n",
       "\n",
       "       [[-0.00000000e+00, -0.00000000e+00,  0.00000000e+00],\n",
       "        [-7.20000000e-10, -2.40000000e-10,  1.00000000e+00]],\n",
       "\n",
       "       [[-0.00000000e+00, -0.00000000e+00,  0.00000000e+00],\n",
       "        [-7.20000000e-10, -1.60000000e-10,  1.00000000e+00]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "        [ 7.20000000e-10,  2.40000000e-10,  1.00000000e+00]],\n",
       "\n",
       "       [[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
       "        [ 7.20000000e-10,  3.20000000e-10,  1.00000000e+00]],\n",
       "\n",
       "       [[ 0.00000000e+00, -0.00000000e+00,  0.00000000e+00],\n",
       "        [ 8.00000000e-10, -1.11022302e-25,  1.00000000e+00]]])"
      ]
     },
     "execution_count": 165,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# so the original has five more points than super_rays...\n",
    "# but which ones?\n",
    "\n",
    "def dump_start(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[0][mc.p], ray[0][mc.d]]\n",
    "    \n",
    "sm.trace_grid(dump_start, 0, form='list', append_if_none=False)[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "25c4bcc4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.0e+00,  0.0e+00,  0.0e+00, -8.0e-10,  0.0e+00,  1.0e+00],\n",
       "       [-0.0e+00, -0.0e+00,  0.0e+00, -7.2e-10, -3.2e-10,  1.0e+00],\n",
       "       [-0.0e+00, -0.0e+00,  0.0e+00, -7.2e-10, -2.4e-10,  1.0e+00],\n",
       "       ...,\n",
       "       [ 0.0e+00,  0.0e+00,  0.0e+00,  7.2e-10,  2.4e-10,  1.0e+00],\n",
       "       [ 0.0e+00,  0.0e+00,  0.0e+00,  7.2e-10,  3.2e-10,  1.0e+00],\n",
       "       [ 0.0e+00,  0.0e+00,  0.0e+00,  8.0e-10,  0.0e+00,  1.0e+00]])"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "super_rays[0,:,0,0:6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "ff8a27fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.00557354, 0.0097539 , 0.01172143]),\n",
       " array([0.00558912, 0.00978306, 0.01175699]))"
      ]
     },
     "execution_count": 167,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# look at that, the last point's different\n",
    "# it looks like it'd be pointed closer to the right edge of the screen than the other ones\n",
    "\n",
    "# let's reduce the objective distance and see if the discrepancy disappears\n",
    "\n",
    "sm.gaps[0].thi=1e9\n",
    "opm.update_model()\n",
    "\n",
    "spot_rms(sm), spot_rms2(sm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "id": "1ac36295",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.44231141, 0.4014741 , 0.38252552]),\n",
       " array([0.44778724, 0.40678566, 0.38775669]),\n",
       " array([0.44231141, 0.4014741 , 0.38252552]),\n",
       " array([0.44778724, 0.40678566, 0.38775669]))"
      ]
     },
     "execution_count": 168,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ugh let's see if it persists with different grid sizes\n",
    "\n",
    "sm.gaps[0].thi=1e10\n",
    "opm.update_model()\n",
    "\n",
    "spot_rms(sm, num_rays=41), spot_rms2(sm, num_rays=41), spot_rms(sm, num_rays=61), spot_rms2(sm, num_rays=61)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "705bc2ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2195380700.2, 65261865.4, 33.63956403550794)"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#  the difference is fairly small, let's see if the speedup is worth trading some accuracy\n",
    "\n",
    "import time\n",
    "ref_start = time.perf_counter_ns()\n",
    "for _ in range(100):\n",
    "    sm.trace_grid(dump_rays, 0, form='list', append_if_none=False)\n",
    "ref_end = time.perf_counter_ns()\n",
    "\n",
    "fast_start = time.perf_counter_ns()\n",
    "for _ in range(100):\n",
    "    super_trace_grid(sm, 0)\n",
    "fast_end = time.perf_counter_ns()\n",
    "\n",
    "ref_elapsed = (ref_end - ref_start)/10.\n",
    "fast_elapsed = (fast_end - fast_start)/10.\n",
    "\n",
    "ref_elapsed, fast_elapsed, ref_elapsed/fast_elapsed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aacb52a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ~30x speedup is probably worth a <1% error, let's just go with it"
   ]
  }
 ],
 "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
}
