Skip to content

Commit

Permalink
nelements fix
Browse files Browse the repository at this point in the history
  • Loading branch information
sarthak-dv committed Jul 11, 2024
1 parent 60c6624 commit 046f3b8
Show file tree
Hide file tree
Showing 2 changed files with 157 additions and 19 deletions.
125 changes: 111 additions & 14 deletions docs/io/visualization/how_to_liv_plot.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,29 @@
"plotter.generate_plot_mpl(packets_mode=\"real\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Plotting only the top contributing elements"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `nelements` option allows you to plot the top contributing elements to the spectrum. Only the top elements are shown in the plot. Please note this works only for elements and not for ions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_mpl(nelements=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand All @@ -135,6 +158,22 @@
"plotter.generate_plot_mpl(species_list = [\"Si I-III\", \"O\", \"Ca\", \"S\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When using both the `nelements` and the `species_list` options, `species_list` takes precedence. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_mpl(species_list = [\"Si I-III\", \"Ca\", \"S\"], nelements=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -190,11 +229,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### Additional plotting options\n",
"The `generate_plot_mpl` method also has options specific to the matplotlib API, thereby providing you with more control over how your last interaction velocity looks. Possible cases where you may use them are:\n",
"- `ax`: To plot on an Axis of a plot you're already working with, e.g. for subplots.\n",
"- `figsize`: To resize the plot as per your requirements.\n",
"- `cmapname`: To use a colormap of your preference, instead of \"jet\"."
"### Additional plotting options"
]
},
{
Expand All @@ -207,6 +242,19 @@
"help(plotter.generate_plot_mpl)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `generate_plot_mpl` method also has options specific to the matplotlib API, thereby providing you with more control over how your last interaction velocity looks. Possible cases where you may use them are:\n",
"\n",
"- `ax`: To plot on an Axis of a plot you're already working with, e.g. for subplots.\n",
"\n",
"- `figsize`: To resize the plot as per your requirements.\n",
"\n",
"- `cmapname`: To use a colormap of your preference, instead of \"jet\"."
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -253,6 +301,29 @@
"plotter.generate_plot_ply(packets_mode=\"real\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Plotting only the top contributing elements"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `nelements` option allows you to plot the top contributing elements to the spectrum. Only the top elements are shown in the plot. Please note this works only for elements and not for ions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(nelements=10)"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand All @@ -268,7 +339,23 @@
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"O\", \"Ca\", \"S\"])"
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"Ca\", \"S\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When using both the `nelements` and the `species_list` options, `species_list` takes precedence. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"Ca\", \"S\"], nelements=3)"
]
},
{
Expand All @@ -285,7 +372,7 @@
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"O\", \"Ca\", \"S\"], num_bins=10)"
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"Ca\", \"S\"], num_bins=10)"
]
},
{
Expand All @@ -302,7 +389,7 @@
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"O\", \"Ca\", \"S\"], xlog_scale=True, ylog_scale=True)"
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"Ca\", \"S\"], xlog_scale=True, ylog_scale=True)"
]
},
{
Expand All @@ -319,18 +406,14 @@
"metadata": {},
"outputs": [],
"source": [
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"O\", \"Ca\", \"S\"], velocity_range=(12500, 15050))"
"plotter.generate_plot_ply(species_list = [\"Si I-III\", \"Ca\", \"S\"], velocity_range=(12500, 15050))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Additional plotting options\n",
"The `generate_plot_ply` method also has options specific to the plotly API, thereby providing you with more control over how your last interaction velocity plot looks. Possible cases where you may use them are:\n",
" - `fig`: To plot the last interaction velocity plot on a figure you are already using e.g. for subplots.\n",
" - `graph_height`: To specify the height of the graph as needed.\n",
" - `cmapname`: To use a colormap of your preference instead of \"jet\"."
"### Additional plotting options"
]
},
{
Expand All @@ -343,6 +426,20 @@
"help(plotter.generate_plot_ply)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"The `generate_plot_ply` method also has options specific to the plotly API, thereby providing you with more control over how your last interaction velocity plot looks. Possible cases where you may use them are:\n",
"\n",
" - `fig`: To plot the last interaction velocity plot on a figure you are already using e.g. for subplots.\n",
"\n",
" - `graph_height`: To specify the height of the graph as needed.\n",
" \n",
" - `cmapname`: To use a colormap of your preference instead of \"jet\"."
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down
51 changes: 46 additions & 5 deletions tardis/visualization/tools/liv_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ def from_hdf(cls, hdf_fpath):
velocity,
)

