Rubin Diffraction Spikes

Diffraction Spikes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0fc96351-c77a-4cb5-94cf-c02a905568b5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-23T20:54:11.732482Z",
     "iopub.status.busy": "2022-06-23T20:54:11.730265Z",
     "iopub.status.idle": "2022-06-23T20:54:12.344896Z",
     "shell.execute_reply": "2022-06-23T20:54:12.344583Z",
     "shell.execute_reply.started": "2022-06-23T20:54:11.732332Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import batoid\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact, FloatSlider, IntText, Checkbox, FloatText\n",
    "from functools import lru_cache\n",
    "# %matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b1e79ab6-28c5-4783-8e68-58d8f15229b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-23T20:54:12.345824Z",
     "iopub.status.busy": "2022-06-23T20:54:12.345700Z",
     "iopub.status.idle": "2022-06-23T20:54:12.391561Z",
     "shell.execute_reply": "2022-06-23T20:54:12.391220Z",
     "shell.execute_reply.started": "2022-06-23T20:54:12.345814Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "telescope = batoid.Optic.fromYaml(\"LSST_r_baffles.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0ecf7486-fbc9-41dd-a31c-c2b7494bfbd8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-23T21:01:35.672282Z",
     "iopub.status.busy": "2022-06-23T21:01:35.671795Z",
     "iopub.status.idle": "2022-06-23T21:01:36.814788Z",
     "shell.execute_reply": "2022-06-23T21:01:36.814504Z",
     "shell.execute_reply.started": "2022-06-23T21:01:35.672240Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2f41384ea3e4c9499f0e23386fa7184",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='thx (deg)', max=2.0, min=-2.0, step=0.01), FloatSlid…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@lru_cache(10)\n",
    "def getwf(thx, thy, wavelength, nx):\n",
    "    return batoid.wavefront(\n",
    "        telescope,\n",
    "        np.deg2rad(thx),\n",
    "        np.deg2rad(thy),\n",
    "        wavelength*1e-9,\n",
    "        nx=nx\n",
    "    ).array.mask\n",
    "\n",
    "@interact(\n",
    "    theta_x=FloatSlider(value=0, min=-2, max=2, step=0.01, description=\"thx (deg)\"),\n",
    "    theta_y=FloatSlider(value=0, min=-2, max=2, step=0.01, description=\"thy (deg)\"),\n",
    "    wavelength=FloatSlider(value=620.0, min=300.0, max=1100.0, step=1.0, description=\"wavelength (nm)\"),\n",
    "    nx=IntText(value=2048, description=\"nx\"),\n",
    "    diff=Checkbox(value=False, description=\"diff\")\n",
    ")\n",
    "def pupil(theta_x, theta_y, wavelength, nx, diff):\n",
    "    mask = getwf(theta_x, theta_y, wavelength, nx)\n",
    "\n",
    "    if diff:\n",
    "        mask = mask | getwf(0, 0, 620, nx=nx)\n",
    "    \n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(mask)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b052b3c5-ec58-4877-a3e8-99295327471d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-06-23T21:01:59.640247Z",
     "iopub.status.busy": "2022-06-23T21:01:59.639343Z",
     "iopub.status.idle": "2022-06-23T21:02:01.749741Z",
     "shell.execute_reply": "2022-06-23T21:02:01.749361Z",
     "shell.execute_reply.started": "2022-06-23T21:01:59.640199Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51798a13c6c54ed0a1e9a967d3eaa655",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='thx (deg)', max=2.0, min=-2.0, step=0.01), FloatSlid…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@lru_cache(10)\n",
    "def getpsf(thx, thy, wavelength, rot, nx):\n",
    "    optic = telescope.withGlobalShift((0,0,0))  # make a copy\n",
    "    optic.stopSurface = optic.stopSurface.withLocalRotation(batoid.RotZ(np.deg2rad(rot)))\n",
    "    return batoid.fftPSF(\n",
    "        optic,\n",
    "        np.deg2rad(thx),\n",
    "        np.deg2rad(thy),\n",
    "        wavelength*1e-9,\n",
    "        nx=nx\n",
    "    )\n",
    "\n",
    "@interact(\n",
    "    theta_x=FloatSlider(value=0, min=-2, max=2, step=0.01, description=\"thx (deg)\"),\n",
    "    theta_y=FloatSlider(value=0, min=-2, max=2, step=0.01, description=\"thy (deg)\"),\n",
    "    wavelength=FloatSlider(value=620.0, min=300.0, max=1100.0, step=1.0, description=\"wavelength (nm)\"),\n",
    "    rot=FloatSlider(value=15.0, min=-360.0, max=360.0, step=15.0, stepdescription=\"rotation (deg)\"),\n",
    "    nx=IntText(value=2048, description=\"nx\"),\n",
    "    log=Checkbox(value=True, description=\"log scale\"),\n",
    "    vmin=FloatText(value=0),\n",
    "    vmax=FloatText(value=12),\n",
    ")\n",
    "\n",
    "def psf(theta_x, theta_y, wavelength, rot, nx, log, vmin, vmax):\n",
    "    lattice = getpsf(theta_x, theta_y, wavelength, rot, nx)\n",
    "    arr = lattice.array\n",
    "    extent0 = 0.5*np.sqrt(np.linalg.det(lattice.primitiveVectors))*np.r_[-1, 1]*arr.shape[0]  # meters\n",
    "    extent0 /= 10e-6  # pixels\n",
    "    extent0 *= 0.2  # arcsec\n",
    "\n",
    "    if log:\n",
    "        arr = np.log10(arr)\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(arr, cmap='inferno', vmin=vmin, vmax=vmax, extent=[*extent0, *extent0])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb592d47-6314-4507-80b2-626c16caf572",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}