{
 "cells": [
  {
   "cell_type": "raw",
   "id": "e12c7d65",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Image Simulation\n",
    "================\n",
    "This notebook demonstrates how to generate synthetic ESIS images\n",
    "using data from the Interface Region Imaging Spectrograph (IRIS) :cite:p:`DePontieu2014`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import named_arrays as na\n",
    "import astropy.units as u\n",
    "import astropy.time\n",
    "import astropy.visualization\n",
    "import optika\n",
    "import iris\n",
    "import esis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44398f271b4d84d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "astropy.visualization.quantity_support();"
   ]
  },
  {
   "cell_type": "raw",
   "id": "bff2e12d",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Download a sequence of 320-step dense rasters, despike, and crop to :math:`\\pm 100 \\text{km/s}`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb8f6532f39b7fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = iris.sg.SpectrographObservation.from_time_range(\n",
    "    time_start=astropy.time.Time(\"2021-09-23T14:00\"),\n",
    "    time_stop=astropy.time.Time(\"2021-09-23T14:01\"),\n",
    ")\n",
    "\n",
    "obs = obs.explicit\n",
    "\n",
    "bg = iris.sg.background.estimate(\n",
    "    obs=obs,\n",
    "    axis_time=obs.axis_time,\n",
    "    axis_wavelength=obs.axis_wavelength,\n",
    "    axis_detector_x=obs.axis_detector_x,\n",
    "    axis_detector_y=obs.axis_detector_y,\n",
    ")\n",
    "\n",
    "obs = obs - bg.outputs\n",
    "\n",
    "obs = na.despike(obs, axis=(obs.axis_wavelength, obs.axis_detector_y))\n",
    "\n",
    "velocity_thresh = 150 * u.km / u.s\n",
    "\n",
    "velocity_centers = obs.velocity_doppler.cell_centers().mean(obs.axis_time)\n",
    "\n",
    "index_lower = np.argmax(-velocity_thresh < velocity_centers)[obs.axis_wavelength]\n",
    "index_upper = np.argmax(velocity_thresh < velocity_centers)[obs.axis_wavelength]\n",
    "\n",
    "crop_wavelength = {obs.axis_wavelength: slice(index_lower.ndarray, index_upper.ndarray)}\n",
    "\n",
    "obs = obs[crop_wavelength]\n",
    "\n",
    "obs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd8bffc63e78ef7",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.to_jshtml(\n",
    "    velocity_min=obs.velocity_doppler.ndarray.min(),\n",
    "    velocity_max=obs.velocity_doppler.ndarray.max(),\n",
    "    fps=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4651db6c35abc5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fwhm_si_iv = (obs.wavelength_center + 0.188 * u.AA).to(u.km/u.s, equivalencies=u.doppler_radio(obs.wavelength_center))\n",
    "fwhm_si_iv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd9eac5a58718f5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength_o_v = 629.73 * u.AA\n",
    "fwhm_o_v = (wavelength_o_v + 0.129 * u.AA).to(u.km/u.s, equivalencies=u.doppler_radio(wavelength_o_v))\n",
    "fwhm_o_v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d80d3a19e1e8604",
   "metadata": {},
   "outputs": [],
   "source": [
    "width_ratio = fwhm_o_v / fwhm_si_iv\n",
    "width_ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae5bd7026f1109a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "velocity = width_ratio * obs.velocity_doppler\n",
    "obs.inputs.wavelength = velocity.to(u.AA, equivalencies=u.doppler_radio(wavelength_o_v))\n",
    "obs.wavelength_center = wavelength_o_v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5bb1b888c0408eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "veranazza_radiance = 690.80 * u.erg / u.cm ** 2 / u.sr / u.s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c13f7dfd514972a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dw = np.diff(obs.inputs.wavelength, axis=obs.axis_wavelength)\n",
    "avg = np.nanmean((dw * obs.outputs).sum(obs.axis_wavelength)[{obs.axis_time: 0}]).ndarray\n",
    "avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "947586c3dbd7b5cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.outputs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a8500c4bcb7a685",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.inputs.wavelength.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "606d9493e4f2f504",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = obs / avg * veranazza_radiance\n",
    "obs = obs.to(u.erg / u.cm ** 2 / u.sr / u.s / u.AA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48c759853a2cde3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.outputs = np.nan_to_num(obs.outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b7c9add9306da69",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.nanmean((dw * obs.outputs).sum(obs.axis_wavelength)).ndarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93bf86680a2cc91",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs.inputs.position.x -= obs.inputs.position.x.mean((obs.axis_detector_x, obs.axis_detector_y))\n",
    "obs.inputs.position.y -= obs.inputs.position.y.mean((obs.axis_detector_x, obs.axis_detector_y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de8dd18e5407a9bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument = esis.flights.f1.optics.design_single(num_distribution=0)\n",
    "instrument.primary_mirror.translation = na.nominal(instrument.primary_mirror.translation)\n",
    "instrument.field_stop.translation = na.nominal(instrument.field_stop.translation)\n",
    "instrument.grating.sag.radius = na.nominal(instrument.grating.sag.radius)\n",
    "instrument.grating.rulings.spacing.coefficients[0] = na.nominal(instrument.grating.rulings.spacing.coefficients[0])\n",
    "instrument.grating.rulings.spacing.coefficients[1] = na.nominal(instrument.grating.rulings.spacing.coefficients[1])\n",
    "instrument.grating.rulings.spacing.coefficients[2] = na.nominal(instrument.grating.rulings.spacing.coefficients[2])\n",
    "instrument.grating.rulings.depth = na.nominal(instrument.grating.rulings.depth)\n",
    "instrument.grating.rulings.ratio_duty = na.nominal(instrument.grating.rulings.ratio_duty)\n",
    "instrument.grating.translation = na.nominal(instrument.grating.translation)\n",
    "instrument.grating.roll = na.nominal(instrument.grating.roll)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc7ce73bda926291",
   "metadata": {},
   "outputs": [],
   "source": [
    "system = instrument.system\n",
    "# system.sensor.material = optika.sensors.IdealImagingSensorMaterial()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8840f76231a34571",
   "metadata": {},
   "outputs": [],
   "source": [
    "pupil = na.Cartesian2dVectorLinearSpace(\n",
    "    start=-1,\n",
    "    stop=1,\n",
    "    axis=na.Cartesian2dVectorArray(\"pupil_x\", \"pupil_y\"),\n",
    "    num=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de8e0d4bac16f9b5",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "image = system.image(\n",
    "    obs,\n",
    "    pupil=pupil,\n",
    "    axis_wavelength=obs.axis_wavelength,\n",
    "    axis_field=(obs.axis_detector_x, obs.axis_detector_y),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d4ea9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = obs.copy()\n",
    "factor = 10000\n",
    "r.outputs = factor * np.nan_to_num(r.outputs)\n",
    "r.inputs.wavelength = (r.velocity_doppler/factor).to(u.AA, equivalencies=u.doppler_radio(wavelength_o_v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "809400cf26ada852",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "image_undispersed = system.image(\n",
    "    r,\n",
    "    pupil=pupil,\n",
    "    axis_wavelength=obs.axis_wavelength,\n",
    "    axis_field=(obs.axis_detector_x, obs.axis_detector_y),\n",
    "    noise=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "357c3dbd7ebb8925",
   "metadata": {},
   "outputs": [],
   "source": [
    "crop = dict(\n",
    "    detector_x=slice(1024 + 256, 1024 + 512 + 256),\n",
    "    detector_y=slice(256, 512 + 256),\n",
    ")\n",
    "image = image[crop]\n",
    "image_undispersed = image_undispersed[crop]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a476333df109263",
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked = na.stack(\n",
    "    arrays=[\n",
    "        image_undispersed,\n",
    "        image,\n",
    "    ],\n",
    "    axis=\"dispersion\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcce6b3ffc3ccb97",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = na.plt.subplots(\n",
    "    axis_cols=\"dispersion\",\n",
    "    ncols=3,\n",
    "    # sharex=True,\n",
    "    # sharey=True,\n",
    "    figsize=(9, 4.3),\n",
    "    constrained_layout=True,\n",
    "    dpi=200,\n",
    "    gridspec_kw=dict(width_ratios=[.9, .9, .1]),\n",
    ")\n",
    "cax = ax.ndarray[~0]\n",
    "cax_twin = cax.twinx()\n",
    "ax = ax[dict(dispersion=slice(0, 2))]\n",
    "cax.xaxis.set_ticks_position(\"top\")\n",
    "cax.xaxis.set_label_position(\"top\")\n",
    "ani, colorbar = na.plt.rgbmovie(\n",
    "    obs.inputs.time,\n",
    "    image.inputs.wavelength,\n",
    "    image.inputs.position.x.to(u.mm),\n",
    "    image.inputs.position.y.to(u.mm),\n",
    "    C=stacked.outputs,\n",
    "    axis_time=obs.axis_time,\n",
    "    axis_wavelength=obs.axis_wavelength,\n",
    "    ax=ax,\n",
    "    vmin=0,\n",
    "    vmax=np.nanpercentile(image.outputs, 99.9, axis=(obs.axis_time, \"detector_x\", \"detector_y\")),\n",
    ")\n",
    "colorbar = colorbar[{obs.axis_time: 0}]\n",
    "na.plt.pcolormesh(\n",
    "    colorbar.inputs.x,\n",
    "    colorbar.inputs.y,\n",
    "    C=colorbar.outputs.value,\n",
    "    axis_rgb=obs.axis_wavelength,\n",
    "    ax=cax,\n",
    ")\n",
    "na.plt.pcolormesh(\n",
    "    colorbar.inputs.x,\n",
    "    colorbar.inputs.y.to(u.km / u.s, equivalencies=u.doppler_radio(wavelength_o_v)),\n",
    "    C=colorbar.outputs.value,\n",
    "    axis_rgb=obs.axis_wavelength,\n",
    "    ax=cax_twin,\n",
    ")\n",
    "na.plt.set_xlim(5.5, 9, ax=ax)\n",
    "na.plt.set_ylim(-2, 2, ax=ax)\n",
    "na.plt.set_aspect(\"equal\", ax=ax)\n",
    "na.plt.set_xlabel(\"detector $x$ (mm)\", ax=ax)\n",
    "na.plt.set_ylabel(\"detector $y$ (mm)\", ax=ax.ndarray[0])\n",
    "na.plt.set_ylabel(\"\", ax=ax.ndarray[~0])\n",
    "ax.ndarray[~0].tick_params(axis=\"y\", labelleft=False)\n",
    "\n",
    "result = ani.to_jshtml(fps=2)\n",
    "result = IPython.display.HTML(result)\n",
    "\n",
    "plt.close(ani._fig)\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4a7ae7c9bb46c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = {obs.axis_time: ~0}\n",
    "s = stacked[t0]\n",
    "\n",
    "fig, ax = plt.subplots(\n",
    "    ncols=2,\n",
    "    constrained_layout=True,\n",
    "    dpi=200,\n",
    "    gridspec_kw=dict(width_ratios=[.9, .1]),\n",
    "    figsize=(7, 5),\n",
    ")\n",
    "ax_twin = ax[1].twinx()\n",
    "ax[1].xaxis.set_ticks_position(\"top\")\n",
    "ax[1].xaxis.set_label_position(\"top\")\n",
    "ani, colorbar = na.plt.rgbmovie(\n",
    "    obs.inputs.time[t0],\n",
    "    image.inputs.wavelength[t0],\n",
    "    s.inputs.position.x.to(u.mm),\n",
    "    s.inputs.position.y.to(u.mm),\n",
    "    C=na.nominal(s.outputs),\n",
    "    axis_time=\"dispersion\",\n",
    "    axis_wavelength=obs.axis_wavelength,\n",
    "    ax=ax[0],\n",
    "    vmin=0,\n",
    "    vmax=np.nanpercentile(image.outputs, 99.9, axis=(obs.axis_time, \"detector_x\", \"detector_y\")),\n",
    ")\n",
    "# colorbar = colorbar[{obs.axis_time: 0}]\n",
    "na.plt.pcolormesh(\n",
    "    colorbar.inputs.x,\n",
    "    colorbar.inputs.y,\n",
    "    C=colorbar.outputs.value,\n",
    "    axis_rgb=obs.axis_wavelength,\n",
    "    ax=ax[1],\n",
    ")\n",
    "na.plt.pcolormesh(\n",
    "    colorbar.inputs.x,\n",
    "    colorbar.inputs.y.to(u.km / u.s, equivalencies=u.doppler_radio(wavelength_o_v)),\n",
    "    C=colorbar.outputs.value,\n",
    "    axis_rgb=obs.axis_wavelength,\n",
    "    ax=ax_twin,\n",
    ")\n",
    "na.plt.set_xlim(5.5, 9, ax=ax[0])\n",
    "na.plt.set_ylim(-2, 2, ax=ax[0])\n",
    "na.plt.set_aspect(\"equal\", ax=ax[0])\n",
    "na.plt.set_xlabel(f\"detector $x$ (mm)\", ax=ax[0])\n",
    "na.plt.set_ylabel(\"detector $y$ (mm)\", ax=ax[0])\n",
    "\n",
    "result = ani.to_jshtml(fps=.5)\n",
    "result = IPython.display.HTML(result)\n",
    "\n",
    "plt.close(ani._fig)\n",
    "\n",
    "result"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "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.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}