def _parse_species_list(self, species_list):
def _parse_species_list(self, species_list, packets_mode, nelements=None):
"""
Parse user requested species list and create list of species ids to be used.
Expand All @@ -109,6 +109,10 @@ def _parse_species_list(self, species_list):
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])
packets_mode : str, optional
Packet mode, either 'virtual' or 'real'. Default is 'virtual'.
nelements : int, optional
Number of elements to include in plot. The most interacting elements are included. If None, displays all elements.
Raises
------
Expand All @@ -121,6 +125,23 @@ def _parse_species_list(self, species_list):
self._species_mapped = self.sdec_plotter._species_mapped
self._keep_colour = self.sdec_plotter._keep_colour

if nelements:
interaction_counts = (
self.data[packets_mode]
.packets_df_line_interaction["last_line_interaction_species"]
.value_counts()
)
interaction_counts.index = interaction_counts.index // 100
element_counts = interaction_counts.groupby(
interaction_counts.index
).sum()
top_elements = element_counts.nlargest(nelements).index
top_species_list = [
atomic_number2element_symbol(element)
for element in top_elements
]
self._parse_species_list(top_species_list, packets_mode)

def _make_colorbar_labels(self):
"""
Generate labels for the colorbar based on species.
Expand Down Expand Up @@ -215,7 +236,7 @@ def _generate_plot_data(self, packets_mode):
return plot_data, plot_colors

def _prepare_plot_data(
self, packets_mode, species_list, cmapname, num_bins
self, packets_mode, species_list, cmapname, num_bins, nelements
):
"""
Prepare data and settings required for generating a plot.
Expand Down Expand Up @@ -265,7 +286,7 @@ def _prepare_plot_data(
f"{atomic_number2element_symbol(specie // 100)}"
for specie in species_in_model
]
self._parse_species_list(species_list)
self._parse_species_list(species_list, packets_mode, nelements)
species_in_model = np.unique(
self.data[packets_mode]
.packets_df_line_interaction["last_line_interaction_species"]
Expand Down Expand Up @@ -328,6 +349,7 @@ def _get_step_plot_data(self, data, bin_edges):
def generate_plot_mpl(
self,
species_list=None,
nelements=None,
packets_mode="virtual",
ax=None,
figsize=(11, 5),
Expand All @@ -344,6 +366,8 @@ def generate_plot_mpl(
----------
species_list : list of str, optional
List of species to plot. Default is None which plots all species in the model.
nelements : int, optional
Number of elements to include in plot. The most interacting elements are included. If None, displays all elements.
packets_mode : str, optional
Packet mode, either 'virtual' or 'real'. Default is 'virtual'.
ax : matplotlib.axes.Axes, optional
Expand All @@ -366,8 +390,15 @@ def generate_plot_mpl(
matplotlib.axes.Axes
Axes object with the plot.
"""
# If species_list and nelements requested, tell user that nelements is ignored
if species_list is not None and nelements is not None:
logger.info(
"Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
)
nelements = None

plot_data, plot_colors, bin_edges = self._prepare_plot_data(
packets_mode, species_list, cmapname, num_bins
packets_mode, species_list, cmapname, num_bins, nelements
)

if ax is None:
Expand Down Expand Up @@ -407,6 +438,7 @@ def generate_plot_mpl(
def generate_plot_ply(
self,
species_list=None,
nelements=None,
packets_mode="virtual",
fig=None,
graph_height=600,
Expand All @@ -423,6 +455,8 @@ def generate_plot_ply(
----------
species_list : list of str, optional
List of species to plot. Default is None which plots all species in the model.
nelements : int, optional
Number of elements to include in plot. The most interacting elements are included. If None, displays all elements.
packets_mode : str, optional
Packet mode, either 'virtual' or 'real'. Default is 'virtual'.
fig : plotly.graph_objects.Figure, optional
Expand All @@ -445,8 +479,15 @@ def generate_plot_ply(
plotly.graph_objects.Figure
Plotly figure object with the plot.
"""
# If species_list and nelements requested, tell user that nelements is ignored
if species_list is not None and nelements is not None:
logger.info(
"Both nelements and species_list were requested. Species_list takes priority; nelements is ignored"
)
nelements = None

plot_data, plot_colors, bin_edges = self._prepare_plot_data(
packets_mode, species_list, cmapname, num_bins
packets_mode, species_list, cmapname, num_bins, nelements
)

if fig is None:
Expand Down

0 comments on commit 046f3b8

Please sign in to comment.