Grotrian diagram IPyWidget functionality (#2328)

* Rename GrotrianWidget to GrotrianPlot. Create skeleton of new wrapper GrotrianWidget class * Update figure on value changes inside widgets * Added widgets for ion, shell, levels and y-scale * Added slider widget for min and max wavelength * Added docstrings for new methods * Handle empty transitions * Remove unnecessary cells from notebook * Fixed refresh logic of wavelength slider * Update docs for GrotrianDiagram * Refactor code and set default y-coordinate scale to log * Updated documentation and minor fixes * Minor changes in documentation (in response to review comments) * Minor change to docs * Update the generating_widgets notebook * Minor markdown edit * Rename model to simulation_state * Fix paths and Simulation object creation * Revert "Fix paths and Simulation object creation" This reverts commit 23afbef. * Revert "Revert "Fix paths and Simulation object creation"" This reverts commit 09b453b. * Render widget in mockup notebook * Revert "Render widget in mockup notebook" This reverts commit e9acd6a. * minor doc change --------- Co-authored-by: Andrew Fullard <[email protected]>
tardis-sn · Apr 12, 2024 · 28caa4e · 28caa4e
1 parent 5bca272
commit 28caa4e
Show file tree

Hide file tree

Showing 7 changed files with 871 additions and 922 deletions.
diff --git a/docs/io/images/energy_level_widget_demo.gif b/docs/io/images/energy_level_widget_demo.gif
diff --git a/docs/io/images/energy_level_widget_options.gif b/docs/io/images/energy_level_widget_options.gif
diff --git a/docs/io/visualization/how_to_generating_widgets.ipynb b/docs/io/visualization/how_to_generating_widgets.ipynb
@@ -14,7 +14,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "First create and run a simulation that we can use to generate widgets (more details about running simulation in [Quickstart](https://tardis-sn.github.io/tardis/quickstart/quickstart.html) section):"
+    "First create and run a simulation that we can use to generate widgets (more details about running simulation in [Quickstart](https://tardis-sn.github.io/tardis/quickstart/quickstart.html) section):\n"
    ]
   },
   {
@@ -33,16 +33,16 @@
     "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",
+    "download_atom_data(\"kurucz_cd23_chianti_H_He\")\n",
     "\n",
-    "sim = run_tardis('tardis_example.yml', virtual_packet_logging=True)"
+    "sim = run_tardis(\"tardis_example.yml\", virtual_packet_logging=True)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now, import functions & class to create widgets from `visualization` subpackage:"
+    "Now, import functions & class to create widgets from `visualization` subpackage:\n"
    ]
   },
   {
@@ -60,6 +60,7 @@
     "    shell_info_from_simulation,\n",
     "    shell_info_from_hdf,\n",
     "    LineInfoWidget,\n",
+    "    GrotrianWidget,\n",
     ")"
    ]
   },
