def test_int_to_roman(test_input, expected_result):
assert int_to_roman(test_input) == expected_result
with pytest.raises(TypeError):
int_to_roman(1.5)
with pytest.raises(ValueError):
int_to_roman(0)
with pytest.raises(ValueError):
int_to_roman(4000)
def test_int_to_roman(test_input, expected_result):
assert int_to_roman(test_input) == expected_result
with pytest.raises(TypeError):
int_to_roman(1.5)
with pytest.raises(ValueError):
int_to_roman(0)
with pytest.raises(ValueError):
int_to_roman(4000)
def _generate_and_add_colormap(self):
"""generate the custom color map, linking colours with atomic
numbers"""
values = [
self.cmap(float(i) / float(self._nelements))
for i in range(self._nelements)
]
custcmap = matplotlib.colors.ListedColormap(values)
bounds = np.arange(self._nelements) + 0.5
norm = matplotlib.colors.Normalize(vmin=0, vmax=self._nelements)
mappable = cm.ScalarMappable(norm=norm, cmap=custcmap)
mappable.set_array(np.linspace(1, self.zmax + 1, 256))
# if a species_list has been specified...
if self._species_list is not None:
labels = []
for zi in self._elements_in_kromer_plot:
ion_number = zi[0] % 100
atomic_number = (zi[0] - ion_number) / 100
ion_numeral = int_to_roman(ion_number + 1)
# using elements dictionary to get atomic symbol for the
# species
atomic_symbol = inv_elements[atomic_number].capitalize()
# 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 labels):
# compiling the label, and adding it to the list
label = f"{atomic_symbol}"
labels.append(label)
elif atomic_number not in self.keep_colour:
# otherwise add the ion to the label list
label = f"{atomic_symbol}$\,${ion_numeral}"
labels.append(label)
else:
# if no species_list specified, generate the labels this way
labels = [
inv_elements[zi].capitalize()
for zi in self._elements_in_kromer_plot[:, 0]
]
mainax = self.ax
cbar = plt.colorbar(mappable, ax=mainax)
cbar.set_ticks(bounds)
cbar.set_ticklabels(labels)
def generate_plot(
self,
ax=None,
cmap=cm.jet,
bins=None,
xlim=None,
ylim=None,
nelements=None,
twinx=False,
species_list=None,
):
"""Generate the actual "Kromer" plot
Parameters
----------
ax : matplotlib.axes or None
axes object into which the emission part of the Kromer plot should
be plotted; if None, a new one is generated (default None)
cmap : matplotlib.cm.ListedColormap or None
color map object used for the illustration of the different atomic
contributions (default matplotlib.cm.jet)
bins : np.ndarray or None
array of the wavelength bins used for the illustration of the
atomic contributions; if None, the same binning as for the stored
virtual spectrum is used (default None)
xlim : tuple or array-like or None
wavelength limits for the display; if None, the x-axis is
automatically scaled (default None)
ylim : tuple or array-like or None
flux limits for the display; if None, the y-axis is automatically
scaled (default None)
nelements: int or None
number of elements that should be included in the Kromer plots.
The top nelements are determined based on those with the most packet
interactions
twinx : boolean
determines where the absorption part of the Kromer plot is placed,
if True, the absorption part is attached at the top of the main
axes box, otherwise it is placed below the emission part (default
False)
species_list: list of strings or None
list of strings containing the names of species that should be included in the Kromer plots,
e.g. ['Si II', 'Ca II']
Returns
-------
fig : matplotlib.figure
figure instance containing the plot
"""
self._ax = None
self._pax = None
self._cmap = cmap
self._ax = ax
self._ylim = ylim
self._twinx = twinx
# the species list can contain either a specific element, a specific
# ion, a range of ions, or any combination of these if the list
# contains a range of ions, separate each one into a new entry in the
# species list
full_species_list = []
if species_list is not None:
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:
element = species.split(" ")[0]
first_ion_numeral = roman_to_int(
species.split(" ")[-1].split("-")[0])
second_ion_numeral = roman_to_int(
species.split(" ")[-1].split("-")[-1])
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)
self._species_list = full_species_list
else:
self._species_list = None
if xlim is None:
self._xlim = [
np.min(self.mdl.spectrum_wave).value,
np.max(self.mdl.spectrum_wave).value,
]
else:
self._xlim = xlim
if bins is None:
self._bins = self.mdl.spectrum_wave[::-1]
else:
self._bins = bins
# get the elements/species to be included in the plot
self._elements_in_kromer_plot = self.line_info
# if no nelements and no species list is specified, then the number of
# elements to be included in the colourbar is determined from the list
# of unique elements that appear in the model
if nelements is None and species_list is None:
self._nelements = len(
np.unique(self.line_in_and_out_infos_within_xlims.
atomic_number.values))
elif nelements is None and species_list is not None:
# if species_list has been specified, then the number of elements
# to be included is set to the length of that list
self._nelements = len(self._species_list)
else:
# if nelements has been specified, then the number of elements to
# be included is set to the length of that list
self._nelements = nelements
# if the length of self._elements_in_kromer_plot exceeds the requested
# number of elements to be included in the colourbar, then this if
# statement applies
if self._species_list is not None:
# if we have specified a species list then only take those species
# that are requested
mask = np.in1d(self._elements_in_kromer_plot[:, 0],
self.requested_species_ids)
self._elements_in_kromer_plot = self._elements_in_kromer_plot[mask]
elif len(self._elements_in_kromer_plot) > self._nelements:
# if nelements is specified, then sort to find the top contributing
# elements, pick the top nelements, and sort back by atomic number
self._elements_in_kromer_plot = self._elements_in_kromer_plot[
np.argsort(self._elements_in_kromer_plot[:, 1])[::-1]]
self._elements_in_kromer_plot = self._elements_in_kromer_plot[:self
.
_nelements]
self._elements_in_kromer_plot = self._elements_in_kromer_plot[
np.argsort(self._elements_in_kromer_plot[:, 0])]
else:
# if the length of self._elements_in_kromer_plot is less than the
# requested number of elements in the model, then this requested
# length is updated to be the length of length of
# self._elements_in_kromer_plot
self._nelements = len(self._elements_in_kromer_plot)
# this will reset nelements if species_list is turned on
# it's possible to request a species that doesn't appear in the plot
# this will ensure that species isn't counted when determining labels
# and colours
if self._species_list is not None:
labels = []
for species in self._species_list:
if " " in species:
atomic_number = species_string_to_tuple(species)[0]
ion_number = species_string_to_tuple(species)[1]
species_id = atomic_number * 100 + ion_number
if species_id in self._elements_in_kromer_plot:
labels.append(species)
else:
labels.append(species)
self._nelements = len(labels)
self._axes_handling_preparation()
self._generate_emission_part()
self._generate_photosphere_part()
self._generate_and_add_colormap()
self._generate_and_add_legend()
self._paxes_handling_preparation()
self._generate_absorption_part()
self._axis_handling_label_rescale()
return plt.gcf()