diff --git a/docs/io/visualization/how_to_liv_plot.ipynb b/docs/io/visualization/how_to_liv_plot.ipynb index 377142e625c..795b9090277 100644 --- a/docs/io/visualization/how_to_liv_plot.ipynb +++ b/docs/io/visualization/how_to_liv_plot.ipynb @@ -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": {}, @@ -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": {}, @@ -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" ] }, { @@ -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": {}, @@ -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": {}, @@ -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)" ] }, { @@ -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)" ] }, { @@ -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)" ] }, { @@ -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" ] }, { @@ -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": {}, diff --git a/tardis/visualization/tools/liv_plot.py b/tardis/visualization/tools/liv_plot.py index e9da85e5b35..2e4c9df87b3 100644 --- a/tardis/visualization/tools/liv_plot.py +++ b/tardis/visualization/tools/liv_plot.py @@ -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. @@ -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 ------ @@ -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. @@ -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. @@ -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"] @@ -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), @@ -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 @@ -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: @@ -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, @@ -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 @@ -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: