{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dark Energy Spectroscopic Instrument"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some calculations to assist with building the DESI model from an existing ZEMAX model and other sources.\n",
    "\n",
    "You can safely ignore this if you just want to use the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.793017Z",
     "start_time": "2019-08-31T19:38:49.540574Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import yaml\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the DESI model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.863606Z",
     "start_time": "2019-08-31T19:38:49.794280Z"
    }
   },
   "outputs": [],
   "source": [
    "fiducial_telescope = batoid.Optic.fromYaml(\"DESI.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Corrector Internal Baffles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup YAML to preserve dictionary order and trunctate distances (in meters) to 5 digits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.866951Z",
     "start_time": "2019-08-31T19:38:49.865088Z"
    }
   },
   "outputs": [],
   "source": [
    "import collections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.872356Z",
     "start_time": "2019-08-31T19:38:49.868689Z"
    }
   },
   "outputs": [],
   "source": [
    "def dict_representer(dumper, data):\n",
    "    return dumper.represent_dict(data.items())\n",
    "\n",
    "yaml.Dumper.add_representer(collections.OrderedDict, dict_representer)\n",
    "\n",
    "def float_representer(dumper, value):\n",
    "    return dumper.represent_scalar(u'tag:yaml.org,2002:float', f'{value:.5f}')\n",
    "\n",
    "yaml.Dumper.add_representer(float, float_representer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the corrector internal baffle apertures, from DESI-4103-v1.  These have been checked against DESI-4037-v6, with the extra baffle between ADC1 and ADC2 added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.877225Z",
     "start_time": "2019-08-31T19:38:49.873658Z"
    }
   },
   "outputs": [],
   "source": [
    "# baffle z-coordinates relative to FP in mm from DESI-4103-v1, checked\n",
    "# against DESI-4037-v6 (and with extra ADC baffle added).\n",
    "ZBAFFLE = np.array([\n",
    "    2302.91, 2230.29, 1916.86, 1823.57, 1617.37, 1586.76, 1457.88, 1349.45, 1314.68,\n",
    "    1232.06, 899.67, 862.08, 568.81, 483.84, 415.22])    \n",
    "# baffle radii in mm from DESI-4103-v1, checked\n",
    "# against DESI-4037-v6 (and with extra ADC baffle added).\n",
    "RBAFFLE = np.array([\n",
    "    558.80, 544.00, 447.75, 417.00, 376.00, 376.00, 378.00, 378.00, 395.00,\n",
    "    403.00, 448.80, 453.70, 492.00, 501.00, 496.00])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate batoid Baffle surfaces for the corrector.  These are mechanically planar, but that would put their (planar) center inside a lens, breaking the sequential tracing model.  We fix this by use spherical baffle surfaces that have the same apertures. This code was originally used to read a batoid model without baffles, but also works if the baffles are already added."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.905965Z",
     "start_time": "2019-08-31T19:38:49.878411Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def baffles(nindent=10):\n",
    "    indent = ' ' * nindent\n",
    "    # Measure z from C1 front face in m.\n",
    "    zbaffle = 1e-3 * (2425.007 - ZBAFFLE)\n",
    "    # Convert r from mm to m.\n",
    "    rbaffle = 1e-3 * RBAFFLE\n",
    "    # By default, all baffles are planar.\n",
    "    nbaffles = len(zbaffle)\n",
    "    baffles = []\n",
    "    for i in range(nbaffles):\n",
    "        baffle = collections.OrderedDict()\n",
    "        baffle['type'] = 'Baffle'\n",
    "        baffle['name'] = f'B{i+1}'\n",
    "        baffle['coordSys'] = {'z': float(zbaffle[i])}\n",
    "        baffle['surface'] = {'type': 'Plane'}\n",
    "        baffle['obscuration'] = {'type': 'ClearCircle', 'radius': float(rbaffle[i])}\n",
    "        baffles.append(baffle)\n",
    "    # Loop over corrector lenses.\n",
    "    corrector = fiducial_telescope['DESI.Hexapod.Corrector']\n",
    "    lenses = 'C1', 'C2', 'ADC1rotator.ADC1', 'ADC2rotator.ADC2', 'C3', 'C4'\n",
    "    for lens in lenses:\n",
    "        obj = corrector['Corrector.' + lens]\n",
    "        assert isinstance(obj, batoid.optic.Lens)\n",
    "        front, back = obj.items[0], obj.items[1]\n",
    "        fTransform = batoid.CoordTransform(front.coordSys, corrector.coordSys)\n",
    "        bTransform = batoid.CoordTransform(back.coordSys, corrector.coordSys)\n",
    "        _, _, zfront = fTransform.applyForwardArray(0, 0, 0)\n",
    "        _, _, zback = bTransform.applyForwardArray(0, 0, 0)\n",
    "        # Find any baffles \"inside\" this lens.\n",
    "        inside = (zbaffle >= zfront) & (zbaffle <= zback)\n",
    "        if not any(inside):\n",
    "            continue\n",
    "        inside = np.where(inside)[0]\n",
    "        for k in inside:\n",
    "            baffle = baffles[k]\n",
    "            r = rbaffle[k]\n",
    "            # Calculate sag at (x,y)=(0,r) to avoid effect of ADC rotation about y.\n",
    "            sagf, sagb = front.surface.sag(0, r), back.surface.sag(0, r)\n",
    "            _, _, zf = fTransform.applyForwardArray(0, r, sagf)\n",
    "            _, _, zb = bTransform.applyForwardArray(0, r, sagb)\n",
    "            if zf > zbaffle[k]:\n",
    "                print(f'{indent}# Move B{k+1} in front of {obj.name} and make spherical to keep model sequential.')\n",
    "                assert isinstance(front.surface, batoid.Sphere)\n",
    "                baffle['surface'] = {'type': 'Sphere', 'R': front.surface.R}\n",
    "                baffle['coordSys']['z'] = float(zfront - (zf - zbaffle[k]))\n",
    "            elif zbaffle[k] > zb:\n",
    "                print(f'{indent}# Move B{k+1} behind {obj.name} and make spherical to keep model sequential.')\n",
    "                assert isinstance(back.surface, batoid.Sphere)\n",
    "                baffle['surface'] = {'type': 'Sphere', 'R': back.surface.R}\n",
    "                baffle['coordSys']['z'] = float(zback + (zbaffle[k] - zb))\n",
    "            else:\n",
    "                print(f'Cannot find a solution for B{k+1} inside {obj.name}!')\n",
    "\n",
    "    lines = yaml.dump(baffles)\n",
    "    for line in lines.split('\\n'):\n",
    "        print(indent + line)\n",
    "\n",
    "baffles()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Validate that the baffle edges in the final model have the correct apertures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:49.911187Z",
     "start_time": "2019-08-31T19:38:49.906961Z"
    }
   },
   "outputs": [],
   "source": [
    "def validate_baffles():\n",
    "    corrector = fiducial_telescope['DESI.Hexapod.Corrector']\n",
    "    for i in range(len(ZBAFFLE)):\n",
    "        baffle = corrector[f'Corrector.B{i+1}']\n",
    "        # Calculate surface z at origin in corrector coordinate system.\n",
    "        _, _, z = batoid.CoordTransform(baffle.coordSys, corrector.coordSys).applyForwardArray(0, 0, 0)\n",
    "        # Calculate surface z at (r,0) in corrector coordinate system.\n",
    "        sag = baffle.surface.sag(1e-3 * RBAFFLE[i], 0)\n",
    "        z += sag\n",
    "        # Measure from FP in mm.\n",
    "        z = np.round(2425.007 - 1e3 * z, 2)\n",
    "        assert z == ZBAFFLE[i], baffle.name\n",
    "        \n",
    "validate_baffles()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Corrector Cage and Spider"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate simplified vane coordinates using parameters from DESI-4110-v1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:50.046236Z",
     "start_time": "2019-08-31T19:38:49.912579Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def spider(dmin=1762, dmax=4940.3, ns_angle=77, widths=[28.5, 28.5, 60., 19.1],\n",
    "           wart_r=958, wart_dth=6, wart_w=300):\n",
    "    # Vane order is [NE, SE, SW, NW], with N along -y and E along +x.\n",
    "    fig, ax = plt.subplots(figsize=(10, 10))\n",
    "    ax.add_artist(plt.Circle((0, 0), 0.5 * dmax, color='yellow'))\n",
    "    ax.add_artist(plt.Circle((0, 0), 0.5 * dmin, color='gray'))\n",
    "    ax.set_xlim(-0.5 * dmax, 0.5 * dmax)\n",
    "    ax.set_ylim(-0.5 * dmax, 0.5 * dmax)\n",
    "    \n",
    "    # Place outer vertices equally along the outer ring at NE, SE, SW, NW.\n",
    "    xymax = 0.5 * dmax * np.array([[1, -1], [1, 1], [-1, 1], [-1, -1]]) / np.sqrt(2)\n",
    "    # Calculate inner vertices so that the planes of the NE and NW vanes intersect\n",
    "    # with an angle of ns_angle (same for the SE and SW planes).\n",
    "    angle = np.deg2rad(ns_angle)\n",
    "    x = xymax[1, 0]\n",
    "    dx = xymax[1, 1] * np.tan(0.5 * angle)\n",
    "    xymin = np.array([[x - dx, 0], [x - dx, 0], [-x+dx, 0], [-x+dx, 0]])\n",
    "    for i in range(4):\n",
    "        plt.plot([xymin[i,0], xymax[i,0]], [xymin[i,1], xymax[i,1]], '-', lw=0.1 * widths[i])\n",
    "\n",
    "    # Calculate batoid rectangle params for the vanes.\n",
    "    xy0 = 0.5 * (xymin + xymax)\n",
    "    heights = np.sqrt(np.sum((xymax - xymin) ** 2, axis=1))\n",
    "\n",
    "    # Calculate wart rectangle coords.\n",
    "    wart_h = 2 * (wart_r - 0.5 * dmin)\n",
    "    wart_dth = np.deg2rad(wart_dth)\n",
    "    wart_xy = 0.5 * dmin * np.array([-np.sin(wart_dth), np.cos(wart_dth)])\n",
    "    plt.plot(*wart_xy, 'rx', ms=25)\n",
    "    # Print batoid config.\n",
    "    indent = ' ' * 10\n",
    "    print(f'{indent}-\\n{indent}  type: ClearAnnulus')\n",
    "    print(f'{indent}  inner: {np.round(0.5e-3 * dmin, 5)}')\n",
    "    print(f'{indent}  outer: {np.round(0.5e-3 * dmax, 5)}')\n",
    "    for i in range(4):\n",
    "        print(f'{indent}-\\n{indent}  type: ObscRectangle')\n",
    "        print(f'{indent}  x: {np.round(1e-3 * xy0[i, 0], 5)}')\n",
    "        print(f'{indent}  y: {np.round(1e-3 * xy0[i, 1], 5)}')\n",
    "        print(f'{indent}  width: {np.round(1e-3 * widths[i], 5)}')\n",
    "        print(f'{indent}  height: {np.round(1e-3 * heights[i], 5)}')\n",
    "        dx, dy = xymax[i] - xymin[i]\n",
    "        angle = np.arctan2(-dx, dy)\n",
    "        print(f'{indent}  theta: {np.round(angle, 5)}')\n",
    "    print(f'-\\n  type: ObscRectangle')\n",
    "    print(f'  x: {np.round(1e-3 * wart_xy[0], 5)}')\n",
    "    print(f'  y: {np.round(1e-3 * wart_xy[1], 5)}')\n",
    "    print(f'  width: {np.round(1e-3 * wart_w, 5)}')\n",
    "    print(f'  height: {np.round(1e-3 * wart_h, 5)}')\n",
    "    print(f'  theta: {np.round(wart_dth, 5)}')\n",
    "    \n",
    "spider()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot \"User Aperture Data\" from the ZEMAX \"spider\" surface 6, as cross check:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-08-31T19:38:50.278622Z",
     "start_time": "2019-08-31T19:38:50.048108Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_obs():\n",
    "    wart1 = np.array([\n",
    "        [ -233.22959, 783.94254],\n",
    "        [-249.32698, 937.09892],\n",
    "        [49.02959, 968.45746],\n",
    "        [ 65.126976, 815.30108],\n",
    "        [ -233.22959, 783.94254],\n",
    "    ])\n",
    "    wart2 = np.array([\n",
    "        [-233.22959, 783.94254],\n",
    "        [ -249.32698, 937.09892],\n",
    "        [49.029593, 968.45746],\n",
    "        [65.126976, 815.30108],\n",
    "        [-233.22959, 783.94254],\n",
    "    ])\n",
    "    vane1 = np.array([\n",
    "        [363.96554,-8.8485008],\n",
    "        [341.66121, 8.8931664],\n",
    "        [1713.4345, 1733.4485],\n",
    "        [1735.7388, 1715.7068],\n",
    "        [363.96554,-8.8485008],\n",
    "    ])\n",
    "    vane2 = np.array([\n",
    "        [-1748.0649, 1705.9022],\n",
    "        [ -1701.1084, 1743.2531],\n",
    "        [ -329.33513, 18.697772],\n",
    "        [ -376.29162, -18.653106],\n",
    "        [-1748.0649, 1705.9022],\n",
    "    ])\n",
    "    vane3 = np.array([\n",
    "        [ -1717.1127, -1730.5227],\n",
    "        [ -1732.0605, -1718.6327],\n",
    "        [ -360.28728, 5.922682],\n",
    "        [-345.33947, -5.9673476],\n",
    "        [ -1717.1127, -1730.5227],\n",
    "    ])\n",
    "    vane4 = np.array([\n",
    "        [ 341.66121, -8.8931664],\n",
    "        [363.96554, 8.8485008],\n",
    "        [1735.7388, -1715.7068],\n",
    "        [1713.4345, -1733.4485],\n",
    "        [ 341.66121, -8.8931664],\n",
    "    ])\n",
    "    extra = np.array([\n",
    "        [           2470 ,               0 ],\n",
    "        [      2422.5396 ,       -481.8731 ],\n",
    "        [      2281.9824 ,      -945.22808 ],\n",
    "        [      2053.7299 ,      -1372.2585 ],\n",
    "        [      1746.5537 ,      -1746.5537 ],\n",
    "        [      1372.2585 ,      -2053.7299 ],\n",
    "        [      945.22808 ,      -2281.9824 ],\n",
    "        [       481.8731 ,      -2422.5396 ],\n",
    "        [  3.0248776e-13 ,           -2470 ],\n",
    "        [      -481.8731 ,      -2422.5396 ],\n",
    "        [     -945.22808 ,      -2281.9824 ],\n",
    "        [     -1372.2585 ,      -2053.7299 ],\n",
    "        [     -1746.5537 ,      -1746.5537 ],\n",
    "        [     -2053.7299 ,      -1372.2585 ],\n",
    "        [     -2281.9824 ,      -945.22808 ],\n",
    "        [     -2422.5396 ,       -481.8731 ],\n",
    "        [          -2470 ,   2.9882133e-12 ],\n",
    "        [     -2422.5396 ,        481.8731 ],\n",
    "        [     -2281.9824 ,       945.22808 ],\n",
    "        [     -2053.7299 ,       1372.2585 ],\n",
    "        [     -1746.5537 ,       1746.5537 ],\n",
    "        [     -1372.2585 ,       2053.7299 ],\n",
    "        [     -945.22808 ,       2281.9824 ],\n",
    "        [      -481.8731 ,       2422.5396 ],\n",
    "        [  5.9764266e-12 ,            2470 ],\n",
    "        [       481.8731 ,       2422.5396 ],\n",
    "        [      945.22808 ,       2281.9824 ],\n",
    "        [      1372.2585 ,       2053.7299 ],\n",
    "        [      1746.5537 ,       1746.5537 ],\n",
    "        [      2053.7299 ,       1372.2585 ],\n",
    "        [      2281.9824 ,       945.22808 ],\n",
    "        [      2422.5396 ,        481.8731 ],\n",
    "        [           2470 ,  -1.0364028e-11 ],\n",
    "        [           2724 ,               0 ],\n",
    "        [      2671.6591 ,      -531.42604 ],\n",
    "        [      2516.6478 ,      -1042.4297 ],\n",
    "        [      2264.9232 ,      -1513.3733 ],\n",
    "        [      1926.1589 ,      -1926.1589 ],\n",
    "        [      1513.3733 ,      -2264.9232 ],\n",
    "        [      1042.4297 ,      -2516.6478 ],\n",
    "        [      531.42604 ,      -2671.6591 ],\n",
    "        [  3.3359379e-13 ,           -2724 ],\n",
    "        [     -531.42604 ,      -2671.6591 ],\n",
    "        [     -1042.4297 ,      -2516.6478 ],\n",
    "        [     -1513.3733 ,      -2264.9232 ],\n",
    "        [     -1926.1589 ,      -1926.1589 ],\n",
    "        [     -2264.9232 ,      -1513.3733 ],\n",
    "        [     -2516.6478 ,      -1042.4297 ],\n",
    "        [     -2671.6591 ,      -531.42604 ],\n",
    "        [          -2724 ,   3.2955032e-12 ],\n",
    "        [     -2671.6591 ,       531.42604 ],\n",
    "        [     -2516.6478 ,       1042.4297 ],\n",
    "        [     -2264.9232 ,       1513.3733 ],\n",
    "        [     -1926.1589 ,       1926.1589 ],\n",
    "        [     -1513.3733 ,       2264.9232 ],\n",
    "        [     -1042.4297 ,       2516.6478 ],\n",
    "        [     -531.42604 ,       2671.6591 ],\n",
    "        [  6.5910065e-12 ,            2724 ],\n",
    "        [      531.42604 ,       2671.6591 ],\n",
    "        [      1042.4297 ,       2516.6478 ],\n",
    "        [      1513.3733 ,       2264.9232 ],\n",
    "        [      1926.1589 ,       1926.1589 ],\n",
    "        [      2264.9232 ,       1513.3733 ],\n",
    "        [      2516.6478 ,       1042.4297 ],\n",
    "        [      2671.6591 ,       531.42604 ],\n",
    "        [           2724 ,  -1.1429803e-11 ],\n",
    "        [           2470 ,               0 ],\n",
    "    ])\n",
    "    plt.figure(figsize=(20, 20))\n",
    "    plt.plot(*wart1.T)\n",
    "    plt.plot(*wart2.T)\n",
    "    plt.plot(*vane1.T)\n",
    "    plt.plot(*vane2.T)\n",
    "    plt.plot(*vane3.T)\n",
    "    plt.plot(*vane4.T)\n",
    "    plt.plot(*extra.T)\n",
    "    w = 1762./2.\n",
    "    plt.gca().add_artist(plt.Circle((0, 0), w, color='gray'))\n",
    "    plt.gca().set_aspect(1.)\n",
    "    \n",
    "plot_obs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.8-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}