diff --git a/PyART/catalogs/sxs.py b/PyART/catalogs/sxs.py index 64e9c7e3..3aea9ed3 100644 --- a/PyART/catalogs/sxs.py +++ b/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 diff --git a/PyART/utils/utils.py b/PyART/utils/utils.py index f61578a9..bb39236a 100644 --- a/PyART/utils/utils.py +++ b/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 diff --git a/tests/test_sxs.py b/tests/test_sxs.py index 16482598..a8c58715 100644 --- a/tests/test_sxs.py +++ b/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)])