diff --git a/tardis/visualization/widgets/custom_abundance.py b/tardis/visualization/widgets/custom_abundance.py index 0d8a9cae23b..7bbc98ae2be 100644 --- a/tardis/visualization/widgets/custom_abundance.py +++ b/tardis/visualization/widgets/custom_abundance.py @@ -29,7 +29,264 @@ BASE_DIR = tardis.__path__[0] YAML_DELIMITER = "---" -COLORMAP = "viridis" +COLORMAP = "jet" + + +class CustomAbundanceWidgetData: + """The model information and data that required in custom + abundance widget. + + Attributes + ---------- + elements : list of str + A list of elements or isotopes' symbols. + """ + + def __init__(self, density_t_0, density, abundance, velocity): + """Initialize CustomAbundanceWidgetData with model information. + + Parameters + ---------- + density_t_0 : astropy.units.quantity.Quantity + Initial time for the density in the model. + density : astropy.units.quantity.Quantity + abundance : pd.DataFrame + velocity : astropy.units.quantity.Quantity + """ + self.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() + + 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) + + @classmethod + def from_csvy(cls, fpath): + """Create a new CustomAbundanceWidgetData instance with data + from CSVY file. + + Parameters + ---------- + fpath : str + the path of CSVY file. + + Returns + ------- + CustomAbundanceWidgetData + """ + 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 CustomAbundanceWidgetData instance with data + from YAML file. + + Parameters + ---------- + fpath : str + The path of YAML file. + + Returns + ------- + CustomAbundanceWidgetData + """ + 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 CustomAbundanceWidgetData instance with data + from HDF file. + + Parameters + ---------- + fpath : str + the path of HDF file. + + Returns + ------- + CustomAbundanceWidgetData + """ + 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 CustomAbundanceWidgetData instance from a + Simulation object. + + Parameters + ---------- + sim : Simulation + + Returns + ------- + CustomAbundanceWidgetData + """ + 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 CustomYAML(yaml.YAMLObject): @@ -42,16 +299,16 @@ def __init__( 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. + 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 @@ -66,8 +323,8 @@ def create_fields_dict(self, elements): Parameters ---------- - elements : list of str - A list of elements or isotopes' symbols. + elements : list of str + A list of elements or isotopes' symbols. """ for i in range(len(elements) + 2): field = {} @@ -90,7 +347,7 @@ class CustomAbundanceWidget: 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 + to visualize the profile, an edit section to allow the user directly edit abundance and density profile, and an output section to output the model to CSVY file. @@ -105,37 +362,27 @@ class CustomAbundanceWidget: 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. + def __init__(self, widget_data): + """Initialize CustomAbundanceWidget with data and generate + the widgets and plot. Parameters ---------- - density : astropy.units.quantity.Quantity - abundance : pd.DataFrame - velocity : astropy.units.quantity.Quantity + widget_data : CustomAbundanceWidgetData """ - 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.data = widget_data self._trigger = True self.create_widgets() self.generate_abundance_density_plot() self.density_editor = DensityEditor( - density_t_0, - self.density, - self.velocity, + self.data, self.fig, self.dpd_shell_no, ) @@ -150,11 +397,11 @@ def shell_no(self, value): @property def no_of_shells(self): - return self.abundance.shape[1] + return self.data.abundance.shape[1] @property def no_of_elements(self): - return self.abundance.shape[0] + return self.data.abundance.shape[0] @property def checked_list(self): # A boolean list to store the value of checkboxes. @@ -164,18 +411,6 @@ def checked_list(self): # A boolean list to store the value of checkboxes. 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( @@ -204,11 +439,11 @@ def create_widgets(self): width="30px", ), ) - for element in self.elements + for element in self.data.elements ] self.input_items = [ ipw.BoundedFloatText(min=0, max=1, step=0.01, description=element) - for element in self.elements + for element in self.data.elements ] for i in range(self.no_of_elements): self.input_items[i].observe(self.input_item_eventhandler, "value") @@ -338,10 +573,10 @@ def update_input_item_value(self, index, value): Parameters ---------- - index : int - The index of the widget in the list of abundance inputs. - value : float - New abundance value. + 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. @@ -353,7 +588,7 @@ def read_abundance(self): shell No. changes. """ for i in range(self.no_of_elements): - value = self.abundance.iloc[i, self.shell_no - 1] + value = self.data.abundance.iloc[i, self.shell_no - 1] self.update_input_item_value(i, value) def bound_locked_sum_to_1(self, index): @@ -363,25 +598,25 @@ def bound_locked_sum_to_1(self, index): Parameters ---------- - index : int - The index of the widget in the list of abundance inputs. + 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[ + back_value = self.data.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() + self.data.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.data.abundance.iloc[index, self.shell_no - 1] = new self.update_input_item_value(index, new) self.update_abundance_plot(index) @@ -390,10 +625,10 @@ def update_abundance_plot(self, index): Parameters ---------- - index : int - The index of the widget in the list of abundance inputs. + index : int + The index of the widget in the list of abundance inputs. """ - y = self.abundance.iloc[index] + y = self.data.abundance.iloc[index] self.fig.data[index + 2].y = np.append(y, y.iloc[-1]) def update_front_end(self): @@ -409,8 +644,8 @@ def update_front_end(self): # 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_inner = self.data.velocity[self.shell_no - 1].value + x_outer = self.data.velocity[self.shell_no].value x[0] = (x_inner + x_outer) / 2 self.fig.data[0].x = x width[0] = x_outer - x_inner @@ -422,39 +657,34 @@ def overwrite_existing_shells(self, v_0, v_1): Parameters ---------- - v_0 : float - The velocity of inner boundary. - v_1 : float - The velocity of outer boundary. + 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. + 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) + v_vals = self.data.velocity.value + position_0 = np.searchsorted(v_vals, v_0) + position_1 = np.searchsorted(v_vals, v_1) index_0 = ( position_0 - 1 - if np.isclose(self.velocity[position_0 - 1].value, v_0) + if np.isclose(v_vals[position_0 - 1], v_0) else position_0 ) index_1 = ( position_1 - 1 - if np.isclose(self.velocity[position_1 - 1].value, v_1) + if np.isclose(v_vals[position_1 - 1], 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) - ) + (index_1 < len(v_vals) and np.isclose(v_vals[index_1], v_1)) + or (index_1 - index_0 == 1 and not np.isclose(v_vals[index_0], v_0)) ): return True else: @@ -466,87 +696,83 @@ def on_btn_add_shell(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ v_start = self.input_v_start.value v_end = self.input_v_end.value + v_vals = self.data.velocity.value - position_0 = np.searchsorted(self.velocity.value, v_start) - position_1 = np.searchsorted(self.velocity.value, v_end) + position_0 = np.searchsorted(v_vals, v_start) + position_1 = np.searchsorted(v_vals, v_end) start_index = ( int(position_0 - 1) - if np.isclose(self.velocity[position_0 - 1].value, v_start) + if np.isclose(v_vals[position_0 - 1], v_start) else int(position_0) ) end_index = ( int(position_1 - 1) - if np.isclose(self.velocity[position_1 - 1].value, v_end) + if np.isclose(v_vals[position_1 - 1], 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 - ): + if end_index < len(v_vals) and np.isclose(v_vals[end_index], 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 + v_vals = np.delete(v_vals, end_index) + self.data.abundance.drop(max(0, end_index - 1), 1, inplace=True) # 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 = ( + v_vals = np.insert(v_vals, [start_index, end_index], [v_start, v_end]) + v_vals = np.delete(v_vals, slice(start_index + 1, end_index + 1)) + + self.data.density = ( np.interp( - v_scalar, - self.velocity[1:].value, - density[1:].value, + v_vals, + self.data.velocity[1:].value, + self.data.density[1:].value, ) - * density.unit + * self.data.density.unit ) - self.velocity = v_scalar * self.velocity.unit - - self.density_editor.v = self.velocity - self.density_editor.d = self.density + self.data.velocity = v_vals * self.data.velocity.unit # 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], + self.data.abundance.insert(end_index - 1, "", 0) + self.data.abundance.drop( + self.data.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( + self.data.abundance.insert(end_index, "new", 0) + self.data.abundance.insert( end_index, "gap", 0 ) # Add a shell to fill the gap. else: - self.abundance.insert(end_index - 1, "new", 0) + self.data.abundance.insert(end_index - 1, "new", 0) if start_index == self.no_of_shells: - self.abundance.insert(end_index - 1, "gap", 0) + self.data.abundance.insert(end_index - 1, "gap", 0) else: - self.abundance.insert( - end_index - 1, "gap", self.abundance.iloc[:, end_index] + self.data.abundance.insert( + end_index - 1, + "gap", + self.data.abundance.iloc[:, end_index], ) # Add a shell to fill the gap with original abundances - self.abundance.columns = range(self.no_of_shells) + self.data.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]) + self.fig.data[1].x = self.data.velocity + self.fig.data[1].y = np.append( + self.data.density[1:], self.data.density[-1] + ) for i in range(self.no_of_elements): - self.fig.data[i + 2].x = self.velocity + self.fig.data[i + 2].x = self.data.velocity self.update_abundance_plot(i) self.dpd_shell_no.options = list(range(1, self.no_of_shells + 1)) @@ -560,8 +786,8 @@ def tbs_scale_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + obj : traitlets.utils.bunch.Bunch + A dictionary holding the information about the change. """ scale_mode = obj.new @@ -588,8 +814,8 @@ def input_item_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + obj : traitlets.utils.bunch.Bunch + A dictionary holding the information about the change. """ if self._trigger: item_index = obj.owner.index @@ -598,12 +824,16 @@ def input_item_eventhandler(self, obj): if is_locked: self.bound_locked_sum_to_1(item_index) - if np.isclose(self.abundance.iloc[:, self.shell_no - 1].sum(), 1): + if np.isclose( + self.data.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 + self.data.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) @@ -615,8 +845,8 @@ def check_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + obj : traitlets.utils.bunch.Bunch + A dictionary holding the information about the change. """ item_index = obj.owner.index @@ -629,8 +859,8 @@ def dpd_shell_no_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + 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: @@ -650,8 +880,8 @@ def on_btn_prev(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ self.shell_no -= 1 @@ -660,8 +890,8 @@ def on_btn_next(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ self.shell_no += 1 @@ -671,18 +901,20 @@ def on_btn_norm(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + 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] + locked_sum = self.data.abundance.loc[ + locked_mask, self.shell_no - 1 + ].sum() + unlocked_arr = self.data.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] = ( + self.data.abundance.loc[~locked_mask, self.shell_no - 1] = ( (1 - locked_sum) * unlocked_arr / unlocked_arr.sum() ) @@ -702,8 +934,8 @@ def input_symb_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + obj : traitlets.utils.bunch.Bunch + A dictionary holding the information about the change. """ element_symbol_string = obj.new.capitalize() @@ -712,7 +944,7 @@ def input_symb_eventhandler(self, obj): self.btn_add_element.disabled = True return - if element_symbol_string in self.elements: + if element_symbol_string in self.data.elements: self.symb_warning.layout.visibility = "visible" self.btn_add_element.disabled = True self.symb_warning.readout = "Already exists!" @@ -739,20 +971,20 @@ def on_btn_add_element(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + 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 + self.data.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.data.abundance.loc[(z, mass_no), :] = 0 - self.abundance.sort_index(inplace=True) + self.data.abundance.sort_index(inplace=True) # Add new BoundedFloatText control and Checkbox control. item = ipw.BoundedFloatText(min=0, max=1, step=0.01) @@ -765,9 +997,9 @@ def on_btn_add_element(self, obj): self.checks.append(check) # Keep the order of description same with atomic number - self.elements = self.get_symbols() + self.data.elements = self.data.get_symbols() for i in range(self.no_of_elements): - self.input_items[i].description = self.elements[i] + self.input_items[i].description = self.data.elements[i] self.box_editor.children = [ ipw.VBox(self.input_items), @@ -777,15 +1009,17 @@ def on_btn_add_element(self, obj): with self.fig.batch_update(): # Add new trace to plot. self.fig.add_scatter( - x=self.velocity, # convert to km/s + x=self.data.velocity, # convert to km/s y=[0] * (self.no_of_shells + 1), mode="lines+markers", name=element_symbol_string, + line=dict(shape="hv"), ) # 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, + np.argwhere(self.data.elements == element_symbol_string)[0][0] + + 2, self.fig.data[-1], ) self.fig.data = fig_data_lst[:-1] @@ -804,16 +1038,16 @@ def apply_to_multiple_shells(self, item_index): Parameters ---------- - item_index : int - The index of the widget in the list of abundance inputs. + 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[ + self.data.abundance.iloc[ item_index, start_index:end_index - ] = self.abundance.iloc[item_index, applied_shell_index] + ] = self.data.abundance.iloc[item_index, applied_shell_index] self.update_abundance_plot(item_index) @@ -823,8 +1057,8 @@ def input_v_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + 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 @@ -846,8 +1080,8 @@ def on_btn_output(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ path = self.input_path.value overwrite = self.ckb_overwrite.value @@ -860,8 +1094,8 @@ def rbs_single_apply_eventhandler(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ self.rbs_multi_apply.unobserve( self.rbs_multi_apply_eventhandler, "value" @@ -876,8 +1110,8 @@ def rbs_multi_apply_eventhandler(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ self.rbs_single_apply.unobserve( self.rbs_single_apply_eventhandler, "value" @@ -895,8 +1129,8 @@ def irs_shell_range_eventhandler(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ x = self.fig.data[0].x width = self.fig.data[0].width @@ -908,8 +1142,8 @@ def irs_shell_range_eventhandler(self, obj): 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_inner = self.data.velocity[start_shell_no - 1].value + x_outer = self.data.velocity[end_shell_no].value x = [x[0], (x_outer + x_inner) / 2] width = [width[0], x_outer - x_inner] y = [1, 1] @@ -924,14 +1158,16 @@ 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 + abundance = self.data.abundance + velocity = self.data.velocity + density = self.data.density # Bar Diagonal self.fig.add_trace( go.Bar( - x=[(self.velocity[0].value + self.velocity[1].value) / 2], + x=[(velocity[0].value + velocity[1].value) / 2], y=[1], - width=[self.velocity[1].value - self.velocity[0].value], + width=[velocity[1].value - velocity[0].value], name="Selected shell", marker=dict( color="rgb(253,205,172)", @@ -942,8 +1178,8 @@ def generate_abundance_density_plot(self): self.fig.add_trace( go.Scatter( - x=self.velocity, - y=np.append(self.density[1:], self.density[-1]), + x=velocity, + y=np.append(density[1:], density[-1]), mode="lines+markers", name="Density", yaxis="y2", @@ -956,11 +1192,11 @@ def generate_abundance_density_plot(self): 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]), + x=velocity, + y=np.append(abundance.iloc[i], abundance.iloc[i, -1]), mode="lines+markers", line=dict(shape="hv", color=colorscale[i]), - name=self.elements[i], + name=self.data.elements[i], ), ) @@ -993,8 +1229,8 @@ def display(self): Returns ------- - ipywidgets.widgets.widget_box.VBox - A box that contains all the widgets in the GUI. + ipywidgets.widgets.widget_box.VBox + A box that contains all the widgets in the GUI. """ self.box_editor = ipw.HBox( [ @@ -1087,10 +1323,10 @@ def to_csvy(self, path, overwrite): Parameters ---------- - path : str - Output path. - overwrite : bool - True if overwriting, False otherwise. + path : str + Output path. + overwrite : bool + True if overwriting, False otherwise. """ posix_path = Path(path) posix_path = posix_path.with_suffix(".csvy") @@ -1109,15 +1345,19 @@ def write_yaml_portion(self, path): Parameters ---------- - path : pathlib.PosixPath + path : pathlib.PosixPath """ name = path.name - d_time_0 = self.density_editor.density_t_0 * u.day + d_time_0 = self.data.density_t_0 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] + name, + d_time_0, + i_time_0, + self.data.velocity[0], + self.data.velocity[-1], ) - custom_yaml.create_fields_dict(self.elements) + custom_yaml.create_fields_dict(self.data.elements) with path.open("w") as f: yaml_output = yaml.dump(custom_yaml, sort_keys=False) @@ -1136,19 +1376,19 @@ def write_csv_portion(self, path): Parameters ---------- - path : pathlib.PosixPath + path : pathlib.PosixPath """ try: - data = self.abundance.T - data.columns = self.elements + data = self.data.abundance.T + data.columns = self.data.elements first_row = pd.DataFrame( - [[0] * self.no_of_elements], columns=self.elements + [[0] * self.no_of_elements], columns=self.data.elements ) data = pd.concat([first_row, data]) - formatted_v = pd.Series(self.velocity.value).apply( + formatted_v = pd.Series(self.data.velocity.value).apply( lambda x: "%.3e" % x ) - density = self.density_editor.d + density = self.data.density data.insert(0, "velocity", formatted_v) data.insert(1, "density", density) @@ -1162,107 +1402,16 @@ def from_csvy(cls, fpath): Parameters ---------- - fpath : str - the path of CSVY file. + fpath : str + the path of CSVY file. Returns ------- - CustomAbundanceWidget + 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) - ) + widget_data = CustomAbundanceWidgetData.from_csvy(fpath) - 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, - ) + return cls(widget_data) @classmethod def from_yml(cls, fpath): @@ -1270,38 +1419,16 @@ def from_yml(cls, fpath): Parameters ---------- - fpath : str - The path of YAML file. + fpath : str + The path of YAML file. Returns ------- - CustomAbundanceWidget + 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 + widget_data = CustomAbundanceWidgetData.from_yml(fpath) - # 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, - ) + return cls(widget_data) @classmethod def from_hdf(cls, fpath): @@ -1309,33 +1436,16 @@ def from_hdf(cls, fpath): Parameters ---------- - fpath : str - the path of HDF file. + fpath : str + the path of HDF file. Returns ------- - CustomAbundanceWidget + 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"] + widget_data = CustomAbundanceWidgetData.from_hdf(fpath) - 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, - ) + return cls(widget_data) @classmethod def from_simulation(cls, sim): @@ -1343,35 +1453,19 @@ def from_simulation(cls, sim): Parameters ---------- - sim : Simulation + sim : Simulation Returns ------- - CustomAbundanceWidget + 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) + widget_data = CustomAbundanceWidgetData.from_simulation(sim) - 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, - ) + return cls(widget_data) class DensityEditor: - """Widget to edit density profile of simulation model. + """Widget to edit density profile of the model. It provides an interface to allow the user directly change the density, or calculate the density with given type and @@ -1381,34 +1475,25 @@ class DensityEditor: ---------- shell_no : int The selected shell number. - trigger : bool + _trigger : bool If False, disable the callback when density input is changed. """ - def __init__(self, density_t_0, density, velocity, fig, shell_no_widget): + def __init__(self, widget_data, 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. + widget_data : CustomAbundanceWidgetData + Data in the custom abundance widget. + 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.data = widget_data 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 @@ -1417,10 +1502,6 @@ def __init__(self, density_t_0, density, velocity, fig, shell_no_widget): 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. @@ -1433,6 +1514,14 @@ def create_widgets(self): ) self.input_d.observe(self.input_d_eventhandler, "value") + self.input_d_time_0 = ipw.FloatText( + value=self.data.density_t_0.value, + description="Density time_0 (day): ", + style={"description_width": "initial"}, + layout=ipw.Layout(margin="0 0 20px 0"), + ) + self.input_d_time_0.observe(self.input_d_time_0_eventhandler, "value") + self.dpd_dtype = ipw.Dropdown( options=["-", "uniform", "exponential", "power_law"], description="Density type: ", @@ -1498,30 +1587,44 @@ def read_density(self): """Read density data in DataFrame to density input box when shell No. changes. """ - dvalue = self.d[self.shell_no].value + dvalue = self.data.density[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]) + y = np.append(self.data.density[1:], self.data.density[-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. + """Update the data and the widgets when the widget gets new + density input. Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + 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.data.density[self.shell_no] = ( + new_value * self.data.density.unit + ) self.update_density_plot() + def input_d_time_0_eventhandler(self, obj): + """Update density time 0 data when the widget gets new input. + + Parameters + ---------- + obj : traitlets.utils.bunch.Bunch + A dictionary holding the information about the change. + """ + new_value = obj.new + self.data.density_t_0 = new_value * self.data.density_t_0.unit + dtype_out = ipw.Output() @dtype_out.capture(clear_output=True) @@ -1531,14 +1634,14 @@ def dpd_dtype_eventhandler(self, obj): Parameters ---------- - obj : traitlets.utils.bunch.Bunch - A dictionary holding the information about the change. + 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. + 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) @@ -1553,10 +1656,12 @@ def on_btn_calculate(self, obj): Parameters ---------- - obj : ipywidgets.widgets.widget_button.Button - The clicked button instance. + obj : ipywidgets.widgets.widget_button.Button + The clicked button instance. """ dtype = self.dpd_dtype.value + density = self.data.density + velocity = self.data.velocity if dtype == "-": return @@ -1565,24 +1670,28 @@ def on_btn_calculate(self, obj): if self.input_value.value == 0: return - self.d = self.input_value.value * self.d.unit * np.ones(len(self.d)) + self.data.density = ( + self.input_value.value * density.unit * np.ones(len(density)) + ) else: if self.input_v_0.value == 0 or self.input_rho_0.value == 0: return - adjusted_velocity = self.v.insert(0, 0) + adjusted_velocity = velocity.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 + v_0 = self.input_v_0.value * velocity.unit + rho_0 = self.input_rho_0.value * density.unit if dtype == "exponential": - self.d = calculate_exponential_density(v_middle, v_0, rho_0) + self.data.density = calculate_exponential_density( + v_middle, v_0, rho_0 + ) elif dtype == "power_law": exponent = self.input_exp.value - self.d = calculate_power_law_density( + self.data.density = calculate_power_law_density( v_middle, v_0, rho_0, exponent ) @@ -1594,8 +1703,8 @@ def display(self): Returns ------- - ipywidgets.widgets.widget_box.VBox - A box that contains all the widgets in the GUI. + 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:"