Source code for tardis.visualization.widgets.line_info

"""Class to create and display Line Info Widget."""

import re

from astropy import units as u
import numpy as np
import pandas as pd
import qgrid
from plotly import graph_objects as go
from plotly.callbacks import BoxSelector
import ipywidgets as ipw

from tardis.analysis import LastLineInteraction
from tardis.util.base import (
    species_tuple_to_string,
    species_string_to_tuple,
    is_notebook,
)
from tardis.visualization.widgets.util import (
    create_table_widget,
    TableSummaryLabel,
)
from tardis.visualization import plot_util as pu


[docs] class LineInfoWidget: """ Widget to explore atomic lines that produced features in the simulated spectrum. It allows selection of a wavelength range in the spectrum to display a table giving the fraction of packets that experienced their last interaction with each species. Using toggle buttons, users can specify whether to filter the selected range by emitted or absorbed wavelengths of packets. Clicking on a row in the fractional species interactions table shows packet counts for each last line interaction of the selected species, which can be grouped in several ways using the dropdown menu. """ FILTER_MODES = ("packet_out_nu", "packet_in_nu") FILTER_MODES_DESC = ("Emitted Wavelength", "Absorbed Wavelength") GROUP_MODES = ("both", "exc", "de-exc") GROUP_MODES_DESC = ( "Both excitation line (absorption) and de-excitation line (emission)", "Only excitation line (absorption)", "Only de-excitation line (emission)", ) COLORS = {"selection_area": "lightpink", "selection_border": "salmon"} def __init__( self, lines_data, line_interaction_analysis, spectrum_wavelength, spectrum_luminosity_density_lambda, virt_spectrum_wavelength, virt_spectrum_luminosity_density_lambda, ): """ Initialize the LineInfoWidget with line interaction and spectrum data. Parameters ---------- lines_data : pd.DataFrame Data about the atomic lines present in simulation model's plasma line_interaction_analysis : dict of tardis.analysis.LastLineInteraction Dictionary in which keys are the FILTER_MODES and values are the LastLineInteraction objects initialized with corresponding modes spectrum_wavelength : astropy.Quantity Wavelength values of a real spectrum, having unit of Angstrom spectrum_luminosity_density_lambda : astropy.Quantity Luminosity density lambda values of a real spectrum, having unit of (erg/s)/Angstrom virt_spectrum_wavelength : astropy.Quantity Wavelength values of a virtual spectrum, having unit of Angstrom virt_spectrum_luminosity_density_lambda : astropy.Quantity Luminosity density lambda values of a virtual spectrum, having unit of (erg/s)/Angstrom """ self.lines_data = lines_data self.line_interaction_analysis = line_interaction_analysis # Widgets ------------------------------------------------ max_rows_option = {"maxVisibleRows": 9} self.species_interactions_table = create_table_widget( data=self.get_species_interactions(None), col_widths=[35, 65], table_options=max_rows_option, ) line_counts_col_widths = [75, 25] self.last_line_counts_table = create_table_widget( data=self.get_last_line_counts(None), col_widths=line_counts_col_widths, table_options=max_rows_option, ) self.total_packets_label = TableSummaryLabel( target_table=self.last_line_counts_table, table_col_widths=line_counts_col_widths, label_key="Total Packets", label_value=0, ) self.figure_widget = self.plot_spectrum( spectrum_wavelength, spectrum_luminosity_density_lambda, virt_spectrum_wavelength, virt_spectrum_luminosity_density_lambda, ) self.filter_mode_buttons = ipw.ToggleButtons( options=self.FILTER_MODES_DESC, index=0 ) self.group_mode_dropdown = ipw.Dropdown( options=self.GROUP_MODES_DESC, index=0 )
[docs] @classmethod def from_simulation(cls, sim): """ Create an instance of LineInfoWidget from a TARDIS simulation object. Parameters ---------- sim : tardis.simulation.Simulation TARDIS Simulation object produced by running a simulation Returns ------- LineInfoWidget object """ spectrum_solver = sim.spectrum_solver return cls( lines_data=sim.plasma.lines.reset_index().set_index("line_id"), line_interaction_analysis={ filter_mode: LastLineInteraction.from_simulation( sim, filter_mode ) for filter_mode in cls.FILTER_MODES }, spectrum_wavelength=spectrum_solver.spectrum_real_packets.wavelength, spectrum_luminosity_density_lambda=spectrum_solver.spectrum_real_packets.luminosity_density_lambda.to( "erg/(s AA)" ), virt_spectrum_wavelength=spectrum_solver.spectrum_virtual_packets.wavelength, virt_spectrum_luminosity_density_lambda=spectrum_solver.spectrum_virtual_packets.luminosity_density_lambda.to( "erg/(s AA)" ), )
[docs] def get_species_interactions( self, wavelength_range, filter_mode=FILTER_MODES[0] ): """ Get fractional species interactions in specified wavelength range. Fractional species interactions means fraction of packets present in the specified wavelength range which experienced their last interaction with a species. The packets to consider are filtered by the specified filter mode. Parameters ---------- wavelength_range : list-like or None A list of two float values to specify the wavelength range - first for the range start and second for the range end. None specifies that no wavelength range is selected and will return empty dataframe filter_mode : str, optional Filter mode of the LastLineInteraction object to use for filtering the selected wavelength range (more details in Notes section). Allowed values are given by the class variable :code:`FILTER_MODES` (default value is :code:`FILTER_MODES[0]`) Returns ------- pandas.DataFrame Dataframe containing species symbols and corresponding fractions of packets interacting with them Notes ----- This method depends on tardis.analysis.LastLineInteraction object for doing computations. So there is a member variable in this class - :code:`line_interaction_analysis` which is a dictionary of such objects (each of them differ in how they filter the selected wavelength range). Thus we have to specify which object to use by specifying the filter_mode parameter. """ if wavelength_range: self.line_interaction_analysis[filter_mode].wavelength_start = ( wavelength_range[0] * u.AA ) self.line_interaction_analysis[filter_mode].wavelength_end = ( wavelength_range[1] * u.AA ) # Obtain species group from last_line_in dataframe selected_species_group = self.line_interaction_analysis[ filter_mode ].last_line_in.groupby(["atomic_number", "ion_number"]) if selected_species_group.groups: selected_species_symbols = [ species_tuple_to_string(item) for item in selected_species_group.groups.keys() ] # Normalize each group's count to find fractions of interactions fractional_species_interactions = ( selected_species_group.size() / self.line_interaction_analysis[ filter_mode ].last_line_in.shape[0] ) else: # No species could be selected in specified wavelength_range # qgrid cannot show empty dataframe properly, # so create one row with empty strings selected_species_symbols = [""] fractional_species_interactions = pd.Series([""]) else: # wavelength_range is None selected_species_symbols = [""] fractional_species_interactions = pd.Series([""]) fractional_species_interactions.index = pd.Index( selected_species_symbols, name="Species" ) fractional_species_interactions.name = "Fraction of packets interacting" return fractional_species_interactions.sort_values( ascending=False ).to_frame()
[docs] def get_last_line_counts( self, selected_species, filter_mode=FILTER_MODES[0], group_mode=GROUP_MODES[0], ): """ Get packet counts of each last line interaction of a species. Parameters ---------- selected_species : str Valid symbol of a species (e.g Si II) selected from the species data returned by :code:`get_species_interactions` (see Notes section) filter_mode : str, optional Filter mode of the LastLineInteraction object to use for fetching the data of last lines interacted (more details in Notes section). Allowed values are given by the class variable :code:`FILTER_MODES` (default value is :code:`FILTER_MODES[0]`) group_mode : str, optional Group mode to use for grouping last line interactions by excitation lines, de-excitation lines or both. Allowed values are given by the class variable :code:`GROUP_MODES` (default value is :code:`GROUP_MODES[0]` i.e. both) Returns ------- pd.DataFrame DataFrame containing last line interactions and corresponding packet counts. Notes ----- This method depends on tardis.analysis.LastLineInteraction object for doing computations. So there is a member variable in this class - :code:`line_interaction_analysis` which is a dictionary of such objects (each of them differ in how they filter the selected wavelength range). Thus we have to specify which object to use by specifying the filter_mode parameter. This method should always be called after calling :code:`get_species_interactions` method which sets a wavelength range on LastLineInteraction object. So selected_species should be one present within that range, otherwise it will result an error. """ if selected_species: selected_species_tuple = species_string_to_tuple(selected_species) try: # Get selected species' rows from last_line_in dataframe current_last_lines_in = ( self.line_interaction_analysis[filter_mode] .last_line_in.xs( key=selected_species_tuple, level=["atomic_number", "ion_number"], drop_level=False, ) .reset_index() ) # Get selected species' rows from last_line_out dataframe current_last_lines_out = ( self.line_interaction_analysis[filter_mode] .last_line_out.xs( key=selected_species_tuple, level=["atomic_number", "ion_number"], drop_level=False, ) .reset_index() ) assert ( current_last_lines_in.empty & current_last_lines_out.empty == False ) except (KeyError, AssertionError): # selected_species is invalid allowed_species = [ species_tuple_to_string(species) for species in self.line_interaction_analysis[filter_mode] .last_line_in.groupby(["atomic_number", "ion_number"]) .groups.keys() ] raise ValueError( "Invalid value of selected_species, it must be one present " "within the currently selected wavelength range in your " f"LineInfoWidget instance, which are {allowed_species}" ) last_line_interaction_string = [] interacting_packets_count = [] if group_mode == "both": # Group by both exc. line ids and de-exc. line ids current_last_lines_in[ "line_id_out" ] = current_last_lines_out.line_id grouped_line_interactions = current_last_lines_in.groupby( ["line_id", "line_id_out"] ) # Iterate over each group's key and size and append them to list for ( line_id, count, ) in grouped_line_interactions.size().items(): current_line_in = self.lines_data.loc[line_id[0]] current_line_out = self.lines_data.loc[line_id[1]] last_line_interaction_string.append( f"exc. {int(current_line_in.level_number_lower):02d}-" f"{int(current_line_in.level_number_upper):02d} " f"({current_line_in.wavelength:.2f} A) " f"de-exc. {int(current_line_out.level_number_upper):02d}-" f"{int(current_line_out.level_number_lower):02d} " f"({current_line_out.wavelength:.2f} A)" ) interacting_packets_count.append(count) elif group_mode == "exc": grouped_line_interactions = current_last_lines_in.groupby( "line_id" ) # Iterate over each group's key and size and append them to list for ( line_id, count, ) in grouped_line_interactions.size().items(): current_line_in = self.lines_data.loc[line_id] last_line_interaction_string.append( f"exc. {int(current_line_in.level_number_lower):02d}-" f"{int(current_line_in.level_number_upper):02d} " f"({current_line_in.wavelength:.2f} A)" ) interacting_packets_count.append(count) elif group_mode == "de-exc": grouped_line_interactions = current_last_lines_out.groupby( "line_id" ) # Iterate over each group's key and size and append them to list for ( line_id, count, ) in grouped_line_interactions.size().items(): current_line_out = self.lines_data.loc[line_id] last_line_interaction_string.append( f"de-exc. {int(current_line_out.level_number_upper):02d}-" f"{int(current_line_out.level_number_lower):02d} " f"({current_line_out.wavelength:.2f} A)" ) interacting_packets_count.append(count) else: raise ValueError( "Invalid value passed to group_mode argument. " f"Allowed values are {self.GROUP_MODES}" ) else: # species_selected is None # qgrid cannot show empty dataframe properly, # so create one row with empty strings interacting_packets_count = [""] last_line_interaction_string = [""] last_line_counts = pd.Series(interacting_packets_count) last_line_counts.name = "No. of packets" last_line_counts.index = pd.Index( last_line_interaction_string, name="Last Line Interaction" ) return last_line_counts.sort_values(ascending=False).to_frame()
[docs] @staticmethod def get_middle_half_edges(arr): """ Get edges of the middle half range of an array. Parameters ---------- arr : np.array Returns ------- list """ arr = np.sort(arr) return [ (arr[-1] - arr[0]) / 4 + arr[1], (arr[-1] - arr[0]) * 3 / 4 + arr[1], ]
[docs] def plot_spectrum( self, wavelength, luminosity_density_lambda, virt_wavelength, virt_luminosity_density_lambda, ): """ Produce a plotly figure widget by plotting the spectrum of model. Parameters ---------- wavelength : astropy.Quantity Wavelength values of a real spectrum, having unit of Angstrom luminosity_density_lambda : astropy.Quantity Luminosity density lambda values of a real spectrum, having unit of (erg/s)/Angstrom virt_wavelength : astropy.Quantity Wavelength values of a virtual spectrum, having unit of Angstrom virt_luminosity_density_lambda : astropy.Quantity Luminosity density lambda values of a virtual spectrum, having unit of (erg/s)/Angstrom Returns ------- plotly.graph_objects.FigureWidget """ # Initially zoomed range in rangeslider should be middle half of spectrum initial_zoomed_range = self.get_middle_half_edges(wavelength.value) # The scatter point should be a middle point in spectrum otherwise # the extra padding around it will be oddly visible when near the edge scatter_point_idx = pu.get_mid_point_idx(wavelength.value) return go.FigureWidget( [ go.Scatter( x=wavelength, y=luminosity_density_lambda, name="Real packets", ), go.Scatter( x=virt_wavelength, y=virt_luminosity_density_lambda, name="Virtual packets", ), # Hide a one point scatter trace, to bring boxselect in modebar go.Scatter( x=wavelength[scatter_point_idx], y=luminosity_density_lambda[scatter_point_idx], mode="markers", marker=dict(opacity=0), showlegend=False, ), ], layout=go.Layout( title="Spectrum", xaxis=dict( title=pu.axis_label_in_latex("Wavelength", wavelength.unit), exponentformat="none", rangeslider=dict(visible=True), range=initial_zoomed_range, ), yaxis=dict( title=pu.axis_label_in_latex( "Luminosity", luminosity_density_lambda.unit, ), exponentformat="e", fixedrange=False, ), dragmode="select", selectdirection="h", height=400, margin=dict(t=50, b=60), ), )
def _update_species_interactions(self, wavelength_range, filter_mode): """ Update data in species_interactions_table. The parameters are exact same as that of :code:`get_species_interactions`. Besides, it also does selection of 1st row in this table to trigger update in last_line_counts_table. """ # Update data in species_interactions_table self.species_interactions_table.df = self.get_species_interactions( wavelength_range, filter_mode ) # Get index of 0th row in species_interactions_table species0 = self.species_interactions_table.df.index[0] # Also update last_line_counts_table by triggering its event listener if self.species_interactions_table.get_selected_rows() == [0]: # Listener won't trigger if last row selected in # species_interactions_table was also 0th, so unselect the rows self.species_interactions_table.change_selection([]) # Select 0th row in this table to trigger _update_last_line_counts self.species_interactions_table.change_selection([species0]) def _add_selection_box(self, selector): """ Draw a shape on plotly figure widget to represent the selection. Parameters ---------- selector : plotly.callbacks.BoxSelector The object containing data about current selection made on plot (x-axis and y-axis range of selection box) """ self.figure_widget.layout.shapes = [ dict( type="rect", xref="x", yref="y", x0=selector.xrange[0], y0=selector.yrange[0], x1=selector.xrange[1], y1=selector.yrange[1], line=dict( color=self.COLORS["selection_border"], width=1, ), fillcolor=self.COLORS["selection_area"], opacity=0.5, ) ] def _update_last_line_counts(self, species, filter_mode, group_mode): """ Update data in last_line_counts_table and associated total_packets_label. The parameters are exact same as that of :code:`get_last_line_counts`. """ # Update data in line counts table self.last_line_counts_table.df = self.get_last_line_counts( species, filter_mode, group_mode ) # Update its corresponding total_packets_label if species: self.total_packets_label.update_and_resize( self.last_line_counts_table.df.iloc[:, 0].sum() ) else: # Line counts table will be empty self.total_packets_label.update_and_resize(0) def _spectrum_selection_handler(self, trace, points, selector): """ Event handler for selection of spectrum in plotly figure widget. This method has the expected signature of the callback function passed to :code:`on_selection` method of a plotly trace as explained in `their docs <https://plotly.com/python-api-reference/generated/plotly.html#plotly.basedatatypes.BaseTraceType.on_selection>`_. """ if isinstance(selector, BoxSelector): self._add_selection_box(selector) self._update_species_interactions( selector.xrange, self.FILTER_MODES[self.filter_mode_buttons.index], ) def _filter_mode_toggle_handler(self, change): """ Event handler for toggle in filter_mode_buttons. This method has the expected signature of the callback function passed to :code:`observe` method of ipywidgets as explained in `their docs <https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Events.html#Signatures>`_. """ try: wavelength_range = [ self.figure_widget.layout.shapes[0][x] for x in ("x0", "x1") ] except IndexError: # No selection is made on figure widget return self._update_species_interactions( wavelength_range, self.FILTER_MODES[self.filter_mode_buttons.index], ) def _species_intrctn_selection_handler(self, event, qgrid_widget): """ Event handler for selection in species_interactions_table. This method has the expected signature of the function passed to :code:`handler` argument of :code:`on_selection` method of qgrid.QgridWidget as explained in `their docs <https://qgrid.readthedocs.io/en/latest/#qgrid.QgridWidget.on>`_. """ # Don't execute function if no row was selected implicitly (by api) if event["new"] == [] and event["source"] == "api": return # Get species from the selected row in species_interactions_table species_selected = self.species_interactions_table.df.index[ event["new"][0] ] if species_selected == "": # when species_interactions_table is empty species_selected = None self._update_last_line_counts( species_selected, self.FILTER_MODES[self.filter_mode_buttons.index], self.GROUP_MODES[self.group_mode_dropdown.index], ) def _group_mode_dropdown_handler(self, change): """ Event handler for selection in group_mode_dropdown. This method has the expected signature of the callback function passed to :code:`observe` method of ipywidgets as explained in `their docs <https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Events.html#Signatures>`_. """ try: selected_row_idx = ( self.species_interactions_table.get_selected_rows()[0] ) species_selected = self.species_interactions_table.df.index[ selected_row_idx ] except IndexError: # No row is selected in species_interactions_table return self._update_last_line_counts( species_selected, self.FILTER_MODES[self.filter_mode_buttons.index], self.GROUP_MODES[self.group_mode_dropdown.index], )
[docs] @staticmethod def ui_control_description(text): """Get description label of a UI control with increased font size.""" return ipw.HTML(f"<span style='font-size: 1.15em;'>{text}:</span>")
[docs] def display(self): """ Display the fully-functional line info widget. It puts together all component widgets nicely together and enables interaction between all the components. Returns ------- ipywidgets.Box Line info widget containing all component widgets """ if not is_notebook(): print("Please use a notebook to display the widget") else: # Set widths of widgets self.species_interactions_table.layout.width = "350px" self.last_line_counts_table.layout.width = "450px" self.total_packets_label.update_and_resize(0) self.group_mode_dropdown.layout.width = "auto" # Attach event listeners to widgets spectrum_trace = self.figure_widget.data[0] spectrum_trace.on_selection(self._spectrum_selection_handler) self.filter_mode_buttons.observe( self._filter_mode_toggle_handler, names="index" ) self.species_interactions_table.on( "selection_changed", self._species_intrctn_selection_handler ) self.group_mode_dropdown.observe( self._group_mode_dropdown_handler, names="index" ) selection_box_symbol = ( "<span style='display: inline-block; " f"background-color: {self.COLORS['selection_area']}; " f"border: 1px solid {self.COLORS['selection_border']}; " "width: 0.8em; height: 1.2em; vertical-align: middle;'></span>" ) table_container_left = ipw.VBox( [ self.ui_control_description( "Filter selected wavelength range " f"( {selection_box_symbol} ) by" ), self.filter_mode_buttons, self.species_interactions_table, ], layout=dict(margin="0px 15px"), ) table_container_right = ipw.VBox( [ self.ui_control_description("Group packet counts by"), self.group_mode_dropdown, self.last_line_counts_table, self.total_packets_label.widget, ], layout=dict(margin="0px 15px"), ) return ipw.VBox( [ self.figure_widget, ipw.Box( [ table_container_left, table_container_right, ], layout=dict( display="flex", align_items="flex-start", justify_content="center", height="420px", ), ), ] )