{
 "cells": [
  {
   "cell_type": "raw",
   "id": "0",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Inverting IRIS Observations with MART\n",
    "=====================================\n",
    "The main target of ESIS-I is the :math:`\\text{O\\,\\textsc{v}}\\;630\\,\\AA` spectral line.\n",
    "To our knowledge, high-quality spectroheliograms of this line are not available to use as a test dataset.\n",
    "The Coronal Diagnostic Spectrometer (CDS) :cite:p:`Harrison1996`\n",
    "and the Solar Ultraviolet Measurements of Emitted Radiation (SUMER) :cite:p:`Wilhelm1995` have both observed this line,\n",
    "but their spatial or spectral resolution is not sufficient to use as a test for ESIS.\n",
    "Instead, we will use observations of the :math:`\\text{Si\\,\\textsc{iv}}\\;1394\\,\\AA` spectral line from the Interface Region Imaging Spectrograph (IRIS) :cite:p:`DePontieu2014` as a proxy for the :math:`\\text{O\\,\\textsc{v}}\\;630\\,\\AA` line.\n",
    "These observations will be used along with a simple model of the ESIS instrument to create synthetic ESIS images.\n",
    "We will then use MART to invert these images and test how well it recovers the original scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import IPython.display\n",
    "import dataclasses\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors\n",
    "import scipy.stats\n",
    "import astropy.units as u\n",
    "import astropy.visualization\n",
    "import named_arrays as na\n",
    "import optika\n",
    "import iris\n",
    "import ctis\n",
    "import esis"
   ]
  },
  {
   "cell_type": "raw",
   "id": "3",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Loading the Test Scene\n",
    "----------------------\n",
    "Load an IRIS observation to use as a test scene to invert.\n",
    "This function shifts and scales the :math:`\\text{Si\\,\\textsc{iv}}\\;1394\\,\\AA` spectral line observed by IRIS to match the :math:`\\text{O\\,\\textsc{v}}\\;630\\,\\AA` line observed by ESIS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "scene = esis.flights.f1.data.synth.scene_iris(\n",
    "    time_start=\"2014-07-04 11:40\",\n",
    "    velocity_max=150 * u.km / u.s,\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "5",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Subtract the average helioprojective :math:`x` and :math:`y` coordinates from the scene to place the observation at disk center."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene.inputs.position = scene.inputs.position - scene.inputs.position.mean()"
   ]
  },
  {
   "cell_type": "raw",
   "id": "7",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Select the first frame of the observation to invert."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "scene = scene[{scene.axis_time: 0}]\n",
    "scene.timedelta = scene.timedelta[{scene.axis_time: 0}]"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Display the test scene as a false-color image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "scene.show();"
   ]
  },
  {
   "cell_type": "raw",
   "id": "11",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Plot the average spectrum of the test scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, ax = plt.subplots(constrained_layout=True)\n",
    "    ax_twin = ax.twiny()\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.wavelength,\n",
    "        scene.outputs.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        ax=ax,\n",
    "        label=\"original\",\n",
    "        axis=scene.axis_wavelength,\n",
    "    )\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.velocity,\n",
    "        scene.outputs.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        ax=ax_twin,\n",
    "        label=\"original\",\n",
    "        axis=scene.axis_wavelength,\n",
    "        linestyle=\"none\",\n",
    "    )\n",
    "    ax.set_xlabel(f\"wavelength ({ax.get_xlabel()})\")\n",
    "    ax.set_ylabel(f\"average radiance ({ax.get_ylabel()})\");"
   ]
  },
  {
   "cell_type": "raw",
   "id": "13",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Defining a Model of the ESIS Instrument\n",
    "---------------------------------------\n",
    "In this tutorial we'll use an idealized version of the ESIS instrument model that considers only dispersion and shot noise.\n",
    "\n",
    "We'll start by saving the rest wavelength of the :math:`\\text{O}\\,\\textsc{v}\\;630\\,\\AA` line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wavelength_rest = scene.inputs.wavelength_rest"
   ]
  },
  {
   "cell_type": "raw",
   "id": "15",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Use the rest wavelength to define an :mod:`astropy.units` equivalency for converting between Doppler velocities and wavelength units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "AA = dict(unit=u.AA, equivalencies=u.doppler_optical(wavelength_rest))"
   ]
  },
  {
   "cell_type": "raw",
   "id": "17",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Use the coordinates of the scene as the coordinates of the reconstructed scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "coordinates_scene = scene.inputs"
   ]
  },