{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spectral element DEComposition (SDEC) Plot\n",
    "The SDEC Plot illustrates the contributions of different chemical elements in the formation of a simulation model's spectrum. It is a spectral diagnostic plot similar to those originally proposed by M. Kromer (see, for example, [Kromer et al. 2013](https://arxiv.org/abs/1311.0310), figure 4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:54:24.525446Z",
     "start_time": "2021-10-01T11:54:24.522894Z"
    },
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# We filter out warnings throughout this notebook\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Due to the large size of the SDEC plots in SVG format, we request output as a\n",
    "# high-resolution PNG\n",
    "%config InlineBackend.figure_formats='png2x'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, create and run a simulation for which you want to generate this plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:28.560853Z",
     "start_time": "2021-10-01T11:54:24.527697Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from tardis import run_tardis\n",
    "from tardis.io.atom_data.util import download_atom_data\n",
    "\n",
    "# We download the atomic data needed to run the simulation\n",
    "download_atom_data('kurucz_cd23_chianti_H_He')\n",
    "\n",
    "sim = run_tardis(\"tardis_example.yml\", virtual_packet_logging=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "Note\n",
    "\n",
    "The virtual packet logging capability must be active in order to produce the SDEC Plot for virtual packets population. Thus, make sure to set `virtual_packet_logging: True` in your configuration file if you want to generate the SDEC Plot with virtual packets. It should be added under the `virtual` property of the `spectrum` property, as described in the [configuration schema](https://tardis-sn.github.io/tardis/io/configuration/components/spectrum.html).\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, import the plotting interface for the SDEC plot, i.e. the `SDECPlotter` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:28.570289Z",
     "start_time": "2021-10-01T11:55:28.565985Z"
    }
   },
   "outputs": [],
   "source": [
    "from tardis.visualization import SDECPlotter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And create a plotter object to process the data of simulation object `sim` for generating the SDEC plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:28.856453Z",
     "start_time": "2021-10-01T11:55:28.573256Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter = SDECPlotter.from_simulation(sim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Static Plot (in matplotlib)\n",
    "You can now call the `generate_plot_mpl()` method on your `plotter` object to get a highly informative, yet beautiful, SDEC plot produced in matplotlib."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Virtual packets mode\n",
    "By default, an SDEC plot is produced for the virtual packet population of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:31.451258Z",
     "start_time": "2021-10-01T11:55:28.858833Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Real packets mode\n",
    "You can produce the SDEC plot for the real packet population of the simulation by setting `packets_mode=\"real\"` which is `\"virtual\"` by default. Since `packets_mode` is the 1st argument, you can simply pass `\"real\"` string only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:32.158763Z",
     "start_time": "2021-10-01T11:55:31.452265Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(\"real\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting a specific wavelength range\n",
    "You can also restrict the wavelength range of escaped packets that you want to plot by specifying `packet_wvl_range`. It should be a quantity in Angstroms, containing two values - lower lambda and upper lambda i.e. `[lower_lambda, upper_lambda] * u.AA`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:32.162522Z",
     "start_time": "2021-10-01T11:55:32.160130Z"
    }
   },
   "outputs": [],
   "source": [
    "from astropy import units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:32.834975Z",
     "start_time": "2021-10-01T11:55:32.165455Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(packet_wvl_range=[3000, 9000] * u.AA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting only the top contributing elements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `nelements` option allows you to plot the top contributing elements to the spectrum. The top elements are shown in unique colors and the rest of the elements are shown in silver. Please note this works only for elements and not for ions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:33.596372Z",
     "start_time": "2021-10-01T11:55:32.836722Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(packet_wvl_range=[2000, 8000] * u.AA, nelements = 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choosing what elements/ions to plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also pass a list of elements/ions of your choice in the `species_list` option and plot them. Valid options include elements (e.g. Si), ions (which must be specified in Roman numeral format, e.g. Si II), a range of ions (e.g. Si I-III), or any combination of these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:34.453347Z",
     "start_time": "2021-10-01T11:55:33.597694Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(packet_wvl_range=[2000, 8000] * u.AA, species_list = ['Si II', 'S I-V', 'Ca'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using both the `nelements` and the `species_list` options, `species_list` takes precedence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:36.222650Z",
     "start_time": "2021-10-01T11:55:34.454511Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(nelements = 3, species_list = ['Si II', 'S I-V', 'Ca'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting flux instead of luminosity\n",
    "You can plot in units of flux on the Y-axis of the SDEC plot, by specifying the `distance` parameter. It should be a quantity with a unit of length like m, Mpc, etc. and must be a positive value. By default, `distance=None` plots luminosity on the Y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:37.709270Z",
     "start_time": "2021-10-01T11:55:36.223839Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(distance=100 * u.Mpc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting an observed spectrum\n",
    "To add an observed spectrum to the SDEC plot, you would need to pass the wavelength and the flux to the `observed_spectrum` parameter. The argument passed should be a tuple/list where the first value is the wavelength and the second value is the flux of the observed spectrum. Please note that these values should be instances of `astropy.Quantity`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:37.763078Z",
     "start_time": "2021-10-01T11:55:37.711523Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "data = np.loadtxt('demo_observed_spectrum.dat')\n",
    "\n",
    "observed_spectrum_wavelength, observed_spectrum_flux = data.T\n",
    "observed_spectrum_wavelength = observed_spectrum_wavelength * u.AA\n",
    "observed_spectrum_flux = observed_spectrum_flux * u.erg / (u.s * u.cm**2 * u.AA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:39.116098Z",
     "start_time": "2021-10-01T11:55:37.764988Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(observed_spectrum = (observed_spectrum_wavelength, observed_spectrum_flux), distance = 6 * u.Mpc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hiding modeled spectrum\n",
    "By default, the modeled spectrum is shown in SDEC plot. You can hide it by setting `show_modeled_spectrum=False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:40.310281Z",
     "start_time": "2021-10-01T11:55:39.117361Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_mpl(show_modeled_spectrum=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional plotting options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:40.316100Z",
     "start_time": "2021-10-01T11:55:40.311570Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# To list all available options (or parameters) with their description\n",
    "help(plotter.generate_plot_mpl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `generate_plot_mpl` method also has options specific to the matplotlib API, thereby providing you with more control over how your SDEC plot looks. Possible cases where you may use them are:\n",
    "\n",
    "- `ax`: To plot SDEC on the Axis of a plot you're already working with, e.g. for subplots.\n",
    "\n",
    "- `figsize`: To resize the SDEC plot as per your requirements.\n",
    "\n",
    "- `cmapname`: To use a colormap of your preference, instead of \"jet\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive Plot (in plotly)\n",
    "If you're using the SDEC plot for exploration purposes, you should plot its interactive version by using `generate_plot_ply()`. This not only allows you to zoom & pan but also to inspect data values by hovering, to resize scale, etc. conveniently (as shown below).\n",
    "\n",
    "![Interactions possible with SDEC plotly plot](../images/sdec_ply_interaction.gif)\n",
    "\n",
    "\n",
    "**This method takes the exact same arguments as `generate_plot_mpl` except a few that are specific to the plotting library.** We can produce all the plots above in plotly, by passing the same arguments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Virtual packets mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:42.308258Z",
     "start_time": "2021-10-01T11:55:40.317173Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_ply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Real packets mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:43.189422Z",
     "start_time": "2021-10-01T11:55:42.310142Z"
    }
   },
   "outputs": [],
   "source": [
    "plotter.generate_plot_ply(packets_mode=\"real\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a similar manner, you can also use the `packet_wvl_range`, `nelements`, `species_list`, `show_modeled_spectrum`, `observed_spectrum` and distance arguments in plotly plots (try it out in interactive mode!)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional plotting options"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `generate_plot_ply` method also has options specific to the plotly API, thereby providing you with more control over how your SDEC plot looks. Possible cases where you may use them are:\n",
    "- `fig`: To plot the SDEC plot on a figure you are already using e.g. for subplots.\n",
    "- `graph_height`: To specify the height of the graph as needed.\n",
    "- `cmapname`: To use a colormap of your preference instead of \"jet\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:47.644716Z",
     "start_time": "2021-10-01T11:55:47.640577Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# To list all available options (or parameters) with their description\n",
    "help(plotter.generate_plot_ply)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using simulation saved as HDF\n",
    "Other than producing the SDEC Plot for simulation objects in runtime, you can also produce it for saved TARDIS simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:47.833327Z",
     "start_time": "2021-10-01T11:55:47.645614Z"
    }
   },
   "outputs": [],
   "source": [
    "# hdf_plotter = SDECPlotter.from_hdf(\"demo.h5\") ## Files is too large - just as an example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This `hdf_plotter` object is similar to the `plotter` object we used above, **so you can use each plotting method demonstrated above with this too.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:48.266791Z",
     "start_time": "2021-10-01T11:55:47.834808Z"
    }
   },
   "outputs": [],
   "source": [
    "# Static plot with virtual packets mode\n",
    "# hdf_plotter.generate_plot_mpl()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:48.659429Z",
     "start_time": "2021-10-01T11:55:48.268776Z"
    }
   },
   "outputs": [],
   "source": [
    "# Static plot with real packets mode\n",
    "# hdf_plotter.generate_plot_mpl(\"real\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-01T11:55:48.854879Z",
     "start_time": "2021-10-01T11:55:48.661398Z"
    }
   },
   "outputs": [],
   "source": [
    "# Interactive plot with virtual packets mode\n",
    "# hdf_plotter.generate_plot_ply()"
   ]
  }
 ],
 "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.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "229.767px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}