From 8b3383658e73557fab32097561c9267b9cada2aa Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 26 Mar 2021 16:29:14 +0000
Subject: [PATCH 01/20] Add catch for missing elements

The elements coloured in the plot are determined via the absorption and emission. It's possible however, that these elements are not the same, e.g. if you specify a limited wavelength range. This could give an error if you try to plot an element not in the emission dataframe, for example. This adds a try statement to catch this case
---
 tardis/visualization/tools/sdec_plot.py | 125 ++++++++++++++----------
 1 file changed, 73 insertions(+), 52 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index c4f22f03af0..9d519b5c4ee 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1016,20 +1016,25 @@ def _plot_emission_mpl(self):
         elements_z = self.elements
         # Contribution from each element
         for i, atomic_number in enumerate(elements_z):
-            lower_level = upper_level
-            upper_level = (
-                lower_level
-                + self.emission_luminosities_df[atomic_number].to_numpy()
-            )
+            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in elements_z is not in the emission df
+            try:
+                lower_level = upper_level
+                upper_level = (
+                    lower_level
+                    + self.emission_luminosities_df[atomic_number].to_numpy()
+                )
 
-            self.ax.fill_between(
-                self.plot_wavelength,
-                lower_level,
-                upper_level,
-                color=self.cmap(i / len(self.elements)),
-                cmap=self.cmap,
-                linewidth=0,
-            )
+                self.ax.fill_between(
+                    self.plot_wavelength,
+                    lower_level,
+                    upper_level,
+                    color=self.cmap(i / len(self.elements)),
+                    cmap=self.cmap,
+                    linewidth=0,
+                )
+            except:
+                print(atomic_number2element_symbol(atomic_number) + " is not in the emitted packets; skipping")
 
     def _plot_absorption_mpl(self):
         """Plot absorption part of the SDEC Plot using matplotlib."""
@@ -1054,20 +1059,25 @@ def _plot_absorption_mpl(self):
 
         elements_z = self.elements
         for i, atomic_number in enumerate(elements_z):
-            upper_level = lower_level
-            lower_level = (
-                upper_level
-                - self.absorption_luminosities_df[atomic_number].to_numpy()
-            )
+            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in elements_z is not in the emission df
+            try:
+                upper_level = lower_level
+                lower_level = (
+                    upper_level
+                    - self.absorption_luminosities_df[atomic_number].to_numpy()
+                )
 
-            self.ax.fill_between(
-                self.plot_wavelength,
-                upper_level,
-                lower_level,
-                color=self.cmap(i / len(elements_z)),
-                cmap=self.cmap,
-                linewidth=0,
-            )
+                self.ax.fill_between(
+                    self.plot_wavelength,
+                    upper_level,
+                    lower_level,
+                    color=self.cmap(i / len(elements_z)),
+                    cmap=self.cmap,
+                    linewidth=0,
+                )
+            except:
+                print(atomic_number2element_symbol(atomic_number) + " is not in the emitted packets; skipping")
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
@@ -1268,19 +1278,25 @@ def _plot_emission_ply(self):
 
         elements_z = self.elements
         for i, atomic_num in enumerate(elements_z):
-            self.fig.add_trace(
-                go.Scatter(
-                    x=self.emission_luminosities_df.index,
-                    y=self.emission_luminosities_df[atomic_num],
-                    mode="none",
-                    name=atomic_number2element_symbol(atomic_num),
-                    fillcolor=self.to_rgb255_string(
-                        self.cmap(i / len(self.elements))
-                    ),
-                    stackgroup="emission",
-                    showlegend=False,
+            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in elements_z is not in the emission df
+            try:
+                self.fig.add_trace(
+                    go.Scatter(
+                        x=self.emission_luminosities_df.index,
+                        y=self.emission_luminosities_df[atomic_num],
+                        mode="none",
+                        name=atomic_number2element_symbol(atomic_num),
+                        fillcolor=self.to_rgb255_string(
+                            self.cmap(i / len(self.elements))
+                        ),
+                        stackgroup="emission",
+                        showlegend=False,
+                    )
                 )
-            )
+            except:
+                print(atomic_number2element_symbol(atomic_num) + " is not in the emitted packets; skipping")
+
 
     def _plot_absorption_ply(self):
         """Plot absorption part of the SDEC Plot using plotly."""
@@ -1290,6 +1306,7 @@ def _plot_absorption_ply(self):
             self.fig.add_trace(
                 go.Scatter(
                     x=self.absorption_luminosities_df.index,
+                    # to plot absorption luminosities along negative y-axis
                     y=self.absorption_luminosities_df.other * -1,
                     mode="none",
                     name="Other elements",
@@ -1300,22 +1317,26 @@ def _plot_absorption_ply(self):
             )
 
         elements_z = self.elements
-
         for i, atomic_num in enumerate(elements_z):
-            self.fig.add_trace(
-                go.Scatter(
-                    x=self.absorption_luminosities_df.index,
-                    # to plot absorption luminosities along negative y-axis
-                    y=self.absorption_luminosities_df[atomic_num] * -1,
-                    mode="none",
-                    name=atomic_number2element_symbol(atomic_num),
-                    fillcolor=self.to_rgb255_string(
-                        self.cmap(i / len(self.elements))
-                    ),
-                    stackgroup="absorption",
-                    showlegend=False,
+            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in elements_z is not in the emission df
+            try:
+                self.fig.add_trace(
+                    go.Scatter(
+                        x=self.absorption_luminosities_df.index,
+                        # to plot absorption luminosities along negative y-axis
+                        y=self.absorption_luminosities_df[atomic_num] * -1,
+                        mode="none",
+                        name=atomic_number2element_symbol(atomic_num),
+                        fillcolor=self.to_rgb255_string(
+                            self.cmap(i / len(self.elements))
+                        ),
+                        stackgroup="absorption",
+                        showlegend=False,
+                    )
                 )
-            )
+            except:
+                print(atomic_number2element_symbol(atomic_num) + " is not in the emitted packets; skipping")
 
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""

From d65ebcb4eb2a7fc8739a94a514b68b01fefd1aba Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 26 Mar 2021 16:34:50 +0000
Subject: [PATCH 02/20] Fixed typo

---
 tardis/visualization/tools/sdec_plot.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 9d519b5c4ee..06f3f72e024 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1077,7 +1077,7 @@ def _plot_absorption_mpl(self):
                     linewidth=0,
                 )
             except:
-                print(atomic_number2element_symbol(atomic_number) + " is not in the emitted packets; skipping")
+                print(atomic_number2element_symbol(atomic_number) + " is not in the absorbed packets; skipping")
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
@@ -1336,7 +1336,7 @@ def _plot_absorption_ply(self):
                     )
                 )
             except:
-                print(atomic_number2element_symbol(atomic_num) + " is not in the emitted packets; skipping")
+                print(atomic_number2element_symbol(atomic_num) + " is not in the absorbed packets; skipping")
 
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""

From 082c0c476d4abb5a8b2b62321d901884e0904ed0 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 26 Mar 2021 16:35:39 +0000
Subject: [PATCH 03/20] Black formatting

---
 tardis/visualization/tools/sdec_plot.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 06f3f72e024..13a0887fff1 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1034,7 +1034,10 @@ def _plot_emission_mpl(self):
                     linewidth=0,
                 )
             except:
-                print(atomic_number2element_symbol(atomic_number) + " is not in the emitted packets; skipping")
+                print(
+                    atomic_number2element_symbol(atomic_number)
+                    + " is not in the emitted packets; skipping"
+                )
 
     def _plot_absorption_mpl(self):
         """Plot absorption part of the SDEC Plot using matplotlib."""
@@ -1077,7 +1080,10 @@ def _plot_absorption_mpl(self):
                     linewidth=0,
                 )
             except:
-                print(atomic_number2element_symbol(atomic_number) + " is not in the absorbed packets; skipping")
+                print(
+                    atomic_number2element_symbol(atomic_number)
+                    + " is not in the absorbed packets; skipping"
+                )
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
@@ -1295,8 +1301,10 @@ def _plot_emission_ply(self):
                     )
                 )
             except:
-                print(atomic_number2element_symbol(atomic_num) + " is not in the emitted packets; skipping")
-
+                print(
+                    atomic_number2element_symbol(atomic_num)
+                    + " is not in the emitted packets; skipping"
+                )
 
     def _plot_absorption_ply(self):
         """Plot absorption part of the SDEC Plot using plotly."""
@@ -1336,7 +1344,10 @@ def _plot_absorption_ply(self):
                     )
                 )
             except:
-                print(atomic_number2element_symbol(atomic_num) + " is not in the absorbed packets; skipping")
+                print(
+                    atomic_number2element_symbol(atomic_num)
+                    + " is not in the absorbed packets; skipping"
+                )
 
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""

From 26078cd83d03118ca9ad168f050f3af737f0833c Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 26 Mar 2021 16:37:25 +0000
Subject: [PATCH 04/20] Fixed typo

---
 tardis/visualization/tools/sdec_plot.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 13a0887fff1..484647fe11f 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1063,7 +1063,7 @@ def _plot_absorption_mpl(self):
         elements_z = self.elements
         for i, atomic_number in enumerate(elements_z):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the emission df
