Make color bar labels method to util function

tardis/visualization/plot_util.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from tardis.util.base import (
+    atomic_number2element_symbol,
     element_symbol2atomic_number,
     int_to_roman,
     roman_to_int,

tardis/visualization/tools/liv_plot.py

@@ -121,25 +121,7 @@ def _make_colorbar_labels(self):
         If a species list is provided, uses that to generate labels.
         Otherwise, generates labels from the species in the model.
         """
-        if self._species_list is None:
-            species_name = [
-                atomic_number2element_symbol(atomic_num)
-                for atomic_num in self.species
-            ]
-        else:
-            species_name = []
-            for species_key, species_ids in self._species_mapped.items():
-                if any(species in self.species for species in species_ids):
-                    if species_key % 100 == 0:
-                        label = atomic_number2element_symbol(species_key // 100)
-                    else:
-                        atomic_number = species_key // 100
-                        ion_number = species_key % 100
-                        ion_numeral = int_to_roman(ion_number + 1)
-                        label = f"{atomic_number2element_symbol(atomic_number)} {ion_numeral}"
-                    species_name.append(label)
-
-        self._species_name = species_name
+        self._species_name = pu.make_colorbar_labels_util(self.species, self._species_mapped, self._keep_colour)
 
     def _make_colorbar_colors(self):
         """

tardis/visualization/tools/sdec_plot.py

@@ -961,36 +961,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
-            ]
-        else:
-            species_name = []
-            for species in self.species:
-                # Go through each species requested
-                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
+        self._species_name = pu.make_colorbar_labels_util(self.species, self._species_mapped, self._keep_colour)
 
     def _make_colorbar_colors(self):
         """Get the colours for the species to be plotted."""