
  {
   "cell_type": "raw",
   "id": "19",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define the vertices of the pixel grid on the sensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "position_sensor = na.Cartesian2dVectorArray(\n",
    "    x=na.arange(0, 300 + 1, axis=\"sensor_x\") * u.pix,\n",
    "    y=na.arange(0, 300 + 1, axis=\"sensor_y\") * u.pix,\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "21",
   "metadata": {
    "editable": true,
    "raw_mimetype": "",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Use these vertices to form the sensor coordinates by combining it with the wavelength coordinates of the scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "coordinates_sensor = na.SpectralPositionalVectorArray(\n",
    "    wavelength=scene.inputs.wavelength,\n",
    "    position=position_sensor,\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "23",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define the dispersion angles of the simulated ESIS instrument and add an arbitrary offset so that the dispersion directions are not aligned with the pixel grid of the scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "angle = na.linspace(0, 180, num=4, axis=\"channel\", endpoint=False) * u.deg\n",
    "angle = angle + 0.4 * u.rad\n",
    "angle.ndarray"
   ]
  },
  {
   "cell_type": "raw",
   "id": "25",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define a unique string for each channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "channel = \"dispersion angle = \" + angle.to_string_array(\"%03d\")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "27",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Compute the effective area of ESIS using values from the instrument paper draft."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ccd = optika.sensors.materials.e2v_ccd97()\n",
    "absorbance = ccd.absorbance(\n",
    "    wavelength=wavelength_rest,\n",
    "    normal=na.Cartesian3dVectorArray(0, 0, -1),\n",
    ").average\n",
    "area_effective = 4 * u.mm ** 2 / 0.39 * absorbance.ndarray\n",
    "area_effective"
   ]
  },
  {
   "cell_type": "raw",
   "id": "29",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Define the name of the logical axis representing the different channels of the ESIS instrument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "axis_channel = \"channel\""
   ]
  },
  {
   "cell_type": "raw",
   "id": "31",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Using the values we computed above, construct a model of the ESIS instrument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "instrument = ctis.instruments.IdealInstrument(\n",
    "    area_effective=area_effective,\n",
    "    timedelta_exposure=10 * u.s,\n",
    "    plate_scale=0.77 * u.arcsec / u.pix,\n",
    "    dispersion=((17.5 * u.km / u.s).to(**AA) - wavelength_rest) / u.pix,\n",
    "    angle=angle,\n",
    "    wavelength_ref=wavelength_rest,\n",
    "    position_ref=na.Cartesian2dVectorArray(150, 150) * u.pix,\n",
    "    coordinates_scene=coordinates_scene,\n",
    "    coordinates_sensor=coordinates_sensor,\n",
    "    channel=channel,\n",
    "    axis_channel=axis_channel,\n",
    "    axis_wavelength=scene.axis_wavelength,\n",
    "    axis_scene_xy=(scene.axis_detector_x, scene.axis_detector_y),\n",
    "    axis_sensor_xy=(\"sensor_x\", \"sensor_y\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "id": "33",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Create synthetic ESIS images\n",
    "----------------------------\n",
    "To calculate synthetic images, all we need to do is apply the :meth:`~ctis.instruments.IdealInstrument.image` method to the IRIS observation we prepared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "images = instrument.image(scene)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "35",
   "metadata": {
    "editable": true,
    "raw_mimetype": "",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Display all four synthetic images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, axs = na.plt.subplots(\n",
    "        axis_rows=\"rows\",\n",
    "        axis_cols=\"cols\",\n",
    "        nrows=2,\n",
    "        ncols=2,\n",
    "        constrained_layout=True,\n",
    "        figsize=(10,9),\n",
    "        sharex=True,\n",
    "        sharey=True,\n",
    "        origin=\"upper\",\n",
    "    )\n",
    "    ax = axs.combine_axes((\"rows\", \"cols\"), axis_channel)\n",
    "    norm = matplotlib.colors.PowerNorm(\n",
    "        gamma=0.5,\n",
    "        vmin=0,\n",
    "        vmax=images.outputs.value.percentile(99.9).ndarray,\n",
    "    )\n",
    "    colorizer = plt.Colorizer(\n",
    "        cmap=\"gray\",\n",
    "        norm=norm,\n",
    "    )\n",
    "    img = na.plt.pcolormesh(\n",
    "        # channel,\n",
    "        images.inputs.position.x,\n",
    "        images.inputs.position.y,\n",
    "        C=images.outputs.value,\n",
    "        ax=ax,\n",
    "        colorizer=colorizer,\n",
    "    )\n",
    "    plt.colorbar(\n",
    "        mappable=plt.cm.ScalarMappable(colorizer=colorizer),\n",
    "        ax=ax.ndarray,\n",
    "        label=f\"signal ({images.outputs.unit:latex_inline})\",\n",
    "    )\n",
    "    na.plt.set_title(channel, ax=ax)\n",
    "    na.plt.set_aspect(\"equal\", ax=ax)\n",
    "    na.plt.set_xlabel(f\"sensor $x$ ({images.inputs.position.x.unit})\", ax=axs[dict(rows=~0)])\n",
    "    na.plt.set_ylabel(f\"sensor $y$ ({images.inputs.position.y.unit})\", ax=axs[dict(cols=0)])"
   ]
  },