diff --git a/tardis/visualization/widgets/custom_abundance.py b/tardis/visualization/widgets/custom_abundance.py index c9e20064ee1..9e2b26aa968 100644 --- a/tardis/visualization/widgets/custom_abundance.py +++ b/tardis/visualization/widgets/custom_abundance.py @@ -1,13 +1,16 @@ """Class to create and display Custom Abundance Widget.""" import os -import math +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 asyncio +import tardis from tardis.util.base import quantity_linspace from tardis.io.config_reader import Configuration from tardis.model import Radial1DModel @@ -24,7 +27,10 @@ from tardis.util.base import atomic_number2element_symbol, quantity_linspace from tardis.visualization.tools.convergence_plot import transition_colors -import asyncio + +BASE_DIR = tardis.__path__[0] +YAML_DELIMITER = "---" +COLORMAP = "viridis" class Timer: @@ -95,11 +101,6 @@ def call_it(): return decorator -import yaml - -YAML_DELIMITER = "---" - - class CustomYAML(yaml.YAMLObject): """A custom YAML object generated by required properties.""" @@ -153,9 +154,6 @@ def create_fields_dict(self, elements): self.datatype["fields"].append(field) -COLORMAP = "viridis" - - class CustomAbundanceWidget: """A widget like object to edit abundances and densities graphically. @@ -318,12 +316,6 @@ def create_widgets(self): ) self.btn_add_element.on_click(self.on_btn_add_element) - # self.btn_apply = ipw.Button( - # description=" Apply", - # icon="pencil-square-o", - # layout = ipw.Layout(width="100px") - # ) - # self.btn_apply.on_click(self.on_btn_apply) self.irs_shell_range = ipw.IntRangeSlider( min=1, max=self.no_of_shells, @@ -513,23 +505,23 @@ def overwrite_existing_shells(self, v_0, v_1): index_0 = ( position_0 - 1 - if math.isclose(self.velocity[position_0 - 1].value, v_0) + if np.isclose(self.velocity[position_0 - 1].value, v_0) else position_0 ) index_1 = ( position_1 - 1 - if math.isclose(self.velocity[position_1 - 1].value, v_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 math.isclose(self.velocity[index_1].value, v_1) + and np.isclose(self.velocity[index_1].value, v_1) ) or ( index_1 - index_0 == 1 - and not math.isclose(self.velocity[index_0].value, v_0) + and not np.isclose(self.velocity[index_0].value, v_0) ) ): return True @@ -551,17 +543,17 @@ def on_btn_add_shell(self, obj): position_1 = np.searchsorted(self.velocity.value, v_end) start_index = ( int(position_0 - 1) - if math.isclose(self.velocity[position_0 - 1].value, v_start) + if np.isclose(self.velocity[position_0 - 1].value, v_start) else int(position_0) ) end_index = ( int(position_1 - 1) - if math.isclose(self.velocity[position_1 - 1].value, v_end) + 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 math.isclose( + 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. @@ -674,7 +666,7 @@ def input_item_eventhandler(self, obj): if is_locked: self.bound_locked_sum_to_1(item_index) - if math.isclose(self.abundance.iloc[:, self.shell_no - 1].sum(), 1): + 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" @@ -1170,16 +1162,16 @@ def to_csvy(self, path, overwrite): overwrite : bool True if overwriting, False otherwise. """ - if not path.endswith(".csvy"): - path += ".csvy" + posix_path = Path(path) + posix_path = posix_path.with_suffix(".csvy") - if os.path.exists(path) and not overwrite: + 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(path) - self.write_csv_portion(path) + self.write_yaml_portion(posix_path) + self.write_csv_portion(posix_path) @error_view.capture(clear_output=True) def write_yaml_portion(self, path): @@ -1187,9 +1179,9 @@ def write_yaml_portion(self, path): Parameters ---------- - path : str + path : pathlib.PosixPath """ - name = path.split("/")[-1] + 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( @@ -1197,13 +1189,13 @@ def write_yaml_portion(self, path): ) custom_yaml.create_fields_dict(self.elements) - with open(path, "w") as f: + with path.open("w") as f: yaml_output = yaml.dump(custom_yaml, sort_keys=False) # Add YAML delimiter - splits = yaml_output.split("\n") - splits[0] = splits[-1] = YAML_DELIMITER - yaml_output = "\n".join(splits) + "\n" + 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) @@ -1214,7 +1206,7 @@ def write_csv_portion(self, path): Parameters ---------- - path : str + path : pathlib.PosixPath """ try: data = self.abundance.T @@ -1248,9 +1240,9 @@ def from_csvy(cls, fpath): CustomAbundanceWidget """ csvy_model_config, csvy_model_data = load_csvy(fpath) - base_dir = os.path.abspath(os.path.dirname(__file__)) - schema_dir = os.path.join(base_dir, "../..", "io", "schemas") - csvy_schema_file = os.path.join(schema_dir, "csvy_model.yml") + 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) )