@@ -68,17 +69,19 @@
    "metadata": {},
    "source": [
     "## Shell Info Widget\n",
+    "\n",
     "This widget allows you to explore chemical abundances of each shell - all the way from elements to ions to levels - by just clicking on the rows you want to explore!\n",
     "\n",
-    "There are two ways in which you can generate the widget:"
+    "There are two ways in which you can generate the widget:\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Using a Simulation object\n",
-    "We will use the simulation object we created in the beginning, `sim` to generate shell info widget. Then simply display it to start using."
+    "\n",
+    "We will use the simulation object we created in the beginning, `sim` to generate shell info widget. Then simply display it to start using.\n"
    ]
   },
   {
@@ -104,15 +107,16 @@
     "\n",
     "![Shell Info Widget Demo](../images/shell_info_widget_demo.gif)\n",
     "\n",
-    "Use the button at the top of this page to run the notebook in interactively to use the widgets!"
+    "Use the button at the top of this page to run the notebook in interactively to use the widgets!\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Using a saved simulation (HDF file)\n",
-    "Alternatively, if you have a TARDIS simulation model saved on your disk as an HDF file, you can also use it to generate the shell info widget."
+    "\n",
+    "Alternatively, if you have a TARDIS simulation model saved on your disk as an HDF file, you can also use it to generate the shell info widget.\n"
    ]
   },
   {
@@ -136,16 +140,17 @@
    "metadata": {},
    "source": [
     "## Line Info Widget\n",
+    "\n",
     "This widget lets you explore the atomic lines responsible for producing features in the simulated spectrum.\n",
     "\n",
-    "You can select any wavelength range in the spectrum interactively to display a table giving the fraction of packets that experienced their last interaction with each species. Using toggle buttons, you can specify whether to filter the selected range by the emitted or absorbed wavelengths of packets. Clicking on a row in the species table, shows packet counts for each last line interaction of the selected species, which can be grouped in several ways."
+    "You can select any wavelength range in the spectrum interactively to display a table giving the fraction of packets that experienced their last interaction with each species. Using toggle buttons, you can specify whether to filter the selected range by the emitted or absorbed wavelengths of packets. Clicking on a row in the species table, shows packet counts for each last line interaction of the selected species, which can be grouped in several ways.\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "To generate line info widget, we will again use the simulation object `sim` and then display the widget:"
+    "To generate line info widget, we will again use the simulation object `sim` and then display the widget:\n"
    ]
   },
   {
@@ -169,7 +174,7 @@
    "source": [
     "You can interact with this widget (which again won't be visible if you're viewing this notebook in our docs as an html page) like this:\n",
     "\n",
-    "![Line Info Widget Demo](../images/line_info_widget_demo.gif)"
+    "![Line Info Widget Demo](../images/line_info_widget_demo.gif)\n"
    ]
   },
   {
@@ -179,11 +184,58 @@
     "<div class=\"alert alert-info\">\n",
     "\n",
     "Note\n",
-    "    \n",
+    "\n",
     "The virtual packet logging capability must be active in order to produce virtual packets' spectrum in `Line Info Widget`. Thus, make sure to set `virtual_packet_logging: True` in your configuration file. It should be added under `virtual` property of `spectrum` property, as described in [configuration schema](https://tardis-sn.github.io/tardis/using/components/configuration/configuration.html#spectrum).\n",
     "\n",
-    "</div>"
+    "</div>\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Energy Level Diagram\n",
+    "\n",
+    "This widget lets you explore and visualize the various level populations and line interactions in a simulation in the form of an Energy Level Diagram.\n",
+    "\n",
+    "You can select any ion present in the simulation and filter the transitions by wavelength or model shell to display an energy level diagram, where:\n",
+    "\n",
+    "- The horizontal lines represent the energy levels. The thickness of each line shows the relative population of that energy level, with thicker lines being more populated.\n",
+    "- The arrows represent the line interactions between levels, with the arrow direction giving the direction of the transition. The thickness of each arrow also shows the number of packets that underwent the transition while the wavelength is given by the color.\n",
+    "\n",
+    "In addition, you can also select between linear- and log-scaling for the y-axis (which represents the energy of each level) and the maximum number of levels to display, beginning from the lowest energy levels.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To generate the energy level diagram, we will again use the simulation object `sim` and then display the widget:\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "energy_level_widget = GrotrianWidget.from_simulation(sim)\n",
+    "energy_level_widget.display()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can interact with this widget (which again won't be visible if you're viewing this notebook in our docs as an html page) like this:\n",
+    "\n",
+    "![Energy Level Diagram Demo](../images/energy_level_widget_options.gif)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

diff --git a/docs/io/visualization/using_widgets.rst b/docs/io/visualization/using_widgets.rst
@@ -178,4 +178,40 @@ There are also several other options in the modebar which we have not explained
 you remember to click back on the **Box Select** option for making selections on
 spectrum.
 
+Energy Level Diagram
+################
+
+This widget lets you visualize the last line interactions
+
+.. image:: ../images/energy_level_widget_demo.gif
+    :alt: Demo of Energy Level Diagram
+
+By selecting an ion on the widget, you can see its energy level diagram, which
+also shows information about the last line interactions experienced by packets
+in the simulation.
+
+The y-axis of the plot represents energy while the horizontal lines show
+discrete energy levels. The thickness of each line represents the level
+population, with thicker lines representing a greater population than the thin lines.
+
+Arrows represent the line interactions experienced by packets. Upwards arrows
+show excitation from lower energy levels to higher levels and downward arrows
+show de-excitation from higher energy levels to lower levels. The thickness of
+each arrow represents the number of packets that underwent that interaction,
+with thicker lines representing more packets than the thin lines.
+The wavelength of the transition is given by the color.
+
+Setting Other Options
+-----------------
+You can select the range on which to filter the wavelength using the slider.
+You can also select the model shell by which to filter the last line interactions
+and the level populations. If no shell is selected, then all the last line
+interactions are plotted and the level populations are averaged across all shells
+in the simulation. You can also set the maximum number of levels to show on the plot.
+
+Lastly, you can also set the scale of the y-axis: Linear or Log.
+
+.. image:: ../images/energy_level_widget_options.gif
+    :alt: Demo of using options
+
 .. Toggle legend
diff --git a/tardis/visualization/__init__.py b/tardis/visualization/__init__.py
@@ -7,5 +7,6 @@
     shell_info_from_hdf,
 )
 from tardis.visualization.widgets.line_info import LineInfoWidget
+from tardis.visualization.widgets.grotrian import GrotrianWidget
 from tardis.visualization.widgets.custom_abundance import CustomAbundanceWidget
 from tardis.visualization.tools.sdec_plot import SDECPlotter