Sxs download

PyART/catalogs/sxs.py

@@ -6,6 +6,7 @@
 import json
 from ..waveform import Waveform
 from ..utils import cat_utils as cat_ut
+from ..utils.utils import LoggerWriter
 
 
 class Waveform_SXS(Waveform):
@@ -110,7 +111,7 @@ def __init__(
                 raise ValueError("basename is None, but unknown src!")
         self.basename = basename
 
-        if self.level is not None and isinstance(self.level, int):
+        if isinstance(self.level, int):
             levpath = f"{self.sxs_data_path}/Lev{self.level}"
         else:
             levpath = self.sxs_data_path
@@ -124,12 +125,31 @@ def __init__(
                     levpath = None
             else:
                 levpath = None
-
         self.check_cut_consistency()
-        if levpath is None or not os.path.exists(levpath):
+
+        needs_download = levpath is None or not os.path.exists(levpath)
+        if not needs_download:
+            # if files are already  downloaded, additionally check that N-order
+            # is there as well. Note: if we enter here, levpath is not None
+            # and the path exists, and thus lev_dirs not empty
+            order_group = f"Extrapolated_N{self.order}.dir"
+            if self.level is None:
+                level = int(lev_dirs[-1].replace("Lev", ""))
+            else:
+                level = self.level
+            fname = self.get_lev_fname(basename=self.basename, level=level)
+            with h5py.File(fname, "r") as f:
+                needs_download = order_group not in f
+            if needs_download:
+                logging.info(
+                    f"{levpath} found, but not the requested N={self.order} order. Download needed."
+                )
+
+        if needs_download:
             if download:
                 logging.info(
-                    f"The path {self.sxs_data_path} does not exist or contains no 'Lev*' directory."
+                    f"The path {self.sxs_data_path} does not exist, contains no 'Lev*'"
+                    + "directory, or does not contain the requested order."
                 )
                 logging.info("Downloading the simulation from the SXS catalog.")
                 self.download_simulation(
@@ -185,6 +205,9 @@ def __init__(
             self.load_horizon()
         if "psi4lm" in load:
             self.load_psi4lm(load_m0=load_m0)
+
+        if self.nr is not None:
+            self.nr.close()
         pass
 
     def check_cut_consistency(self):
@@ -282,11 +305,9 @@ def download_simulation(
             name_level = name
 
         # based on the logging level, redirect stdout to null
-        # This is because the sxs module prints a lot of information to stdout
         original_stdout = sys.stdout
-        original_stderr = sys.stderr
-        sys.stdout = open(os.devnull, "w")
-        sys.stderr = open(os.devnull, "w")
+        sys.stdout = LoggerWriter(logging.getLogger(__name__))
+
         sxs_sim = sxsmod.load(
             name_level,
             extrapolation_order=extrapolation_order,
@@ -300,17 +321,8 @@ def download_simulation(
         self.level = self.level or int(
             sxs_sim.Lev.replace("Lev", "")
         )  # Guarantees int(self.level)
-        lev = f"Lev{self.level}"
-
-        # Create the output directory
-        sxs_dir = f"SXS_{self.src}_{ID}"
-        # Only add sxs_dir if it's not already the last part of the path
-        if not path.endswith(sxs_dir):
-            full_path = os.path.join(path, sxs_dir)
-        else:
-            full_path = path
 
-        out_dir = os.path.join(full_path, lev)
+        out_dir = self.get_lev_fname(level=self.level, basename="")
         os.makedirs(out_dir, exist_ok=True)
 
         # Save hlm data if requested
@@ -337,11 +349,15 @@ def download_simulation(
                         f"Mode Y_l{ell}_m{m} not found in the waveform data! Skipping."
                     )
                     continue
-            # create the h5 file
-            h5file = h5py.File(
-                os.path.join(out_dir, f"rhOverM_Asymptotic_GeometricUnits_CoM.h5"), "w"
+            # create/udpate the h5 file
+            filename = os.path.join(
+                out_dir, f"rhOverM_Asymptotic_GeometricUnits_CoM.h5"
             )
-            save_dict_to_h5(h5file, to_h5file)
+            with h5py.File(filename, "a") as h5file:
+                if extp in h5file:
+                    logging.info(f"{extp} already present, skipping.")
+                else:
+                    save_dict_to_h5(h5file, {extp: to_h5file[extp]})
             h5file.close()
             logging.info("Saved hlm data.")
 
@@ -361,11 +377,13 @@ def download_simulation(
                 if mode_string in wav:
                     to_h5file[extp][mode_string] = wav[mode_string]
 
-            # create the h5 file
-            h5file = h5py.File(
-                os.path.join(out_dir, f"rMPsi4_Asymptotic_GeometricUnits_CoM.h5"), "w"
-            )
-            save_dict_to_h5(h5file, to_h5file)
+            # create/udpdate the h5 file
+            filename = os.path.join(out_dir, f"rMPsi4_Asymptotic_GeometricUnits_CoM.h5")
+            with h5py.File(filename, "a") as h5file:
+                if extp in h5file:
+                    logging.info(f"{extp} already present, skipping.")
+                else:
+                    save_dict_to_h5(h5file, {extp: to_h5file[extp]})
             h5file.close()
             logging.info("Saved psi4lm data.")
 
@@ -408,9 +426,8 @@ def download_simulation(
             if ":" in fld:
                 shutil.rmtree(os.path.join(os.environ["SXSCACHEDIR"], fld))
 
-        # Restore stdout/stderr
+        # Restore stdout
         sys.stdout = original_stdout
-        sys.stderr = original_stderr
 
         pass
 

PyART/utils/utils.py

@@ -791,7 +791,7 @@ def D02(xp, yp, pad=True):
     return dyp
 
 
-def D1(f, x, order=4, uniform_check=True):
+def D1(f, x, order=4, uniform_check=True, uc_atol=1e-4, uc_rtol=1e-4):
     """
     Computes the first derivative of function f(x)
 
@@ -804,9 +804,12 @@ def D1(f, x, order=4, uniform_check=True):
     order : int, optional
        finite differencing order (default is 4)
     uniform_check: bool, optional
-       check that the arrayr has uniform spacing
+       check that the array has uniform spacing
        (default is true)
-
+    uc_atol : float, optional
+        atol to use in np.allclose if uniform_check is True
+    uc_rtol : float, optional
+        rtol to use in np.allclose if uniform_check is True
     Returns
     -------
     df : list (or numpy array)
@@ -815,7 +818,7 @@ def D1(f, x, order=4, uniform_check=True):
 
     if uniform_check:
         dx = np.diff(x)
-        is_constant = np.allclose(dx, dx[0])
+        is_constant = np.allclose(dx, dx[0], rtol=uc_rtol, atol=uc_atol)
         if not is_constant:
             raise RuntimeError("Array not uniformly spaced")
 
@@ -1322,3 +1325,23 @@ def refine_local(idx, find_max=True):
     tap = refine_local(idxap, find_max=True)
     tpe = refine_local(idxpe, find_max=False)
     return tap, tpe
+
+
+class LoggerWriter:
+    """
+    Small class to redirect stdout to logging like:
+    original_stdout = sys.stdout
+    sys.stdout = LoggerWriter(logger)
+    """
+
+    def __init__(self, logger, level=logging.INFO):
+        self.logger = logger
+        self.level = level
+
+    def write(self, message):
+        message = message.rstrip()
+        if message:
+            self.logger.log(self.level, message)
+
+    def flush(self):
+        pass  # Needed because sys.stdout expects it

tests/test_sxs.py

@@ -12,17 +12,19 @@ def test_sxs():
     """
     Test the SXS download function.
     """
-    wf = sxs.Waveform_SXS(
-        ID="0180",
-        path="./",
-        download=True,
-        downloads=["hlm", "metadata", "horizons"],
-        load=["hlm", "metadata", "horizons"],
-        ignore_deprecation=True,
-        level=4,
-        order=2,
-        nu_rescale=False,
-    )
+    opts = {
+        "ID": "0180",
+        "path": "./",
+        "download": True,
+        "downloads": ["hlm", "metadata", "horizons"],
+        "load": ["hlm", "metadata", "horizons"],
+        "level": 4,
+        "order": 2,
+        "nu_rescale": False,
+        "ignore_deprecation": True,
+    }
+    wf = sxs.Waveform_SXS(**opts)
+
     # check attributes
     assert wf.ID == "0180"
     assert wf.level == 4
@@ -54,5 +56,9 @@ def test_sxs():
         # check length
         assert len(wf.hlm[mode]["A"]) == len(wf.u)
 
+    # get also order=3
+    opts["order"] = 3
+    wf = sxs.Waveform_SXS(**opts)
+
     # check that conversion to LVKNR works
     wf.to_lvk(modes=[(2, 2)])