
  {
   "cell_type": "raw",
   "id": "37",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/x-rst",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Invert synthetic ESIS images\n",
    "----------------------------\n",
    "Now for the tricky part, inverting these images into a reconstructed scene using MART.\n",
    "An important part of using MART effectively is the initial guess supplied to the algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "image_separated = instrument.image(scene, integrate=False, noise=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "backprojected = instrument.backproject(image_separated, integrate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "scene_degraded = np.mean(backprojected, axis=axis_channel, keepdims=False)\n",
    "scene_degraded.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_degraded = scene_degraded.replace(inputs=scene.inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "backprojected.outputs.sum().to(scene.outputs.unit) / 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "coords_scene_single = coordinates_scene.replace(wavelength=images.inputs.wavelength)\n",
    "coords_sensor_single = images.inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument_single = dataclasses.replace(\n",
    "    instrument,\n",
    "    coordinates_scene=coords_scene_single,\n",
    "    coordinates_sensor=coords_sensor_single,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = instrument_single.backproject(images) * scene.outputs.shape[scene.axis_wavelength]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "smin = np.min(s, axis=axis_channel)[{axis_channel: 0}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {},
   "outputs": [],
   "source": [
    "s.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "smin.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "scene.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, ax = plt.subplots()\n",
    "    na.plt.pcolormesh(\n",
    "        smin.inputs.position.x,\n",
    "        smin.inputs.position.y,\n",
    "        C=smin.outputs[{axis_channel: 0, scene.axis_wavelength:0 }].value,\n",
    "        vmin=0,\n",
    "        vmax=np.percentile(smin.outputs, 99.5).value,\n",
    "    )\n",
    "    ax.set_aspect(\"equal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "mart = ctis.inverters.MartInverter(\n",
    "    instrument=instrument,\n",
    "    intermediate=True,\n",
    "    num_iteration=20,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "spectrum_avg = scene.outputs.mean((scene.axis_detector_x, scene.axis_detector_y))\n",
    "spectrum_avg = spectrum_avg / spectrum_avg.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53",
   "metadata": {},
   "outputs": [],
   "source": [
    "std_OV = esis.flights.f1.spectrum.O_V.width_doppler\n",
    "v = scene.inputs.velocity.cell_centers(scene.axis_wavelength)\n",
    "# guess_spectral = np.exp(-np.square(v / std_OV) / 2)\n",
    "std_OV = 37.5  * 1.5 * u.km / u.s\n",
    "kappa = 0.2\n",
    "v0 = 0 * u.km / u.s\n",
    "guess_spectral = (1 + np.square((v - v0) / std_OV) / kappa)**(-kappa - 1)\n",
    "guess_spectral = guess_spectral / guess_spectral.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "guess_spatial = smin.outputs\n",
    "# guess_spatial = guess_spatial.broadcast_to(scene.outputs.shape)\n",
    "guess_spatial = guess_spatial / guess_spatial.sum()\n",
    "guess_spatial.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "guess_total = backprojected.outputs.sum().to(scene.outputs.unit) / 4\n",
    "guess_total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "guess = guess_spectral * guess_spatial * guess_total\n",
    "# guess = spectrum_avg * scene.outputs.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, ax = plt.subplots(constrained_layout=True)\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.velocity,\n",
    "        scene.outputs.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        ax=ax,\n",
    "        label=\"original\",\n",
    "        axis=scene.axis_wavelength,\n",
    "    )\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.velocity,\n",
    "        guess.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        label=\"guess\",\n",
    "        axis=scene.axis_wavelength,\n",
    "    )\n",
    "    ax.set_xlabel(f\"wavelength ({ax.get_xlabel()})\")\n",
    "    ax.set_ylabel(f\"average radiance ({ax.get_ylabel()})\");\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "guess.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "scene.outputs.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "inversion = mart(images, guess=guess, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61",
   "metadata": {},
   "outputs": [],
   "source": [
    "axis_iter = inversion.inverter.axis_iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(\n",
    "    nrows=2,\n",
    "    sharex=True,\n",
    "    constrained_layout=True,\n",
    ")\n",
    "na.plt.plot(\n",
    "    inversion.iteration,\n",
    "    inversion.mean_chi_squared,\n",
    "    ax=ax[0],\n",
    "    axis=axis_iter,\n",
    "    label=channel,\n",
    ")\n",
    "na.plt.plot(\n",
    "    inversion.iteration,\n",
    "    inversion.correlation_residual,\n",
    "    ax=ax[1],\n",
    "    axis=axis_iter,\n",
    "    label=channel,\n",
    ")\n",
    "ax[0].set_ylabel(r\"$\\langle \\chi^2 \\rangle$\")\n",
    "ax[1].set_xlabel(\"iteration\")\n",
    "ax[1].set_ylabel(\"signal-correlated residual\")\n",
    "ax[0].set_yscale(\"log\")\n",
    "ax[0].legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "solution = inversion.solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, ax = plt.subplots(constrained_layout=True)\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.wavelength,\n",
    "        scene.outputs.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        ax=ax,\n",
    "        label=\"original\",\n",
    "    )\n",
    "    na.plt.stairs(\n",
    "        scene.inputs.wavelength,\n",
    "        guess.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        label=\"guess\",\n",
    "        axis=scene.axis_wavelength,\n",
    "    )\n",
    "    na.plt.stairs(\n",
    "        solution.inputs.wavelength,\n",
    "        solution.outputs.mean((scene.axis_detector_x, scene.axis_detector_y)),\n",
    "        ax=ax,\n",
    "        label=\"reconstructed\",\n",
    "    )\n",
    "    ax.set_xlabel(f\"wavelength ({ax.get_xlabel()})\")\n",
    "    # ax2.set_xlabel(f\"wavelength ({ax2.get_xlabel()})\")\n",
    "    ax.set_ylabel(f\"average radiance ({ax.get_ylabel()})\")\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {},
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, axs = plt.subplots(\n",
    "        ncols=3,\n",
    "        gridspec_kw=dict(width_ratios=[.5, .5, .1]),\n",
    "        constrained_layout=True,\n",
    "        figsize=(10, 6),\n",
    "    )\n",
    "    ax1, ax2, cax = axs\n",
    "    ax2.set_yticklabels([])\n",
    "    vmin = np.nanpercentile(\n",
    "        a=scene.outputs,\n",
    "        q=0.5,\n",
    "        axis=(scene.axis_detector_x, scene.axis_detector_y),\n",
    "    )\n",
    "    vmax = np.nanpercentile(\n",
    "        a=scene.outputs,\n",
    "        q=99.5,\n",
    "        axis=(scene.axis_detector_x, scene.axis_detector_y),\n",
    "    )\n",
    "    wavelength_min=-100 * u.km / u.s,\n",
    "    wavelength_max=+100 * u.km / u.s,\n",
    "    na.plt.rgbmesh(\n",
    "        scene.inputs.velocity.cell_centers(scene.axis_wavelength),\n",
    "        scene.inputs.position.x,\n",
    "        scene.inputs.position.y,\n",
    "        C=scene_degraded.outputs,\n",
    "        axis_wavelength=scene.axis_wavelength,\n",
    "        ax=ax1,\n",
    "        vmin=vmin,\n",
    "        vmax=vmax,\n",
    "        wavelength_min=-100 * u.km / u.s,\n",
    "        wavelength_max=+100 * u.km / u.s,\n",
    "    )\n",
    "    colorbar = na.plt.rgbmesh(\n",
    "        scene.inputs.velocity.cell_centers(scene.axis_wavelength),\n",
    "        scene.inputs.position.x,\n",
    "        scene.inputs.position.y,\n",
    "        C=solution.outputs,\n",
    "        axis_wavelength=scene.axis_wavelength,\n",
    "        ax=ax2,\n",
    "        vmin=vmin,\n",
    "        vmax=vmax,\n",
    "        wavelength_min=-100 * u.km / u.s,\n",
    "        wavelength_max=+100 * u.km / u.s,\n",
    "    )\n",
    "    na.plt.pcolormesh(\n",
    "        C=colorbar,\n",
    "        axis_rgb=scene.axis_wavelength,\n",
    "        ax=cax,\n",
    "    )\n",
    "    ax1.set_title(\"degraded\")\n",
    "    ax2.set_title(\"reconstructed\")\n",
    "    unit_x = scene.inputs.position.x.unit\n",
    "    unit_y = scene.inputs.position.y.unit\n",
    "    ax1.set_xlabel(f\"scene $x$ ({unit_x:latex_inline})\")\n",
    "    ax2.set_xlabel(f\"scene $x$ ({unit_x:latex_inline})\")\n",
    "    ax1.set_ylabel(f\"scene $y$ ({unit_y:latex_inline})\")\n",
    "    cax.xaxis.set_ticks_position(\"top\")\n",
    "    cax.xaxis.set_label_position(\"top\")\n",
    "    cax.yaxis.tick_right()\n",
    "    cax.yaxis.set_label_position(\"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with astropy.visualization.quantity_support():\n",
    "    fig, ax = plt.subplots(\n",
    "        ncols=2,\n",
    "        constrained_layout=True,\n",
    "        sharex=True,\n",
    "        sharey=True,\n",
    "        figsize=(6, 5),\n",
    "    )\n",
    "    i = {scene.axis_detector_x: 210}\n",
    "    vmax = scene_degraded.outputs.percentile(99.9).ndarray.value\n",
    "    norm = matplotlib.colors.PowerNorm(gamma=0.5, vmin=0, vmax=vmax)\n",
    "    na.plt.pcolormesh(\n",
    "        scene.inputs.velocity[{scene.axis_detector_x: 0, scene.axis_detector_y: 0}],\n",
    "        scene.inputs.position.y[i],\n",
    "        C=scene_degraded.outputs[i].value,\n",
    "        ax=ax[0],\n",
    "        norm=norm,\n",
    "    )\n",
    "    ani = na.plt.pcolormovie(\n",
    "        inversion.iteration,\n",
    "        scene.inputs.velocity[{scene.axis_detector_x: 0, scene.axis_detector_y: 0}],\n",
    "        inversion.solutions.inputs.position.y[i],\n",
    "        C=inversion.solutions.outputs[i].to_value(scene_degraded.outputs.unit),\n",
    "        ax=ax[1],\n",
    "        axis_time=inversion.inverter.axis_iteration,\n",
    "        norm=norm,\n",
    "    )\n",
    "    ax[0].set_title(\"original\")\n",
    "    ax[1].set_title(\"reconstruction\")\n",
    "\n",
    "result = ani.to_jshtml(fps=10)\n",
    "result = IPython.display.HTML(result)\n",
    "\n",
    "ani.save(\"mart-iris-spectra.gif\", dpi=200, fps=30)\n",
    "\n",
    "plt.close(ani._fig)\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "inversion.plot_moments(\n",
    "    truth=scene, \n",
    "    axis=scene.axis_wavelength,\n",
    "    range_radiance=[0, 3e3] * u.erg / u.s / u.sr / u.cm**2,\n",
    "    range_median=[-60, 60] * u.km / u.s,\n",
    "    range_iqr=[0, 100] * u.km / u.s,\n",
    "    percentile_radiance=25,\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "predictions = instrument.image(solution, noise=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "residual = images - predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = na.plt.subplots(\n",
    "    axis_rows=\"rows\",\n",
    "    axis_cols=\"cols\",\n",
    "    nrows=2,\n",
    "    ncols=2,\n",
    "    constrained_layout=True,\n",
    "    figsize=(8,7),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    "    origin=\"upper\",\n",
    ")\n",
    "ax = ax.reshape(dict(channel=-1))\n",
    "vmin = -35\n",
    "vmax = +35\n",
    "img = na.plt.pcolormesh(\n",
    "    residual.inputs.position.x,\n",
    "    residual.inputs.position.y,\n",
    "    C=residual.outputs.value,\n",
    "    ax=ax,\n",
    "    vmin=vmin,\n",
    "    vmax=vmax,\n",
    "    cmap=\"gray\",\n",
    ")\n",
    "na.plt.set_title\n",
    "na.plt.set_aspect(\"equal\", ax=ax)\n",
    "plt.colorbar(\n",
    "    ax=ax.ndarray,\n",
    "    mappable=img.ndarray[0],\n",
    "    label=f\"residual ({residual.outputs.unit:latex_inline})\",\n",
    ");\n",
    "na.plt.set_title(channel, ax=ax);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "float(scipy.stats.spearmanr(predictions.outputs.ndarray.reshape(-1), residual.outputs.ndarray.reshape(-1)).statistic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72",
   "metadata": {},
   "outputs": [],
   "source": [
    "float(scipy.stats.pearsonr(predictions.outputs.ndarray.reshape(-1), residual.outputs.ndarray.reshape(-1)).statistic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "images_perfect = instrument.image(scene, noise=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "inversion.inverter.mean_chi_squared(\n",
    "    images.outputs,\n",
    "    images_perfect.outputs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "inversion.inverter.correlation_residual(\n",
    "    images.outputs,\n",
    "    predictions.outputs,\n",
    ").ptp()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# with astropy.visualization.quantity_support():\n",
    "#     fig, axs = plt.subplots(\n",
    "#         ncols=3,\n",
    "#         gridspec_kw=dict(width_ratios=[.5, .5, .1]),\n",
    "#         constrained_layout=True,\n",
    "#         figsize=(10, 6),\n",
    "#     )\n",
    "#     ax1, ax2, cax = axs\n",
    "#     ax2.set_yticklabels([])\n",
    "#     vmax = np.nanpercentile(\n",
    "#         a=scene.outputs,\n",
    "#         q=99.5,\n",
    "        \n",
    "#         axis=(scene.axis_detector_x, scene.axis_detector_y),\n",
    "#     )\n",
    "\n",
    "#     na.plt.rgbmesh(\n",
    "#         C=scene_degraded,\n",
    "#         axis_wavelength=\"wavelength\",\n",
    "#         ax=ax1,\n",
    "#         vmin=0,\n",
    "#         vmax=vmax,\n",
    "#     )\n",
    "#     label = \"iteration = \" + inversion.iteration.to_string_array(\"%d\")\n",
    "#     chisq_str = r\"$\\langle \\chi^2 \\rangle$\"\n",
    "#     label = label + f\"\\n{chisq_str} = \" + inversion.mean_chi_squared.mean(\"channel\").to_string_array(\"%.03f\")\n",
    "#     ani, colorbar = na.plt.rgbmovie(\n",
    "#         label,\n",
    "#         scene.velocity_doppler,\n",
    "#         scene.inputs.position.x,\n",
    "#         scene.inputs.position.y,\n",
    "#         C=inversion.solutions.outputs,\n",
    "#         axis_time=inversion.inverter.axis_iteration,\n",
    "#         axis_wavelength=\"wavelength\",\n",
    "#         ax=ax2,\n",
    "#         vmin=0,\n",
    "#         vmax=vmax,\n",
    "#     )\n",
    "#     na.plt.pcolormesh(\n",
    "#         C=colorbar,\n",
    "#         axis_rgb=\"wavelength\",\n",
    "#         ax=cax,\n",
    "#     )\n",
    "#     ax1.set_title(\"original\")\n",
    "#     ax2.set_title(\"reconstructed\")\n",
    "#     unit_x = scene.inputs.position.x.unit\n",
    "#     unit_y = scene.inputs.position.y.unit\n",
    "#     ax1.set_xlabel(f\"scene $x$ ({unit_x:latex_inline})\")\n",
    "#     ax2.set_xlabel(f\"scene $x$ ({unit_x:latex_inline})\")\n",
    "#     ax1.set_ylabel(f\"scene $y$ ({unit_y:latex_inline})\")\n",
    "#     cax.xaxis.set_ticks_position(\"top\")\n",
    "#     cax.xaxis.set_label_position(\"top\")\n",
    "#     cax.yaxis.tick_right()\n",
    "#     cax.yaxis.set_label_position(\"right\")\n",
    "\n",
    "# result = ani.to_jshtml(fps=30)\n",
    "# result = IPython.display.HTML(result)\n",
    "\n",
    "# ani.save(\"mart-iris-rgb.mp4\")\n",
    "\n",
    "# plt.close(ani._fig)\n",
    "\n",
    "# result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77",
   "metadata": {},
   "source": [
    " - mulitply counts by 10 to see how it changes the results\n",
    " - smooth the outputs before taking the residual (and computing the moments)\n",
    " - negative correlation coefficient implies that we went a little too far (crossed zero)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "$d \\chi = \\frac{d' - d}{\\sqrt{d}}$ contribution of the residual to the total chi square (SNR-weighted residual)"
   ]
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}