+            # Therefore it's possible that something in elements_z is not in the absorption df
             try:
                 upper_level = lower_level
                 lower_level = (
@@ -1327,7 +1327,7 @@ def _plot_absorption_ply(self):
         elements_z = self.elements
         for i, atomic_num in enumerate(elements_z):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the emission df
+            # Therefore it's possible that something in elements_z is not in the absorption df
             try:
                 self.fig.add_trace(
                     go.Scatter(

From 7cf1464a32da2c858aec14d7dd3a5f47c20caf54 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Sun, 4 Apr 2021 22:19:57 +0100
Subject: [PATCH 05/20] Add species column to packet df

Added a new column to the packet dataframe that gives the species id. This combines the atomic and ion numbers into a unique id that will be used to group species later.
---
 tardis/visualization/tools/sdec_plot.py | 27 ++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 484647fe11f..d71e3dcefc9 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -14,8 +14,15 @@
 import matplotlib.colors as clr
 import plotly.graph_objects as go
 
-from tardis.util.base import atomic_number2element_symbol
-from tardis.visualization import plot_util as pu
+from tardis.util.base import (
+    atomic_number2element_symbol,
+    element_symbol2atomic_number,
+    species_string_to_tuple,
+    species_tuple_to_string,
+    roman_to_int,
+    int_to_roman,
+)
+from tardis.widgets import plot_util as pu
 
 
 class SDECData:
@@ -111,17 +118,23 @@ def __init__(
         self.packets_df_line_interaction = self.packets_df.loc[line_mask].copy()
 
         # Add columns for atomic number of last interaction in/out
-        self.packets_df_line_interaction["last_line_interaction_out_atom"] = (
+        self.packets_df_line_interaction["last_line_interaction_atom"] = (
             self.lines_df["atomic_number"]
             .iloc[
                 self.packets_df_line_interaction["last_line_interaction_out_id"]
             ]
             .to_numpy()
         )
-        self.packets_df_line_interaction["last_line_interaction_in_atom"] = (
+        self.packets_df_line_interaction["last_line_interaction_species"] = (
             self.lines_df["atomic_number"]
             .iloc[
-                self.packets_df_line_interaction["last_line_interaction_in_id"]
+                self.packets_df_line_interaction["last_line_interaction_out_id"]
+            ]
+            .to_numpy() * 100
+            +
+            self.lines_df["ion_number"]
+            .iloc[
+                self.packets_df_line_interaction["last_line_interaction_out_id"]
             ]
             .to_numpy()
         )
@@ -725,7 +738,7 @@ def _calculate_emission_luminosities(self, packets_mode, packet_wvl_range):
         packets_df_grouped = (
             self.data[packets_mode]
             .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
-            .groupby(by="last_line_interaction_out_atom")
+            .groupby(by="last_line_interaction_atom")
         )
 
         # Contribution of each element with which packets interacted ----------
@@ -806,7 +819,7 @@ def _calculate_absorption_luminosities(
         packets_df_grouped = (
             self.data[packets_mode]
             .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
-            .groupby(by="last_line_interaction_in_atom")
+            .groupby(by="last_line_interaction_atom")
         )
         for atomic_number, group in packets_df_grouped:
             # Histogram of specific element

From 7429924ae5bf0b7dbc9a80c25df11f0249d2a500 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Sun, 4 Apr 2021 22:55:29 +0100
Subject: [PATCH 06/20] Added species list parser

Added parser to convert request species list into ids to be used by plotter
---
 tardis/visualization/tools/sdec_plot.py | 85 ++++++++++++++++++++++++-
 1 file changed, 84 insertions(+), 1 deletion(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index d71e3dcefc9..c4ecfcd7f7b 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -437,6 +437,89 @@ def from_hdf(cls, hdf_fpath):
             )
         )
 
+
+    def _parse_species_list(
+        self, species_list
+    ):
+        """
+        Parse user requested species list and create list of species ids to be used.
+
+        Parameters
+        ----------
+        species_list : list of species to plot
+            List of species (e.g. Si II, Ca II, etc.) that the user wants to show as unique colours.
+            Species can be given as an ion (e.g. Si II), an element (e.g. Si), a range of ions
+            (e.g. Si I - V), or any combination of these (e.g. species_list = [Si II, Fe I-V, Ca])
+
+        """
+
+        # check if there are any digits in the species list. If there are then exit.
+        # species_list should only contain species in the Roman numeral
+        # format, e.g. Si II, and each ion must contain a space
+        if any(char.isdigit() for char in " ".join(species_list)) == True:
+            raise ValueError(
+            "All species must be in Roman numeral form, e.g. Si II"
+           )
+        else:
+            if species_list is not None:
+                full_species_list = []
+                for species in species_list:
+                    # check if a hyphen is present. If it is, then it indicates a
+                    # range of ions. Add each ion in that range to the list
+                    if "-" in species:
+                        # split the string on spaces. First thing in the list is then the element
+                        element = species.split(" ")[0]
+                        # Next thing is the ion range
+                        # convert the requested ions into numerals
+                        first_ion_numeral = roman_to_int(
+                            species.split(" ")[-1].split("-")[0]
+                            )
+                        second_ion_numeral = roman_to_int(
+                            species.split(" ")[-1].split("-")[-1]
+                            )
+                            # add each ion between the two requested into the species list
+                        for i in np.arange(first_ion_numeral, second_ion_numeral + 1):
+                            full_species_list.append(element + " " + int_to_roman(i))
+                    else:
+                        full_species_list.append(species)
+
+                # full_species_list is now a list containing each individual species requested
+                # e.g. it parses species_list = [Si I - V] in to species_list = [Si I, Si II, Si III, Si IV, Si V]
+                self._full_species_list = full_species_list
+                requested_species_ids = []
+                keep_colour = []
+
+                # go through each of the request species. Check whether it is
+                # an element or ion (ions have spaces). If it is an element,
+                # add all possible ions to the ions list. Otherwise just add
+                # the requested ion
+                for species in full_species_list:
+                    if " " in species:
+                        requested_species_ids.append(
+                            [
+                                species_string_to_tuple(species)[0] * 100
+                                + species_string_to_tuple(species)[1]
+                            ]
+                        )
+                    else:
+                        atomic_number = element_symbol2atomic_number(species)
+                        requested_species_ids.append(
+                            [atomic_number * 100 + i for i in np.arange(atomic_number)]
+                        )
+                        # add the atomic number to a list so you know that this element should
+                        # have all species in the same colour, i.e. it was requested like
+                        # species_list = [Si]
+                        keep_colour.append(atomic_number)
+                requested_species_ids = [
+                    species_id for list in requested_species_ids for species_id in list
+                ]
+
+
+                self._species_list = requested_species_ids
+                self._keep_colour = keep_colour
+            else:
+                self._species_list = None
+
     def _calculate_plotting_data(
         self, packets_mode, packet_wvl_range, distance, nelements
     ):
@@ -558,7 +641,7 @@ def _calculate_plotting_data(
         )
 
         # If nelements is not included, the list of elements is just all elements
-        if nelements is None:
+        if nelements is None and self._species_list is None:
             self.elements = np.array(list(self.total_luminosities_df.keys()))
         else:
             # If nelements is included then create a new column which is the sum

From 6952a190b8def33f668a2631365d4ba320b9e7ac Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Sun, 4 Apr 2021 23:02:23 +0100
Subject: [PATCH 07/20] Implement parser call

Call species list parser function
---
 tardis/visualization/tools/sdec_plot.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index c4ecfcd7f7b..1b70344e15e 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -517,6 +517,7 @@ def _parse_species_list(
 
                 self._species_list = requested_species_ids
                 self._keep_colour = keep_colour
+
             else:
                 self._species_list = None
 
@@ -969,6 +970,7 @@ def generate_plot_mpl(
         figsize=(12, 7),
         cmapname="jet",
         nelements=None,
+        species_list = None
     ):
         """
         Generate Spectral element DEComposition (SDEC) Plot using matplotlib.
@@ -1008,6 +1010,10 @@ def generate_plot_mpl(
         matplotlib.axes._subplots.AxesSubplot
             Axis on which SDEC Plot is created
         """
+
+        if species_list is not None:
+            self._parse_species_list(species_list = species_list)
+
         # Calculate data attributes required for plotting
         # and save them in instance itself
         self._calculate_plotting_data(

From cac5a30e51ec09358e80382ac77a6f2ed96e6d04 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Sun, 4 Apr 2021 23:29:48 +0100
Subject: [PATCH 08/20] Fixed parser

Rearranged parser so it didn't crash if species_list was none
---
 tardis/visualization/tools/sdec_plot.py | 26 +++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 1b70344e15e..f8180090030 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -453,15 +453,15 @@ def _parse_species_list(
 
         """
 
+        if species_list is not None:
         # check if there are any digits in the species list. If there are then exit.
         # species_list should only contain species in the Roman numeral
         # format, e.g. Si II, and each ion must contain a space
-        if any(char.isdigit() for char in " ".join(species_list)) == True:
-            raise ValueError(
-            "All species must be in Roman numeral form, e.g. Si II"
-           )
-        else:
-            if species_list is not None:
+            if any(char.isdigit() for char in " ".join(species_list)) == True:
+                raise ValueError(
+                "All species must be in Roman numeral form, e.g. Si II"
+                )
+            else:
                 full_species_list = []
                 for species in species_list:
                     # check if a hyphen is present. If it is, then it indicates a
@@ -514,12 +514,10 @@ def _parse_species_list(
                     species_id for list in requested_species_ids for species_id in list
                 ]
 
-
                 self._species_list = requested_species_ids
                 self._keep_colour = keep_colour
-
-            else:
-                self._species_list = None
+        else:
+            self._species_list = None
 
     def _calculate_plotting_data(
         self, packets_mode, packet_wvl_range, distance, nelements
@@ -644,6 +642,11 @@ def _calculate_plotting_data(
         # If nelements is not included, the list of elements is just all elements
         if nelements is None and self._species_list is None:
             self.elements = np.array(list(self.total_luminosities_df.keys()))
+        elif self._species_list is not None:
+            mask = np.in1d(
+                np.array(list(temp.keys())), self._species_list
+            )
+            print(mask)
         else:
             # If nelements is included then create a new column which is the sum
             # of all other elements, i.e. those that aren't in the top contributing nelements
@@ -1011,8 +1014,7 @@ def generate_plot_mpl(
             Axis on which SDEC Plot is created
         """
 
-        if species_list is not None:
-            self._parse_species_list(species_list = species_list)
+        self._parse_species_list(species_list = species_list)
 
         # Calculate data attributes required for plotting
         # and save them in instance itself

From 494e615396153afcc06ee54048d970a7b55b1a50 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Mon, 5 Apr 2021 20:11:15 +0100
Subject: [PATCH 09/20] Implement MPL ion contributions

Added in the ion contribution stuff for the mpl plotting.
Also added a new function to get the colorbar labels. This is necessary to determine how many unique colours there should be
---
 tardis/visualization/tools/sdec_plot.py | 313 +++++++++++++++++++-----
 1 file changed, 246 insertions(+), 67 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index f8180090030..7779297f27b 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -484,12 +484,12 @@ def _parse_species_list(
                         full_species_list.append(species)
 
                 # full_species_list is now a list containing each individual species requested
-                # e.g. it parses species_list = [Si I - V] in to species_list = [Si I, Si II, Si III, Si IV, Si V]
+                # e.g. it parses species_list = [Si I - V] into species_list = [Si I, Si II, Si III, Si IV, Si V]
                 self._full_species_list = full_species_list
                 requested_species_ids = []
                 keep_colour = []
 
-                # go through each of the request species. Check whether it is
+                # go through each of the requested species. Check whether it is
                 # an element or ion (ions have spaces). If it is an element,
                 # add all possible ions to the ions list. Otherwise just add
                 # the requested ion
@@ -614,13 +614,13 @@ def _calculate_plotting_data(
         # Calculate luminosities to be shown in plot
         (
             self.emission_luminosities_df,
-            self.emission_elements,
+            self.emission_species,
         ) = self._calculate_emission_luminosities(
             packets_mode=packets_mode, packet_wvl_range=packet_wvl_range
         )
         (
             self.absorption_luminosities_df,
-            self.absorption_elements,
+            self.absorption_species,
         ) = self._calculate_absorption_luminosities(
             packets_mode=packets_mode, packet_wvl_range=packet_wvl_range
         )
@@ -639,14 +639,55 @@ def _calculate_plotting_data(
             ascending=False
         )
 
-        # If nelements is not included, the list of elements is just all elements
+        # If nelements and species_list are not included, the list of elements is just all elements
         if nelements is None and self._species_list is None:
-            self.elements = np.array(list(self.total_luminosities_df.keys()))
+            self.species = np.array(list(self.total_luminosities_df.keys()))
         elif self._species_list is not None:
+            # Compare the species present in the model to those in the requested list
+            # Mask out those that aren't in the model
             mask = np.in1d(
-                np.array(list(temp.keys())), self._species_list
+                np.array(list(sorted_list.keys())), self._species_list
             )
-            print(mask)
+            # If species_list is included then create a new column which is the sum
+            # of all other species i.e. those that aren't in the requested list
+            self.total_luminosities_df.insert(
+                loc=0,
+                column="other",
+                value=self.total_luminosities_df[
+                    sorted_list.keys()[~mask]
+                ].sum(axis=1),
+            )
+            # Then drop all of the individual columns for species included in 'other'
+            self.total_luminosities_df.drop(
+                sorted_list.keys()[~mask], inplace=True, axis=1
+            )
+            # Repeat this for the emission and absorption dfs
+            self.emission_luminosities_df.insert(
+                loc=0,
+                column="other",
+                value=self.emission_luminosities_df[
+                    sorted_list.keys()[~mask]
+                ].sum(axis=1),
+            )
+            self.emission_luminosities_df.drop(
+                sorted_list.keys()[~mask], inplace=True, axis=1
+            )
+
+            self.absorption_luminosities_df.insert(
+                loc=0,
+                column="other",
+                value=self.absorption_luminosities_df[
+                    sorted_list.keys()[~mask]
+                ].sum(axis=1),
+            )
+            self.absorption_luminosities_df.drop(
+                sorted_list.keys()[~mask], inplace=True, axis=1
+            )
+
+            # Get the list of species in the model
+            # Index from 1: to avoid the 'other' column
+            self.species = np.sort(self.total_luminosities_df.keys()[1:])
+
         else:
             # If nelements is included then create a new column which is the sum
             # of all other elements, i.e. those that aren't in the top contributing nelements
@@ -687,9 +728,9 @@ def _calculate_plotting_data(
             self.absorption_luminosities_df.drop(
                 sorted_list.keys()[nelements:], inplace=True, axis=1
             )
-
+            # Get the list of species in the model
             # Index from 1: to avoid the 'other' column
-            self.elements = np.sort(self.total_luminosities_df.keys()[1:])
+            self.species = np.sort(self.total_luminosities_df.keys()[1:])
 
         self.photosphere_luminosity = self._calculate_photosphere_luminosity(
             packets_mode=packets_mode
@@ -822,14 +863,22 @@ def _calculate_emission_luminosities(self, packets_mode, packet_wvl_range):
 
         # Group packets_df by atomic number of elements with which packets
         # had their last emission (interaction out)
-        packets_df_grouped = (
-            self.data[packets_mode]
-            .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
-            .groupby(by="last_line_interaction_atom")
-        )
+        # or if species_list is requested then group by species id
+        if self._species_list is None:
+            packets_df_grouped = (
+                self.data[packets_mode]
+                .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
+                .groupby(by="last_line_interaction_atom")
+                )
+        else:
+            packets_df_grouped = (
+                self.data[packets_mode]
+                .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
+                .groupby(by="last_line_interaction_species")
+                )
 
-        # Contribution of each element with which packets interacted ----------
-        for atomic_number, group in packets_df_grouped:
+        # Contribution of each species with which packets interacted ----------
+        for identifier, group in packets_df_grouped:
             # Histogram of specific element
             hist_el = np.histogram(
                 group["nus"],
@@ -848,14 +897,14 @@ def _calculate_emission_luminosities(self, packets_mode, packet_wvl_range):
             )
             L_lambda_el = L_nu_el * self.plot_frequency / self.plot_wavelength
 
-            luminosities_df[atomic_number] = L_lambda_el.value
+            luminosities_df[identifier] = L_lambda_el.value
 
-        # Create an array of the elements with which packets interacted
-        emission_elements_present = np.array(
+        # Create an array of the species with which packets interacted
+        emission_species_present = np.array(
             list(packets_df_grouped.groups.keys())
         )
 
-        return luminosities_df, emission_elements_present
+        return luminosities_df, emission_species_present
 
     def _calculate_absorption_luminosities(
         self, packets_mode, packet_wvl_range
@@ -903,12 +952,20 @@ def _calculate_absorption_luminosities(
 
         # Group packets_df by atomic number of elements with which packets
         # had their last absorption (interaction in)
-        packets_df_grouped = (
-            self.data[packets_mode]
-            .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
-            .groupby(by="last_line_interaction_atom")
-        )
-        for atomic_number, group in packets_df_grouped:
+        if self._species_list is None:
+            packets_df_grouped = (
+                self.data[packets_mode]
+                .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
+                .groupby(by="last_line_interaction_atom")
+                )
+        else:
+            packets_df_grouped = (
+                self.data[packets_mode]
+                .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
+                .groupby(by="last_line_interaction_species")
+                )
+
+        for identifier, group in packets_df_grouped:
             # Histogram of specific element
             hist_el = np.histogram(
                 group["last_line_interaction_in_nu"],
@@ -927,13 +984,13 @@ def _calculate_absorption_luminosities(
             )
             L_lambda_el = L_nu_el * self.plot_frequency / self.plot_wavelength
 
-            luminosities_df[atomic_number] = L_lambda_el.value
+            luminosities_df[identifier] = L_lambda_el.value
 
-        absorption_elements_present = np.array(
+        absorption_species_present = np.array(
             list(packets_df_grouped.groups.keys())
         )
 
-        return luminosities_df, absorption_elements_present
+        return luminosities_df, absorption_species_present
 
     def _calculate_photosphere_luminosity(self, packets_mode):
         """
@@ -1007,6 +1064,11 @@ def generate_plot_mpl(
             largest contribution to total luminosity absorbed and emitted.
             Other elements are shown in silver. Default value is
             None, which displays all elements
+        species_list: list of strings or None
+            list of strings containing the names of species that should be included in the Kromer plots.
+            Must be given in Roman numeral format. Can include specific ions, a range of ions,
+            individual elements, or any combination of these:
+            e.g. ['Si II', 'Ca II', 'C', 'Fe I-V']
 
         Returns
         -------
@@ -1031,7 +1093,10 @@ def generate_plot_mpl(
             self.ax = ax
 
         # Set colormap to be used in elements of emission and absorption plots
-        self.cmap = cm.get_cmap(cmapname, self.elements.size)
+        self._make_colorbar_labels_mpl()
+        self.cmap = cm.get_cmap(cmapname, len(self._species_name))
+        self._show_colorbar_mpl()
+
 
         self._plot_emission_mpl()
         self._plot_absorption_mpl()
@@ -1054,7 +1119,6 @@ def generate_plot_mpl(
             label="Blackbody Photosphere",
         )
 
-        self._show_colorbar_mpl()
 
         # Set legends and labels
         self.ax.legend(fontsize=12)
@@ -1117,31 +1181,76 @@ def _plot_emission_mpl(self):
                 label="Other elements",
             )
 
-        elements_z = self.elements
         # Contribution from each element
-        for i, atomic_number in enumerate(elements_z):
+        # Create new variables to keep track of the last atomic number that was plotted
+        # This is used when plotting species incase an element was given in the list
+        # This is to ensure that all ions of that element are grouped together
+        # ii is to track the colour index
+        ii = 0
+        previous_atomic_number = 0
+        for i, identifier in enumerate(self.species):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the emission df
             try:
                 lower_level = upper_level
                 upper_level = (
                     lower_level
-                    + self.emission_luminosities_df[atomic_number].to_numpy()
+                    + self.emission_luminosities_df[identifier].to_numpy()
                 )
 
+                if self._species_list is not None:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+                    if previous_atomic_number == 0:
+                        # If this is the first species being plotted, then take note of the atomic number
+                        # don't update the colour index
+                        ii = ii
+                        previous_atomic_number = atomic_number
+                    elif previous_atomic_number in self._keep_colour:
+                        # If the atomic number is in the list of elements that should all be plotted in the same colour
+                        # then dont update the colour index
+                        if previous_atomic_number == atomic_number:
+                            ii = ii
+                            previous_atomic_number = atomic_number
+                        else:
+                            # Otherwise, increase the colour counter by one
+                            ii = ii + 1
+                            previous_atomic_number = atomic_number
+                    else:
+                        # If this is just a normal species that was requested then increment the colour index
+                        ii = ii + 1
+                        previous_atomic_number = atomic_number
+                    # Calculate the colour of this species
+                    color = self.cmap(ii / len(self._species_name))
+                    print(ii, identifier, self._species_name, color)
+
+                else:
+                    # If you're not using species list then this is just a fraction based on the total
+                    # number of columns in the dataframe
+                    color = self.cmap(i / len(self.species))
+                    print(i, identifier, self.species, color)
+
                 self.ax.fill_between(
                     self.plot_wavelength,
                     lower_level,
                     upper_level,
-                    color=self.cmap(i / len(self.elements)),
+                    color=color,
                     cmap=self.cmap,
                     linewidth=0,
                 )
             except:
-                print(
-                    atomic_number2element_symbol(atomic_number)
-                    + " is not in the emitted packets; skipping"
-                )
+                if self._species_list is None:
+                    print(
+                        atomic_number2element_symbol(identifier)
+                        + " is not in the emitted packets; skipping"
+                        )
+                else:
+                    print(
+                        atomic_number2element_symbol(atomic_number)
+                        + int_to_roman(ion_number + 1)
+                        + " is not in the emitted packets; skipping"
+                        )
 
     def _plot_absorption_mpl(self):
         """Plot absorption part of the SDEC Plot using matplotlib."""
@@ -1164,50 +1273,120 @@ def _plot_absorption_mpl(self):
                 color="silver",
             )
 
-        elements_z = self.elements
-        for i, atomic_number in enumerate(elements_z):
+        ii = 0
+        previous_atomic_number = 0
+        for i, identifier in enumerate(self.species):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the absorption df
             try:
                 upper_level = lower_level
                 lower_level = (
                     upper_level
-                    - self.absorption_luminosities_df[atomic_number].to_numpy()
+                    - self.absorption_luminosities_df[identifier].to_numpy()
                 )
 
+                if self._species_list is not None:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+                    if previous_atomic_number == 0:
+                        # If this is the first species being plotted, then take note of the atomic number
+                        # don't update the colour index
+                        ii = ii
+                        previous_atomic_number = atomic_number
+                    elif previous_atomic_number in self._keep_colour:
+                        # If the atomic number is in the list of elements that should all be plotted in the same colour
+                        # then dont update the colour index
+                        if previous_atomic_number == atomic_number:
+                            ii = ii
+                            previous_atomic_number = atomic_number
+                        else:
+                            # Otherwise, increase the colour counter by one
+                            ii = ii + 1
+                            previous_atomic_number = atomic_number
+                    else:
+                        # If this is just a normal species that was requested then increment the colour index
+                        ii = ii + 1
+                        previous_atomic_number = atomic_number
+                    # Calculate the colour of this species
+                    color = self.cmap(ii / len(self._species_name))
+                    print(ii, identifier, self._species_name, color)
+
+                else:
+                    # If you're not using species list then this is just a fraction based on the total
+                    # number of columns in the dataframe
+                    color = self.cmap(i / len(self.species))
+                    print(i, identifier, self.species, color)
+
                 self.ax.fill_between(
                     self.plot_wavelength,
                     upper_level,
                     lower_level,
-                    color=self.cmap(i / len(elements_z)),
+                    color=color,
                     cmap=self.cmap,
                     linewidth=0,
                 )
             except:
-                print(
-                    atomic_number2element_symbol(atomic_number)
-                    + " is not in the absorbed packets; skipping"
-                )
+                if self._species_list is None:
+                    print(
+                        atomic_number2element_symbol(identifier)
+                        + " is not in the emitted packets; skipping"
+                        )
+                else:
+                    print(
+                        atomic_number2element_symbol(atomic_number)
+                        + int_to_roman(ion_number + 1)
+                        + " is not in the emitted packets; skipping"
+                        )
+
+    def _make_colorbar_labels_mpl(self):
+        """Get the labels for the species in the colorbar."""
+        if self._species_list is None:
+            species_name = [
+                atomic_number2element_symbol(atomic_num)
+                for atomic_num in self.species
+            ]
+        else:
+            species_name = []
+            for species in self.species:
+                ion_number = species % 100
+                atomic_number = (species - ion_number) / 100
+
+                ion_numeral = int_to_roman(ion_number + 1)
+                atomic_symbol = atomic_number2element_symbol(atomic_number)
+
+                # if the element was requested, and not a specific ion, then
+                # add the element symbol to the label list
+                if (atomic_number in self._keep_colour) & (atomic_symbol not in species_name):
+                    # compiling the label, and adding it to the list
+                    label = f"{atomic_symbol}"
+                    species_name.append(label)
+                elif atomic_number not in self._keep_colour:
+                    # otherwise add the ion to the label list
+                    label = f"{atomic_symbol}$\,${ion_numeral}"
+                    species_name.append(label)
+
+        self._species_name = species_name
+
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
+
         color_values = [
-            self.cmap(i / self.elements.size) for i in range(self.elements.size)
+            self.cmap(i / len(self._species_name)) for i in range(len(self._species_name))
         ]
+
         custcmap = clr.ListedColormap(color_values)
-        norm = clr.Normalize(vmin=0, vmax=self.elements.size)
+        norm = clr.Normalize(vmin=0, vmax=len(self._species_name))
         mappable = cm.ScalarMappable(norm=norm, cmap=custcmap)
-        mappable.set_array(np.linspace(1, self.elements.size + 1, 256))
+        mappable.set_array(np.linspace(1, len(self._species_name) + 1, 256))
         cbar = plt.colorbar(mappable, ax=self.ax)
 
-        bounds = np.arange(self.elements.size) + 0.5
+        bounds = np.arange(len(self._species_name)) + 0.5
         cbar.set_ticks(bounds)
 
-        elements_name = [
-            atomic_number2element_symbol(atomic_num)
-            for atomic_num in self.elements
-        ]
-        cbar.set_ticklabels(elements_name)
+        cbar.set_ticklabels(self._species_name)
+
 
     def generate_plot_ply(
         self,
@@ -1274,7 +1453,7 @@ def generate_plot_ply(
             self.fig = fig
 
         # Set colormap to be used in elements of emission and absorption plots
-        self.cmap = cm.get_cmap(cmapname, self.elements.size)
+        self.cmap = cm.get_cmap(cmapname, self.species.size)
 
         self._plot_emission_ply()
         self._plot_absorption_ply()
@@ -1386,7 +1565,7 @@ def _plot_emission_ply(self):
                 )
             )
 
-        elements_z = self.elements
+        elements_z = self.species
         for i, atomic_num in enumerate(elements_z):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the emission df
@@ -1398,7 +1577,7 @@ def _plot_emission_ply(self):
                         mode="none",
                         name=atomic_number2element_symbol(atomic_num),
                         fillcolor=self.to_rgb255_string(
-                            self.cmap(i / len(self.elements))
+                            self.cmap(i / len(self.species))
                         ),
                         stackgroup="emission",
                         showlegend=False,
@@ -1428,7 +1607,7 @@ def _plot_absorption_ply(self):
                 )
             )
 
-        elements_z = self.elements
+        elements_z = self.species
         for i, atomic_num in enumerate(elements_z):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the absorption df
@@ -1441,7 +1620,7 @@ def _plot_absorption_ply(self):
                         mode="none",
                         name=atomic_number2element_symbol(atomic_num),
                         fillcolor=self.to_rgb255_string(
-                            self.cmap(i / len(self.elements))
+                            self.cmap(i / len(self.species))
                         ),
                         stackgroup="absorption",
                         showlegend=False,
@@ -1456,13 +1635,13 @@ def _plot_absorption_ply(self):
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""
         # Interpolate [0, 1] range to create bins equal to number of elements
-        colorscale_bins = np.linspace(0, 1, num=self.elements.size + 1)
+        colorscale_bins = np.linspace(0, 1, num=self.species.size + 1)
 
         # Create a categorical colorscale [a list of (reference point, color)]
         # by mapping same reference points (excluding 1st and last bin edge)
         # twice in a row (https://plotly.com/python/colorscales/#constructing-a-discrete-or-discontinuous-color-scale)
         categorical_colorscale = []
-        for i in range(self.elements.size):
+        for i in range(self.species.size):
             color = self.to_rgb255_string(self.cmap(colorscale_bins[i]))
             categorical_colorscale.append((colorscale_bins[i], color))
             categorical_colorscale.append((colorscale_bins[i + 1], color))
@@ -1471,13 +1650,13 @@ def _show_colorbar_ply(self):
             colorscale=categorical_colorscale,
             showscale=True,
             cmin=0,
-            cmax=self.elements.size,
+            cmax=self.species.size,
             colorbar=dict(
                 title="Elements",
-                tickvals=np.arange(0, self.elements.size) + 0.5,
+                tickvals=np.arange(0, self.species.size) + 0.5,
                 ticktext=[
                     atomic_number2element_symbol(atomic_num)
-                    for atomic_num in self.elements
+                    for atomic_num in self.species
                 ],
                 # to change length and position of colorbar
                 len=0.75,

From 0a9c708e14c532b8431c589f99f7991ecf2073e9 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Mon, 5 Apr 2021 20:12:50 +0100
Subject: [PATCH 10/20] Changed order of functions

---
 tardis/visualization/tools/sdec_plot.py | 42 ++++++++++++-------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 7779297f27b..e908ab0db66 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1093,7 +1093,7 @@ def generate_plot_mpl(
             self.ax = ax
 
         # Set colormap to be used in elements of emission and absorption plots
-        self._make_colorbar_labels_mpl()
+        self._make_colorbar_labels()
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
         self._show_colorbar_mpl()
 
@@ -1339,7 +1339,26 @@ def _plot_absorption_mpl(self):
                         + " is not in the emitted packets; skipping"
                         )
 
-    def _make_colorbar_labels_mpl(self):
+    def _show_colorbar_mpl(self):
+        """Show matplotlib colorbar with labels of elements mapped to colors."""
+
+        color_values = [
+            self.cmap(i / len(self._species_name)) for i in range(len(self._species_name))
+        ]
+
+        custcmap = clr.ListedColormap(color_values)
+        norm = clr.Normalize(vmin=0, vmax=len(self._species_name))
+        mappable = cm.ScalarMappable(norm=norm, cmap=custcmap)
+        mappable.set_array(np.linspace(1, len(self._species_name) + 1, 256))
+        cbar = plt.colorbar(mappable, ax=self.ax)
+
+        bounds = np.arange(len(self._species_name)) + 0.5
+        cbar.set_ticks(bounds)
+
+        cbar.set_ticklabels(self._species_name)
+
+
+    def _make_colorbar_labels(self):
         """Get the labels for the species in the colorbar."""
         if self._species_list is None:
             species_name = [
@@ -1369,25 +1388,6 @@ def _make_colorbar_labels_mpl(self):
         self._species_name = species_name
 
 
-    def _show_colorbar_mpl(self):
-        """Show matplotlib colorbar with labels of elements mapped to colors."""
-
-        color_values = [
-            self.cmap(i / len(self._species_name)) for i in range(len(self._species_name))
-        ]
-
-        custcmap = clr.ListedColormap(color_values)
-        norm = clr.Normalize(vmin=0, vmax=len(self._species_name))
-        mappable = cm.ScalarMappable(norm=norm, cmap=custcmap)
-        mappable.set_array(np.linspace(1, len(self._species_name) + 1, 256))
-        cbar = plt.colorbar(mappable, ax=self.ax)
-
-        bounds = np.arange(len(self._species_name)) + 0.5
-        cbar.set_ticks(bounds)
-
-        cbar.set_ticklabels(self._species_name)
-
-
     def generate_plot_ply(
         self,
         packets_mode="virtual",

From 6c56d4c719b59ec562dac77dd5662427365a374c Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Mon, 5 Apr 2021 21:37:14 +0100
Subject: [PATCH 11/20] Implement plotly ion contributions

Added colouring for ion contributions in plotly figures
---
 tardis/visualization/tools/sdec_plot.py | 303 ++++++++++++++++--------
 1 file changed, 198 insertions(+), 105 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index e908ab0db66..7ecdeb4ba25 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1030,7 +1030,7 @@ def generate_plot_mpl(
         figsize=(12, 7),
         cmapname="jet",
         nelements=None,
-        species_list = None
+        species_list=None
     ):
         """
         Generate Spectral element DEComposition (SDEC) Plot using matplotlib.
@@ -1092,8 +1092,9 @@ def generate_plot_mpl(
         else:
             self.ax = ax
 
-        # Set colormap to be used in elements of emission and absorption plots
+        # Get the labels in the color bar. This determines the number of unique colors
         self._make_colorbar_labels()
+        # Set colormap to be used in elements of emission and absorption plots
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
         self._show_colorbar_mpl()
 
@@ -1191,45 +1192,45 @@ def _plot_emission_mpl(self):
         for i, identifier in enumerate(self.species):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the emission df
-            try:
-                lower_level = upper_level
-                upper_level = (
-                    lower_level
-                    + self.emission_luminosities_df[identifier].to_numpy()
-                )
 
-                if self._species_list is not None:
-                    # Get the ion number and atomic number for each species
-                    ion_number = identifier % 100
-                    atomic_number = (identifier - ion_number) / 100
-                    if previous_atomic_number == 0:
-                        # If this is the first species being plotted, then take note of the atomic number
-                        # don't update the colour index
+
+            if self._species_list is not None:
+                # Get the ion number and atomic number for each species
+                ion_number = identifier % 100
+                atomic_number = (identifier - ion_number) / 100
+                if previous_atomic_number == 0:
+                    # If this is the first species being plotted, then take note of the atomic number
+                    # don't update the colour index
+                    ii = ii
+                    previous_atomic_number = atomic_number
+                elif previous_atomic_number in self._keep_colour:
+                    # If the atomic number is in the list of elements that should all be plotted in the same colour
+                    # then dont update the colour index
+                    if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
-                    elif previous_atomic_number in self._keep_colour:
-                        # If the atomic number is in the list of elements that should all be plotted in the same colour
-                        # then dont update the colour index
-                        if previous_atomic_number == atomic_number:
-                            ii = ii
-                            previous_atomic_number = atomic_number
-                        else:
-                            # Otherwise, increase the colour counter by one
-                            ii = ii + 1
-                            previous_atomic_number = atomic_number
                     else:
-                        # If this is just a normal species that was requested then increment the colour index
+                        # Otherwise, increase the colour counter by one
                         ii = ii + 1
                         previous_atomic_number = atomic_number
-                    # Calculate the colour of this species
-                    color = self.cmap(ii / len(self._species_name))
-                    print(ii, identifier, self._species_name, color)
-
                 else:
-                    # If you're not using species list then this is just a fraction based on the total
-                    # number of columns in the dataframe
-                    color = self.cmap(i / len(self.species))
-                    print(i, identifier, self.species, color)
+                    # If this is just a normal species that was requested then increment the colour index
+                    ii = ii + 1
+                    previous_atomic_number = atomic_number
+                # Calculate the colour of this species
+                color = self.cmap(ii / len(self._species_name))
+
+            else:
+                # If you're not using species list then this is just a fraction based on the total
+                # number of columns in the dataframe
+                color = self.cmap(i / len(self.species))
+
+            try:
+                lower_level = upper_level
+                upper_level = (
+                    lower_level
+                    + self.emission_luminosities_df[identifier].to_numpy()
+                )
 
                 self.ax.fill_between(
                     self.plot_wavelength,
@@ -1278,45 +1279,44 @@ def _plot_absorption_mpl(self):
         for i, identifier in enumerate(self.species):
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the absorption df
-            try:
-                upper_level = lower_level
-                lower_level = (
-                    upper_level
-                    - self.absorption_luminosities_df[identifier].to_numpy()
-                )
 
-                if self._species_list is not None:
-                    # Get the ion number and atomic number for each species
-                    ion_number = identifier % 100
-                    atomic_number = (identifier - ion_number) / 100
-                    if previous_atomic_number == 0:
-                        # If this is the first species being plotted, then take note of the atomic number
-                        # don't update the colour index
+            if self._species_list is not None:
+                # Get the ion number and atomic number for each species
+                ion_number = identifier % 100
+                atomic_number = (identifier - ion_number) / 100
+                if previous_atomic_number == 0:
+                    # If this is the first species being plotted, then take note of the atomic number
+                    # don't update the colour index
+                    ii = ii
+                    previous_atomic_number = atomic_number
+                elif previous_atomic_number in self._keep_colour:
+                    # If the atomic number is in the list of elements that should all be plotted in the same colour
+                    # then dont update the colour index
+                    if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
-                    elif previous_atomic_number in self._keep_colour:
-                        # If the atomic number is in the list of elements that should all be plotted in the same colour
-                        # then dont update the colour index
-                        if previous_atomic_number == atomic_number:
-                            ii = ii
-                            previous_atomic_number = atomic_number
-                        else:
-                            # Otherwise, increase the colour counter by one
-                            ii = ii + 1
-                            previous_atomic_number = atomic_number
                     else:
-                        # If this is just a normal species that was requested then increment the colour index
+                        # Otherwise, increase the colour counter by one
                         ii = ii + 1
                         previous_atomic_number = atomic_number
-                    # Calculate the colour of this species
-                    color = self.cmap(ii / len(self._species_name))
-                    print(ii, identifier, self._species_name, color)
-
                 else:
-                    # If you're not using species list then this is just a fraction based on the total
-                    # number of columns in the dataframe
-                    color = self.cmap(i / len(self.species))
-                    print(i, identifier, self.species, color)
+                    # If this is just a normal species that was requested then increment the colour index
+                    ii = ii + 1
+                    previous_atomic_number = atomic_number
+                # Calculate the colour of this species
+                color = self.cmap(ii / len(self._species_name))
+
+            else:
+                # If you're not using species list then this is just a fraction based on the total
+                # number of columns in the dataframe
+                color = self.cmap(i / len(self.species))
+
+            try:
+                upper_level = lower_level
+                lower_level = (
+                    upper_level
+                    - self.absorption_luminosities_df[identifier].to_numpy()
+                )
 
                 self.ax.fill_between(
                     self.plot_wavelength,
@@ -1326,6 +1326,7 @@ def _plot_absorption_mpl(self):
                     cmap=self.cmap,
                     linewidth=0,
                 )
+
             except:
                 if self._species_list is None:
                     print(
@@ -1382,7 +1383,7 @@ def _make_colorbar_labels(self):
                     species_name.append(label)
                 elif atomic_number not in self._keep_colour:
                     # otherwise add the ion to the label list
-                    label = f"{atomic_symbol}$\,${ion_numeral}"
+                    label = f"{atomic_symbol} {ion_numeral}"
                     species_name.append(label)
 
         self._species_name = species_name
@@ -1399,6 +1400,7 @@ def generate_plot_ply(
         graph_height=600,
         cmapname="jet",
         nelements=None,
+        species_list=None,
     ):
         """
         Generate interactive Spectral element DEComposition (SDEC) Plot using plotly.
@@ -1432,12 +1434,19 @@ def generate_plot_ply(
             largest contribution to total luminosity absorbed and emitted.
             Other elements are shown in silver. Default value is
             None, which displays all elements
-
+        species_list: list of strings or None
+            list of strings containing the names of species that should be included in the Kromer plots.
+            Must be given in Roman numeral format. Can include specific ions, a range of ions,
+            individual elements, or any combination of these:
+            e.g. ['Si II', 'Ca II', 'C', 'Fe I-V']
         Returns
         -------
         plotly.graph_objs._figure.Figure
             Figure object on which SDEC Plot is created
         """
+
+        self._parse_species_list(species_list = species_list)
+
         # Calculate data attributes required for plotting
         # and save them in instance itself
         self._calculate_plotting_data(
@@ -1452,12 +1461,17 @@ def generate_plot_ply(
         else:
             self.fig = fig
 
+        # Get the labels in the color bar. This determines the number of unique colors
+        self._make_colorbar_labels()
         # Set colormap to be used in elements of emission and absorption plots
-        self.cmap = cm.get_cmap(cmapname, self.species.size)
+        self.cmap = cm.get_cmap(cmapname, len(self._species_name))
 
         self._plot_emission_ply()
         self._plot_absorption_ply()
 
+        print(self.emission_luminosities_df)
+        print(self.absorption_luminosities_df)
+
         # Plot modeled spectrum
         if show_modeled_spectrum:
             self.fig.add_trace(
@@ -1511,7 +1525,6 @@ def generate_plot_ply(
     def to_rgb255_string(color_tuple):
         """
         Convert a matplotlib RGBA tuple to a generic RGB 255 string.
-
         Parameters
         ----------
         color_tuple : tuple
@@ -1565,29 +1578,72 @@ def _plot_emission_ply(self):
                 )
             )
 
-        elements_z = self.species
-        for i, atomic_num in enumerate(elements_z):
+        # Contribution from each element
+        # Create new variables to keep track of the last atomic number that was plotted
+        # This is used when plotting species incase an element was given in the list
+        # This is to ensure that all ions of that element are grouped together
+        # ii is to track the colour index
+        ii = 0
+        previous_atomic_number = 0
+        for i, identifier in enumerate(self.species):
+
+            if self._species_list is not None:
+                # Get the ion number and atomic number for each species
+                ion_number = identifier % 100
+                atomic_number = (identifier - ion_number) / 100
+                if previous_atomic_number == 0:
+                    # If this is the first species being plotted, then take note of the atomic number
+                    # don't update the colour index
+                    ii = ii
+                    previous_atomic_number = atomic_number
+                elif previous_atomic_number in self._keep_colour:
+                    # If the atomic number is in the list of elements that should all be plotted in the same colour
+                   # then dont update the colour index
+                    if previous_atomic_number == atomic_number:
+                        ii = ii
+                        previous_atomic_number = atomic_number
+                    else:
+                        # Otherwise, increase the colour counter by one
+                        ii = ii + 1
+                        previous_atomic_number = atomic_number
+                else:
+                    # If this is just a normal species that was requested then increment the colour index
+                    ii = ii + 1
+                    previous_atomic_number = atomic_number
+                # Calculate the colour of this species
+                color = self.cmap(ii / len(self._species_name))
+            else:
+                # If you're not using species list then this is just a fraction based on the total
+                # number of columns in the dataframe
+                color = self.cmap(i / len(self.species))
+
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the emission df
             try:
                 self.fig.add_trace(
-                    go.Scatter(
-                        x=self.emission_luminosities_df.index,
-                        y=self.emission_luminosities_df[atomic_num],
-                        mode="none",
-                        name=atomic_number2element_symbol(atomic_num),
-                        fillcolor=self.to_rgb255_string(
-                            self.cmap(i / len(self.species))
-                        ),
-                        stackgroup="emission",
-                        showlegend=False,
+                        go.Scatter(
+                            x=self.emission_luminosities_df.index,
+                            y=self.emission_luminosities_df[identifier],
+                            mode="none",
+                            name="none",
+                            fillcolor=self.to_rgb255_string(
+                                color),
+                            stackgroup="emission",
+                            showlegend=False,
+                        )
                     )
-                )
             except:
-                print(
-                    atomic_number2element_symbol(atomic_num)
-                    + " is not in the emitted packets; skipping"
-                )
+                if self._species_list is None:
+                    print(
+                        atomic_number2element_symbol(identifier)
+                        + " is not in the emitted packets; skipping"
+                        )
+                else:
+                    print(
+                        atomic_number2element_symbol(atomic_number)
+                        + int_to_roman(ion_number + 1)
+                        + " is not in the emitted packets; skipping"
+                        )
 
     def _plot_absorption_ply(self):
         """Plot absorption part of the SDEC Plot using plotly."""
@@ -1607,8 +1663,39 @@ def _plot_absorption_ply(self):
                 )
             )
 
-        elements_z = self.species
-        for i, atomic_num in enumerate(elements_z):
+        ii = 0
+        previous_atomic_number = 0
+        for i, identifier in enumerate(self.species):
+            if self._species_list is not None:
+                # Get the ion number and atomic number for each species
+                ion_number = identifier % 100
+                atomic_number = (identifier - ion_number) / 100
+                if previous_atomic_number == 0:
+                    # If this is the first species being plotted, then take note of the atomic number
+                    # don't update the colour index
+                    ii = ii
+                    previous_atomic_number = atomic_number
+                elif previous_atomic_number in self._keep_colour:
+                    # If the atomic number is in the list of elements that should all be plotted in the same colour
+                   # then dont update the colour index
+                    if previous_atomic_number == atomic_number:
+                        ii = ii
+                        previous_atomic_number = atomic_number
+                    else:
+                        # Otherwise, increase the colour counter by one
+                        ii = ii + 1
+                        previous_atomic_number = atomic_number
+                else:
+                    # If this is just a normal species that was requested then increment the colour index
+                    ii = ii + 1
+                    previous_atomic_number = atomic_number
+                # Calculate the colour of this species
+                color = self.cmap(ii / len(self._species_name))
+            else:
+                # If you're not using species list then this is just a fraction based on the total
+                # number of columns in the dataframe
+                color = self.cmap(i / len(self.species))
+
             # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
             # Therefore it's possible that something in elements_z is not in the absorption df
             try:
@@ -1616,32 +1703,39 @@ def _plot_absorption_ply(self):
                     go.Scatter(
                         x=self.absorption_luminosities_df.index,
                         # to plot absorption luminosities along negative y-axis
-                        y=self.absorption_luminosities_df[atomic_num] * -1,
+                        y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
-                        name=atomic_number2element_symbol(atomic_num),
+                        name="none",
                         fillcolor=self.to_rgb255_string(
-                            self.cmap(i / len(self.species))
-                        ),
+                            color),
                         stackgroup="absorption",
                         showlegend=False,
                     )
                 )
+
             except:
-                print(
-                    atomic_number2element_symbol(atomic_num)
-                    + " is not in the absorbed packets; skipping"
-                )
+                if self._species_list is None:
+                    print(
+                        atomic_number2element_symbol(identifier)
+                        + " is not in the emitted packets; skipping"
+                        )
+                else:
+                    print(
+                        atomic_number2element_symbol(atomic_number)
+                        + int_to_roman(ion_number + 1)
+                        + " is not in the emitted packets; skipping"
+                        )
 
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""
         # Interpolate [0, 1] range to create bins equal to number of elements
-        colorscale_bins = np.linspace(0, 1, num=self.species.size + 1)
+        colorscale_bins = np.linspace(0, 1, num=len(self._species_name) + 1)
 
         # Create a categorical colorscale [a list of (reference point, color)]
         # by mapping same reference points (excluding 1st and last bin edge)
         # twice in a row (https://plotly.com/python/colorscales/#constructing-a-discrete-or-discontinuous-color-scale)
         categorical_colorscale = []
-        for i in range(self.species.size):
+        for i in range(len(self._species_name)):
             color = self.to_rgb255_string(self.cmap(colorscale_bins[i]))
             categorical_colorscale.append((colorscale_bins[i], color))
             categorical_colorscale.append((colorscale_bins[i + 1], color))
@@ -1650,20 +1744,19 @@ def _show_colorbar_ply(self):
             colorscale=categorical_colorscale,
             showscale=True,
             cmin=0,
-            cmax=self.species.size,
+            cmax=len(self._species_name),
             colorbar=dict(
                 title="Elements",
-                tickvals=np.arange(0, self.species.size) + 0.5,
-                ticktext=[
-                    atomic_number2element_symbol(atomic_num)
-                    for atomic_num in self.species
-                ],
+                tickvals=np.arange(0, len(self._species_name)) + 0.5,
+                ticktext=self._species_name,
                 # to change length and position of colorbar
                 len=0.75,
                 yanchor="top",
                 y=0.75,
             ),
         )
+        print(self._species_name)
+        print(coloraxis_options)
 
         # Plot an invisible one point scatter trace, to make colorbar show up
         scatter_point_idx = pu.get_mid_point_idx(self.plot_wavelength)

From 607f1557692fa7cdaf4c87e350cf7bbbee00eeb2 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Mon, 5 Apr 2021 21:38:28 +0100
Subject: [PATCH 12/20] Remove print statements

---
 tardis/visualization/tools/sdec_plot.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 7ecdeb4ba25..ebba6c4525a 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -22,7 +22,7 @@
     roman_to_int,
     int_to_roman,
 )
-from tardis.widgets import plot_util as pu
+from tardis.visualization import plot_util as pu
 
 
 class SDECData:
@@ -1469,9 +1469,6 @@ def generate_plot_ply(
         self._plot_emission_ply()
         self._plot_absorption_ply()
 
-        print(self.emission_luminosities_df)
-        print(self.absorption_luminosities_df)
-
         # Plot modeled spectrum
         if show_modeled_spectrum:
             self.fig.add_trace(
@@ -1755,8 +1752,6 @@ def _show_colorbar_ply(self):
                 y=0.75,
             ),
         )
-        print(self._species_name)
-        print(coloraxis_options)
 
         # Plot an invisible one point scatter trace, to make colorbar show up
         scatter_point_idx = pu.get_mid_point_idx(self.plot_wavelength)

From 385c5b55a320d4ab7dcf8cbeb4773188fee198ec Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Wed, 14 Apr 2021 23:46:16 +0100
Subject: [PATCH 13/20] Fixed df bug and added comments

Fixed bug in the absorption and emission df calculations. Now they elements should be handled correctly so there won't be errors if a species appears in emission and not absorption, and vice versa

Added comments throughout
---
 tardis/visualization/tools/sdec_plot.py | 98 +++++++++++++++++--------
 1 file changed, 67 insertions(+), 31 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index ebba6c4525a..60dcadca7d4 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -117,7 +117,7 @@ def __init__(
         )  # & operator is quite faster than np.logical_and on pd.Series
         self.packets_df_line_interaction = self.packets_df.loc[line_mask].copy()
 
-        # Add columns for atomic number of last interaction in/out
+        # Add columns for atomic number of last interaction out
         self.packets_df_line_interaction["last_line_interaction_atom"] = (
             self.lines_df["atomic_number"]
             .iloc[
@@ -125,6 +125,8 @@ def __init__(
             ]
             .to_numpy()
         )
+        # Add columns for the species id of last interaction
+        # Species id is given by 100 * Z + X, where Z is atomic number and X is ion number
         self.packets_df_line_interaction["last_line_interaction_species"] = (
             self.lines_df["atomic_number"]
             .iloc[
@@ -454,7 +456,7 @@ def _parse_species_list(
         """
 
         if species_list is not None:
-        # check if there are any digits in the species list. If there are then exit.
+        # check if there are any digits in the species list. If there are, then exit.
         # species_list should only contain species in the Roman numeral
         # format, e.g. Si II, and each ion must contain a space
             if any(char.isdigit() for char in " ".join(species_list)) == True:
@@ -465,7 +467,7 @@ def _parse_species_list(
                 full_species_list = []
                 for species in species_list:
                     # check if a hyphen is present. If it is, then it indicates a
-                    # range of ions. Add each ion in that range to the list
+                    # range of ions. Add each ion in that range to the list as a new entry
                     if "-" in species:
                         # split the string on spaces. First thing in the list is then the element
                         element = species.split(" ")[0]
@@ -481,6 +483,7 @@ def _parse_species_list(
                         for i in np.arange(first_ion_numeral, second_ion_numeral + 1):
                             full_species_list.append(element + " " + int_to_roman(i))
                     else:
+                        # Otherwise it's either an element or ion so just add to the list
                         full_species_list.append(species)
 
                 # full_species_list is now a list containing each individual species requested
@@ -662,26 +665,47 @@ def _calculate_plotting_data(
                 sorted_list.keys()[~mask], inplace=True, axis=1
             )
             # Repeat this for the emission and absorption dfs
+            # This will require creating a temporary list that includes 'noint' and 'escatter'
+            # packets, because you don't want them dropped or included in 'other'
+            temp = [species for species in self._species_list]
+            temp.append('noint')
+            temp.append('escatter')
+            mask = np.in1d(
+                np.array(list(self.emission_luminosities_df.keys())), temp
+            )
+            # If species_list is included then create a new column which is the sum
+            # of all other species i.e. those that aren't in the requested list
             self.emission_luminosities_df.insert(
                 loc=0,
                 column="other",
                 value=self.emission_luminosities_df[
-                    sorted_list.keys()[~mask]
+                    self.emission_luminosities_df.keys()[~mask]
                 ].sum(axis=1),
             )
+            # Need to add a new value to the mask array for the 'other' column just added
+            mask = np.insert(mask, 0, True)
+            # Then drop all of the individual columns for species included in 'other'
             self.emission_luminosities_df.drop(
-                sorted_list.keys()[~mask], inplace=True, axis=1
+                self.emission_luminosities_df.keys()[~mask], inplace=True, axis=1
             )
 
+            mask = np.in1d(
+                np.array(list(self.absorption_luminosities_df.keys())), temp
+            )
+            # If species_list is included then create a new column which is the sum
+            # of all other species i.e. those that aren't in the requested list
             self.absorption_luminosities_df.insert(
                 loc=0,
                 column="other",
                 value=self.absorption_luminosities_df[
-                    sorted_list.keys()[~mask]
+                    self.absorption_luminosities_df.keys()[~mask]
                 ].sum(axis=1),
             )
+            # Need to add a new value to the mask array for the 'other' column just added
+            mask = np.insert(mask, 0, True)
+            # Then drop all of the individual columns for species included in 'other'
             self.absorption_luminosities_df.drop(
-                sorted_list.keys()[~mask], inplace=True, axis=1
+                self.absorption_luminosities_df.keys()[~mask], inplace=True, axis=1
             )
 
             # Get the list of species in the model
@@ -879,7 +903,7 @@ def _calculate_emission_luminosities(self, packets_mode, packet_wvl_range):
 
         # Contribution of each species with which packets interacted ----------
         for identifier, group in packets_df_grouped:
-            # Histogram of specific element
+            # Histogram of specific species
             hist_el = np.histogram(
                 group["nus"],
                 bins=self.plot_frequency_bins,
@@ -952,6 +976,7 @@ def _calculate_absorption_luminosities(
 
         # Group packets_df by atomic number of elements with which packets
         # had their last absorption (interaction in)
+        # or if species_list is requested then group by species id
         if self._species_list is None:
             packets_df_grouped = (
                 self.data[packets_mode]
@@ -966,7 +991,7 @@ def _calculate_absorption_luminosities(
                 )
 
         for identifier, group in packets_df_grouped:
-            # Histogram of specific element
+            # Histogram of specific species
             hist_el = np.histogram(
                 group["last_line_interaction_in_nu"],
                 bins=self.plot_frequency_bins,
@@ -1076,6 +1101,11 @@ def generate_plot_mpl(
             Axis on which SDEC Plot is created
         """
 
+        # If species_list and nelements requested, tell user that nelements is ignored
+        if species_list is not None and nelements is not None:
+            print("Both nelements and species_list were requested. Species_list takes priority; nelements is ignored")
+
+        # Parse the requested species list
         self._parse_species_list(species_list = species_list)
 
         # Calculate data attributes required for plotting
@@ -1098,7 +1128,7 @@ def generate_plot_mpl(
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
         self._show_colorbar_mpl()
 
-
+        # Plot emission and absorption components
         self._plot_emission_mpl()
         self._plot_absorption_mpl()
 
@@ -1187,13 +1217,10 @@ def _plot_emission_mpl(self):
         # This is used when plotting species incase an element was given in the list
         # This is to ensure that all ions of that element are grouped together
         # ii is to track the colour index
+        # e.g. if Si is given in species_list, this is to ensure Si I, Si II, etc. all have the same colour
         ii = 0
         previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the emission df
-
-
             if self._species_list is not None:
                 # Get the ion number and atomic number for each species
                 ion_number = identifier % 100
@@ -1205,12 +1232,12 @@ def _plot_emission_mpl(self):
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index
+                    # then dont update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
                     else:
-                        # Otherwise, increase the colour counter by one
+                        # Otherwise, increase the colour counter by one, because this is a new element
                         ii = ii + 1
                         previous_atomic_number = atomic_number
                 else:
@@ -1224,7 +1251,8 @@ def _plot_emission_mpl(self):
                 # If you're not using species list then this is just a fraction based on the total
                 # number of columns in the dataframe
                 color = self.cmap(i / len(self.species))
-
+            # Add a try catch because the identifier comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in list is not in the emission df
             try:
                 lower_level = upper_level
                 upper_level = (
@@ -1241,6 +1269,7 @@ def _plot_emission_mpl(self):
                     linewidth=0,
                 )
             except:
+                # Add notifications that this species was not in the emission df
                 if self._species_list is None:
                     print(
                         atomic_number2element_symbol(identifier)
@@ -1277,9 +1306,6 @@ def _plot_absorption_mpl(self):
         ii = 0
         previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the absorption df
-
             if self._species_list is not None:
                 # Get the ion number and atomic number for each species
                 ion_number = identifier % 100
@@ -1291,12 +1317,12 @@ def _plot_absorption_mpl(self):
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index
+                    # then dont update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
                     else:
-                        # Otherwise, increase the colour counter by one
+                        # Otherwise, increase the colour counter by one, because this is a new element
                         ii = ii + 1
                         previous_atomic_number = atomic_number
                 else:
@@ -1310,7 +1336,8 @@ def _plot_absorption_mpl(self):
                 # If you're not using species list then this is just a fraction based on the total
                 # number of columns in the dataframe
                 color = self.cmap(i / len(self.species))
-
+            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in elements_z is not in the absorption df
             try:
                 upper_level = lower_level
                 lower_level = (
@@ -1328,6 +1355,7 @@ def _plot_absorption_mpl(self):
                 )
 
             except:
+                # Add notifications that this species was not in the emission df
                 if self._species_list is None:
                     print(
                         atomic_number2element_symbol(identifier)
@@ -1362,6 +1390,7 @@ def _show_colorbar_mpl(self):
     def _make_colorbar_labels(self):
         """Get the labels for the species in the colorbar."""
         if self._species_list is None:
+            # If species_list is none then the labels are just elements
             species_name = [
                 atomic_number2element_symbol(atomic_num)
                 for atomic_num in self.species
@@ -1369,6 +1398,7 @@ def _make_colorbar_labels(self):
         else:
             species_name = []
             for species in self.species:
+                # Go through each species requested
                 ion_number = species % 100
                 atomic_number = (species - ion_number) / 100
 
@@ -1445,6 +1475,11 @@ def generate_plot_ply(
             Figure object on which SDEC Plot is created
         """
 
+        # If species_list and nelements requested, tell user that nelements is ignored
+        if species_list is not None and nelements is not None:
+            print("Both nelements and species_list were requested. Species_list takes priority; nelements is ignored")
+
+        # Parse the requested species list
         self._parse_species_list(species_list = species_list)
 
         # Calculate data attributes required for plotting
@@ -1466,6 +1501,7 @@ def generate_plot_ply(
         # Set colormap to be used in elements of emission and absorption plots
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
 
+        # Plot absorption and emission components
         self._plot_emission_ply()
         self._plot_absorption_ply()
 
@@ -1595,12 +1631,12 @@ def _plot_emission_ply(self):
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
-                   # then dont update the colour index
+                    # then dont update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
                     else:
-                        # Otherwise, increase the colour counter by one
+                        # Otherwise, increase the colour counter by one, because this is a new element
                         ii = ii + 1
                         previous_atomic_number = atomic_number
                 else:
@@ -1614,8 +1650,8 @@ def _plot_emission_ply(self):
                 # number of columns in the dataframe
                 color = self.cmap(i / len(self.species))
 
-            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the emission df
+            # Add a try catch because identifier comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in the list is not in the emission df
             try:
                 self.fig.add_trace(
                         go.Scatter(
@@ -1674,12 +1710,12 @@ def _plot_absorption_ply(self):
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
-                   # then dont update the colour index
+                    # then dont update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
                         ii = ii
                         previous_atomic_number = atomic_number
                     else:
-                        # Otherwise, increase the colour counter by one
+                        # Otherwise, increase the colour counter by one, because this is a new element
                         ii = ii + 1
                         previous_atomic_number = atomic_number
                 else:
@@ -1693,8 +1729,8 @@ def _plot_absorption_ply(self):
                 # number of columns in the dataframe
                 color = self.cmap(i / len(self.species))
 
-            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the absorption df
+            # Add a try catch because the identifier comes from the total contribution of absorption and emission.
+            # Therefore it's possible that something in list is not in the absorption df
             try:
                 self.fig.add_trace(
                     go.Scatter(

From f97307b265d5eafe53f01dcec764f4c4d2eab403 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Thu, 15 Apr 2021 00:07:23 +0100
Subject: [PATCH 14/20] Black formatting

---
 tardis/visualization/tools/sdec_plot.py | 132 +++++++++++++-----------
 1 file changed, 71 insertions(+), 61 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 1d1d444ba86..42cb380101c 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -132,9 +132,9 @@ def __init__(
             .iloc[
                 self.packets_df_line_interaction["last_line_interaction_out_id"]
             ]
-            .to_numpy() * 100
-            +
-            self.lines_df["ion_number"]
+            .to_numpy()
+            * 100
+            + self.lines_df["ion_number"]
             .iloc[
                 self.packets_df_line_interaction["last_line_interaction_out_id"]
             ]
@@ -439,10 +439,7 @@ def from_hdf(cls, hdf_fpath):
             )
         )
 
-
-    def _parse_species_list(
-        self, species_list
-    ):
+    def _parse_species_list(self, species_list):
         """
         Parse user requested species list and create list of species ids to be used.
 
@@ -456,12 +453,12 @@ def _parse_species_list(
         """
 
         if species_list is not None:
-        # check if there are any digits in the species list. If there are, then exit.
-        # species_list should only contain species in the Roman numeral
-        # format, e.g. Si II, and each ion must contain a space
+            # check if there are any digits in the species list. If there are, then exit.
+            # species_list should only contain species in the Roman numeral
+            # format, e.g. Si II, and each ion must contain a space
             if any(char.isdigit() for char in " ".join(species_list)) == True:
                 raise ValueError(
-                "All species must be in Roman numeral form, e.g. Si II"
+                    "All species must be in Roman numeral form, e.g. Si II"
                 )
             else:
                 full_species_list = []
@@ -475,13 +472,17 @@ def _parse_species_list(
                         # convert the requested ions into numerals
                         first_ion_numeral = roman_to_int(
                             species.split(" ")[-1].split("-")[0]
-                            )
+                        )
                         second_ion_numeral = roman_to_int(
                             species.split(" ")[-1].split("-")[-1]
+                        )
+                        # add each ion between the two requested into the species list
+                        for i in np.arange(
+                            first_ion_numeral, second_ion_numeral + 1
+                        ):
+                            full_species_list.append(
+                                element + " " + int_to_roman(i)
                             )
-                            # add each ion between the two requested into the species list
-                        for i in np.arange(first_ion_numeral, second_ion_numeral + 1):
-                            full_species_list.append(element + " " + int_to_roman(i))
                     else:
                         # Otherwise it's either an element or ion so just add to the list
                         full_species_list.append(species)
@@ -507,14 +508,19 @@ def _parse_species_list(
                     else:
                         atomic_number = element_symbol2atomic_number(species)
                         requested_species_ids.append(
-                            [atomic_number * 100 + i for i in np.arange(atomic_number)]
+                            [
+                                atomic_number * 100 + i
+                                for i in np.arange(atomic_number)
+                            ]
                         )
                         # add the atomic number to a list so you know that this element should
                         # have all species in the same colour, i.e. it was requested like
                         # species_list = [Si]
                         keep_colour.append(atomic_number)
                 requested_species_ids = [
-                    species_id for list in requested_species_ids for species_id in list
+                    species_id
+                    for list in requested_species_ids
+                    for species_id in list
                 ]
 
                 self._species_list = requested_species_ids
@@ -656,9 +662,9 @@ def _calculate_plotting_data(
             self.total_luminosities_df.insert(
                 loc=0,
                 column="other",
-                value=self.total_luminosities_df[
-                    sorted_list.keys()[~mask]
-                ].sum(axis=1),
+                value=self.total_luminosities_df[sorted_list.keys()[~mask]].sum(
+                    axis=1
+                ),
             )
             # Then drop all of the individual columns for species included in 'other'
             self.total_luminosities_df.drop(
@@ -668,8 +674,8 @@ def _calculate_plotting_data(
             # This will require creating a temporary list that includes 'noint' and 'escatter'
             # packets, because you don't want them dropped or included in 'other'
             temp = [species for species in self._species_list]
-            temp.append('noint')
-            temp.append('escatter')
+            temp.append("noint")
+            temp.append("escatter")
             mask = np.in1d(
                 np.array(list(self.emission_luminosities_df.keys())), temp
             )
@@ -686,7 +692,9 @@ def _calculate_plotting_data(
             mask = np.insert(mask, 0, True)
             # Then drop all of the individual columns for species included in 'other'
             self.emission_luminosities_df.drop(
-                self.emission_luminosities_df.keys()[~mask], inplace=True, axis=1
+                self.emission_luminosities_df.keys()[~mask],
+                inplace=True,
+                axis=1,
             )
 
             mask = np.in1d(
@@ -705,7 +713,9 @@ def _calculate_plotting_data(
             mask = np.insert(mask, 0, True)
             # Then drop all of the individual columns for species included in 'other'
             self.absorption_luminosities_df.drop(
-                self.absorption_luminosities_df.keys()[~mask], inplace=True, axis=1
+                self.absorption_luminosities_df.keys()[~mask],
+                inplace=True,
+                axis=1,
             )
 
             # Get the list of species in the model
@@ -893,13 +903,13 @@ def _calculate_emission_luminosities(self, packets_mode, packet_wvl_range):
                 self.data[packets_mode]
                 .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
                 .groupby(by="last_line_interaction_atom")
-                )
+            )
         else:
             packets_df_grouped = (
                 self.data[packets_mode]
                 .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
                 .groupby(by="last_line_interaction_species")
-                )
+            )
 
         # Contribution of each species with which packets interacted ----------
         for identifier, group in packets_df_grouped:
@@ -982,13 +992,13 @@ def _calculate_absorption_luminosities(
                 self.data[packets_mode]
                 .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
                 .groupby(by="last_line_interaction_atom")
-                )
+            )
         else:
             packets_df_grouped = (
                 self.data[packets_mode]
                 .packets_df_line_interaction.loc[self.packet_nu_line_range_mask]
                 .groupby(by="last_line_interaction_species")
-                )
+            )
 
         for identifier, group in packets_df_grouped:
             # Histogram of specific species
@@ -1055,7 +1065,7 @@ def generate_plot_mpl(
         figsize=(12, 7),
         cmapname="jet",
         nelements=None,
-        species_list=None
+        species_list=None,
     ):
         """
         Generate Spectral element DEComposition (SDEC) Plot using matplotlib.
@@ -1103,10 +1113,12 @@ def generate_plot_mpl(
 
         # If species_list and nelements requested, tell user that nelements is ignored
         if species_list is not None and nelements is not None:
-            print("Both nelements and species_list were requested. Species_list takes priority; nelements is ignored")
+            print(
+                "Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
+            )
 
         # Parse the requested species list
-        self._parse_species_list(species_list = species_list)
+        self._parse_species_list(species_list=species_list)
 
         # Calculate data attributes required for plotting
         # and save them in instance itself
@@ -1150,7 +1162,6 @@ def generate_plot_mpl(
             label="Blackbody Photosphere",
         )
 
-
         # Set legends and labels
         self.ax.legend(fontsize=12)
         self.ax.set_xlabel(r"Wavelength $[\mathrm{\AA}]$", fontsize=12)
@@ -1274,13 +1285,13 @@ def _plot_emission_mpl(self):
                     print(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
-                        )
+                    )
                 else:
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
-                        )
+                    )
 
     def _plot_absorption_mpl(self):
         """Plot absorption part of the SDEC Plot using matplotlib."""
@@ -1360,20 +1371,20 @@ def _plot_absorption_mpl(self):
                     print(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
-                        )
+                    )
                 else:
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
-                        )
-
+                    )
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
 
         color_values = [
-            self.cmap(i / len(self._species_name)) for i in range(len(self._species_name))
+            self.cmap(i / len(self._species_name))
+            for i in range(len(self._species_name))
         ]
 
         custcmap = clr.ListedColormap(color_values)
@@ -1387,7 +1398,6 @@ def _show_colorbar_mpl(self):
 
         cbar.set_ticklabels(self._species_name)
 
-
     def _make_colorbar_labels(self):
         """Get the labels for the species in the colorbar."""
         if self._species_list is None:
@@ -1408,7 +1418,9 @@ def _make_colorbar_labels(self):
 
                 # if the element was requested, and not a specific ion, then
                 # add the element symbol to the label list
-                if (atomic_number in self._keep_colour) & (atomic_symbol not in species_name):
+                if (atomic_number in self._keep_colour) & (
+                    atomic_symbol not in species_name
+                ):
                     # compiling the label, and adding it to the list
                     label = f"{atomic_symbol}"
                     species_name.append(label)
@@ -1419,7 +1431,6 @@ def _make_colorbar_labels(self):
 
         self._species_name = species_name
 
-
     def generate_plot_ply(
         self,
         packets_mode="virtual",
@@ -1478,10 +1489,12 @@ def generate_plot_ply(
 
         # If species_list and nelements requested, tell user that nelements is ignored
         if species_list is not None and nelements is not None:
-            print("Both nelements and species_list were requested. Species_list takes priority; nelements is ignored")
+            print(
+                "Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
+            )
 
         # Parse the requested species list
-        self._parse_species_list(species_list = species_list)
+        self._parse_species_list(species_list=species_list)
 
         # Calculate data attributes required for plotting
         # and save them in instance itself
@@ -1655,30 +1668,29 @@ def _plot_emission_ply(self):
             # Therefore it's possible that something in the list is not in the emission df
             try:
                 self.fig.add_trace(
-                        go.Scatter(
-                            x=self.emission_luminosities_df.index,
-                            y=self.emission_luminosities_df[identifier],
-                            mode="none",
-                            name="none",
-                            fillcolor=self.to_rgb255_string(
-                                color),
-                            stackgroup="emission",
-                            showlegend=False,
-                        )
+                    go.Scatter(
+                        x=self.emission_luminosities_df.index,
+                        y=self.emission_luminosities_df[identifier],
+                        mode="none",
+                        name="none",
+                        fillcolor=self.to_rgb255_string(color),
+                        stackgroup="emission",
+                        showlegend=False,
                     )
+                )
             except:
                 if self._species_list is None:
                     print(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
-                        )
+                    )
                 else:
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
-                        )
-                    
+                    )
+
     def _plot_absorption_ply(self):
         """Plot absorption part of the SDEC Plot using plotly."""
 
@@ -1740,8 +1752,7 @@ def _plot_absorption_ply(self):
                         y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(
-                            color),
+                        fillcolor=self.to_rgb255_string(color),
                         stackgroup="absorption",
                         showlegend=False,
                     )
@@ -1752,14 +1763,13 @@ def _plot_absorption_ply(self):
                     print(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
-                        )
+                    )
                 else:
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
-                        )
-
+                    )
 
     def _show_colorbar_ply(self):
         """Show plotly colorbar with labels of elements mapped to colors."""

From aaa0af87f822b24b9216c966c46566fb9997f269 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Thu, 22 Apr 2021 01:34:32 +0100
Subject: [PATCH 15/20] Added colour function

Added a new function to work out what colours will be plotted. This was originally repeated multiple times in the code, but now it's a single function that's called multiple times

Also fixed a bug in the absorption luminosities dataframe
---
 tardis/visualization/tools/sdec_plot.py | 237 ++++++++----------------
 1 file changed, 82 insertions(+), 155 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 42cb380101c..ec002c80225 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -22,7 +22,7 @@
     roman_to_int,
     int_to_roman,
 )
-from tardis.visualization import plot_util as pu
+from tardis.widgets import plot_util as pu
 
 
 class SDECData:
@@ -697,6 +697,7 @@ def _calculate_plotting_data(
                 axis=1,
             )
 
+            temp = [species for species in self._species_list]
             mask = np.in1d(
                 np.array(list(self.absorption_luminosities_df.keys())), temp
             )
@@ -1138,6 +1139,8 @@ def generate_plot_mpl(
         self._make_colorbar_labels()
         # Set colormap to be used in elements of emission and absorption plots
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
+        # Get the number of unqie colors
+        self._make_colorbar_colors()
         self._show_colorbar_mpl()
 
         # Plot emission and absorption components
@@ -1224,46 +1227,7 @@ def _plot_emission_mpl(self):
             )
 
         # Contribution from each element
-        # Create new variables to keep track of the last atomic number that was plotted
-        # This is used when plotting species incase an element was given in the list
-        # This is to ensure that all ions of that element are grouped together
-        # ii is to track the colour index
-        # e.g. if Si is given in species_list, this is to ensure Si I, Si II, etc. all have the same colour
-        ii = 0
-        previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-            if self._species_list is not None:
-                # Get the ion number and atomic number for each species
-                ion_number = identifier % 100
-                atomic_number = (identifier - ion_number) / 100
-                if previous_atomic_number == 0:
-                    # If this is the first species being plotted, then take note of the atomic number
-                    # don't update the colour index
-                    ii = ii
-                    previous_atomic_number = atomic_number
-                elif previous_atomic_number in self._keep_colour:
-                    # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index if this element has been plotted already
-                    if previous_atomic_number == atomic_number:
-                        ii = ii
-                        previous_atomic_number = atomic_number
-                    else:
-                        # Otherwise, increase the colour counter by one, because this is a new element
-                        ii = ii + 1
-                        previous_atomic_number = atomic_number
-                else:
-                    # If this is just a normal species that was requested then increment the colour index
-                    ii = ii + 1
-                    previous_atomic_number = atomic_number
-                # Calculate the colour of this species
-                color = self.cmap(ii / len(self._species_name))
-
-            else:
-                # If you're not using species list then this is just a fraction based on the total
-                # number of columns in the dataframe
-                color = self.cmap(i / len(self.species))
-            # Add a try catch because the identifier comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in list is not in the emission df
             try:
                 lower_level = upper_level
                 upper_level = (
@@ -1275,7 +1239,7 @@ def _plot_emission_mpl(self):
                     self.plot_wavelength,
                     lower_level,
                     upper_level,
-                    color=color,
+                    color=self._color_list[i],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1287,6 +1251,10 @@ def _plot_emission_mpl(self):
                         + " is not in the emitted packets; skipping"
                     )
                 else:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
@@ -1314,41 +1282,7 @@ def _plot_absorption_mpl(self):
                 color="silver",
             )
 
-        ii = 0
-        previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-            if self._species_list is not None:
-                # Get the ion number and atomic number for each species
-                ion_number = identifier % 100
-                atomic_number = (identifier - ion_number) / 100
-                if previous_atomic_number == 0:
-                    # If this is the first species being plotted, then take note of the atomic number
-                    # don't update the colour index
-                    ii = ii
-                    previous_atomic_number = atomic_number
-                elif previous_atomic_number in self._keep_colour:
-                    # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index if this element has been plotted already
-                    if previous_atomic_number == atomic_number:
-                        ii = ii
-                        previous_atomic_number = atomic_number
-                    else:
-                        # Otherwise, increase the colour counter by one, because this is a new element
-                        ii = ii + 1
-                        previous_atomic_number = atomic_number
-                else:
-                    # If this is just a normal species that was requested then increment the colour index
-                    ii = ii + 1
-                    previous_atomic_number = atomic_number
-                # Calculate the colour of this species
-                color = self.cmap(ii / len(self._species_name))
-
-            else:
-                # If you're not using species list then this is just a fraction based on the total
-                # number of columns in the dataframe
-                color = self.cmap(i / len(self.species))
-            # Add a try catch because elemets_z comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in elements_z is not in the absorption df
             try:
                 upper_level = lower_level
                 lower_level = (
@@ -1360,7 +1294,7 @@ def _plot_absorption_mpl(self):
                     self.plot_wavelength,
                     upper_level,
                     lower_level,
-                    color=color,
+                    color=self._color_list[i],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1370,13 +1304,17 @@ def _plot_absorption_mpl(self):
                 if self._species_list is None:
                     print(
                         atomic_number2element_symbol(identifier)
-                        + " is not in the emitted packets; skipping"
+                        + " is not in the absorbed packets; skipping"
                     )
                 else:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
-                        + " is not in the emitted packets; skipping"
+                        + " is not in the absorbed packets; skipping"
                     )
 
     def _show_colorbar_mpl(self):
@@ -1431,6 +1369,58 @@ def _make_colorbar_labels(self):
 
         self._species_name = species_name
 
+    def _make_colorbar_colors(self):
+        """Get the colours for the species to be plotted."""
+        # the colours depends on the species present in the model and what's requested
+        # some species need to be shown in the same colour, so the exact colours have to be
+        # worked out
+
+        color_list = []
+
+        # Colors for each element
+        # Create new variables to keep track of the last atomic number that was plotted
+        # This is used when plotting species incase an element was given in the list
+        # This is to ensure that all ions of that element are grouped together
+        # ii is to track the colour index
+        # e.g. if Si is given in species_list, this is to ensure Si I, Si II, etc. all have the same colour
+        ii = 0
+        previous_atomic_number = 0
+        for i, identifier in enumerate(self.species):
+            if self._species_list is not None:
+                # Get the ion number and atomic number for each species
+                ion_number = identifier % 100
+                atomic_number = (identifier - ion_number) / 100
+                if previous_atomic_number == 0:
+                    # If this is the first species being plotted, then take note of the atomic number
+                    # don't update the colour index
+                    ii = ii
+                    previous_atomic_number = atomic_number
+                elif previous_atomic_number in self._keep_colour:
+                    # If the atomic number is in the list of elements that should all be plotted in the same colour
+                    # then dont update the colour index if this element has been plotted already
+                    if previous_atomic_number == atomic_number:
+                        ii = ii
+                        previous_atomic_number = atomic_number
+                    else:
+                        # Otherwise, increase the colour counter by one, because this is a new element
+                        ii = ii + 1
+                        previous_atomic_number = atomic_number
+                else:
+                    # If this is just a normal species that was requested then increment the colour index
+                    ii = ii + 1
+                    previous_atomic_number = atomic_number
+                # Calculate the colour of this species
+                color = self.cmap(ii / len(self._species_name))
+
+            else:
+                # If you're not using species list then this is just a fraction based on the total
+                # number of columns in the dataframe
+                color = self.cmap(i / len(self.species))
+
+            color_list.append(color)
+
+        self._color_list = color_list
+
     def generate_plot_ply(
         self,
         packets_mode="virtual",
@@ -1514,6 +1504,8 @@ def generate_plot_ply(
         self._make_colorbar_labels()
         # Set colormap to be used in elements of emission and absorption plots
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
+        # Get the number of unqie colors
+        self._make_colorbar_colors()
 
         # Plot absorption and emission components
         self._plot_emission_ply()
@@ -1626,46 +1618,7 @@ def _plot_emission_ply(self):
             )
 
         # Contribution from each element
-        # Create new variables to keep track of the last atomic number that was plotted
-        # This is used when plotting species incase an element was given in the list
-        # This is to ensure that all ions of that element are grouped together
-        # ii is to track the colour index
-        ii = 0
-        previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-
-            if self._species_list is not None:
-                # Get the ion number and atomic number for each species
-                ion_number = identifier % 100
-                atomic_number = (identifier - ion_number) / 100
-                if previous_atomic_number == 0:
-                    # If this is the first species being plotted, then take note of the atomic number
-                    # don't update the colour index
-                    ii = ii
-                    previous_atomic_number = atomic_number
-                elif previous_atomic_number in self._keep_colour:
-                    # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index if this element has been plotted already
-                    if previous_atomic_number == atomic_number:
-                        ii = ii
-                        previous_atomic_number = atomic_number
-                    else:
-                        # Otherwise, increase the colour counter by one, because this is a new element
-                        ii = ii + 1
-                        previous_atomic_number = atomic_number
-                else:
-                    # If this is just a normal species that was requested then increment the colour index
-                    ii = ii + 1
-                    previous_atomic_number = atomic_number
-                # Calculate the colour of this species
-                color = self.cmap(ii / len(self._species_name))
-            else:
-                # If you're not using species list then this is just a fraction based on the total
-                # number of columns in the dataframe
-                color = self.cmap(i / len(self.species))
-
-            # Add a try catch because identifier comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in the list is not in the emission df
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1673,7 +1626,7 @@ def _plot_emission_ply(self):
                         y=self.emission_luminosities_df[identifier],
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(color),
+                        fillcolor=self.to_rgb255_string(self._color_list[i]),
                         stackgroup="emission",
                         showlegend=False,
                     )
@@ -1685,6 +1638,10 @@ def _plot_emission_ply(self):
                         + " is not in the emitted packets; skipping"
                     )
                 else:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
@@ -1709,41 +1666,7 @@ def _plot_absorption_ply(self):
                 )
             )
 
-        ii = 0
-        previous_atomic_number = 0
         for i, identifier in enumerate(self.species):
-            if self._species_list is not None:
-                # Get the ion number and atomic number for each species
-                ion_number = identifier % 100
-                atomic_number = (identifier - ion_number) / 100
-                if previous_atomic_number == 0:
-                    # If this is the first species being plotted, then take note of the atomic number
-                    # don't update the colour index
-                    ii = ii
-                    previous_atomic_number = atomic_number
-                elif previous_atomic_number in self._keep_colour:
-                    # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index if this element has been plotted already
-                    if previous_atomic_number == atomic_number:
-                        ii = ii
-                        previous_atomic_number = atomic_number
-                    else:
-                        # Otherwise, increase the colour counter by one, because this is a new element
-                        ii = ii + 1
-                        previous_atomic_number = atomic_number
-                else:
-                    # If this is just a normal species that was requested then increment the colour index
-                    ii = ii + 1
-                    previous_atomic_number = atomic_number
-                # Calculate the colour of this species
-                color = self.cmap(ii / len(self._species_name))
-            else:
-                # If you're not using species list then this is just a fraction based on the total
-                # number of columns in the dataframe
-                color = self.cmap(i / len(self.species))
-
-            # Add a try catch because the identifier comes from the total contribution of absorption and emission.
-            # Therefore it's possible that something in list is not in the absorption df
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1752,7 +1675,7 @@ def _plot_absorption_ply(self):
                         y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(color),
+                        fillcolor=self.to_rgb255_string(self._color_list[i]),
                         stackgroup="absorption",
                         showlegend=False,
                     )
@@ -1762,13 +1685,17 @@ def _plot_absorption_ply(self):
                 if self._species_list is None:
                     print(
                         atomic_number2element_symbol(identifier)
-                        + " is not in the emitted packets; skipping"
+                        + " is not in the absorbed packets; skipping"
                     )
                 else:
+                    # Get the ion number and atomic number for each species
+                    ion_number = identifier % 100
+                    atomic_number = (identifier - ion_number) / 100
+
                     print(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
-                        + " is not in the emitted packets; skipping"
+                        + " is not in the absorbed packets; skipping"
                     )
 
     def _show_colorbar_ply(self):

From a83fa819c342e19e86e6b8f230471008fc4347d7 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Thu, 22 Apr 2021 01:45:03 +0100
Subject: [PATCH 16/20] Address review comments

Change name of some variables and fixed some typos.
---
 tardis/visualization/tools/sdec_plot.py | 60 ++++++++++++-------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index ec002c80225..92ac80063eb 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -477,11 +477,11 @@ def _parse_species_list(self, species_list):
                             species.split(" ")[-1].split("-")[-1]
                         )
                         # add each ion between the two requested into the species list
-                        for i in np.arange(
+                        for ii in np.arange(
                             first_ion_numeral, second_ion_numeral + 1
                         ):
                             full_species_list.append(
-                                element + " " + int_to_roman(i)
+                                element + " " + int_to_roman(ii)
                             )
                     else:
                         # Otherwise it's either an element or ion so just add to the list
@@ -509,8 +509,8 @@ def _parse_species_list(self, species_list):
                         atomic_number = element_symbol2atomic_number(species)
                         requested_species_ids.append(
                             [
-                                atomic_number * 100 + i
-                                for i in np.arange(atomic_number)
+                                atomic_number * 100 + ii
+                                for ii in np.arange(atomic_number)
                             ]
                         )
                         # add the atomic number to a list so you know that this element should
@@ -1101,7 +1101,7 @@ def generate_plot_mpl(
             Other elements are shown in silver. Default value is
             None, which displays all elements
         species_list: list of strings or None
-            list of strings containing the names of species that should be included in the Kromer plots.
+            list of strings containing the names of species that should be included in the SDEC plots.
             Must be given in Roman numeral format. Can include specific ions, a range of ions,
             individual elements, or any combination of these:
             e.g. ['Si II', 'Ca II', 'C', 'Fe I-V']
@@ -1227,7 +1227,7 @@ def _plot_emission_mpl(self):
             )
 
         # Contribution from each element
-        for i, identifier in enumerate(self.species):
+        for ii, identifier in enumerate(self.species):
             try:
                 lower_level = upper_level
                 upper_level = (
@@ -1239,7 +1239,7 @@ def _plot_emission_mpl(self):
                     self.plot_wavelength,
                     lower_level,
                     upper_level,
-                    color=self._color_list[i],
+                    color=self._color_list[ii],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1282,7 +1282,7 @@ def _plot_absorption_mpl(self):
                 color="silver",
             )
 
-        for i, identifier in enumerate(self.species):
+        for ii, identifier in enumerate(self.species):
             try:
                 upper_level = lower_level
                 lower_level = (
@@ -1294,7 +1294,7 @@ def _plot_absorption_mpl(self):
                     self.plot_wavelength,
                     upper_level,
                     lower_level,
-                    color=self._color_list[i],
+                    color=self._color_list[ii],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1321,8 +1321,8 @@ def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
 
         color_values = [
-            self.cmap(i / len(self._species_name))
-            for i in range(len(self._species_name))
+            self.cmap(ii / len(self._species_name))
+            for ii in range(len(self._species_name))
         ]
 
         custcmap = clr.ListedColormap(color_values)
@@ -1379,13 +1379,13 @@ def _make_colorbar_colors(self):
 
         # Colors for each element
         # Create new variables to keep track of the last atomic number that was plotted
-        # This is used when plotting species incase an element was given in the list
+        # This is used when plotting species in case an element was given in the list
         # This is to ensure that all ions of that element are grouped together
         # ii is to track the colour index
         # e.g. if Si is given in species_list, this is to ensure Si I, Si II, etc. all have the same colour
-        ii = 0
+        counter = 0
         previous_atomic_number = 0
-        for i, identifier in enumerate(self.species):
+        for ii, identifier in enumerate(self.species):
             if self._species_list is not None:
                 # Get the ion number and atomic number for each species
                 ion_number = identifier % 100
@@ -1393,29 +1393,29 @@ def _make_colorbar_colors(self):
                 if previous_atomic_number == 0:
                     # If this is the first species being plotted, then take note of the atomic number
                     # don't update the colour index
-                    ii = ii
+                    counter = counter
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
-                    # then dont update the colour index if this element has been plotted already
+                    # then don't update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
-                        ii = ii
+                        counter = counter
                         previous_atomic_number = atomic_number
                     else:
                         # Otherwise, increase the colour counter by one, because this is a new element
-                        ii = ii + 1
+                        counter = counter + 1
                         previous_atomic_number = atomic_number
                 else:
                     # If this is just a normal species that was requested then increment the colour index
-                    ii = ii + 1
+                    counter = counter + 1
                     previous_atomic_number = atomic_number
                 # Calculate the colour of this species
-                color = self.cmap(ii / len(self._species_name))
+                color = self.cmap(counter / len(self._species_name))
 
             else:
                 # If you're not using species list then this is just a fraction based on the total
                 # number of columns in the dataframe
-                color = self.cmap(i / len(self.species))
+                color = self.cmap(ii / len(self.species))
 
             color_list.append(color)
 
@@ -1467,7 +1467,7 @@ def generate_plot_ply(
             Other elements are shown in silver. Default value is
             None, which displays all elements
         species_list: list of strings or None
-            list of strings containing the names of species that should be included in the Kromer plots.
+            list of strings containing the names of species that should be included in the SDEC plots.
             Must be given in Roman numeral format. Can include specific ions, a range of ions,
             individual elements, or any combination of these:
             e.g. ['Si II', 'Ca II', 'C', 'Fe I-V']
@@ -1618,7 +1618,7 @@ def _plot_emission_ply(self):
             )
 
         # Contribution from each element
-        for i, identifier in enumerate(self.species):
+        for ii, identifier in enumerate(self.species):
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1626,7 +1626,7 @@ def _plot_emission_ply(self):
                         y=self.emission_luminosities_df[identifier],
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[i]),
+                        fillcolor=self.to_rgb255_string(self._color_list[ii]),
                         stackgroup="emission",
                         showlegend=False,
                     )
@@ -1666,7 +1666,7 @@ def _plot_absorption_ply(self):
                 )
             )
 
-        for i, identifier in enumerate(self.species):
+        for ii, identifier in enumerate(self.species):
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1675,7 +1675,7 @@ def _plot_absorption_ply(self):
                         y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[i]),
+                        fillcolor=self.to_rgb255_string(self._color_list[ii]),
                         stackgroup="absorption",
                         showlegend=False,
                     )
@@ -1707,10 +1707,10 @@ def _show_colorbar_ply(self):
         # by mapping same reference points (excluding 1st and last bin edge)
         # twice in a row (https://plotly.com/python/colorscales/#constructing-a-discrete-or-discontinuous-color-scale)
         categorical_colorscale = []
-        for i in range(len(self._species_name)):
-            color = self.to_rgb255_string(self.cmap(colorscale_bins[i]))
-            categorical_colorscale.append((colorscale_bins[i], color))
-            categorical_colorscale.append((colorscale_bins[i + 1], color))
+        for ii in range(len(self._species_name)):
+            color = self.to_rgb255_string(self.cmap(colorscale_bins[ii]))
+            categorical_colorscale.append((colorscale_bins[ii], color))
+            categorical_colorscale.append((colorscale_bins[ii + 1], color))
 
         coloraxis_options = dict(
             colorscale=categorical_colorscale,

From ccb9507cd05462825e0c08c98cf76e78c5f5d19d Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Thu, 22 Apr 2021 19:00:58 +0100
Subject: [PATCH 17/20] Change print to logger

---
 tardis/visualization/tools/sdec_plot.py | 31 ++++++++++++++-----------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 92ac80063eb..167063ddd5b 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -22,8 +22,10 @@
     roman_to_int,
     int_to_roman,
 )
-from tardis.widgets import plot_util as pu
+from tardis.visualization import plot_util as pu
 
+import logging
+logger = logging.getLogger(__name__)
 
 class SDECData:
     """The data of simulation model used by Spectral element DEComposition (SDEC) Plot.
@@ -519,8 +521,8 @@ def _parse_species_list(self, species_list):
                         keep_colour.append(atomic_number)
                 requested_species_ids = [
                     species_id
-                    for list in requested_species_ids
-                    for species_id in list
+                    for temp_list in requested_species_ids
+                    for species_id in temp_list
                 ]
 
                 self._species_list = requested_species_ids
@@ -1114,7 +1116,7 @@ def generate_plot_mpl(
 
         # If species_list and nelements requested, tell user that nelements is ignored
         if species_list is not None and nelements is not None:
-            print(
+            logger.info(
                 "Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
             )
 
@@ -1246,7 +1248,7 @@ def _plot_emission_mpl(self):
             except:
                 # Add notifications that this species was not in the emission df
                 if self._species_list is None:
-                    print(
+                    longer.info(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
                     )
@@ -1255,7 +1257,7 @@ def _plot_emission_mpl(self):
                     ion_number = identifier % 100
                     atomic_number = (identifier - ion_number) / 100
 
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
@@ -1302,7 +1304,7 @@ def _plot_absorption_mpl(self):
             except:
                 # Add notifications that this species was not in the emission df
                 if self._species_list is None:
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(identifier)
                         + " is not in the absorbed packets; skipping"
                     )
@@ -1311,7 +1313,7 @@ def _plot_absorption_mpl(self):
                     ion_number = identifier % 100
                     atomic_number = (identifier - ion_number) / 100
 
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the absorbed packets; skipping"
@@ -1479,7 +1481,7 @@ def generate_plot_ply(
 
         # If species_list and nelements requested, tell user that nelements is ignored
         if species_list is not None and nelements is not None:
-            print(
+            logger.info(
                 "Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
             )
 
@@ -1504,7 +1506,7 @@ def generate_plot_ply(
         self._make_colorbar_labels()
         # Set colormap to be used in elements of emission and absorption plots
         self.cmap = cm.get_cmap(cmapname, len(self._species_name))
-        # Get the number of unqie colors
+        # Get the number of unique colors
         self._make_colorbar_colors()
 
         # Plot absorption and emission components
@@ -1564,6 +1566,7 @@ def generate_plot_ply(
     def to_rgb255_string(color_tuple):
         """
         Convert a matplotlib RGBA tuple to a generic RGB 255 string.
+
         Parameters
         ----------
         color_tuple : tuple
@@ -1633,7 +1636,7 @@ def _plot_emission_ply(self):
                 )
             except:
                 if self._species_list is None:
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(identifier)
                         + " is not in the emitted packets; skipping"
                     )
@@ -1642,7 +1645,7 @@ def _plot_emission_ply(self):
                     ion_number = identifier % 100
                     atomic_number = (identifier - ion_number) / 100
 
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the emitted packets; skipping"
@@ -1683,7 +1686,7 @@ def _plot_absorption_ply(self):
 
             except:
                 if self._species_list is None:
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(identifier)
                         + " is not in the absorbed packets; skipping"
                     )
@@ -1692,7 +1695,7 @@ def _plot_absorption_ply(self):
                     ion_number = identifier % 100
                     atomic_number = (identifier - ion_number) / 100
 
-                    print(
+                    logger.info(
                         atomic_number2element_symbol(atomic_number)
                         + int_to_roman(ion_number + 1)
                         + " is not in the absorbed packets; skipping"

From ac88aa3ce1e268491cf078828a2a917343731e9e Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Thu, 22 Apr 2021 19:03:15 +0100
Subject: [PATCH 18/20] Black formatting

---
 tardis/visualization/tools/sdec_plot.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 167063ddd5b..005646357dc 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -25,8 +25,10 @@
 from tardis.visualization import plot_util as pu
 
 import logging
+
 logger = logging.getLogger(__name__)
 
+
 class SDECData:
     """The data of simulation model used by Spectral element DEComposition (SDEC) Plot.
 

From 3ed59339adbab7fc9340ea643ca4caf87153e0a8 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 23 Apr 2021 11:43:16 +0100
Subject: [PATCH 19/20] Updated variable names and string formatting

---
 tardis/visualization/tools/sdec_plot.py | 100 ++++++++++++------------
 1 file changed, 51 insertions(+), 49 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 005646357dc..55d11b38f97 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -481,11 +481,11 @@ def _parse_species_list(self, species_list):
                             species.split(" ")[-1].split("-")[-1]
                         )
                         # add each ion between the two requested into the species list
-                        for ii in np.arange(
+                        for ion_number in np.arange(
                             first_ion_numeral, second_ion_numeral + 1
                         ):
                             full_species_list.append(
-                                element + " " + int_to_roman(ii)
+                                f"{element} {int_to_roman(ion_number)}"
                             )
                     else:
                         # Otherwise it's either an element or ion so just add to the list
@@ -513,8 +513,8 @@ def _parse_species_list(self, species_list):
                         atomic_number = element_symbol2atomic_number(species)
                         requested_species_ids.append(
                             [
-                                atomic_number * 100 + ii
-                                for ii in np.arange(atomic_number)
+                                atomic_number * 100 + ion_number
+                                for ion_number in np.arange(atomic_number)
                             ]
                         )
                         # add the atomic number to a list so you know that this element should
@@ -1231,7 +1231,7 @@ def _plot_emission_mpl(self):
             )
 
         # Contribution from each element
-        for ii, identifier in enumerate(self.species):
+        for species_counter, identifier in enumerate(self.species):
             try:
                 lower_level = upper_level
                 upper_level = (
@@ -1243,7 +1243,7 @@ def _plot_emission_mpl(self):
                     self.plot_wavelength,
                     lower_level,
                     upper_level,
-                    color=self._color_list[ii],
+                    color=self._color_list[species_counter],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1251,8 +1251,8 @@ def _plot_emission_mpl(self):
                 # Add notifications that this species was not in the emission df
                 if self._species_list is None:
                     longer.info(
-                        atomic_number2element_symbol(identifier)
-                        + " is not in the emitted packets; skipping"
+                        f"{atomic_number2element_symbol(identifier)}"
+                        f" is not in the emitted packets; skipping"
                     )
                 else:
                     # Get the ion number and atomic number for each species
@@ -1260,9 +1260,9 @@ def _plot_emission_mpl(self):
                     atomic_number = (identifier - ion_number) / 100
 
                     logger.info(
-                        atomic_number2element_symbol(atomic_number)
-                        + int_to_roman(ion_number + 1)
-                        + " is not in the emitted packets; skipping"
+                        f"{atomic_number2element_symbol(atomic_number)}"
+                        f"{int_to_roman(ion_number + 1)}"
+                        f" is not in the emitted packets; skipping"
                     )
 
     def _plot_absorption_mpl(self):
@@ -1286,7 +1286,7 @@ def _plot_absorption_mpl(self):
                 color="silver",
             )
 
-        for ii, identifier in enumerate(self.species):
+        for species_counter, identifier in enumerate(self.species):
             try:
                 upper_level = lower_level
                 lower_level = (
@@ -1298,7 +1298,7 @@ def _plot_absorption_mpl(self):
                     self.plot_wavelength,
                     upper_level,
                     lower_level,
-                    color=self._color_list[ii],
+                    color=self._color_list[species_counter],
                     cmap=self.cmap,
                     linewidth=0,
                 )
@@ -1306,9 +1306,9 @@ def _plot_absorption_mpl(self):
             except:
                 # Add notifications that this species was not in the emission df
                 if self._species_list is None:
-                    logger.info(
-                        atomic_number2element_symbol(identifier)
-                        + " is not in the absorbed packets; skipping"
+                    longer.info(
+                        f"{atomic_number2element_symbol(identifier)}"
+                        f" is not in the absorbed packets; skipping"
                     )
                 else:
                     # Get the ion number and atomic number for each species
@@ -1316,17 +1316,17 @@ def _plot_absorption_mpl(self):
                     atomic_number = (identifier - ion_number) / 100
 
                     logger.info(
-                        atomic_number2element_symbol(atomic_number)
-                        + int_to_roman(ion_number + 1)
-                        + " is not in the absorbed packets; skipping"
+                        f"{atomic_number2element_symbol(atomic_number)}"
+                        f"{int_to_roman(ion_number + 1)}"
+                        f" is not in the absorbed packets; skipping"
                     )
 
     def _show_colorbar_mpl(self):
         """Show matplotlib colorbar with labels of elements mapped to colors."""
 
         color_values = [
-            self.cmap(ii / len(self._species_name))
-            for ii in range(len(self._species_name))
+            self.cmap(species_counter / len(self._species_name))
+            for species_counter in range(len(self._species_name))
         ]
 
         custcmap = clr.ListedColormap(color_values)
@@ -1387,9 +1387,9 @@ def _make_colorbar_colors(self):
         # This is to ensure that all ions of that element are grouped together
         # ii is to track the colour index
         # e.g. if Si is given in species_list, this is to ensure Si I, Si II, etc. all have the same colour
-        counter = 0
+        color_counter = 0
         previous_atomic_number = 0
-        for ii, identifier in enumerate(self.species):
+        for species_counter, identifier in enumerate(self.species):
             if self._species_list is not None:
                 # Get the ion number and atomic number for each species
                 ion_number = identifier % 100
@@ -1397,29 +1397,29 @@ def _make_colorbar_colors(self):
                 if previous_atomic_number == 0:
                     # If this is the first species being plotted, then take note of the atomic number
                     # don't update the colour index
-                    counter = counter
+                    color_counter = color_counter
                     previous_atomic_number = atomic_number
                 elif previous_atomic_number in self._keep_colour:
                     # If the atomic number is in the list of elements that should all be plotted in the same colour
                     # then don't update the colour index if this element has been plotted already
                     if previous_atomic_number == atomic_number:
-                        counter = counter
+                        color_counter = color_counter
                         previous_atomic_number = atomic_number
                     else:
                         # Otherwise, increase the colour counter by one, because this is a new element
-                        counter = counter + 1
+                        color_counter = color_counter + 1
                         previous_atomic_number = atomic_number
                 else:
                     # If this is just a normal species that was requested then increment the colour index
-                    counter = counter + 1
+                    color_counter = color_counter + 1
                     previous_atomic_number = atomic_number
                 # Calculate the colour of this species
-                color = self.cmap(counter / len(self._species_name))
+                color = self.cmap(color_counter / len(self._species_name))
 
             else:
                 # If you're not using species list then this is just a fraction based on the total
                 # number of columns in the dataframe
-                color = self.cmap(ii / len(self.species))
+                color = self.cmap(species_counter / len(self.species))
 
             color_list.append(color)
 
@@ -1623,7 +1623,7 @@ def _plot_emission_ply(self):
             )
 
         # Contribution from each element
-        for ii, identifier in enumerate(self.species):
+        for species_counter, identifier in enumerate(self.species):
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1631,16 +1631,17 @@ def _plot_emission_ply(self):
                         y=self.emission_luminosities_df[identifier],
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[ii]),
+                        fillcolor=self.to_rgb255_string(self._color_list[species_counter]),
                         stackgroup="emission",
                         showlegend=False,
                     )
                 )
             except:
+                # Add notifications that this species was not in the emission df
                 if self._species_list is None:
-                    logger.info(
-                        atomic_number2element_symbol(identifier)
-                        + " is not in the emitted packets; skipping"
+                    longer.info(
+                        f"{atomic_number2element_symbol(identifier)}"
+                        f" is not in the emitted packets; skipping"
                     )
                 else:
                     # Get the ion number and atomic number for each species
@@ -1648,9 +1649,9 @@ def _plot_emission_ply(self):
                     atomic_number = (identifier - ion_number) / 100
 
                     logger.info(
-                        atomic_number2element_symbol(atomic_number)
-                        + int_to_roman(ion_number + 1)
-                        + " is not in the emitted packets; skipping"
+                        f"{atomic_number2element_symbol(atomic_number)}"
+                        f"{int_to_roman(ion_number + 1)}"
+                        f" is not in the emitted packets; skipping"
                     )
 
     def _plot_absorption_ply(self):
@@ -1671,7 +1672,7 @@ def _plot_absorption_ply(self):
                 )
             )
 
-        for ii, identifier in enumerate(self.species):
+        for species_counter, identifier in enumerate(self.species):
             try:
                 self.fig.add_trace(
                     go.Scatter(
@@ -1680,17 +1681,18 @@ def _plot_absorption_ply(self):
                         y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[ii]),
+                        fillcolor=self.to_rgb255_string(self._color_list[species_counter]),
                         stackgroup="absorption",
                         showlegend=False,
                     )
                 )
 
             except:
+                # Add notifications that this species was not in the emission df
                 if self._species_list is None:
-                    logger.info(
-                        atomic_number2element_symbol(identifier)
-                        + " is not in the absorbed packets; skipping"
+                    longer.info(
+                        f"{atomic_number2element_symbol(identifier)}"
+                        f" is not in the absorbed packets; skipping"
                     )
                 else:
                     # Get the ion number and atomic number for each species
@@ -1698,9 +1700,9 @@ def _plot_absorption_ply(self):
                     atomic_number = (identifier - ion_number) / 100
 
                     logger.info(
-                        atomic_number2element_symbol(atomic_number)
-                        + int_to_roman(ion_number + 1)
-                        + " is not in the absorbed packets; skipping"
+                        f"{atomic_number2element_symbol(atomic_number)}"
+                        f"{int_to_roman(ion_number + 1)}"
+                        f" is not in the absorbed packets; skipping"
                     )
 
     def _show_colorbar_ply(self):
@@ -1712,10 +1714,10 @@ def _show_colorbar_ply(self):
         # by mapping same reference points (excluding 1st and last bin edge)
         # twice in a row (https://plotly.com/python/colorscales/#constructing-a-discrete-or-discontinuous-color-scale)
         categorical_colorscale = []
-        for ii in range(len(self._species_name)):
-            color = self.to_rgb255_string(self.cmap(colorscale_bins[ii]))
-            categorical_colorscale.append((colorscale_bins[ii], color))
-            categorical_colorscale.append((colorscale_bins[ii + 1], color))
+        for species_counter in range(len(self._species_name)):
+            color = self.to_rgb255_string(self.cmap(colorscale_bins[species_counter]))
+            categorical_colorscale.append((colorscale_bins[species_counter], color))
+            categorical_colorscale.append((colorscale_bins[species_counter + 1], color))
 
         coloraxis_options = dict(
             colorscale=categorical_colorscale,

From c5362f72f5ce71ae6798ac1e6d735b366b23d272 Mon Sep 17 00:00:00 2001
From: Mark Magee <mrmagee.astro@gmail.com>
Date: Fri, 23 Apr 2021 11:44:41 +0100
Subject: [PATCH 20/20] Black formatting

---
 tardis/visualization/tools/sdec_plot.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py
index 55d11b38f97..f633ceedf59 100644
--- a/tardis/visualization/tools/sdec_plot.py
+++ b/tardis/visualization/tools/sdec_plot.py
@@ -1631,7 +1631,9 @@ def _plot_emission_ply(self):
                         y=self.emission_luminosities_df[identifier],
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[species_counter]),
+                        fillcolor=self.to_rgb255_string(
+                            self._color_list[species_counter]
+                        ),
                         stackgroup="emission",
                         showlegend=False,
                     )
@@ -1681,7 +1683,9 @@ def _plot_absorption_ply(self):
                         y=self.absorption_luminosities_df[identifier] * -1,
                         mode="none",
                         name="none",
-                        fillcolor=self.to_rgb255_string(self._color_list[species_counter]),
+                        fillcolor=self.to_rgb255_string(
+                            self._color_list[species_counter]
+                        ),
                         stackgroup="absorption",
                         showlegend=False,
                     )
@@ -1715,9 +1719,15 @@ def _show_colorbar_ply(self):
         # twice in a row (https://plotly.com/python/colorscales/#constructing-a-discrete-or-discontinuous-color-scale)
         categorical_colorscale = []
         for species_counter in range(len(self._species_name)):
-            color = self.to_rgb255_string(self.cmap(colorscale_bins[species_counter]))
-            categorical_colorscale.append((colorscale_bins[species_counter], color))
-            categorical_colorscale.append((colorscale_bins[species_counter + 1], color))
+            color = self.to_rgb255_string(
+                self.cmap(colorscale_bins[species_counter])
+            )
+            categorical_colorscale.append(
+                (colorscale_bins[species_counter], color)
+            )
+            categorical_colorscale.append(
+                (colorscale_bins[species_counter + 1], color)
+            )
 
         coloraxis_options = dict(
             colorscale=categorical_colorscale,