Use pathlib and made minor changes

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,40 +1162,40 @@ 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):
         """Write the YAML portion of the output file.
 
         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(
             name, d_time_0, i_time_0, self.velocity[0], self.velocity[-1]
         )
         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)
         )