`_ section to see the
notebook in action.
Currently, TARDIS supports the following widgets:
+Custom Abundance Widget
+#######################
+This widget (or GUI) allows you to develop custom model compositions
+graphically and output the model to a new file.
+
+.. image:: ../images/custom_abundance_widget.png
+ :alt: Demo of Custom Abundance Widget
+
+The GUI consists of three parts:
+
+1. **Visualization plot** - An interactive step graph that shows abundances
+and densities as a function of velocity.
+
+2. **Data editor** - An interface to edit abundances, densities and velocity
+shells.
+
+3. **File output** - An output module to save the model compositions as a CSVY
+file.
+
+Interacting with the GUI
+========================
+
+You can interact with this GUI in the following ways:
+
+Edit Abundances
+---------------
+There are two radio buttons which allow you to edit either single shell or
+mutiple shells. The new input will be applied to selected shell(s) immediately
+and the plot is updated at the same time. If you want to edit multiple shells
+at a time, remember to choose the second radio button and set the range of
+shell number using int slider. The selected shell(s) is highlighted in the
+plot.
+
+.. image:: ../images/cus_abund_edit_abundance.gif
+ :alt: Demo of editing abundances
+
+Normalize Abundances
+--------------------
+Click `Normalize` button to normalize the abundances on selected shell(s) to
+1. If you wish to keep the abundance of a certain element unchanged during the
+normalization, you can select the checkbox near that element to lock it.
+The normalization will be applied to other elements and the sum of the
+abundances still remains at 1.
+
+.. image:: ../images/cus_abund_norm.gif
+ :alt: Demo of normalization
+
+Edit Densities
+--------------
+In `Edit Density` tab, you also can edit either single shell or the whole
+density profile. To calculate a new profile, you need to choose the density
+type and input required parameters.
+
+.. image:: ../images/cus_abund_edit_density.gif
+ :alt: Demo of editing densities
+
+Add New Element
+---------------
+At the bottom of the editor, the symbol input box allows you to add new element
+or isotope to the model. It automatically recognizes whether the symbol exists
+or whether the element is already in the model.
+
+.. image:: ../images/cus_abund_add_element.gif
+ :alt: Demo of adding new element
+
+Add New Shell
+-------------
+Another important functionality is to add new shell to the model. You need to
+specify the velocity range of the new shell and abundances at that new shell
+will be set to 0. Note that the existing shell(s) will be resized smartly if
+the new shell overlaps with it.
+
+.. image:: ../images/cus_abund_add_shell.gif
+ :alt: Demo of adding new shell
+
Shell Info Widget
#################
diff --git a/tardis/visualization/__init__.py b/tardis/visualization/__init__.py
index bad425bd8d9..4b806fdd147 100644
--- a/tardis/visualization/__init__.py
+++ b/tardis/visualization/__init__.py
@@ -7,4 +7,5 @@
shell_info_from_hdf,
)
from tardis.visualization.widgets.line_info import LineInfoWidget
+from tardis.visualization.widgets.custom_abundance import CustomAbundanceWidget
from tardis.visualization.tools.sdec_plot import SDECPlotter
diff --git a/tardis/visualization/widgets/custom_abundance.py b/tardis/visualization/widgets/custom_abundance.py
new file mode 100644
index 00000000000..0d8a9cae23b
--- /dev/null
+++ b/tardis/visualization/widgets/custom_abundance.py
@@ -0,0 +1,1620 @@
+"""Class to create and display Custom Abundance Widget."""
+import os
+import yaml
+import numpy as np
+import pandas as pd
+import ipywidgets as ipw
+import plotly.graph_objects as go
+from astropy import units as u
+from pyne import nucname
+from pathlib import Path
+
+import tardis
+from tardis.util.base import quantity_linspace
+from tardis.io.config_reader import Configuration
+from tardis.model import Radial1DModel
+from tardis.model.density import (
+ calculate_power_law_density,
+ calculate_exponential_density,
+)
+from tardis.io.config_validator import validate_dict
+from tardis.io.parsers.csvy import load_csvy
+from tardis.io.model_reader import (
+ read_uniform_abundances,
+ parse_csv_abundances,
+)
+from tardis.util.base import atomic_number2element_symbol, quantity_linspace
+from tardis.visualization.tools.convergence_plot import transition_colors
+from tardis.visualization.widgets.util import debounce
+
+BASE_DIR = tardis.__path__[0]
+YAML_DELIMITER = "---"
+COLORMAP = "viridis"
+
+
+class CustomYAML(yaml.YAMLObject):
+ """A custom YAML object generated by required properties."""
+
+ def __init__(
+ self, name, d_time_0, i_time_0, v_inner_boundary, v_outer_boundary
+ ):
+ """Initialize CustomYAML object with given properties.
+
+ Parameters
+ ----------
+ name : str
+ Name of the YAML file.
+ d_time_0 : astropy.units.quantity.Quantity
+ Initial time for the density in the model.
+ i_time_0 : astropy.units.quantity.Quantity
+ Initial time for isotope decay. Set to 0 for no isotopes.
+ v_inner_boundary : astropy.units.quantity.Quantity
+ Velocity of the inner boundary.
+ v_outer_boundary : astropy.units.quantity.Quantity
+ Velocity of the outer boundary.
+ """
+ self.name = name
+ self.model_density_time_0 = d_time_0
+ self.model_isotope_time_0 = i_time_0
+ self.datatype = {}
+ self.datatype["fields"] = []
+ self.v_inner_boundary = v_inner_boundary
+ self.v_outer_boundary = v_outer_boundary
+
+ def create_fields_dict(self, elements):
+ """Create a dictionary to store the items in 'fields' part.
+
+ Parameters
+ ----------
+ elements : list of str
+ A list of elements or isotopes' symbols.
+ """
+ for i in range(len(elements) + 2):
+ field = {}
+
+ if i == 0:
+ field["name"] = "velocity"
+ field["unit"] = "km/s"
+ elif i == 1:
+ field["name"] = "density"
+ field["unit"] = "g/cm^3"
+ else:
+ field["name"] = elements[i - 2]
+ field["desc"] = f"fractional {elements[i-2]} abundance"
+
+ self.datatype["fields"].append(field)
+
+
+class CustomAbundanceWidget:
+ """Widget to edit abundances and densities of simulation model
+ graphically.
+
+ It generates a GUI based on input data. The GUI has a plot section
+ to visualize the profile, a edit section to allow the user directly
+ edit abundance and density profile, and an output section to output
+ the model to CSVY file.
+
+ Attributes
+ ----------
+ shell_no : int
+ The selected shell number.
+ no_of_shells : int
+ The number of shells in the model.
+ no_of_elements : int
+ The number of elements in the model.
+ checked_list : list of bool
+ A list of bool to record whether the checkbox is checked.
+ The index of the bool corresponds to the index of checkbox.
+ elements : list of str
+ A list of elements or isotopes' symbols.
+ _trigger : bool
+ If False, disable the callback when abundance input is changed.
+ """
+
+ error_view = ipw.Output()
+
+ def __init__(self, density_t_0, density, abundance, velocity):
+ """Initialize CustomAbundanceWidget with density, abundance and
+ velocity data.
+
+ Parameters
+ ----------
+ density : astropy.units.quantity.Quantity
+ abundance : pd.DataFrame
+ velocity : astropy.units.quantity.Quantity
+ """
+ density_t_0 = density_t_0.to("day")
+ self.density = density.to("g cm^-3")
+ self.abundance = abundance
+ self.velocity = velocity.to("km/s")
+ self.elements = self.get_symbols()
+ self._trigger = True
+
+ self.create_widgets()
+ self.generate_abundance_density_plot()
+ self.density_editor = DensityEditor(
+ density_t_0,
+ self.density,
+ self.velocity,
+ self.fig,
+ self.dpd_shell_no,
+ )
+
+ @property
+ def shell_no(self):
+ return self.dpd_shell_no.value
+
+ @shell_no.setter
+ def shell_no(self, value):
+ self.dpd_shell_no.value = value
+
+ @property
+ def no_of_shells(self):
+ return self.abundance.shape[1]
+
+ @property
+ def no_of_elements(self):
+ return self.abundance.shape[0]
+
+ @property
+ def checked_list(self): # A boolean list to store the value of checkboxes.
+ _checked_list = []
+ for check in self.checks:
+ _checked_list.append(check.value)
+
+ return _checked_list
+
+ def get_symbols(self):
+ """Get symbol string from atomic number and mass number."""
+ str_symbols = np.array(
+ self.abundance.index.get_level_values(0).map(
+ atomic_number2element_symbol
+ )
+ )
+ str_mass = np.array(
+ self.abundance.index.get_level_values(1), dtype="str"
+ )
+ return np.add(str_symbols, str_mass)
+
+ def create_widgets(self):
+ """Create widget components in GUI and register callbacks for widgets."""
+ self.dpd_shell_no = ipw.Dropdown(
+ options=list(range(1, self.no_of_shells + 1)),
+ description="Shell No. ",
+ value=1,
+ layout=ipw.Layout(width="160px"),
+ )
+ self.dpd_shell_no.observe(self.dpd_shell_no_eventhandler, "value")
+ self.btn_prev = ipw.Button(
+ icon="chevron-left",
+ disabled=True,
+ layout=ipw.Layout(width="30px", height="30px"),
+ )
+ self.btn_prev.on_click(self.on_btn_prev)
+ self.btn_next = ipw.Button(
+ icon="chevron-right", layout=ipw.Layout(width="30px", height="30px")
+ )
+ self.btn_next.on_click(self.on_btn_next)
+
+ self.checks = [
+ ipw.Checkbox(
+ indent=False,
+ icon="lock",
+ layout=ipw.Layout(
+ width="30px",
+ ),
+ )
+ for element in self.elements
+ ]
+ self.input_items = [
+ ipw.BoundedFloatText(min=0, max=1, step=0.01, description=element)
+ for element in self.elements
+ ]
+ for i in range(self.no_of_elements):
+ self.input_items[i].observe(self.input_item_eventhandler, "value")
+ self.input_items[i].index = i
+ self.checks[i].observe(self.check_eventhandler, "value")
+ self.checks[i].index = i
+
+ self.btn_norm = ipw.Button(
+ description="Normalize",
+ icon="cog",
+ layout=ipw.Layout(width="100px", margin="0 0 0 50px"),
+ )
+ self.btn_norm.on_click(self.on_btn_norm)
+ self.norm_warning = ipw.Valid(
+ value=False,
+ readout="Unnormalized",
+ style={"description_width": "initial"},
+ layout=ipw.Layout(visibility="hidden"),
+ )
+
+ self.symb_warning = ipw.Valid(
+ value=False, layout=ipw.Layout(visibility="hidden")
+ )
+ self.input_symb = ipw.Text(
+ description="Element: ",
+ style={"description_width": "initial"},
+ placeholder="symbol",
+ layout=ipw.Layout(width="125px"),
+ )
+ self.input_symb.observe(self.input_symb_eventhandler, "value")
+ self.btn_add_element = ipw.Button(
+ icon="plus-square",
+ description="Add",
+ disabled=True,
+ layout=ipw.Layout(width="60px"),
+ )
+ self.btn_add_element.on_click(self.on_btn_add_element)
+
+ self.irs_shell_range = ipw.IntRangeSlider(
+ min=1,
+ max=self.no_of_shells,
+ step=1,
+ description="Shell No. ",
+ disabled=True,
+ style={"description_width": "initial"},
+ continuous_update=False,
+ layout=ipw.Layout(margin="5px 0 0 0"),
+ )
+ self.irs_shell_range.observe(self.irs_shell_range_eventhandler)
+
+ self.btn_add_shell = ipw.Button(
+ icon="plus-square",
+ description="Add",
+ disabled=True,
+ layout=ipw.Layout(
+ width="80px",
+ ),
+ )
+ self.btn_add_shell.on_click(self.on_btn_add_shell)
+ self.input_v_start = ipw.FloatText(
+ min=0,
+ description="Add shell(s) with velocity range (km/s): ",
+ style={"description_width": "initial"},
+ layout=ipw.Layout(
+ width="320px",
+ ),
+ )
+ self.input_v_end = ipw.FloatText(
+ min=0,
+ description="to",
+ style={"description_width": "initial"},
+ layout=ipw.Layout(width="110px"),
+ )
+ self.input_v_start.observe(self.input_v_eventhandler, "value")
+ self.input_v_end.observe(self.input_v_eventhandler, "value")
+ self.overwrite_warning = ipw.HTML(
+ value="Warning: Existing shell(s) will be overwritten!",
+ layout=ipw.Layout(visibility="hidden", margin="0 0 0 10px"),
+ )
+
+ self.btn_output = ipw.Button(
+ description="Output CSVY File",
+ )
+ self.btn_output.on_click(self.on_btn_output)
+
+ self.input_path = ipw.Text(
+ description="File path: ", placeholder="Input file name or path"
+ )
+
+ self.input_i_time_0 = ipw.FloatText(
+ description="Isotope time_0 (day): ",
+ style={"description_width": "initial"},
+ )
+
+ self.ckb_overwrite = ipw.Checkbox(
+ description="overwrite",
+ indent=False,
+ )
+
+ self.tbs_scale = ipw.ToggleButtons(
+ options=["Linear", "Log"],
+ description="Scale of yaxes: ",
+ style={"description_width": "initial"},
+ value="Linear",
+ )
+ self.tbs_scale.observe(self.tbs_scale_eventhandler, "value")
+
+ self.rbs_single_apply = ipw.RadioButtons(
+ options=["Only selected shell"],
+ layout=ipw.Layout(margin="10px 0 0 0"),
+ )
+ self.rbs_single_apply.observe(
+ self.rbs_single_apply_eventhandler, "value"
+ )
+ self.rbs_multi_apply = ipw.RadioButtons(
+ options=["A range of shells: "],
+ index=None,
+ layout=ipw.Layout(width="130px", margin="10px 0 10px 0"),
+ )
+ self.rbs_multi_apply.observe(self.rbs_multi_apply_eventhandler, "value")
+
+ def update_input_item_value(self, index, value):
+ """Update the value of the widget in the list of abundance inputs.
+
+ Keep two decimal places for displayed value and disable
+ `input_item_eventhandler` while changing the value.
+
+ Parameters
+ ----------
+ index : int
+ The index of the widget in the list of abundance inputs.
+ value : float
+ New abundance value.
+ """
+ self._trigger = False
+ # `input_items` is the list of abundance input widgets.
+ self.input_items[index].value = float("{:.2e}".format(value))
+ self._trigger = True
+
+ def read_abundance(self):
+ """Read abundances data in DataFrame to input items box when
+ shell No. changes.
+ """
+ for i in range(self.no_of_elements):
+ value = self.abundance.iloc[i, self.shell_no - 1]
+ self.update_input_item_value(i, value)
+
+ def bound_locked_sum_to_1(self, index):
+ """Ensure the sum of locked abundances is no more than 1. If the
+ locked sum is more than 1, calculate the maximum with the sum no
+ more than 1 and return it.
+
+ Parameters
+ ----------
+ index : int
+ The index of the widget in the list of abundance inputs.
+ """
+ locked_mask = np.array(self.checked_list)
+ back_value = self.abundance.iloc[
+ index, self.shell_no - 1
+ ] # abundance value in back end (DataFrame)
+ front_value = self.input_items[
+ index
+ ].value # abundance value in front end (widget)
+ locked_sum = (
+ self.abundance.loc[locked_mask, self.shell_no - 1].sum()
+ - back_value
+ + front_value
+ )
+
+ if locked_sum > 1:
+ new = 1 - (locked_sum - front_value)
+ self.abundance.iloc[index, self.shell_no - 1] = new
+ self.update_input_item_value(index, new)
+ self.update_abundance_plot(index)
+
+ def update_abundance_plot(self, index):
+ """Update the abundance line of certain element in the plot.
+
+ Parameters
+ ----------
+ index : int
+ The index of the widget in the list of abundance inputs.
+ """
+ y = self.abundance.iloc[index]
+ self.fig.data[index + 2].y = np.append(y, y.iloc[-1])
+
+ def update_front_end(self):
+ """Update checkbox widgets, input widgets and plot in the front
+ end when selected shell No. is changed.
+ """
+ # Update checkboxes, input boxes and plot.
+ for check in self.checks:
+ check.value = False
+ self.read_abundance()
+ self.density_editor.read_density()
+ with self.fig.batch_update():
+ # Change bar diagonal
+ x = list(self.fig.data[0].x)
+ width = list(self.fig.data[0].width)
+ x_inner = self.velocity[self.shell_no - 1].value
+ x_outer = self.velocity[self.shell_no].value
+ x[0] = (x_inner + x_outer) / 2
+ self.fig.data[0].x = x
+ width[0] = x_outer - x_inner
+ self.fig.data[0].width = width
+
+ def overwrite_existing_shells(self, v_0, v_1):
+ """Judge whether the existing shell(s) will be overwritten when
+ inserting a new shell within the entered velocity range.
+
+ Parameters
+ ----------
+ v_0 : float
+ The velocity of inner boundary.
+ v_1 : float
+ The velocity of outer boundary.
+
+ Returns
+ -------
+ bool
+ True if the existing shell will be overwritten, False otherwise.
+ """
+ position_0 = np.searchsorted(self.velocity.value, v_0)
+ position_1 = np.searchsorted(self.velocity.value, v_1)
+
+ index_0 = (
+ position_0 - 1
+ if np.isclose(self.velocity[position_0 - 1].value, v_0)
+ else position_0
+ )
+ index_1 = (
+ position_1 - 1
+ if np.isclose(self.velocity[position_1 - 1].value, v_1)
+ else position_1
+ )
+
+ if (index_1 - index_0 > 1) or (
+ (
+ index_1 < len(self.velocity)
+ and np.isclose(self.velocity[index_1].value, v_1)
+ )
+ or (
+ index_1 - index_0 == 1
+ and not np.isclose(self.velocity[index_0].value, v_0)
+ )
+ ):
+ return True
+ else:
+ return False
+
+ def on_btn_add_shell(self, obj):
+ """Add new shell with given boundary velocities. Triggered if
+ the button is clicked.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ v_start = self.input_v_start.value
+ v_end = self.input_v_end.value
+
+ position_0 = np.searchsorted(self.velocity.value, v_start)
+ position_1 = np.searchsorted(self.velocity.value, v_end)
+ start_index = (
+ int(position_0 - 1)
+ if np.isclose(self.velocity[position_0 - 1].value, v_start)
+ else int(position_0)
+ )
+ end_index = (
+ int(position_1 - 1)
+ if np.isclose(self.velocity[position_1 - 1].value, v_end)
+ else int(position_1)
+ )
+
+ # Delete the overwritten shell (abundances and velocities).
+ if end_index < len(self.velocity) and np.isclose(
+ self.velocity[end_index].value, v_end
+ ):
+ # New shell will overwrite the original shell that ends at v_end.
+ v_scalar = np.delete(self.velocity, end_index).value
+ self.abundance.drop(max(0, end_index - 1), 1, inplace=True)
+ else:
+ v_scalar = self.velocity.value
+
+ # Insert new velocities calculate new densities according
+ # to new velocities through interpolation.
+ v_scalar = np.insert(
+ v_scalar, [start_index, end_index], [v_start, v_end]
+ )
+ v_scalar = np.delete(v_scalar, slice(start_index + 1, end_index + 1))
+ density = self.density_editor.d
+ self.density = (
+ np.interp(
+ v_scalar,
+ self.velocity[1:].value,
+ density[1:].value,
+ )
+ * density.unit
+ )
+ self.velocity = v_scalar * self.velocity.unit
+
+ self.density_editor.v = self.velocity
+ self.density_editor.d = self.density
+
+ # Change abundances after adding new shell.
+ if start_index != end_index:
+ self.abundance.insert(end_index - 1, "", 0)
+ self.abundance.drop(
+ self.abundance.iloc[:, start_index : end_index - 1],
+ 1,
+ inplace=True,
+ )
+ else:
+ if start_index == 0:
+ self.abundance.insert(end_index, "new", 0)
+ self.abundance.insert(
+ end_index, "gap", 0
+ ) # Add a shell to fill the gap.
+ else:
+ self.abundance.insert(end_index - 1, "new", 0)
+ if start_index == self.no_of_shells:
+ self.abundance.insert(end_index - 1, "gap", 0)
+ else:
+ self.abundance.insert(
+ end_index - 1, "gap", self.abundance.iloc[:, end_index]
+ ) # Add a shell to fill the gap with original abundances
+
+ self.abundance.columns = range(self.no_of_shells)
+
+ # Update data and x axis in plot.
+ with self.fig.batch_update():
+ self.fig.layout.xaxis.autorange = True
+ self.fig.data[1].x = self.velocity
+ self.fig.data[1].y = np.append(self.density[1:], self.density[-1])
+ for i in range(self.no_of_elements):
+ self.fig.data[i + 2].x = self.velocity
+ self.update_abundance_plot(i)
+
+ self.dpd_shell_no.options = list(range(1, self.no_of_shells + 1))
+ self.shell_no = start_index + 1
+ self.update_front_end()
+ self.irs_shell_range.max = self.no_of_shells
+
+ def tbs_scale_eventhandler(self, obj):
+ """Switch the scale type of y axis between linear mode and log
+ mode. Triggered if the toggle button is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ scale_mode = obj.new
+
+ if scale_mode == "Linear":
+ self.fig.update_layout(
+ yaxis=dict(
+ type="linear",
+ range=[0, 1],
+ ),
+ yaxis2=dict(type="linear"),
+ )
+ else:
+ self.fig.update_layout(
+ yaxis=dict(
+ type="log",
+ range=[-8, 0],
+ ),
+ yaxis2=dict(type="log"),
+ )
+
+ def input_item_eventhandler(self, obj):
+ """Update the data and the widget when it gets new abundance
+ input. Triggered if the abundance input is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ if self._trigger:
+ item_index = obj.owner.index
+ is_locked = self.checks[item_index]
+
+ if is_locked:
+ self.bound_locked_sum_to_1(item_index)
+
+ if np.isclose(self.abundance.iloc[:, self.shell_no - 1].sum(), 1):
+ self.norm_warning.layout.visibility = "hidden"
+ else:
+ self.norm_warning.layout.visibility = "visible"
+
+ self.abundance.iloc[item_index, self.shell_no - 1] = obj.owner.value
+
+ if self.rbs_multi_apply.index is None:
+ self.update_abundance_plot(item_index)
+ else:
+ self.apply_to_multiple_shells(item_index)
+
+ def check_eventhandler(self, obj):
+ """Triggered if the checkbox is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ item_index = obj.owner.index
+
+ if obj.new == True:
+ self.bound_locked_sum_to_1(item_index)
+
+ def dpd_shell_no_eventhandler(self, obj):
+ """Make the data in widgets correspond with the selected shell.
+ Triggered if the dropdown value is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ # Disable "previous" and "next" buttons when shell no comes to boundaries.
+ if obj.new == 1:
+ self.btn_prev.disabled = True
+ else:
+ self.btn_prev.disabled = False
+
+ if obj.new == self.no_of_shells:
+ self.btn_next.disabled = True
+ else:
+ self.btn_next.disabled = False
+
+ self.update_front_end()
+
+ def on_btn_prev(self, obj):
+ """Move to previous shell.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ self.shell_no -= 1
+
+ def on_btn_next(self, obj):
+ """Move to next shell.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ self.shell_no += 1
+
+ def on_btn_norm(self, obj):
+ """Normalize unlocked abundances to 1. Triggered if the
+ normalize button is clicked.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ locked_mask = np.array(self.checked_list)
+ locked_sum = self.abundance.loc[locked_mask, self.shell_no - 1].sum()
+ unlocked_arr = self.abundance.loc[~locked_mask, self.shell_no - 1]
+
+ # if abundances are all zero
+ if unlocked_arr.sum() == 0:
+ return
+
+ self.abundance.loc[~locked_mask, self.shell_no - 1] = (
+ (1 - locked_sum) * unlocked_arr / unlocked_arr.sum()
+ )
+
+ self.read_abundance()
+
+ for i in range(self.no_of_elements):
+ if self.rbs_multi_apply.index is None:
+ # Only normalize selected shell
+ self.update_abundance_plot(i)
+ else:
+ self.apply_to_multiple_shells(i)
+
+ @debounce(0.5)
+ def input_symb_eventhandler(self, obj):
+ """Judge whether the input symbol is valid. Triggered after 0.5s
+ when the symbol input is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ element_symbol_string = obj.new.capitalize()
+
+ if element_symbol_string == "":
+ self.symb_warning.layout.visibility = "hidden"
+ self.btn_add_element.disabled = True
+ return
+
+ if element_symbol_string in self.elements:
+ self.symb_warning.layout.visibility = "visible"
+ self.btn_add_element.disabled = True
+ self.symb_warning.readout = "Already exists!"
+ return
+
+ try:
+ if nucname.iselement(element_symbol_string) or nucname.isnuclide(
+ element_symbol_string
+ ):
+ self.symb_warning.layout.visibility = "hidden"
+ self.btn_add_element.disabled = False
+ return
+
+ except RuntimeError:
+ pass
+
+ self.symb_warning.layout.visibility = "visible"
+ self.btn_add_element.disabled = True
+ self.symb_warning.readout = "invalid"
+
+ def on_btn_add_element(self, obj):
+ """Add new element and update the display in the front end.
+ Triggered if the add button is clicked.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ element_symbol_string = self.input_symb.value.capitalize()
+
+ if element_symbol_string in nucname.name_zz:
+ z = nucname.name_zz[element_symbol_string]
+ self.abundance.loc[(z, ""), :] = 0
+ else:
+ mass_no = nucname.anum(element_symbol_string)
+ z = nucname.znum(element_symbol_string)
+ self.abundance.loc[(z, mass_no), :] = 0
+
+ self.abundance.sort_index(inplace=True)
+
+ # Add new BoundedFloatText control and Checkbox control.
+ item = ipw.BoundedFloatText(min=0, max=1, step=0.01)
+ check = ipw.Checkbox(indent=False, layout=ipw.Layout(width="30px"))
+ item.index = self.no_of_elements - 1
+ check.index = self.no_of_elements - 1
+ item.observe(self.input_item_eventhandler, "value")
+ check.observe(self.check_eventhandler, "value")
+ self.input_items.append(item)
+ self.checks.append(check)
+
+ # Keep the order of description same with atomic number
+ self.elements = self.get_symbols()
+ for i in range(self.no_of_elements):
+ self.input_items[i].description = self.elements[i]
+
+ self.box_editor.children = [
+ ipw.VBox(self.input_items),
+ ipw.VBox(self.checks),
+ ]
+
+ with self.fig.batch_update():
+ # Add new trace to plot.
+ self.fig.add_scatter(
+ x=self.velocity, # convert to km/s
+ y=[0] * (self.no_of_shells + 1),
+ mode="lines+markers",
+ name=element_symbol_string,
+ )
+ # Sort the legend in atomic order.
+ fig_data_lst = list(self.fig.data)
+ fig_data_lst.insert(
+ np.argwhere(self.elements == element_symbol_string)[0][0] + 2,
+ self.fig.data[-1],
+ )
+ self.fig.data = fig_data_lst[:-1]
+
+ colorscale = transition_colors(self.no_of_elements, COLORMAP)
+ for i in range(self.no_of_elements):
+ self.fig.data[2 + i].line.color = colorscale[i]
+
+ self.read_abundance()
+
+ # Clear symbol input box.
+ self.input_symb.value = ""
+
+ def apply_to_multiple_shells(self, item_index):
+ """Apply the changed abundances to specified range of shell(s).
+
+ Parameters
+ ----------
+ item_index : int
+ The index of the widget in the list of abundance inputs.
+ """
+ start_index = self.irs_shell_range.value[0] - 1
+ end_index = self.irs_shell_range.value[1]
+ applied_shell_index = self.shell_no - 1
+
+ self.abundance.iloc[
+ item_index, start_index:end_index
+ ] = self.abundance.iloc[item_index, applied_shell_index]
+
+ self.update_abundance_plot(item_index)
+
+ def input_v_eventhandler(self, obj):
+ """Judge whether the input velocity range is valid. Triggered if the
+ velocity input is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ v_start = self.input_v_start.value
+ v_end = self.input_v_end.value
+
+ if v_start >= v_end or v_start < 0 or v_end < 0:
+ self.btn_add_shell.disabled = True
+ return
+ else:
+ self.btn_add_shell.disabled = False
+
+ # Whether overwrite existing shells
+ if self.overwrite_existing_shells(v_start, v_end):
+ self.overwrite_warning.layout.visibility = "visible"
+ else:
+ self.overwrite_warning.layout.visibility = "hidden"
+
+ def on_btn_output(self, obj):
+ """Triggered if the output button is clicked.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ path = self.input_path.value
+ overwrite = self.ckb_overwrite.value
+
+ self.to_csvy(path, overwrite)
+
+ def rbs_single_apply_eventhandler(self, obj):
+ """Switch to single shell editing mode. Triggered if the
+ first radio button is selected.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ self.rbs_multi_apply.unobserve(
+ self.rbs_multi_apply_eventhandler, "value"
+ )
+ self.rbs_multi_apply.index = None
+ self.rbs_multi_apply.observe(self.rbs_multi_apply_eventhandler, "value")
+ self.irs_shell_range.disabled = True
+
+ def rbs_multi_apply_eventhandler(self, obj):
+ """Switch to multi-shells editing mode. Triggered if the
+ second radio button is selected.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ self.rbs_single_apply.unobserve(
+ self.rbs_single_apply_eventhandler, "value"
+ )
+ self.rbs_single_apply.index = None
+ self.rbs_single_apply.observe(
+ self.rbs_single_apply_eventhandler, "value"
+ )
+ # self.irs_shell_range.value = [self.shell_no, self.shell_no]
+ self.irs_shell_range.disabled = False
+
+ def irs_shell_range_eventhandler(self, obj):
+ """Select the velocity range of new shell and highlight the range
+ in the plot. Triggered if the shell range slider is changed.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ x = self.fig.data[0].x
+ width = self.fig.data[0].width
+ range = self.fig.layout.xaxis.range
+
+ if self.irs_shell_range.disabled:
+ x = [x[0]]
+ width = [width[0]]
+ y = [1]
+ else:
+ (start_shell_no, end_shell_no) = self.irs_shell_range.value
+ x_inner = self.velocity[start_shell_no - 1].value
+ x_outer = self.velocity[end_shell_no].value
+ x = [x[0], (x_outer + x_inner) / 2]
+ width = [width[0], x_outer - x_inner]
+ y = [1, 1]
+
+ with self.fig.batch_update():
+ self.fig.data[0].x = x
+ self.fig.data[0].width = width
+ self.fig.data[0].y = y
+ self.fig.layout.xaxis.range = range
+
+ def generate_abundance_density_plot(self):
+ """Generate abundance and density plot in different shells."""
+ self.fig = go.FigureWidget()
+ title = "Abundance/Density vs Velocity"
+ data = self.abundance
+
+ # Bar Diagonal
+ self.fig.add_trace(
+ go.Bar(
+ x=[(self.velocity[0].value + self.velocity[1].value) / 2],
+ y=[1],
+ width=[self.velocity[1].value - self.velocity[0].value],
+ name="Selected shell",
+ marker=dict(
+ color="rgb(253,205,172)",
+ ),
+ hoverinfo="none",
+ )
+ )
+
+ self.fig.add_trace(
+ go.Scatter(
+ x=self.velocity,
+ y=np.append(self.density[1:], self.density[-1]),
+ mode="lines+markers",
+ name="Density",
+ yaxis="y2",
+ line=dict(color="black", shape="hv"),
+ marker_symbol="square",
+ ),
+ )
+
+ colorscale = transition_colors(self.no_of_elements, COLORMAP)
+ for i in range(self.no_of_elements):
+ self.fig.add_trace(
+ go.Scatter(
+ x=self.velocity,
+ y=np.append(data.iloc[i], data.iloc[i, -1]),
+ mode="lines+markers",
+ line=dict(shape="hv", color=colorscale[i]),
+ name=self.elements[i],
+ ),
+ )
+
+ self.fig.update_layout(
+ xaxis=dict(
+ title="Velocity (km/s)",
+ tickformat="f",
+ ),
+ yaxis=dict(
+ title="Fractional Abundance",
+ exponentformat="e",
+ range=[0, 1],
+ ),
+ yaxis2=dict(
+ title="Density (g/cm^3)",
+ exponentformat="e",
+ overlaying="y",
+ side="right",
+ ),
+ height=500,
+ title=title,
+ hovermode="closest",
+ legend=dict(
+ x=1.15,
+ ),
+ )
+
+ def display(self):
+ """Display the GUI.
+
+ Returns
+ -------
+ ipywidgets.widgets.widget_box.VBox
+ A box that contains all the widgets in the GUI.
+ """
+ self.box_editor = ipw.HBox(
+ [
+ ipw.VBox(self.input_items),
+ ipw.VBox(self.checks, layout=ipw.Layout(margin="0 0 0 10px")),
+ ]
+ )
+
+ box_add_shell = ipw.HBox(
+ [
+ self.input_v_start,
+ self.input_v_end,
+ self.btn_add_shell,
+ self.overwrite_warning,
+ ],
+ layout=ipw.Layout(margin="0 0 0 50px"),
+ )
+
+ box_head = ipw.HBox(
+ [self.dpd_shell_no, self.btn_prev, self.btn_next, box_add_shell]
+ )
+
+ box_add_element = ipw.HBox(
+ [self.input_symb, self.btn_add_element, self.symb_warning],
+ layout=ipw.Layout(margin="0 0 0 80px"),
+ )
+
+ help_note = ipw.HTML(
+ value='* Select a checkbox to lock the '
+ "abundance of corresponding element.
"
+ ' On clicking the "Normalize" button, '
+ "the locked abundance(s) will not be normalized.
",
+ indent=True,
+ )
+
+ box_norm = ipw.HBox([self.btn_norm, self.norm_warning])
+
+ box_apply = ipw.VBox(
+ [
+ ipw.Label(value="Apply abundance(s) to:"),
+ self.rbs_single_apply,
+ ipw.HBox([self.rbs_multi_apply, self.irs_shell_range]),
+ ],
+ layout=ipw.Layout(margin="0 0 15px 50px"),
+ )
+
+ box_features = ipw.VBox([box_norm, help_note])
+ box_abundance = ipw.VBox(
+ [
+ box_apply,
+ ipw.HBox([self.box_editor, box_features]),
+ box_add_element,
+ ]
+ )
+ box_density = self.density_editor.display()
+
+ main_tab = ipw.Tab([box_abundance, box_density])
+ main_tab.set_title(0, "Edit Abundance")
+ main_tab.set_title(1, "Edit Density")
+
+ hint = ipw.HTML(
+ value="Save model as file: "
+ )
+ box_output = ipw.VBox(
+ [
+ hint,
+ self.input_i_time_0,
+ ipw.HBox(
+ [self.input_path, self.btn_output, self.ckb_overwrite]
+ ),
+ ]
+ )
+ self.read_abundance()
+ self.density_editor.read_density()
+
+ return ipw.VBox(
+ [
+ self.tbs_scale,
+ self.fig,
+ box_head,
+ main_tab,
+ box_output,
+ self.error_view,
+ ]
+ )
+
+ @error_view.capture(clear_output=True)
+ def to_csvy(self, path, overwrite):
+ """Output CSVY file on the specified path.
+
+ Parameters
+ ----------
+ path : str
+ Output path.
+ overwrite : bool
+ True if overwriting, False otherwise.
+ """
+ posix_path = Path(path)
+ posix_path = posix_path.with_suffix(".csvy")
+
+ if posix_path.exists() and not overwrite:
+ raise FileExistsError(
+ "The file already exists. Click the 'overwrite' checkbox to overwrite it."
+ )
+ else:
+ self.write_yaml_portion(posix_path)
+ self.write_csv_portion(posix_path)
+
+ @error_view.capture(clear_output=True)
+ def write_yaml_portion(self, path):
+ """Write the YAML portion of the output file.
+
+ Parameters
+ ----------
+ path : pathlib.PosixPath
+ """
+ name = path.name
+ d_time_0 = self.density_editor.density_t_0 * u.day
+ i_time_0 = self.input_i_time_0.value * u.day
+ custom_yaml = CustomYAML(
+ name, d_time_0, i_time_0, self.velocity[0], self.velocity[-1]
+ )
+ custom_yaml.create_fields_dict(self.elements)
+
+ with path.open("w") as f:
+ yaml_output = yaml.dump(custom_yaml, sort_keys=False)
+
+ # Add YAML delimiter
+ yaml_output_snippets = yaml_output.split("\n")
+ yaml_output_snippets[0] = yaml_output_snippets[-1] = YAML_DELIMITER
+ yaml_output = "\n".join(yaml_output_snippets) + "\n"
+
+ f.write(yaml_output)
+
+ print("Saved Successfully!")
+
+ def write_csv_portion(self, path):
+ """Write the CSV portion of the output file.
+
+ Parameters
+ ----------
+ path : pathlib.PosixPath
+ """
+ try:
+ data = self.abundance.T
+ data.columns = self.elements
+ first_row = pd.DataFrame(
+ [[0] * self.no_of_elements], columns=self.elements
+ )
+ data = pd.concat([first_row, data])
+ formatted_v = pd.Series(self.velocity.value).apply(
+ lambda x: "%.3e" % x
+ )
+ density = self.density_editor.d
+ data.insert(0, "velocity", formatted_v)
+ data.insert(1, "density", density)
+
+ data.to_csv(path, mode="a", index=False)
+ except pd.errors.EmptyDataError:
+ data = None
+
+ @classmethod
+ def from_csvy(cls, fpath):
+ """Create a new CustomAbundanceWidget instance with data from CSVY file.
+
+ Parameters
+ ----------
+ fpath : str
+ the path of CSVY file.
+
+ Returns
+ -------
+ CustomAbundanceWidget
+ """
+ csvy_model_config, csvy_model_data = load_csvy(fpath)
+ csvy_schema_file = os.path.join(
+ BASE_DIR, "../..", "io", "schemas", "csvy_model.yml"
+ )
+ csvy_model_config = Configuration(
+ validate_dict(csvy_model_config, schemapath=csvy_schema_file)
+ )
+
+ if hasattr(csvy_model_config, "velocity"):
+ velocity = quantity_linspace(
+ csvy_model_config.velocity.start,
+ csvy_model_config.velocity.stop,
+ csvy_model_config.velocity.num + 1,
+ ).cgs
+ else:
+ velocity_field_index = [
+ field["name"] for field in csvy_model_config.datatype.fields
+ ].index("velocity")
+ velocity_unit = u.Unit(
+ csvy_model_config.datatype.fields[velocity_field_index]["unit"]
+ )
+ velocity = csvy_model_data["velocity"].values * velocity_unit
+
+ no_of_shells = len(velocity) - 1
+
+ if hasattr(csvy_model_config, "density"):
+ adjusted_velocity = velocity.insert(0, 0)
+ v_middle = (
+ adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5
+ )
+ no_of_shells = len(adjusted_velocity) - 1
+
+ d_conf = csvy_model_config.density
+ density_type = d_conf.type
+ if density_type == "branch85_w7":
+ density_0 = calculate_power_law_density(
+ v_middle, d_conf.w7_v_0, d_conf.w7_rho_0, -7
+ )
+ time_0 = d_conf.w7_time_0
+ elif density_type == "uniform":
+ density_0 = d_conf.value.to("g cm^-3") * np.ones(no_of_shells)
+ time_0 = d_conf.get("time_0", 0 * u.day)
+ elif density_type == "power_law":
+ density_0 = calculate_power_law_density(
+ v_middle, d_conf.v_0, d_conf.rho_0, d_conf.exponent
+ )
+ time_0 = d_conf.get("time_0", 0 * u.day)
+ elif density_type == "exponential":
+ density_0 = calculate_exponential_density(
+ v_middle, d_conf.v_0, d_conf.rho_0
+ )
+ time_0 = d_conf.get("time_0", 0 * u.day)
+ else:
+ raise ValueError(f"Unrecognized density type " f"{d_conf.type}")
+ else:
+ density_field_index = [
+ field["name"] for field in csvy_model_config.datatype.fields
+ ].index("density")
+ density_unit = u.Unit(
+ csvy_model_config.datatype.fields[density_field_index]["unit"]
+ )
+ density_0 = csvy_model_data["density"].values * density_unit
+
+ if hasattr(csvy_model_config, "abundance"):
+ abundances_section = csvy_model_config.abundance
+ abundance, isotope_abundance = read_uniform_abundances(
+ abundances_section, no_of_shells
+ )
+ else:
+ _, abundance, isotope_abundance = parse_csv_abundances(
+ csvy_model_data
+ )
+ abundance = abundance.loc[:, 1:]
+ abundance.columns = np.arange(abundance.shape[1])
+ isotope_abundance = isotope_abundance.loc[:, 1:]
+ isotope_abundance.columns = np.arange(isotope_abundance.shape[1])
+
+ abundance = abundance.replace(np.nan, 0.0)
+ abundance = abundance[abundance.sum(axis=1) > 0]
+ isotope_abundance = isotope_abundance.replace(np.nan, 0.0)
+ isotope_abundance = isotope_abundance[isotope_abundance.sum(axis=1) > 0]
+
+ # Combine elements and isotopes to one DataFrame
+ abundance["mass_number"] = ""
+ abundance.set_index("mass_number", append=True, inplace=True)
+ abundance = pd.concat([abundance, isotope_abundance])
+ abundance.sort_index(inplace=True)
+
+ return cls(
+ density_t_0=time_0,
+ density=density_0,
+ abundance=abundance,
+ velocity=velocity,
+ )
+
+ @classmethod
+ def from_yml(cls, fpath):
+ """Create a new CustomAbundanceWidget instance with data from YAML file.
+
+ Parameters
+ ----------
+ fpath : str
+ The path of YAML file.
+
+ Returns
+ -------
+ CustomAbundanceWidget
+ """
+ config = Configuration.from_yaml(fpath)
+
+ if hasattr(config, "csvy_model"):
+ model = Radial1DModel.from_csvy(config)
+ else:
+ model = Radial1DModel.from_config(config)
+
+ velocity = model.velocity
+ density_t_0 = model.homologous_density.time_0
+ density = model.homologous_density.density_0
+ abundance = model.raw_abundance
+ isotope_abundance = model.raw_isotope_abundance
+
+ # Combine elements and isotopes to one DataFrame
+ abundance["mass_number"] = ""
+ abundance.set_index("mass_number", append=True, inplace=True)
+ abundance = pd.concat([abundance, isotope_abundance])
+ abundance.sort_index(inplace=True)
+
+ return cls(
+ density_t_0=density_t_0,
+ density=density,
+ abundance=abundance,
+ velocity=velocity,
+ )
+
+ @classmethod
+ def from_hdf(cls, fpath):
+ """Create a new CustomAbundanceWidget instance with data from HDF file.
+
+ Parameters
+ ----------
+ fpath : str
+ the path of HDF file.
+
+ Returns
+ -------
+ CustomAbundanceWidget
+ """
+ with pd.HDFStore(fpath, "r") as hdf:
+ abundance = hdf["/simulation/plasma/abundance"]
+ _density_t_0 = hdf["/simulation/model/homologous_density/scalars"]
+ _density = hdf["/simulation/model/homologous_density/density_0"]
+ v_inner = hdf["/simulation/model/v_inner"]
+ v_outer = hdf["/simulation/model/v_outer"]
+
+ density_t_0 = float(_density_t_0) * u.s
+ density = np.array(_density) * u.g / (u.cm) ** 3
+ velocity = np.append(v_inner, v_outer[len(v_outer) - 1]) * u.cm / u.s
+
+ abundance["mass_number"] = ""
+ abundance.set_index("mass_number", append=True, inplace=True)
+
+ return cls(
+ density_t_0=density_t_0,
+ density=density,
+ abundance=abundance,
+ velocity=velocity,
+ )
+
+ @classmethod
+ def from_simulation(cls, sim):
+ """Create a new CustomAbundanceWidget instance from a Simulation object.
+
+ Parameters
+ ----------
+ sim : Simulation
+
+ Returns
+ -------
+ CustomAbundanceWidget
+ """
+ abundance = sim.model.raw_abundance.copy()
+ isotope_abundance = sim.model.raw_isotope_abundance.copy()
+
+ # integrate element and isotope to one DataFrame
+ abundance["mass_number"] = ""
+ abundance.set_index("mass_number", append=True, inplace=True)
+ abundance = pd.concat([abundance, isotope_abundance])
+ abundance.sort_index(inplace=True)
+
+ velocity = sim.model.velocity
+ density_t_0 = sim.model.homologous_density.time_0
+ density = sim.model.homologous_density.density_0
+
+ return cls(
+ density_t_0=density_t_0,
+ density=density,
+ abundance=abundance,
+ velocity=velocity,
+ )
+
+
+class DensityEditor:
+ """Widget to edit density profile of simulation model.
+
+ It provides an interface to allow the user directly change
+ the density, or calculate the density with given type and
+ parameters.
+
+ Attributes
+ ----------
+ shell_no : int
+ The selected shell number.
+ trigger : bool
+ If False, disable the callback when density input is changed.
+ """
+
+ def __init__(self, density_t_0, density, velocity, fig, shell_no_widget):
+ """Initialize DensityEditor with data and widget components.
+
+ Parameters
+ ----------
+ density : astropy.units.quantity.Quantity
+ Density data.
+ velocity : astropy.units.quantity.Quantity
+ Velocity data.
+ fig : plotly.graph_objs._figurewidget.FigureWidget
+ The figure object of density plot.
+ shell_no_widget : ipywidgets.widgets.widget_selection.Dropdown
+ A widget to record the selected shell number.
+ """
+ self.d = density
+ self.v = velocity
+ self.fig = fig
+ self.shell_no_widget = shell_no_widget
+ self.input_d_time_0 = ipw.FloatText(
+ value=density_t_0.value,
+ description="Density time_0 (day): ",
+ style={"description_width": "initial"},
+ layout=ipw.Layout(margin="0 0 20px 0"),
+ )
+
+ self.create_widgets()
+ self._trigger = True
+
+ @property
+ def shell_no(self):
+ return self.shell_no_widget.value
+
+ @property
+ def density_t_0(self):
+ return self.input_d_time_0.value
+
+ def create_widgets(self):
+ """Create widget components in density editor GUI and register
+ callbacks for widgets.
+ """
+ self.input_d = ipw.FloatText(
+ description="Density",
+ layout=ipw.Layout(
+ width="230px",
+ ),
+ )
+ self.input_d.observe(self.input_d_eventhandler, "value")
+
+ self.dpd_dtype = ipw.Dropdown(
+ options=["-", "uniform", "exponential", "power_law"],
+ description="Density type: ",
+ style={"description_width": "initial"},
+ layout=ipw.Layout(width="300px"),
+ )
+ self.dpd_dtype.observe(self.dpd_dtype_eventhandler, "value")
+
+ self.input_rho_0 = ipw.FloatText(
+ description="rho_0", layout=ipw.Layout(width="300px")
+ )
+
+ self.input_exp = ipw.FloatText(
+ description="exponent", layout=ipw.Layout(width="300px")
+ )
+
+ self.input_v_0 = ipw.FloatText(
+ description="v_0", layout=ipw.Layout(width="300px")
+ )
+
+ self.input_value = ipw.FloatText(
+ description="value", layout=ipw.Layout(width="300px")
+ )
+
+ self.btn_calculate = ipw.Button(
+ icon="calculator",
+ description="Calculate Density",
+ layout=ipw.Layout(width="300px"),
+ )
+ self.btn_calculate.on_click(self.on_btn_calculate)
+
+ self.uniform_box = ipw.HBox(
+ [self.input_value, ipw.Label(value="g cm^3")]
+ )
+
+ # Formula to compute density profile
+ form_exp = ipw.HTMLMath(
+ value=r"$$\rho = \rho_0 \times \exp \left( -\frac{v}{v_0} \right)$$",
+ layout=ipw.Layout(margin="0 0 0 100px"),
+ )
+ form_pow = ipw.HTMLMath(
+ value=r"$$\rho = \rho_0 \times \left( \frac{v}{v_0} \right)^n$$",
+ layout=ipw.Layout(margin="0 0 0 100px"),
+ )
+
+ self.exp_box = ipw.VBox(
+ [
+ form_exp,
+ ipw.HBox([self.input_rho_0, ipw.Label(value="g cm^3")]),
+ ipw.HBox([self.input_v_0, ipw.Label(value="km/s")]),
+ ]
+ )
+ self.pow_box = ipw.VBox(
+ [
+ form_pow,
+ ipw.HBox([self.input_rho_0, ipw.Label(value="g cm^3")]),
+ ipw.HBox([self.input_v_0, ipw.Label(value="km/s")]),
+ self.input_exp,
+ ]
+ )
+
+ def read_density(self):
+ """Read density data in DataFrame to density input box when
+ shell No. changes.
+ """
+ dvalue = self.d[self.shell_no].value
+ self._trigger = False
+ self.input_d.value = float("{:.3e}".format(dvalue))
+ self._trigger = True
+
+ def update_density_plot(self):
+ """Update the density line in the plot."""
+ y = np.append(self.d[1:], self.d[-1])
+ self.fig.data[1].y = y
+
+ def input_d_eventhandler(self, obj):
+ """Update the data and the widgets when it gets new density input.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+ """
+ if self._trigger:
+ new_value = obj.new
+ self.d[self.shell_no] = new_value * self.d.unit
+
+ self.update_density_plot()
+
+ dtype_out = ipw.Output()
+
+ @dtype_out.capture(clear_output=True)
+ def dpd_dtype_eventhandler(self, obj):
+ """Display the input boxes of the specified density type.
+ Triggered if the density type dropdown is changed.
+
+ Parameters
+ ----------
+ obj : traitlets.utils.bunch.Bunch
+ A dictionary holding the information about the change.
+
+ Returns
+ -------
+ ipywidgets.widgets.widget_box.VBox
+ A box widget that contains the input boxes of certain density
+ type parameters.
+ """
+ if obj.new == "uniform":
+ display(self.uniform_box)
+ elif obj.new == "exponential":
+ display(self.exp_box)
+ elif obj.new == "power_law":
+ display(self.pow_box)
+
+ def on_btn_calculate(self, obj):
+ """Calculate density according to density parameters input.
+ Triggered if the density calculate button is clicked.
+
+ Parameters
+ ----------
+ obj : ipywidgets.widgets.widget_button.Button
+ The clicked button instance.
+ """
+ dtype = self.dpd_dtype.value
+
+ if dtype == "-":
+ return
+
+ if dtype == "uniform":
+ if self.input_value.value == 0:
+ return
+
+ self.d = self.input_value.value * self.d.unit * np.ones(len(self.d))
+ else:
+ if self.input_v_0.value == 0 or self.input_rho_0.value == 0:
+ return
+
+ adjusted_velocity = self.v.insert(0, 0)
+ v_middle = (
+ adjusted_velocity[1:] * 0.5 + adjusted_velocity[:-1] * 0.5
+ )
+ v_0 = self.input_v_0.value * self.v.unit
+ rho_0 = self.input_rho_0.value * self.d.unit
+
+ if dtype == "exponential":
+ self.d = calculate_exponential_density(v_middle, v_0, rho_0)
+
+ elif dtype == "power_law":
+ exponent = self.input_exp.value
+ self.d = calculate_power_law_density(
+ v_middle, v_0, rho_0, exponent
+ )
+
+ self.read_density()
+ self.update_density_plot()
+
+ def display(self):
+ """Display the GUI.
+
+ Returns
+ -------
+ ipywidgets.widgets.widget_box.VBox
+ A box that contains all the widgets in the GUI.
+ """
+ hint1 = ipw.HTML(
+ value="1) Edit density of the selected shell:"
+ )
+ hint2 = ipw.HTML(
+ value="2) Edit densities for all shells:"
+ )
+ d_box = ipw.HBox(
+ [self.input_d, ipw.Label(value="g/cm^3")],
+ layout=ipw.Layout(margin="0 0 20px 0"),
+ )
+ return ipw.VBox(
+ [
+ self.input_d_time_0,
+ hint1,
+ d_box,
+ hint2,
+ self.dpd_dtype,
+ self.dtype_out,
+ self.btn_calculate,
+ ]
+ )
diff --git a/tardis/visualization/widgets/util.py b/tardis/visualization/widgets/util.py
index 05d6af9da6c..5592c71b09b 100644
--- a/tardis/visualization/widgets/util.py
+++ b/tardis/visualization/widgets/util.py
@@ -3,6 +3,7 @@
import logging
import qgrid
import ipywidgets as ipw
+import asyncio
logger = logging.getLogger(__name__)
@@ -226,3 +227,71 @@ def _create(self, key, value):
justify_content="flex-start",
),
)
+
+
+class Timer:
+ """Timer to implement debouncing using an asynchronous loop.
+
+ Notes
+ -----
+ This class is reproduced from ipywidgets documentation, for more information
+ please see https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Events.html#debounce
+ """
+
+ def __init__(self, timeout, callback):
+ """Initialize the Timer with delay time and delayed function.
+
+ Parameters
+ ----------
+ timeout : float
+ callback : function
+ """
+ self._timeout = timeout
+ self._callback = callback
+
+ async def _job(self):
+ await asyncio.sleep(self._timeout)
+ self._callback()
+
+ def start(self):
+ self._task = asyncio.ensure_future(self._job())
+
+ def cancel(self):
+ self._task.cancel()
+
+
+def debounce(wait):
+ """Decorator that will postpone a function's execution until after
+ `wait` seconds have elapsed since the last time it was invoked.
+
+ Parameters
+ ----------
+ wait : float
+
+ Returns
+ -------
+ function
+
+ Notes
+ -----
+ This decorator is reproduced from ipywidgets documentation, for more information
+ please see https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Events.html#debounce
+ """
+
+ def decorator(fn):
+ timer = None
+
+ def debounced(*args, **kwargs):
+ nonlocal timer
+
+ def call_it():
+ fn(*args, **kwargs)
+
+ if timer is not None:
+ timer.cancel()
+ timer = Timer(wait, call_it)
+ timer.start()
+
+ return debounced
+
+ return decorator