{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:04:04.077691Z",
     "start_time": "2019-09-17T06:04:03.112109Z"
    }
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import os\n",
    "import yaml\n",
    "import numpy as np\n",
    "from scipy.interpolate import griddata\n",
    "from sklearn.neighbors import KDTree\n",
    "import matplotlib.pyplot as plt\n",
    "import galsim\n",
    "from tqdm import tqdm\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:04:04.101473Z",
     "start_time": "2019-09-17T06:04:04.079105Z"
    }
   },
   "outputs": [],
   "source": [
    "telescope = batoid.Optic.fromYaml(\"LSST_r.yaml\")\n",
    "# fiducial_telescope.clearObscuration()\n",
    "# Re-insert M1 outer edge obscuration though.\n",
    "# fiducial_telescope.itemDict['LSST.M1'].obscuration = batoid.ObscNegation(batoid.ObscCircle(8.360001/2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:04:04.922928Z",
     "start_time": "2019-09-17T06:04:04.919876Z"
    }
   },
   "outputs": [],
   "source": [
    "def simpleGridData(inputCoords, inputValue, outputGrid):\n",
    "    output = griddata(\n",
    "        inputCoords,\n",
    "        inputValue,\n",
    "        outputGrid,\n",
    "        method='cubic'\n",
    "    )\n",
    "    nearest = griddata(\n",
    "        inputCoords,\n",
    "        inputValue,\n",
    "        outputGrid,\n",
    "        method='nearest'\n",
    "    )\n",
    "    w = np.isfinite(output)\n",
    "    # Fill in cubic NaNs with nearest\n",
    "    output[~w] = nearest[~w]\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:04:09.408188Z",
     "start_time": "2019-09-17T06:04:09.403696Z"
    }
   },
   "outputs": [],
   "source": [
    "def betterGridData(inputCoords, inputValue, outputGrid, k=20):\n",
    "    output = griddata(inputCoords, inputValue, outputGrid, method='cubic')\n",
    "    tree = KDTree(inputCoords)\n",
    "    _, neighbors = tree.query(outputGrid, k=k)\n",
    "    toFix = np.nonzero(~np.isfinite(output))[0]\n",
    "    for idx in toFix:\n",
    "        neighbor = neighbors[idx]\n",
    "        basis = galsim.zernike.zernikeBasis(\n",
    "            10, \n",
    "            inputCoords[neighbor, 0], \n",
    "            inputCoords[neighbor, 1]\n",
    "        )\n",
    "        coefs, _, _, _ = np.linalg.lstsq(\n",
    "            basis.T,\n",
    "            inputValue[neighbor],\n",
    "            rcond=None\n",
    "        )\n",
    "        value = galsim.zernike.Zernike(coefs).evalCartesian(\n",
    "            *outputGrid[idx]\n",
    "        )\n",
    "        output[idx] = value\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:04:14.434132Z",
     "start_time": "2019-09-17T06:04:14.039989Z"
    }
   },
   "outputs": [],
   "source": [
    "M1M3dir = \"/Users/josh/src/M1M3_ML/data/\"\n",
    "M1M3fn = os.path.join(M1M3dir, \"M1M3_1um_156_grid.txt\")\n",
    "M1M3data = np.loadtxt(M1M3fn).T\n",
    "\n",
    "w1 = np.where(M1M3data[0] == 1)[0]\n",
    "w3 = np.where(M1M3data[0] == 3)[0]\n",
    "\n",
    "M1x = M1M3data[1][w1]\n",
    "M1y = M1M3data[2][w1]\n",
    "\n",
    "M3x = M1M3data[1][w3]\n",
    "M3y = M1M3data[2][w3]\n",
    "\n",
    "M1modes = M1M3data[3:23, w1]  # First 20 mirror modes\n",
    "M3modes = M1M3data[3:23, w3]\n",
    "\n",
    "# Modes are currently in the direction normal to surface.  \n",
    "# We want perturbations in the direction of the optic axis.\n",
    "# So divide by the z-component of the surface normal.\n",
    "M1modes /= telescope['M1'].surface.normal(M1x, M1y)[:,2]\n",
    "M3modes /= telescope['M3'].surface.normal(M3x, M3y)[:,2]\n",
    "\n",
    "M1points = np.column_stack([M1x, M1y])\n",
    "M3points = np.column_stack([M3x, M3y])\n",
    "\n",
    "xgrid = np.linspace(-4.2, 4.2, 100)\n",
    "xgrid, ygrid = np.meshgrid(xgrid, xgrid)\n",
    "xgrid = xgrid.ravel()\n",
    "ygrid = ygrid.ravel()\n",
    "rgrid = np.hypot(xgrid, ygrid)\n",
    "\n",
    "centerX, centerY = np.array(np.meshgrid(\n",
    "    np.linspace(-0.15, 0.15, 10), \n",
    "    np.linspace(-0.15, 0.15, 10)\n",
    "))\n",
    "centerPoints = np.array(np.stack([centerX.ravel(), centerY.ravel()])).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:06:58.906118Z",
     "start_time": "2019-09-17T06:04:18.872375Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "M1ModeBetterGrid = []\n",
    "M3ModeBetterGrid = []\n",
    "for i in tqdm(range(20)):\n",
    "    M1cubic = betterGridData(\n",
    "        np.vstack([M1points, centerPoints]),\n",
    "        np.hstack([M1modes[i], np.zeros(len(centerPoints))]),\n",
    "        np.column_stack([xgrid, ygrid])\n",
    "    )\n",
    "    M1ModeBetterGrid.append(M1cubic)\n",
    "    \n",
    "    M3cubic = betterGridData(\n",
    "        np.vstack([M3points, centerPoints]),\n",
    "        np.hstack([M3modes[i], np.zeros(len(centerPoints))]),\n",
    "        np.column_stack([xgrid, ygrid])\n",
    "    )\n",
    "    M3ModeBetterGrid.append(M3cubic)\n",
    "\n",
    "M1ModeBetterGrid = np.array(M1ModeBetterGrid)\n",
    "M3ModeBetterGrid = np.array(M3ModeBetterGrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:00.326172Z",
     "start_time": "2019-09-17T06:06:58.907658Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "M1ModeSimpleGrid = []\n",
    "M3ModeSimpleGrid = []\n",
    "for i in tqdm(range(20)):\n",
    "    M1cubic = simpleGridData(\n",
    "        np.vstack([M1points, centerPoints]),\n",
    "        np.hstack([M1modes[i], np.zeros(len(centerPoints))]),\n",
    "        np.column_stack([xgrid, ygrid])\n",
    "    )\n",
    "    M1ModeSimpleGrid.append(M1cubic)\n",
    "    \n",
    "    M3cubic = simpleGridData(\n",
    "        np.vstack([M3points, centerPoints]),\n",
    "        np.hstack([M3modes[i], np.zeros(len(centerPoints))]),\n",
    "        np.column_stack([xgrid, ygrid])\n",
    "    )\n",
    "    M3ModeSimpleGrid.append(M3cubic)\n",
    "\n",
    "M1ModeSimpleGrid = np.array(M1ModeSimpleGrid)\n",
    "M3ModeSimpleGrid = np.array(M3ModeSimpleGrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:05.185515Z",
     "start_time": "2019-09-17T06:07:00.327975Z"
    }
   },
   "outputs": [],
   "source": [
    "r = np.hypot(xgrid, ygrid)\n",
    "wM1 = (r<4.18) & (r > 2.558)\n",
    "wM3 = (r<2.508) & (r > 0.55)\n",
    "vmin, vmax = -2, 2\n",
    "\n",
    "fig, axes = plt.subplots(nrows=4, ncols=5, figsize=(12, 8))\n",
    "for j, ax in enumerate(axes.ravel()):\n",
    "    bend = [0]*j+[1]+[0]*(19-j)\n",
    "    bend = np.array(bend)*1e6\n",
    "    ax.scatter(\n",
    "        xgrid, \n",
    "        ygrid, \n",
    "        c=np.ma.masked_array(np.dot(bend, M3ModeSimpleGrid), mask=~wM3),\n",
    "        s=5,\n",
    "        vmin=vmin, vmax=vmax\n",
    "    )\n",
    "    ax.scatter(\n",
    "        xgrid, \n",
    "        ygrid, \n",
    "        c=np.ma.masked_array(np.dot(bend, M1ModeSimpleGrid), mask=~wM1),\n",
    "        s=5,\n",
    "        vmin=vmin, vmax=vmax\n",
    "    )\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.286021Z",
     "start_time": "2019-09-17T06:07:05.186618Z"
    }
   },
   "outputs": [],
   "source": [
    "r = np.hypot(xgrid, ygrid)\n",
    "wM1 = (r<4.18) & (r > 2.558)\n",
    "wM3 = (r<2.508) & (r > 0.55)\n",
    "vmin, vmax = -0.1, 0.1\n",
    "\n",
    "fig, axes = plt.subplots(nrows=4, ncols=5, figsize=(12, 8))\n",
    "for j, ax in enumerate(axes.ravel()):\n",
    "    bend = [0]*j+[1]+[0]*(19-j)\n",
    "    bend = np.array(bend)*1e6\n",
    "    ax.scatter(\n",
    "        xgrid, \n",
    "        ygrid, \n",
    "        c=np.ma.masked_array(np.dot(bend, M3ModeSimpleGrid-M3ModeBetterGrid), mask=~wM3),\n",
    "        s=5,\n",
    "        cmap='seismic',\n",
    "        vmin=vmin, vmax=vmax\n",
    "    )\n",
    "    ax.scatter(\n",
    "        xgrid, \n",
    "        ygrid, \n",
    "        c=np.ma.masked_array(np.dot(bend, M1ModeSimpleGrid-M1ModeBetterGrid), mask=~wM1),\n",
    "        s=5,\n",
    "        cmap='seismic',\n",
    "        vmin=vmin, vmax=vmax\n",
    "    )\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.294026Z",
     "start_time": "2019-09-17T06:07:10.287255Z"
    }
   },
   "outputs": [],
   "source": [
    "def rot(thx, thy):\n",
    "    return np.dot(batoid.RotX(thx), batoid.RotY(thy))\n",
    "\n",
    "def perturbed_telescope(M2Shift, M2Tilt, cameraShift, cameraTilt, M1M3Bend, M2Bend):\n",
    "    modes=slice(0,20)\n",
    "    M1grid = np.dot(M1M3Bend[modes], M1ModeBetterGrid[modes])\n",
    "    M3grid = np.dot(M1M3Bend[modes], M3ModeBetterGrid[modes])\n",
    "    M1_delta = batoid.Bicubic(xgrid[0:100], ygrid[::100], M1grid.reshape(100, 100))\n",
    "    M3_delta = batoid.Bicubic(xgrid[0:100], ygrid[::100], M3grid.reshape(100, 100))\n",
    "    perturbedM1 = batoid.Sum([telescope['M1'].surface, M1_delta])\n",
    "    perturbedM3 = batoid.Sum([telescope['M3'].surface, M3_delta])\n",
    "\n",
    "    return (telescope\n",
    "            .withGloballyShiftedOptic('M2', M2Shift)\n",
    "            .withLocallyRotatedOptic('M2', rot(*M2Tilt))\n",
    "            .withGloballyShiftedOptic('LSSTCamera', cameraShift)\n",
    "            .withLocallyRotatedOptic('LSSTCamera', rot(*cameraTilt))\n",
    "            .withSurface('M1', perturbedM1)\n",
    "            .withSurface('M3', perturbedM3)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.300031Z",
     "start_time": "2019-09-17T06:07:10.295371Z"
    }
   },
   "outputs": [],
   "source": [
    "def tracePerturbed(thx, thy, *args, **kwargs):\n",
    "    ptelescope = perturbed_telescope(*args, **kwargs)\n",
    "    rays = batoid.RayVector.asPolar(\n",
    "        optic=telescope, wavelength=625e-9,\n",
    "        theta_x=np.deg2rad(thx), theta_y=np.deg2rad(thy),\n",
    "        nrad=100, naz=300,\n",
    "    )\n",
    "    ptelescope.trace(rays)\n",
    "    w = ~rays.vignetted\n",
    "\n",
    "    plt.scatter(\n",
    "        rays.x[w]-np.mean(rays.x[w]), \n",
    "        rays.y[w]-np.mean(rays.y[w]), \n",
    "        s=0.1, alpha=0.3\n",
    "    )\n",
    "    plt.xlim(1*np.array([-10e-6, 10e-6]))\n",
    "    plt.ylim(1*np.array([-10e-6, 10e-6]))\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.496566Z",
     "start_time": "2019-09-17T06:07:10.301741Z"
    }
   },
   "outputs": [],
   "source": [
    "ibend = 4\n",
    "bend = [0]*(ibend)+[1]+[0]*(19-ibend)\n",
    "tracePerturbed(\n",
    "    0.0, 0.0, \n",
    "    M2Shift=(0.0, 0.0, 0.0),\n",
    "    M2Tilt=(0.0, 0.0),\n",
    "    cameraShift=(0.0, 0.0, 0.0),\n",
    "    cameraTilt=(0.0, 0.0),\n",
    "    M1M3Bend=np.array(bend)*1.4e-1,\n",
    "    M2Bend=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.503380Z",
     "start_time": "2019-09-17T06:07:10.498630Z"
    }
   },
   "outputs": [],
   "source": [
    "def zernikeField(thxs, thys, *args, **kwargs):\n",
    "    ptelescope = perturbed_telescope(*args, **kwargs)\n",
    "\n",
    "    jmax = 28\n",
    "    zs = np.zeros([len(thxs), jmax+1], dtype=float)\n",
    "    for i, (thx, thy) in enumerate(zip(thxs, thys)):\n",
    "        if np.hypot(thx, thy) > 1.76:\n",
    "            zs[i] = float('nan')\n",
    "            continue\n",
    "        zs[i] = batoid.analysis.zernikeGQ(\n",
    "            ptelescope, np.deg2rad(thx), np.deg2rad(thy), 625e-9, \n",
    "            jmax=jmax, rings=10,\n",
    "            reference='chief'\n",
    "        )\n",
    "    return zs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:10.508936Z",
     "start_time": "2019-09-17T06:07:10.505167Z"
    }
   },
   "outputs": [],
   "source": [
    "thx = np.linspace(-1.75, 1.75, 30)\n",
    "thxs, thys = np.meshgrid(thx, thx)\n",
    "thxs = thxs.ravel()\n",
    "thys = thys.ravel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:07:13.143687Z",
     "start_time": "2019-09-17T06:07:10.510528Z"
    }
   },
   "outputs": [],
   "source": [
    "z0s = zernikeField(\n",
    "    thxs, thys,\n",
    "    M2Shift=(0.0, 0.0, 0.0),\n",
    "    M2Tilt=(0.0, 0.0),\n",
    "    cameraShift=(0.0, 0.0, 0.0),\n",
    "    cameraTilt=(0.0, 0.0),\n",
    "    M1M3Bend=np.array([0]*20),\n",
    "    M2Bend=None    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:08:07.312620Z",
     "start_time": "2019-09-17T06:07:13.144537Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "zss = []\n",
    "for ibend in tqdm(range(20)):\n",
    "    bend = [0]*ibend+[1]+[0]*(19-ibend)\n",
    "    zss.append(zernikeField(\n",
    "        thxs, thys,    \n",
    "        M2Shift=(0.0, 0.0, 0.0),\n",
    "        M2Tilt=(0.0, 0.0),\n",
    "        cameraShift=(0.0, 0.0, 0.0),\n",
    "        cameraTilt=(0.0, 0.0),\n",
    "        M1M3Bend=np.array(bend)*1e-1,\n",
    "        M2Bend=None    \n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:08:55.614035Z",
     "start_time": "2019-09-17T06:08:07.313969Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for j in range(20):\n",
    "    fig = plt.figure(figsize=(13, 8))\n",
    "    batoid.plotUtils.zernikePyramid(\n",
    "        thxs, thys, zss[j][:, 4:].T-z0s[:,4:].T, \n",
    "        title=\"mode {}\".format(j+1), \n",
    "        vmin=-2, vmax=2,\n",
    "        fig=fig\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-09-17T06:09:43.120307Z",
     "start_time": "2019-09-17T06:08:55.615069Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for zj in range(4, 29):\n",
    "    fig, axes = plt.subplots(nrows=4, ncols=5, figsize=(12, 8))\n",
    "    for j, ax in enumerate(axes.ravel()):\n",
    "        ax.scatter(\n",
    "            thxs,\n",
    "            thys,\n",
    "            c=zss[j][:, zj] - z0s[:, zj],\n",
    "            s=5,\n",
    "            vmin=-0.5,\n",
    "            vmax=0.5,\n",
    "            cmap='Spectral_r'\n",
    "        )\n",
    "        ax.set_aspect('equal')\n",
    "        ax.set_title(f\"Mode {j+1}\")\n",
    "    fig.suptitle(\"Z{}\".format(zj))\n",
    "    fig.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "    plt.show()"
   ]
  },
  {
   "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
}