diff --git a/PyART/analysis/match.py b/PyART/analysis/match.py index b4f62b1f..209ff033 100644 --- a/PyART/analysis/match.py +++ b/PyART/analysis/match.py @@ -122,8 +122,12 @@ def __init__( ) if self.settings["cut_second_waveform"] or self.settings["cut_longer"]: - tmrg1, _, _, _ = WaveForm1.find_max() - WaveForm1.u[0] - tmrg2, _, _, _ = WaveForm2.find_max() - WaveForm2.u[0] + # merger time measured from the start of each waveform. find_max + # returns (t_mrg, A_mrg, omg_mrg, domg_mrg), so only the first entry + # is wanted: subtracting u[0] from the whole tuple relies on numpy + # broadcasting it, which holds only while the entries are np.float64 + tmrg1 = WaveForm1.find_max()[0] - WaveForm1.u[0] + tmrg2 = WaveForm2.find_max()[0] - WaveForm2.u[0] DeltaT = tmrg2 - tmrg1 if DeltaT > 0: WaveForm2.cut(DeltaT) @@ -529,7 +533,9 @@ def _compute_mm_single_mode(self, wf1, wf2, settings): out = { "h1f": h1f, "h2f": h2f, - "j_shift": j_shift * h2.delta_t, + # take the time step from h2f: h2 only exists when wf2 is in the + # time domain and was not served from the cache + "j_shift": j_shift * h2f.delta_t, "ph_shift": ph_shift, } return m, out @@ -801,7 +807,9 @@ def _compute_mm_skymax(self, wf1, wf2, settings): ) mms.append(mm) out = {} # FIXME: just for consistency with compute_mm_single_mode - return np.average(mm), out + # average over the whole coa_phase x eff_pols grid, not just the last + # entry, which is what the loop above accumulates mms for + return np.average(mms), out def skymax_match( self, s, wf, inc, psd, modes, dT=1.0 / 4096, fmin_mm=20.0, fmax=2048.0 diff --git a/PyART/catalogs/cataloger.py b/PyART/catalogs/cataloger.py index 83497626..abed5efb 100644 --- a/PyART/catalogs/cataloger.py +++ b/PyART/catalogs/cataloger.py @@ -48,7 +48,7 @@ def __init__( name = wave.metadata["name"] self.data[name] = {"Waveform": wave, "Optimizer": None} except Exception as e: - logging.error(f"Issues with {ID:04}: {e}") + logging.error(f"Issues with {ID}: {e}") self.nsims = len(self.data) pass @@ -57,7 +57,7 @@ def get_Waveform(self, ID, add_opts={}, verbose=None): if verbose is None: verbose = self.verbose if verbose: - logging.info(f"Loading {self.catalog} waveform with ID:{ID:04}") + logging.info(f"Loading {self.catalog} waveform with ID:{ID}") if self.catalog == "sxs": from .sxs import Waveform_SXS @@ -148,6 +148,45 @@ def optimize_mismatches_batch(self, batch, optimizer_opts={}, verbose=None): Optimizer(self.data[name]["Waveform"], **optimizer_opts) pass + def collect_mismatch_jsons(self, json_tmp_list): + """ + Merge the per-process temporary JSON files produced by a parallel run + into self.json_file, then delete them. + + Mismatches already known to the base file are kept: each process starts + from a copy of it and only adds its own batch, so the temporary files + overlap on everything but the entries they computed themselves. + + Parameters + ---------- + json_tmp_list: list of str + paths of the temporary JSON files to merge + + Returns + ------- + json_data: dict + the merged content, as written to self.json_file + """ + # 1) load original json file if it exists, otherwise load first temp file + if os.path.exists(self.json_file): + json_base = self.json_file + else: + json_base = json_tmp_list[0] + with open(json_base, "r") as file: + json_data = json.loads(file.read()) + # 2) add info from temporary json, delete temporary json + for json_tmp in json_tmp_list: + with open(json_tmp, "r") as file: + json_data_tmp = json.loads(file.read()) + mismatches = json_data_tmp["mismatches"] + for key in mismatches: + if key not in json_data["mismatches"]: + json_data["mismatches"][key] = mismatches[key] + os.remove(json_tmp) + with open(self.json_file, "w") as file: + file.write(json.dumps(json_data, indent=2)) + return json_data + def process_with_redirect(self, process_id, nproc, opts={}): """ If we are running in parallel, use log files. @@ -238,24 +277,7 @@ def optimize_mismatches( process.join() # info stored in the temporary JSONs, collect info - # 1) load original json file if it exists, otherwise load first temp file - if os.path.exists(self.json_file): - json_base = self.json_file - else: - json_base = json_tmp_list[0] - with open(json_base, "r") as file: - json_data = json.loads(file.read()) - # 2) add info from temporary json, delete temporary json - for json_tmp in json_tmp_list: - with open(json_tmp, "r") as file: - json_data_tmp = json.loads(file.read()) - mismatches = json_data_tmp["mismatches"] - for key in mismatches: - if key not in json_data: - json_data["mismatches"][key] = mismatches[key] - os.remove(json_tmp) - with open(self.json_file, "w") as file: - file.write(json.dumps(json_data, indent=2)) + self.collect_mismatch_jsons(json_tmp_list) # read collated json optimizer_opts["json_file"] = self.json_file @@ -403,6 +425,9 @@ def mm_at_M(self, name, M, mm_settings=None, model_opts={}): eob = self.get_model_waveform(name, model_opts=model_opts) nr = self.data[name]["Waveform"] + # copy: the mass is specific to this call, and the caller usually reuses + # the same settings dict across masses + mm_settings = {} if mm_settings is None else dict(mm_settings) mm_settings["M"] = M matcher = Matcher(nr, eob, settings=mm_settings) return matcher.mismatch @@ -423,9 +448,6 @@ def mm_vs_M( ylim=None, figname=None, # save if not None ): - if figname is not None: - savepng = True - # select waveforms and get colors subset = self.find_subset(ranges=ranges) colors_dict = self.get_colors_for_subset( @@ -524,8 +546,7 @@ def mm_vs_M( if ylim is not None: plt.ylim(ylim) - if figname is None: - figname = f"mismatches_{self.catalog}_{cmap_var}.png" + if figname is not None: plt.savefig(figname, dpi=200, bbox_inches="tight") logging.info(f"Figure saved: {figname}") plt.show() diff --git a/PyART/catalogs/sxs.py b/PyART/catalogs/sxs.py index 3aea9ed3..f4c250e3 100644 --- a/PyART/catalogs/sxs.py +++ b/PyART/catalogs/sxs.py @@ -7,6 +7,7 @@ from ..waveform import Waveform from ..utils import cat_utils as cat_ut from ..utils.utils import LoggerWriter +from ..utils.wf_utils import get_multipole_dict class Waveform_SXS(Waveform): @@ -293,9 +294,19 @@ def download_simulation( import sxs as sxsmod import shutil - if path is not None: - logging.info(f"Setting the download (cache) directory to {path}") - os.environ["SXSCACHEDIR"] = path + if path is None: + raise ValueError( + "download_simulation needs a path: it is both where the data is " + "written and the cache directory handed to the sxs module." + ) + + # The sxs module reads SXSCACHEDIR to decide where to download. Keep the + # directory in a local variable too, and use that below: reading the + # environment back would couple this call to whatever a previous one + # left there, and rmtree is run against it. + cache_dir = path + logging.info(f"Setting the download (cache) directory to {cache_dir}") + os.environ["SXSCACHEDIR"] = cache_dir # Define the simulation ID and load it name = f"SXS:{self.src}:{ID}" @@ -304,130 +315,139 @@ def download_simulation( else: name_level = name - # based on the logging level, redirect stdout to null + # based on the logging level, redirect stdout to the logger. This is + # because the sxs module prints a lot of information to stdout. + # try/finally: an exception in between must not leave stdout redirected + # for good. original_stdout = sys.stdout sys.stdout = LoggerWriter(logging.getLogger(__name__)) - - sxs_sim = sxsmod.load( - name_level, - extrapolation_order=extrapolation_order, - extrapolation=f"N{extrapolation_order}", - ignore_deprecation=ignore_deprecation, - progress=True, - ) - logging.info(f"Loaded SXS simulation {name_level}.") - - # Set Level if not already set - self.level = self.level or int( - sxs_sim.Lev.replace("Lev", "") - ) # Guarantees int(self.level) - - out_dir = self.get_lev_fname(level=self.level, basename="") - os.makedirs(out_dir, exist_ok=True) - - # Save hlm data if requested - if "hlm" in downloads: - wav = sxs_sim.h - extp = f"Extrapolated_N{extrapolation_order}.dir" - to_h5file = {extp: {}} - ellmax = 8 # wav.ellmax - modes = [ - (l, m) - for l, m in product(range(2, ellmax + 1), range(-ellmax, ellmax + 1)) - if l >= np.abs(m) - ] - for ell, m in modes: - try: - idx = wav.index(ell, m) - mode_string = f"Y_l{ell}_m{m}.dat" - data = np.column_stack( - (wav.time, wav[:, idx].real, wav[:, idx].imag) - ) - to_h5file[extp][mode_string] = data - except ValueError: - logging.warning( - f"Mode Y_l{ell}_m{m} not found in the waveform data! Skipping." - ) - continue - # create/udpate the h5 file - filename = os.path.join( - out_dir, f"rhOverM_Asymptotic_GeometricUnits_CoM.h5" + try: + sxs_sim = sxsmod.load( + name_level, + extrapolation_order=extrapolation_order, + extrapolation=f"N{extrapolation_order}", + ignore_deprecation=ignore_deprecation, + progress=True, ) - 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.") - - # Save psi4lm data if requested - if "psi4lm" in downloads: - wav = sxs_sim.psi4 - extp = f"Extrapolated_N{extrapolation_order}.dir" - to_h5file = {extp: {}} - ellmax = 8 # wav.ellmax - modes = [ - (l, m) - for l, m in product(range(2, ellmax + 1), range(-ellmax, ellmax + 1)) - if l >= np.abs(m) - ] - for mode in modes: - mode_string = "Y_l" + str(mode[0]) + "_m" + str(mode[1]) + ".dat" - if mode_string in wav: - to_h5file[extp][mode_string] = wav[mode_string] - - # 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.") - - # Save horizons if requested - if "horizons" in downloads: - hrz = sxs_sim.horizons - to_h5file = {} - for object in ["AhA.dir", "AhB.dir", "AhC.dir"]: - to_h5file[object] = {} - for key in [ - "CoordCenterInertial.dat", - "ChristodoulouMass.dat", - "DimensionfulInertialSpinMag.dat", - "chiInertial.dat", - ]: + logging.info(f"Loaded SXS simulation {name_level}.") + + # Set Level if not already set + self.level = self.level or int( + sxs_sim.Lev.replace("Lev", "") + ) # Guarantees int(self.level) + + out_dir = self.get_lev_fname(level=self.level, basename="") + os.makedirs(out_dir, exist_ok=True) + + # Save hlm data if requested + if "hlm" in downloads: + wav = sxs_sim.h + extp = f"Extrapolated_N{extrapolation_order}.dir" + to_h5file = {extp: {}} + ellmax = 8 # wav.ellmax + modes = [ + (l, m) + for l, m in product( + range(2, ellmax + 1), range(-ellmax, ellmax + 1) + ) + if l >= np.abs(m) + ] + for ell, m in modes: try: - to_h5file[object][key] = hrz[f"{object}/{key}"] - except KeyError: + idx = wav.index(ell, m) + mode_string = f"Y_l{ell}_m{m}.dat" + data = np.column_stack( + (wav.time, wav[:, idx].real, wav[:, idx].imag) + ) + to_h5file[extp][mode_string] = data + except ValueError: logging.warning( - f"{object}/{key} not found in horizons data! Skipping." + 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"Horizons.h5"), "w") - save_dict_to_h5(h5file, to_h5file) - h5file.close() - logging.info("Saved horizons data.") - - # Save metadata if requested - if "metadata" in downloads: - import json - - with open(os.path.join(out_dir, "metadata.json"), "w") as file: - json.dump(sxs_sim.metadata, file, indent=2) - logging.info("Saved metadata.") - - # find old SXS download foders and remove them - flds = [f for f in os.listdir(os.environ["SXSCACHEDIR"]) if ID in f] - for fld in flds: - if ":" in fld: - shutil.rmtree(os.path.join(os.environ["SXSCACHEDIR"], fld)) - - # Restore stdout - sys.stdout = original_stdout + # create/update the h5 file + filename = os.path.join( + out_dir, f"rhOverM_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]}) + logging.info("Saved hlm data.") + + # Save psi4lm data if requested + if "psi4lm" in downloads: + wav = sxs_sim.psi4 + extp = f"Extrapolated_N{extrapolation_order}.dir" + to_h5file = {extp: {}} + ellmax = 8 # wav.ellmax + modes = [ + (l, m) + for l, m in product( + range(2, ellmax + 1), range(-ellmax, ellmax + 1) + ) + if l >= np.abs(m) + ] + for mode in modes: + mode_string = "Y_l" + str(mode[0]) + "_m" + str(mode[1]) + ".dat" + if mode_string in wav: + to_h5file[extp][mode_string] = wav[mode_string] + + # create/update 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]}) + logging.info("Saved psi4lm data.") + + # Save horizons if requested + if "horizons" in downloads: + hrz = sxs_sim.horizons + to_h5file = {} + for object in ["AhA.dir", "AhB.dir", "AhC.dir"]: + to_h5file[object] = {} + for key in [ + "CoordCenterInertial.dat", + "ChristodoulouMass.dat", + "DimensionfulInertialSpinMag.dat", + "chiInertial.dat", + ]: + try: + to_h5file[object][key] = hrz[f"{object}/{key}"] + except KeyError: + logging.warning( + f"{object}/{key} not found in horizons data! Skipping." + ) + continue + # create the h5 file + h5file = h5py.File(os.path.join(out_dir, f"Horizons.h5"), "w") + save_dict_to_h5(h5file, to_h5file) + h5file.close() + logging.info("Saved horizons data.") + + # Save metadata if requested + if "metadata" in downloads: + import json + + with open(os.path.join(out_dir, "metadata.json"), "w") as file: + json.dump(sxs_sim.metadata, file, indent=2) + logging.info("Saved metadata.") + + # find old SXS download folders and remove them. Only the + # colon-named ones the sxs module creates, and only in the directory + # this call actually downloaded into. + flds = [f for f in os.listdir(cache_dir) if ID in f] + for fld in flds: + if ":" in fld: + shutil.rmtree(os.path.join(cache_dir, fld)) + finally: + # Restore stdout + sys.stdout = original_stdout pass @@ -481,13 +501,18 @@ def read_spin_variable(spin_idx): key = attempt + str(spin_idx) if is_valid(key): hS = np.array(ometa[key]) - if attempt == "reference_": + # 'reference_spin' is dimensionful, unlike the two + # '*_dimensionless_spin' entries: normalize it by the mass + # squared, as done for the remnant spin below + if attempt == "reference_spin": if spin_idx == 1: hS = hS / M1**2 elif spin_idx == 2: hS = hS / M2**2 - break - return hS, attempt + return hS, attempt + raise KeyError( + f"No valid spin entry found for body {spin_idx}, tried: {attempts}" + ) hS1, skey1 = read_spin_variable(1) hS2, skey2 = read_spin_variable(2) @@ -525,7 +550,9 @@ def read_spin_variable(spin_idx): ecc = ometa["reference_eccentricity"] if isinstance(ecc, str): - if "<" in ecc and "e+00": # there are things like '<1.7e+00' in meta + # there are things like '<1.7e+00' in meta: an upper bound of order + # unity carries no information, while a small one means circular + if "<" in ecc and "e+00" in ecc: ecc = None else: ecc = 1e-5 @@ -635,24 +662,71 @@ def read_spin_variable(spin_idx): def load_horizon(self): """ Load the horizon data from Horizons.h5 file. - Store the data in self._dyn dictionary + Store the data in self._dyn dictionary. + + The datasets are stored with the time in the first column: + ChristodoulouMass.dat and DimensionfulInertialSpinMag.dat are (N,2), + while CoordCenterInertial.dat and chiInertial.dat are (N,4). Here the + time is stored once, in dyn["t"], and the vectors are stored without + it, i.e. with shape (N,3). + + AhA/AhB are the two individual horizons; AhC is the common horizon, + which only forms at merger and therefore lives on its own, shorter time + array (dyn["t_remnant"]) and is absent for runs that do not merge. """ horizon = h5py.File(self.get_lev_fname(basename="Horizons.h5")) - mA = horizon["AhA.dir"]["ChristodoulouMass.dat"] - mB = horizon["AhB.dir"]["ChristodoulouMass.dat"] - chiA = horizon["AhA.dir/DimensionfulInertialSpinMag.dat"] - chiB = horizon["AhB.dir/DimensionfulInertialSpinMag.dat"] - xA = horizon["AhA.dir/CoordCenterInertial.dat"] - xB = horizon["AhB.dir/CoordCenterInertial.dat"] - - self._dyn["t"] = chiA[:, 0] - self._dyn["m1"] = mA[:, 1] - self._dyn["m2"] = mB[:, 1] - self._dyn["chi1"] = chiA[:, 1] - self._dyn["chi2"] = chiB[:, 1] - self._dyn["x1"] = xA[:, 1:] - self._dyn["x2"] = xB[:, 1:] + def read_horizon(obj): + """ + Read one apparent horizon, dropping the leading time column from + the vectors. Datasets that the download skipped are returned as + None rather than raising. + """ + grp = horizon[obj] + + def dset(name, vector=False): + if name not in grp: + logging.warning(f"{obj}/{name} not found in horizons data!") + return None + return grp[name][:, 1:] if vector else grp[name][:, 1] + + return { + "t": grp["ChristodoulouMass.dat"][:, 0], + "m": dset("ChristodoulouMass.dat"), + # chiInertial is the dimensionless spin vector, which is what + # dyn["chi"] is meant to hold. DimensionfulInertialSpinMag is + # |S|: dimensionful (chi = S/m^2) and a magnitude, so it has no + # direction to project on L. + "chi": dset("chiInertial.dat", vector=True), + "S_mag": dset("DimensionfulInertialSpinMag.dat"), + "x": dset("CoordCenterInertial.dat", vector=True), + } + + A = read_horizon("AhA.dir") + B = read_horizon("AhB.dir") + + self._dyn["t"] = A["t"] + self._dyn["m1"] = A["m"] + self._dyn["m2"] = B["m"] + self._dyn["chi1"] = A["chi"] + self._dyn["chi2"] = B["chi"] + self._dyn["S1_mag"] = A["S_mag"] + self._dyn["S2_mag"] = B["S_mag"] + self._dyn["x1"] = A["x"] + self._dyn["x2"] = B["x"] + + if "AhC.dir" in horizon: + C = read_horizon("AhC.dir") + self._dyn["t_remnant"] = C["t"] + self._dyn["m_remnant"] = C["m"] + self._dyn["chi_remnant"] = C["chi"] + self._dyn["S_remnant_mag"] = C["S_mag"] + self._dyn["x_remnant"] = C["x"] + else: + logging.info( + "No common horizon (AhC.dir) in the horizons data: " + "remnant quantities not loaded." + ) pass @@ -673,19 +747,20 @@ def compute_spins_at_tref(self, tref): """ d = self.dyn - # find the index of the reference time + # find the index of the reference time. The dyn vectors carry no time + # column, so they are indexed by time only. idx = np.argmin(np.abs(d["t"] - tref)) - chi1_ref = d["chi1"][idx][1:] - chi2_ref = d["chi2"][idx][1:] - x1_ref = d["x1"][idx][1:] - x2_ref = d["x2"][idx][1:] + chi1_ref = d["chi1"][idx] + chi2_ref = d["chi2"][idx] + x1_ref = d["x1"][idx] + x2_ref = d["x2"][idx] # time derivative of x1 and x2 - x1_dot = np.transpose([np.gradient(d["x1"][:, i], d["t"]) for i in range(1, 4)]) - x2_dot = np.transpose([np.gradient(d["x2"][:, i], d["t"]) for i in range(1, 4)]) + x1_dot = np.transpose([np.gradient(d["x1"][:, i], d["t"]) for i in range(3)]) + x2_dot = np.transpose([np.gradient(d["x2"][:, i], d["t"]) for i in range(3)]) x = x1_ref - x2_ref - x_dot = [x1_dot[idx][i] - x2_dot[idx][i] for i in range(3)] + x_dot = x1_dot[idx] - x2_dot[idx] L_hat_ref = np.cross(x, x_dot) / np.linalg.norm(np.cross(x, x_dot)) # compute the spins projected on L_hat_ref @@ -740,17 +815,13 @@ def load_hlm(self, ellmax=None, load_m0=False): h = hlm[:, 1] + 1j * hlm[:, 2] if self.nu_rescale: h /= self.metadata["nu"] - # amp and phase - Alm = abs(h)[self.cut_N :] - plm = -np.unwrap(np.angle(h))[self.cut_N :] - # save in dictionary + # Build the mode dict with the shared helper, so that the sign + # conventions match every other producer. It is applied to the whole + # mode and the junk is cut afterwards: the phase must be unwrapped + # before the cut, or it would be offset by a multiple of 2pi. key = (l, m) dict_hlm[key] = { - "real": Alm * np.cos(plm), - "imag": Alm * np.sin(plm), - "A": Alm, - "p": plm, - "z": h[self.cut_N :], + ky: val[self.cut_N :] for ky, val in get_multipole_dict(h).items() } self._hlm = dict_hlm pass @@ -768,8 +839,9 @@ def load_psi4lm(self, ellmax=None, load_m0=False): """ psi4_basename = self.basename.replace("rhOverM", "rMPsi4") fname = self.get_lev_fname(level=self.level, basename=psi4_basename) - if os.path.exists(fname): - self.nr_psi = h5py.File(fname) + if not os.path.exists(fname): + raise FileNotFoundError(f"psi4 file not found: {fname}") + self.nr_psi = h5py.File(fname) if ellmax == None: ellmax = self.ellmax @@ -793,6 +865,11 @@ def load_psi4lm(self, ellmax=None, load_m0=False): if self.cut_U is None: self.cut_U = tmp_u[self.cut_N] + if self._u is None: + raise RuntimeError( + "psi4 times are taken from the hlm time array, but hlm was " + "never loaded: add 'hlm' to the load list." + ) self._t_psi4 = self._u # FIXME: should we use another time? dict_psi4lm = {} @@ -804,17 +881,10 @@ def load_psi4lm(self, ellmax=None, load_m0=False): psi4 = psi4lm[:, 1] + 1j * psi4lm[:, 2] if self.nu_rescale: psi4 /= self.metadata["nu"] - # amp and phase - Alm = abs(psi4)[self.cut_N :] - plm = -np.unwrap(np.angle(psi4))[self.cut_N :] - # save in dictionary + # see load_hlm: shared helper first, junk cut afterwards key = (l, m) dict_psi4lm[key] = { - "real": Alm * np.cos(plm), - "imag": Alm * np.sin(plm), - "A": Alm, - "p": plm, - "z": psi4[self.cut_N :], + ky: val[self.cut_N :] for ky, val in get_multipole_dict(psi4).items() } self._psi4lm = dict_psi4lm pass diff --git a/PyART/models/teob.py b/PyART/models/teob.py index 7335c78f..fb749bfc 100644 --- a/PyART/models/teob.py +++ b/PyART/models/teob.py @@ -271,7 +271,7 @@ def CreateDict( "r_hyp": r_hyp, "H_hyp": H_hyp, "j_hyp": J_hyp, - "prs_sign_hyp": -1, + "prs_sign_hyp": prs_sign_hyp, "coalescence_angle": phi_ref, "df": df, "anomaly": anomaly, diff --git a/PyART/utils/utils.py b/PyART/utils/utils.py index bb39236a..63dd3cbe 100644 --- a/PyART/utils/utils.py +++ b/PyART/utils/utils.py @@ -670,7 +670,8 @@ def upoly_fits( nmax: int maximum polynomial order n_extract: int or None - if not None, extract the extrapolated value from this fit order + fit order to extract the extrapolated value from; must lie in + [nmin, nmax]. If None, defaults to nmax. r_cutoff_low: float or None if not None, only use data with r0 >= r_cutoff_low r_cutoff_high: float or None @@ -684,7 +685,7 @@ def upoly_fits( ------- out: dict dictionary with the following keys - 'extrap': extrapolated value (mean over fit orders or from n_extract) + 'extrap': extrapolated value, taken from the n_extract fit order 'extrap_vec': extrapolated values for each fit order 'fit_orders': list of fit orders 'coeffs': polynomial coefficients for each fit order @@ -700,6 +701,10 @@ def upoly_fits( r_cutoff_high = max(r0) if nmin > nmax: raise ValueError(f"nmin>nmax: {nmin}>{nmax} !") + if not nmin <= n_extract <= nmax: + raise ValueError( + f"n_extract must lie in [nmin, nmax]=[{nmin}, {nmax}], got {n_extract}!" + ) i_rmin = np.argmin(r0) if direction == "in": @@ -730,10 +735,8 @@ def upoly_fits( b[:, i] = zero_pad_before(b_tmp, nmax + 1) p[:, i] = np.polyval(b_tmp, u) ye_vec[i] = b[-1, i] - if n_extract is not None and fit_order == n_extract: + if fit_order == n_extract: ye = ye_vec[i] - if n_extract is None: - ye = np.mean(ye_vec) out = { "extrap": ye, "extrap_vec": ye_vec, @@ -836,7 +839,7 @@ def D1(f, x, order=4, uniform_check=True, uc_atol=1e-4, uc_rtol=1e-4): i = np.arange(Nmin, Nmax) df[i] = (f[i + 1] - f[i]) * oodx - f[Nmax] = f[Nmax - 1] + df[Nmax] = df[Nmax - 1] elif order == 2: i = np.arange(Nmin + 1, Nmax) df[i] = 0.5 * (f[i + 1] - f[i - 1]) * oodx @@ -1177,9 +1180,9 @@ def list_to_str(mylist): if isinstance(value1, dict): dbool = are_dictionaries_equal(value1, value2, verbose=True) if dbool: - logging.info(f">> issues with {key:16s} (dictionary)") + logging.info(f">> {key:16s} is dict: no differences") else: - logging.info(f">> {key:16s} is dict:") + logging.info(f">> issues with {key:16s} (dictionary)") else: if value1 is None: value1 = "None" diff --git a/PyART/waveform.py b/PyART/waveform.py index 75c7e049..01229b15 100644 --- a/PyART/waveform.py +++ b/PyART/waveform.py @@ -232,13 +232,15 @@ def find_max( height = np.mean(Alm) / 2 peaks, props = find_peaks(Alm, height=height) + if len(peaks) == 0: + raise ValueError("No peaks found") + if kind == "first-max-after-t": - for i in range(len(peaks)): - if t[peaks[i]] > umin: - break + after_umin = np.where(t[peaks] > umin)[0] + if len(after_umin) == 0: + raise ValueError(f"No peak found after umin={umin}") + i = after_umin[0] elif kind == "last-peak": - if len(peaks) == 0: - raise ValueError("No peaks found") i = len(peaks) - 1 elif kind == "global": Alms = props["peak_heights"] @@ -551,25 +553,33 @@ def to_frequency(self, taper=True, pad=True): if True, apply a tapering window to the time-domain waveform pad: bool if True, pad the time-domain waveform to the next power of 2 + Returns ------- - out: (f, hp, hc) + out: None + sets the f, hp, hc and domain attributes in place """ dt = self.u[1] - self.u[0] # window if taper: - self._hp = ut.windowing(self.hp, alpha=0.1) - self._hc = ut.windowing(self.hc, alpha=0.1) + self._hp, _ = ut.windowing(self.hp, alpha=0.1) + self._hc, _ = ut.windowing(self.hc, alpha=0.1) if pad: - srate = 1.0 / dt seglen = ut.nextpow2(self.u[-1]) - dN = int((seglen - len(self.u) * srate) / srate) self._u = np.arange(0.0, seglen, dt) self._t = np.arange(0.0, seglen, dt) - self._hp = ut.zero_pad_before(self.hp, dN, return_column=False) - self._hc = ut.zero_pad_before(self.hp, dN, return_column=False) + # zero_pad_before wants the final number of samples, not the + # number of zeroes to prepend + N = len(self.u) + if N < len(self.hp): + raise ValueError( + f"Cannot pad to {N} samples: waveform is already {len(self.hp)} " + "samples long. Is the time array starting at t=0?" + ) + self._hp = ut.zero_pad_before(self.hp, N, return_column=False) + self._hc = ut.zero_pad_before(self.hc, N, return_column=False) assert len(self.u) == len(self.hp) assert len(self.u) == len(self.hc) @@ -604,7 +614,7 @@ def integrate_data( self, t_psi4, radius, - integr_opts={"method": "FFI", "f0": 0.01, "integrand": "psi4"}, + integr_opts=None, modes=None, M=1.0, ): @@ -617,12 +627,13 @@ def integrate_data( time array for psi4lm radius: float extraction radius - integr_opts: dict + integr_opts: dict or None dictionary with integration options method: 'FFI' or 'trapezoid' f0: frequency cutoff for FFI deg: degree of the polynomial for trapezoid integrand: 'psi4' or 'news' + if None, use {'method': 'FFI', 'f0': 0.01, 'integrand': 'psi4'} modes: list list of (l,m) modes to integrate M: float @@ -632,6 +643,11 @@ def integrate_data( out: dict dictionary with integration options used """ + if integr_opts is None: + integr_opts = {"method": "FFI", "f0": 0.01} + else: + # copy: the defaults filled in below must not leak back to the caller + integr_opts = dict(integr_opts) if modes is None: modes = self.psi4lm.keys() if "integrand" not in integr_opts: @@ -677,7 +693,7 @@ def to_geom(self, M, distance): D_sec = distance * ut.consts["DMpc"] M_sec = M * ut.consts["Msun"] - time_attrs = ["_u", "_t", "t_psi4", "u_pc"] + time_attrs = ["_u", "_t", "_t_psi4", "u_pc"] for time_attr in time_attrs: val = getattr(self, time_attr, None) if val is not None: @@ -728,7 +744,7 @@ def to_SI(self, M, distance): D_sec = distance * ut.consts["DMpc"] M_sec = M * ut.consts["Msun"] - time_attrs = ["_u", "_t", "t_psi4", "u_pc"] + time_attrs = ["_u", "_t", "_t_psi4", "u_pc"] for time_attr in time_attrs: val = getattr(self, time_attr, None) if val is not None: @@ -1065,7 +1081,7 @@ def __init__( r_extr=1, M=1, modes=[(2, 2)], - integr_opts={}, + integr_opts=None, fmt="etk", fname="mp_psi4_l@L@_m@M@_r100.00.asc", integrand="psi4", @@ -1086,8 +1102,9 @@ def __init__( total mass modes: list list of (l,m) modes to load - integr_opts: dict - dictionary with integration options + integr_opts: dict or None + dictionary with integration options; missing entries are filled + in with defaults, and the caller's dictionary is left untouched method: 'FFI' or 'trapezoid' f0: frequency cutoff for FFI deg: degree of the polynomial for trapezoid @@ -1114,6 +1131,8 @@ def __init__( self.integrand = integrand.lower() self.norm = norm + # copy: the defaults filled in below must not leak back to the caller + integr_opts = {} if integr_opts is None else dict(integr_opts) if "method" not in integr_opts: integr_opts["method"] = "FFI" if "f0" not in integr_opts: diff --git a/docs/source/tutorials/intro_to_waveforms.ipynb b/docs/source/tutorials/intro_to_waveforms.ipynb index 87492aec..d87a2520 100644 --- a/docs/source/tutorials/intro_to_waveforms.ipynb +++ b/docs/source/tutorials/intro_to_waveforms.ipynb @@ -189,9 +189,7 @@ "cell_type": "markdown", "id": "bfc68330", "metadata": {}, - "source": [ - "We will now look at the dynamics. This is not available for all catalogs/models. In the case of SXS waveforms, we do load Horizon data when we download the waveform and can access it." - ] + "source": "We will now look at the dynamics. This is not available for all catalogs/models. In the case of SXS waveforms, we do load Horizon data when we download the waveform and can access it.\n\nThe horizon data is stored in `waveform.dyn`. The time is stored once, in `dyn['t']`, and the vectors are stored without it, with shape `(N, 3)`:\n\n| key | meaning |\n| --- | --- |\n| `t` | time |\n| `m1`, `m2` | Christodoulou masses |\n| `chi1`, `chi2` | dimensionless spin vectors $\\vec{\\chi}_i = \\vec{S}_i/m_i^2$ |\n| `S1_mag`, `S2_mag` | dimensionful spin magnitudes $\\|\\vec{S}_i\\|$ |\n| `x1`, `x2` | coordinate positions in the inertial frame |\n\nIf the binary merges, the common horizon is loaded too. It only forms at merger, so it lives on its own, shorter time array `t_remnant`, alongside `m_remnant`, `chi_remnant`, `S_remnant_mag` and `x_remnant`." }, { "cell_type": "code", @@ -199,32 +197,7 @@ "id": "d74e74ec", "metadata": {}, "outputs": [], - "source": [ - "# plot the mass / spin evolution\n", - "fig, axs = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n", - "axs[0].semilogy(sxs_waveform.dyn['t'], abs(sxs_waveform.dyn['m1']-m1), label=r'$m_1$')\n", - "axs[0].semilogy(sxs_waveform.dyn['t'], abs(sxs_waveform.dyn['m2']-m2), label=r'$m_2$', linestyle='--')\n", - "axs[1].semilogy(sxs_waveform.dyn['t'], sxs_waveform.dyn['chi1'], label=r'S_1')\n", - "axs[1].semilogy(sxs_waveform.dyn['t'], sxs_waveform.dyn['chi2'], label=r'S_2', linestyle='--')\n", - "axs[0].set_ylabel(r'$\\Delta m_i$')\n", - "axs[1].set_ylabel(r'$|S_i|$')\n", - "axs[1].set_xlabel('t/M')\n", - "axs[0].legend()\n", - "plt.show()\n", - "\n", - "\n", - "# plot the trajectories\n", - "xA = sxs_waveform.dyn['x1']\n", - "xB = sxs_waveform.dyn['x2']\n", - "fig, ax = plt.subplots(figsize=(8, 8))\n", - "ax.plot(xA[:, 0], xA[:, 1], label='BH 1')\n", - "ax.plot(xB[:, 0], xB[:, 1], label='BH 2')\n", - "ax.set_xlabel('x/M')\n", - "ax.set_ylabel('y/M')\n", - "ax.set_title('Black Hole Trajectories')\n", - "ax.legend()\n", - "plt.show()" - ] + "source": "# plot the mass / spin evolution\nfig, axs = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\naxs[0].semilogy(sxs_waveform.dyn['t'], abs(sxs_waveform.dyn['m1']-m1), label=r'$m_1$')\naxs[0].semilogy(sxs_waveform.dyn['t'], abs(sxs_waveform.dyn['m2']-m2), label=r'$m_2$', linestyle='--')\naxs[1].semilogy(sxs_waveform.dyn['t'], sxs_waveform.dyn['S1_mag'], label=r'$S_1$')\naxs[1].semilogy(sxs_waveform.dyn['t'], sxs_waveform.dyn['S2_mag'], label=r'$S_2$', linestyle='--')\naxs[0].set_ylabel(r'$\\Delta m_i$')\naxs[1].set_ylabel(r'$|S_i|$')\naxs[1].set_xlabel('t/M')\naxs[0].legend()\naxs[1].legend()\nplt.show()\n\n# the dimensionless spins are vectors: chi_i = S_i / m_i^2\nfig, ax = plt.subplots(figsize=(12, 4))\nfor i, comp in enumerate(['x', 'y', 'z']):\n ax.plot(sxs_waveform.dyn['t'], sxs_waveform.dyn['chi1'][:, i], label=rf'$\\chi_1^{comp}$')\nax.set_xlabel('t/M')\nax.set_ylabel(r'$\\chi_1$')\nax.set_title('Dimensionless spin components of BH 1')\nax.legend()\nplt.show()\n\n# the common horizon forms at merger, so it has its own (shorter) time array\nif 'm_remnant' in sxs_waveform.dyn:\n chi_rem = np.linalg.norm(sxs_waveform.dyn['chi_remnant'][-1])\n print(f\"remnant mass: {sxs_waveform.dyn['m_remnant'][-1]:.6f} \"\n f\"(metadata: {sxs_waveform.metadata['Mf']:.6f})\")\n print(f\"remnant spin: {chi_rem:.6f} \"\n f\"(metadata: {sxs_waveform.metadata['af']:.6f})\")\n\n# plot the trajectories\nxA = sxs_waveform.dyn['x1']\nxB = sxs_waveform.dyn['x2']\nfig, ax = plt.subplots(figsize=(8, 8))\nax.plot(xA[:, 0], xA[:, 1], label='BH 1')\nax.plot(xB[:, 0], xB[:, 1], label='BH 2')\nax.set_xlabel('x/M')\nax.set_ylabel('y/M')\nax.set_title('Black Hole Trajectories')\nax.legend()\nplt.show()" }, { "cell_type": "markdown", @@ -250,16 +223,7 @@ "id": "df054c5c", "metadata": {}, "outputs": [], - "source": [ - "# we compute the polarizations & r for default plots\n", - "sxs_waveform.compute_hphc(phi=0., i=np.pi/3)\n", - "sxs_waveform.dyn['r'] = np.sqrt((xA[:, 0]-xB[:, 0])**2 + (xA[:, 1]-xB[:, 1])**2 + (xA[:, 2]-xB[:, 2])**2)\n", - "\n", - "# plot pol, modes and r(t)\n", - "fix, ax = plt.subplots(4, 1, figsize=(6, 10), sharex=True)s\n", - "for i, quantity in enumerate(['hp', 'hc', 'hlm', 'dyn']):\n", - " sxs_waveform.plot(quantity, show=False, ax=ax[i])" - ] + "source": "# we compute the polarizations & r for default plots\nsxs_waveform.compute_hphc(phi=0., i=np.pi/3)\nsxs_waveform.dyn['r'] = np.sqrt((xA[:, 0]-xB[:, 0])**2 + (xA[:, 1]-xB[:, 1])**2 + (xA[:, 2]-xB[:, 2])**2)\n\n# plot pol, modes and r(t)\nfig, ax = plt.subplots(4, 1, figsize=(6, 10), sharex=True)\nfor i, quantity in enumerate(['hp', 'hc', 'hlm', 'dyn']):\n sxs_waveform.plot(quantity, show=False, ax=ax[i])" }, { "cell_type": "code", @@ -291,4 +255,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} +} \ No newline at end of file diff --git a/tests/test_cataloger.py b/tests/test_cataloger.py new file mode 100644 index 00000000..067eed35 --- /dev/null +++ b/tests/test_cataloger.py @@ -0,0 +1,283 @@ +""" +Test Cataloger bookkeeping that does not require catalog data: +- ID handling in log messages +- merging of the per-process JSONs written by a parallel run +- mm_at_M settings handling +- figure saving in mm_vs_M +""" + +import json +import logging + +import matplotlib +import numpy as np +import pytest + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from PyART.catalogs.cataloger import Cataloger + + +@pytest.fixture(autouse=True) +def no_usetex(): + """ + cataloger.py sets usetex at import; the labels it draws cannot be rendered + without a LaTeX installation, which is unrelated to what we test here. + """ + old = matplotlib.rcParams["text.usetex"] + matplotlib.rc("text", usetex=False) + yield + matplotlib.rc("text", usetex=old) + + +def make_cataloger(**attrs): + """ + Build a Cataloger without running __init__, which would need real catalog + data on disk. + """ + cat = Cataloger.__new__(Cataloger) + cat.path = "./" + cat.catalog = "rit" + cat.verbose = True + cat.sim_list = [] + cat.data = {} + cat.json_file = "mismatches.json" + for key, value in attrs.items(): + setattr(cat, key, value) + return cat + + +############################## +# ID handling in logs (#16) +############################## + + +@pytest.mark.parametrize("ID", ["q1", "0180", "ab", 1362, "SXS:BBH:0180"]) +def test_get_Waveform_logs_id_verbatim(ID, caplog): + """ + IDs may be strings; a format spec of :04 pads short strings on the right + ('q1' -> 'q100'), reporting an ID that does not exist. + """ + cat = make_cataloger(catalog="not-a-catalog") + + with caplog.at_level(logging.INFO): + with pytest.raises(ValueError, match="Unknown catalog"): + cat.get_Waveform(ID) + + assert f"ID:{ID}" in caplog.text + + +def test_get_Waveform_logs_short_string_id_without_padding(caplog): + """Regression: the ':04' spec turned the RIT-style ID 'q1' into 'q100'.""" + cat = make_cataloger(catalog="not-a-catalog") + + with caplog.at_level(logging.INFO): + with pytest.raises(ValueError): + cat.get_Waveform("q1") + + assert "q100" not in caplog.text + + +############################## +# parallel JSON merge (#17) +############################## + + +def write_json(path, mismatches): + with open(path, "w") as f: + json.dump({"mismatches": mismatches}, f) + + +def test_collect_mismatch_jsons_keeps_all_process_results(tmp_path): + """ + Each process writes only its own batch; the merge must retain every + process's results rather than letting later files overwrite earlier ones. + """ + base = tmp_path / "mismatches.json" + tmp0 = tmp_path / "mismatches_0.json" + tmp1 = tmp_path / "mismatches_1.json" + + # process 0 computed simA, process 1 computed simB + write_json(tmp0, {"simA": {"mm_min": 0.1}}) + write_json(tmp1, {"simB": {"mm_min": 0.2}}) + + cat = make_cataloger(json_file=str(base)) + merged = cat.collect_mismatch_jsons([str(tmp0), str(tmp1)]) + + assert set(merged["mismatches"]) == {"simA", "simB"} + assert merged["mismatches"]["simA"]["mm_min"] == 0.1 + assert merged["mismatches"]["simB"]["mm_min"] == 0.2 + + # and it is what landed on disk + with open(base) as f: + on_disk = json.load(f) + assert on_disk == merged + + # temporary files are cleaned up + assert not tmp0.exists() + assert not tmp1.exists() + + +def test_collect_mismatch_jsons_does_not_overwrite_existing_entries(tmp_path): + """ + Entries already in the base file win: 'if key not in json_data' was always + true (json_data holds 'mismatches', never the sim names), so a stale value + from a temporary file could clobber a fresh one. + """ + base = tmp_path / "mismatches.json" + write_json(base, {"simA": {"mm_min": 999.0}}) + + tmp0 = tmp_path / "mismatches_0.json" + write_json(tmp0, {"simA": {"mm_min": 0.1}, "simB": {"mm_min": 0.2}}) + + cat = make_cataloger(json_file=str(base)) + merged = cat.collect_mismatch_jsons([str(tmp0)]) + + assert merged["mismatches"]["simA"]["mm_min"] == 999.0 + assert merged["mismatches"]["simB"]["mm_min"] == 0.2 + + +def test_collect_mismatch_jsons_without_existing_base(tmp_path): + """With no base file, the first temporary file seeds the merge.""" + base = tmp_path / "mismatches.json" + tmp0 = tmp_path / "mismatches_0.json" + tmp1 = tmp_path / "mismatches_1.json" + write_json(tmp0, {"simA": {"mm_min": 0.1}}) + write_json(tmp1, {"simB": {"mm_min": 0.2}}) + + cat = make_cataloger(json_file=str(base)) + merged = cat.collect_mismatch_jsons([str(tmp0), str(tmp1)]) + + assert set(merged["mismatches"]) == {"simA", "simB"} + assert base.exists() + + +############################## +# mm_at_M settings (#18) +############################## + + +def test_mm_at_M_does_not_mutate_caller_settings(monkeypatch): + """mm_at_M sets 'M' for its own call; the caller's dict must be untouched.""" + cat = make_cataloger(data={"sim": {"Waveform": object(), "Optimizer": None}}) + monkeypatch.setattr( + Cataloger, "get_model_waveform", lambda self, name, **kw: object() + ) + + seen = {} + + def fake_matcher(nr, eob, settings=None): + seen.update(settings) + return type("M", (), {"mismatch": 0.5})() + + monkeypatch.setattr("PyART.catalogs.cataloger.Matcher", fake_matcher) + + settings = {"kind": "single-mode"} + mm = cat.mm_at_M("sim", 120.0, mm_settings=settings) + + assert mm == 0.5 + assert seen["M"] == 120.0 + assert "M" not in settings, "mm_at_M mutated the caller's settings dict" + + +def test_mm_at_M_accepts_default_settings(monkeypatch): + """mm_settings defaults to None; that used to raise TypeError.""" + cat = make_cataloger(data={"sim": {"Waveform": object(), "Optimizer": None}}) + monkeypatch.setattr( + Cataloger, "get_model_waveform", lambda self, name, **kw: object() + ) + + seen = {} + + def fake_matcher(nr, eob, settings=None): + seen.update(settings) + return type("M", (), {"mismatch": 0.25})() + + monkeypatch.setattr("PyART.catalogs.cataloger.Matcher", fake_matcher) + + mm = cat.mm_at_M("sim", 90.0) + + assert mm == 0.25 + assert seen == {"M": 90.0} + + +############################## +# mm_vs_M figure saving (#15) +############################## + + +@pytest.fixture +def cataloger_for_mm_vs_M(monkeypatch, tmp_path): + """ + A Cataloger whose mismatches are pre-loaded from JSON, so mm_vs_M plots + without generating any waveform. + """ + masses = list(np.linspace(100, 200, num=5)) + mm_json = tmp_path / "loaded.json" + with open(mm_json, "w") as f: + json.dump( + { + "masses": masses, + "options": {"mm_settings": {}}, + "mismatches": { + "sim": { + "mm_vs_M": [1e-3] * len(masses), + "mm_max": 1e-3, + "mm_min": 1e-3, + } + }, + }, + f, + ) + + cat = make_cataloger() + monkeypatch.setattr(Cataloger, "find_subset", lambda self, ranges=None: ["sim"]) + monkeypatch.setattr( + Cataloger, + "get_colors_for_subset", + lambda self, subset, cmap_var=None, cmap_name=None: { + "colors": ["C0"], + "indices": [0], + "range": (0.0, 1.0), + "cmap": plt.get_cmap("jet"), + }, + ) + monkeypatch.setattr(Cataloger, "tex_label_from_key", lambda self, key: key) + monkeypatch.setattr(plt, "show", lambda *a, **k: None) + return cat, str(mm_json), len(masses) + + +def test_mm_vs_M_saves_user_supplied_figname(cataloger_for_mm_vs_M, tmp_path): + """A caller-supplied figname must actually be written.""" + cat, mm_json, n = cataloger_for_mm_vs_M + figname = tmp_path / "myfig.png" + + cat.mm_vs_M(N=n, json_load=mm_json, figname=str(figname)) + + assert figname.exists(), "mm_vs_M did not save the user-supplied figname" + + +def test_mm_vs_M_saves_nothing_when_figname_is_none( + cataloger_for_mm_vs_M, tmp_path, monkeypatch +): + """figname=None means 'do not save', matching the json_save convention.""" + cat, mm_json, n = cataloger_for_mm_vs_M + monkeypatch.chdir(tmp_path) + + cat.mm_vs_M(N=n, json_load=mm_json, figname=None) + + assert ( + list(tmp_path.glob("*.png")) == [] + ), "mm_vs_M saved a figure when figname was None" + + +def test_mm_vs_M_saves_json_when_requested(cataloger_for_mm_vs_M, tmp_path): + cat, mm_json, n = cataloger_for_mm_vs_M + out = tmp_path / "saved.json" + + cat.mm_vs_M(N=n, json_load=mm_json, json_save=str(out)) + + assert out.exists() + with open(out) as f: + assert "sim" in json.load(f)["mismatches"] diff --git a/tests/test_match.py b/tests/test_match.py index 7ca6dbaf..06980c7c 100644 --- a/tests/test_match.py +++ b/tests/test_match.py @@ -3,6 +3,7 @@ """ import numpy as np +import pytest import matplotlib.pyplot as plt from PyART.catalogs import sxs from PyART.analysis.match import Matcher @@ -89,3 +90,150 @@ def test_self_match_pol_helper(cp, ep): test_self_match_pol_helper(cp, ep) pass + + +def test_skymax_averages_over_the_whole_grid(monkeypatch): + """ + _compute_mm_skymax loops over coa_phase x eff_pols and accumulates the + matches in `mms`, but used to return np.average(mm) -- the last scalar -- + throwing the loop away. The result must be the average over the grid. + """ + fake_matches = [0.90, 0.92, 0.94, 0.98] + calls = [] + + def fake_skymax_match(self, s, wf, inc, psd, modes, **kwargs): + calls.append(1) + return fake_matches[len(calls) - 1] + + monkeypatch.setattr(Matcher, "skymax_match", fake_skymax_match) + + settings = { + "kind": "hm", + "initial_frequency_mm": fmin, + "final_frequency_mm": fmax, + "tlen": len(nr.u), + "dt": 1 / srate, + "M": M, + "resize_factor": 4, + "pad_end_frac": 0.5, + "taper_alpha": 0.2, + "taper_start": 0.05, + "taper": "sigmoid", + "debug": False, + "coa_phase": [0.0, np.pi / 2], + "eff_pols": [0.0, np.pi / 3], + } + + m = Matcher(nr, nr_2, settings=settings) + + assert len(calls) == 4, "the grid should be 2 coa_phase x 2 eff_pols" + # mismatch = 1 - , averaged over the whole grid + assert m.mismatch == pytest.approx(1 - np.mean(fake_matches)) + # and specifically not just the last entry + assert m.mismatch != pytest.approx(1 - fake_matches[-1]) + + +def test_skymax_single_point_grid(monkeypatch): + """With one coa_phase and one eff_pol the average is that single value.""" + + def fake_skymax_match(self, s, wf, inc, psd, modes, **kwargs): + return 0.75 + + monkeypatch.setattr(Matcher, "skymax_match", fake_skymax_match) + + settings = { + "kind": "hm", + "initial_frequency_mm": fmin, + "final_frequency_mm": fmax, + "tlen": len(nr.u), + "dt": 1 / srate, + "M": M, + "resize_factor": 4, + "pad_end_frac": 0.5, + "taper_alpha": 0.2, + "taper_start": 0.05, + "taper": "sigmoid", + "debug": False, + "coa_phase": [0.3], + "eff_pols": [0.4], + } + + m = Matcher(nr, nr_2, settings=settings) + assert m.mismatch == pytest.approx(0.25) + + +def base_single_mode_settings(**overrides): + settings = { + "kind": "single-mode", + "modes-or-pol": "modes", + "modes": [(2, 2)], + "initial_frequency_mm": fmin, + "final_frequency_mm": fmax, + "tlen": len(nr.u), + "dt": 1 / srate, + "M": M, + "resize_factor": 4, + "pad_end_frac": 0.5, + "taper_alpha": 0.2, + "taper_start": 0.05, + "taper": "sigmoid", + "debug": False, + } + settings.update(overrides) + return settings + + +@pytest.mark.parametrize("option", ["cut_longer", "cut_second_waveform"]) +def test_cut_options_run(option): + """ + These options had no coverage at all. + + They used to read 'tmrg, _, _, _ = WaveForm.find_max() - WaveForm.u[0]', + which survives only because find_max returns np.float64 entries: numpy + broadcasts the tuple and subtracts u[0] from all four, the last three being + discarded. The merger time is now taken explicitly, which is equivalent for + a numpy time array and does not depend on that accident. + """ + wf1 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf2 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf1.cut(300) + wf2.cut(500) # different start -> non-zero DeltaT + + m = Matcher(wf1, wf2, settings=base_single_mode_settings(**{option: True})) + assert np.isfinite(m.mismatch) + + +def test_cut_options_are_mutually_exclusive(): + wf1 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf2 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf1.cut(300) + wf2.cut(300) + + with pytest.raises(RuntimeError, match="cannot be used together"): + Matcher( + wf1, + wf2, + settings=base_single_mode_settings( + cut_longer=True, cut_second_waveform=True + ), + ) + + +def test_single_mode_with_cached_h2f(): + """ + The output dict referenced h2.delta_t, but h2 only exists when wf2 is in the + time domain and was not served from the cache -> NameError. + """ + settings = base_single_mode_settings() + + # first pass: populate the cache the way callers are meant to + m0 = Matcher(nr, nr_2, settings=settings) + cache = {"h1f": m0.h1f, "h2f": m0.h2f, "M": M} + + wf1 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf2 = sxs.Waveform_SXS(ID=sxs_id, download=False, ignore_deprecation=True) + wf1.cut(300) + wf2.cut(300) + + m = Matcher(wf1, wf2, settings=base_single_mode_settings(), cache=cache) + assert np.isfinite(m.mismatch) diff --git a/tests/test_sxs.py b/tests/test_sxs.py index a8c58715..a5d7a0ce 100644 --- a/tests/test_sxs.py +++ b/tests/test_sxs.py @@ -3,7 +3,13 @@ """ from PyART.catalogs import sxs +from PyART.waveform import Waveform +import json import os +import sys + +import numpy as np +import pytest mode_keys = ["A", "p", "real", "imag", "z"] @@ -37,13 +43,15 @@ def test_sxs(): assert os.path.exists(f"SXS_BBH_0180/Lev{wf.level}/metadata.json") assert os.path.exists(f"SXS_BBH_0180/Lev{wf.level}/Horizons.h5") - # check that the old folder was removed - cache_dir = os.environ.get("SXSCACHEDIR") - assert cache_dir, "SXSCACHEDIR environment variable is not set." - flds = os.listdir(cache_dir) - for fld in flds: - if fld.startswith("SXS:BBH:0180"): - assert False, f"Old folder {fld} still exists." + # Check that the colon-named folder the sxs module downloads into was + # cleaned up, leaving only SXS_BBH_0180. Look in the directory we asked to + # download into rather than at SXSCACHEDIR: that variable is only set when a + # download actually happens, so asserting on it made this test pass on a + # clean tree and fail on every rerun, once the data is already there. + for fld in os.listdir(opts["path"]): + assert not fld.startswith( + "SXS:BBH:0180" + ), f"Old folder {fld} still exists in {opts['path']}." # check that the modes loaded make sense for mode in wf.hlm.keys(): @@ -62,3 +70,404 @@ def test_sxs(): # check that conversion to LVKNR works wf.to_lvk(modes=[(2, 2)]) + + +############################## +# metadata parsing +############################## +# +# The bundled SXS:BBH:0180 carries dimensionless spins and a float +# eccentricity, so it exercises none of the branches below. These tests drive +# load_metadata from synthetic metadata.json files instead. + +M1 = 0.6 +M2 = 0.4 + + +def base_metadata(**overrides): + """A minimal SXS-like metadata dict, with the keys load_metadata needs.""" + meta = { + "reference_mass1": M1, + "reference_mass2": M2, + "reference_dimensionless_spin1": [0.0, 0.0, 0.3], + "reference_dimensionless_spin2": [0.0, 0.0, -0.2], + "reference_position1": [5.0, 0.0, 0.0], + "reference_position2": [-5.0, 0.0, 0.0], + "remnant_mass": 0.95, + "remnant_dimensionless_spin": [0.0, 0.0, 0.68], + "alternative_names": ["SXS:BBH:0001", "SXS:BBH:0001"], + "reference_eccentricity": 1e-4, + "reference_orbital_frequency": [0.0, 0.0, 0.005], + "reference_time": 250.0, + "initial_ADM_energy": 0.99, + "initial_ADM_linear_momentum": [0.0, 0.0, 0.0], + "initial_ADM_angular_momentum": [0.0, 0.0, 1.1], + } + meta.update(overrides) + return meta + + +def make_sxs_stub(tmp_path, meta): + """A Waveform_SXS pointing at a synthetic metadata.json, without __init__.""" + lev = tmp_path / "Lev4" + lev.mkdir(parents=True, exist_ok=True) + (lev / "metadata.json").write_text(json.dumps(meta)) + + wf = sxs.Waveform_SXS.__new__(sxs.Waveform_SXS) + Waveform.__init__(wf) + wf.sxs_data_path = str(tmp_path) + wf.level = 4 + wf.ID = "0001" + wf.src = "BBH" + wf.order = 2 + wf.ellmax = 4 + wf.nu_rescale = False + wf.cut_N = None + wf.cut_U = 0.0 + wf.basename = "rhOverM_Asymptotic_GeometricUnits_CoM.h5" + return wf + + +def test_dimensionless_spins_are_used_as_is(tmp_path): + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_metadata() + + assert wf.metadata["chi1z"] == pytest.approx(0.3) + assert wf.metadata["chi2z"] == pytest.approx(-0.2) + + +def test_dimensionful_reference_spin_is_normalized_by_mass_squared(tmp_path): + """ + 'reference_spin' is dimensionful and must be divided by the mass squared; + the 'attempt == "reference_"' check could never match the actual key, so + these spins were silently left un-normalized. + """ + S1z, S2z = 0.3 * M1**2, -0.2 * M2**2 + meta = base_metadata( + reference_spin1=[0.0, 0.0, S1z], + reference_spin2=[0.0, 0.0, S2z], + ) + # remove the dimensionless entries so the fallback is used + del meta["reference_dimensionless_spin1"] + del meta["reference_dimensionless_spin2"] + + wf = make_sxs_stub(tmp_path, meta) + wf.load_metadata() + + assert wf.metadata["chi1z"] == pytest.approx(0.3) + assert wf.metadata["chi2z"] == pytest.approx(-0.2) + + +def test_initial_dimensionless_spin_fallback(tmp_path): + meta = base_metadata( + initial_dimensionless_spin1=[0.0, 0.0, 0.5], + initial_dimensionless_spin2=[0.0, 0.0, -0.4], + ) + del meta["reference_dimensionless_spin1"] + del meta["reference_dimensionless_spin2"] + + wf = make_sxs_stub(tmp_path, meta) + wf.load_metadata() + + assert wf.metadata["chi1z"] == pytest.approx(0.5) + assert wf.metadata["chi2z"] == pytest.approx(-0.4) + + +def test_missing_spin_entries_raise(tmp_path): + """hS used to be unbound when no attempt matched -> UnboundLocalError.""" + meta = base_metadata() + del meta["reference_dimensionless_spin1"] + del meta["reference_dimensionless_spin2"] + + wf = make_sxs_stub(tmp_path, meta) + with pytest.raises(KeyError, match="No valid spin entry"): + wf.load_metadata() + + +@pytest.mark.parametrize( + "ecc_str, expected", + [ + ("<1.7e+00", None), # bound of order unity: no information + ("<1e-04", 1e-5), # small bound: effectively circular + ("<3.2e-03", 1e-5), + ], +) +def test_string_eccentricity(tmp_path, ecc_str, expected): + """ + 'if "<" in ecc and "e+00"' had a constant truthy second operand, so every + string eccentricity became None -- including the small bounds that should + be read as quasi-circular. + """ + wf = make_sxs_stub(tmp_path, base_metadata(reference_eccentricity=ecc_str)) + wf.load_metadata() + + assert wf.metadata["e0"] == expected + + +def test_float_eccentricity_is_kept(tmp_path): + wf = make_sxs_stub(tmp_path, base_metadata(reference_eccentricity=5.11e-05)) + wf.load_metadata() + + assert wf.metadata["e0"] == pytest.approx(5.11e-05) + + +############################## +# psi4 loading +############################## + + +def test_load_psi4lm_missing_file_raises_filenotfound(tmp_path): + """ + nr_psi was only assigned when the file existed, but used unconditionally, + giving AttributeError instead of naming the missing file. + """ + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_metadata() + + with pytest.raises(FileNotFoundError, match="psi4 file not found"): + wf.load_psi4lm() + + +def test_load_psi4lm_without_hlm_raises_runtime(tmp_path): + """psi4 times are copied from the hlm time array, which must exist.""" + import h5py + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_metadata() + + # a psi4 file that exists, so we get past the FileNotFoundError above + psi4_name = wf.basename.replace("rhOverM", "rMPsi4") + fname = tmp_path / "Lev4" / psi4_name + with h5py.File(fname, "w") as f: + grp = f.create_group(f"Extrapolated_N{wf.order}.dir") + t = np.linspace(0.0, 100.0, 50) + grp["Y_l2_m2.dat"] = np.column_stack((t, np.sin(t), np.cos(t))) + + assert wf.u is None # hlm never loaded + with pytest.raises(RuntimeError, match="hlm was never loaded"): + wf.load_psi4lm() + + +############################## +# download stdout redirect +############################## + + +def test_download_simulation_restores_stdout_on_failure(tmp_path, monkeypatch): + """ + download_simulation redirects stdout/stderr to devnull around the noisy sxs + calls. Without try/finally, any failure in between left the *whole process* + writing to devnull, silencing everything afterwards. + """ + import sxs as sxsmod + + def boom(*args, **kwargs): + raise RuntimeError("simulated sxs.load failure") + + monkeypatch.setattr(sxsmod, "load", boom) + + wf = sxs.Waveform_SXS.__new__(sxs.Waveform_SXS) + wf.src = "BBH" + wf.level = 4 + wf.order = 2 + + saved_stdout, saved_stderr = sys.stdout, sys.stderr + + with pytest.raises(RuntimeError, match="simulated sxs.load failure"): + wf.download_simulation(ID="0001", path=str(tmp_path), extrapolation_order=2) + + assert sys.stdout is saved_stdout, "stdout left redirected to devnull" + assert sys.stderr is saved_stderr, "stderr left redirected to devnull" + + +############################## +# horizon loading +############################## + + +REMNANT_MASS = 0.95 +REMNANT_CHI = [0.0, 0.0, 0.68] + + +def write_horizons( + tmp_path, + chi1, + chi2, + mass=0.5, + omega=0.05, + radius=5.0, + nt=400, + remnant=True, +): + """ + A synthetic Horizons.h5 for a circular orbit in the xy-plane, so that the + orbital angular momentum points along +z and the spin projections are known + analytically. + + The columns follow the SXS layout: time first, so ChristodoulouMass.dat and + DimensionfulInertialSpinMag.dat are (N,2) while CoordCenterInertial.dat and + chiInertial.dat are (N,4). The spin magnitudes are kept consistent with the + dimensionless spins, |S| = |chi| m^2. + + AhC, the common horizon, exists only after merger, so it is written on its + own later and shorter time array. remnant=False omits it, as happens for + runs that do not merge. + """ + import h5py + + def add(f, obj, t, chi, x, m): + chi = np.tile(np.asarray(chi, dtype=float), (len(t), 1)) + grp = f.create_group(obj) + grp["ChristodoulouMass.dat"] = np.column_stack((t, np.full_like(t, m))) + grp["CoordCenterInertial.dat"] = np.column_stack((t, x)) + grp["chiInertial.dat"] = np.column_stack((t, chi)) + grp["DimensionfulInertialSpinMag.dat"] = np.column_stack( + (t, np.linalg.norm(chi, axis=1) * m**2) + ) + + t = np.linspace(0.0, 100.0, nt) + x1 = np.column_stack( + (radius * np.cos(omega * t), radius * np.sin(omega * t), np.zeros_like(t)) + ) + x2 = -x1 + + fname = tmp_path / "Lev4" / "Horizons.h5" + fname.parent.mkdir(parents=True, exist_ok=True) + with h5py.File(fname, "w") as f: + add(f, "AhA.dir", t, chi1, x1, mass) + add(f, "AhB.dir", t, chi2, x2, mass) + if remnant: + t_rem = np.linspace(100.0, 150.0, nt // 4) + x_rem = np.zeros((len(t_rem), 3)) + add(f, "AhC.dir", t_rem, REMNANT_CHI, x_rem, REMNANT_MASS) + return t + + +def test_load_horizon_reads_dimensionless_spin_vector(tmp_path): + """ + dyn["chi"] must come from chiInertial (the dimensionless spin vector), not + from DimensionfulInertialSpinMag, which is |S|: dimensionful by m^2 and a + magnitude, so it carries no direction. + """ + chi1 = [0.3, 0.0, 0.4] + chi2 = [0.0, 0.1, -0.2] + write_horizons(tmp_path, chi1, chi2) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + d = wf.dyn + + # full 3-vectors, time stripped and kept once in dyn["t"] + assert np.shape(d["chi1"])[1] == 3 + assert np.shape(d["chi2"])[1] == 3 + assert np.shape(d["x1"])[1] == 3 + + assert np.allclose(d["chi1"], chi1) + assert np.allclose(d["chi2"], chi2) + + +def test_load_horizon_time_and_masses(tmp_path): + t = write_horizons(tmp_path, [0.0, 0.0, 0.1], [0.0, 0.0, 0.2], mass=0.4) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + d = wf.dyn + + assert np.allclose(d["t"], t) + assert len(d["m1"]) == len(t) + assert np.allclose(d["m1"], 0.4) + assert np.allclose(d["m2"], 0.4) + + +def test_compute_spins_at_tref_circular_orbit(tmp_path): + """ + For a circular orbit in the xy-plane L is along +z, so the parallel spin + component is chi_z and the perpendicular one is the in-plane norm. + """ + chi1 = [0.3, 0.0, 0.4] + chi2 = [0.0, 0.1, -0.2] + write_horizons(tmp_path, chi1, chi2) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + + chi1_L, chi1_perp, chi2_L, chi2_perp = wf.compute_spins_at_tref(50.0) + + assert chi1_L == pytest.approx(0.4, abs=1e-8) + assert chi1_perp == pytest.approx(0.3, abs=1e-8) + assert chi2_L == pytest.approx(-0.2, abs=1e-8) + assert chi2_perp == pytest.approx(0.1, abs=1e-8) + + +def test_load_horizon_spin_magnitudes(tmp_path): + """ + The dimensionful spin magnitude is loaded alongside the dimensionless spin + vector; the two are related by chi = S/m^2. + """ + chi1 = [0.3, 0.0, 0.4] # |chi1| = 0.5 + chi2 = [0.0, 0.1, -0.2] + mass = 0.5 + write_horizons(tmp_path, chi1, chi2, mass=mass) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + d = wf.dyn + + # a magnitude: one scalar per time + assert np.shape(d["S1_mag"]) == np.shape(d["t"]) + + assert np.allclose(d["S1_mag"], np.linalg.norm(chi1) * mass**2) + assert np.allclose(d["S1_mag"] / d["m1"] ** 2, np.linalg.norm(d["chi1"], axis=1)) + assert np.allclose(d["S2_mag"] / d["m2"] ** 2, np.linalg.norm(d["chi2"], axis=1)) + + +def test_load_horizon_remnant(tmp_path): + """ + AhC is the common horizon: it forms at merger, so it carries its own time + array rather than sharing dyn["t"]. + """ + write_horizons(tmp_path, [0.0, 0.0, 0.1], [0.0, 0.0, 0.2]) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + d = wf.dyn + + assert np.allclose(d["m_remnant"], REMNANT_MASS) + assert np.allclose(d["chi_remnant"], REMNANT_CHI) + assert np.allclose( + d["S_remnant_mag"], np.linalg.norm(REMNANT_CHI) * REMNANT_MASS**2 + ) + assert np.shape(d["chi_remnant"])[1] == 3 + assert np.shape(d["x_remnant"])[1] == 3 + + # its own, later time array + assert len(d["t_remnant"]) != len(d["t"]) + assert d["t_remnant"][0] >= d["t"][-1] + + +def test_load_horizon_without_remnant(tmp_path): + """Runs that do not merge have no common horizon; that is not an error.""" + write_horizons(tmp_path, [0.0, 0.0, 0.1], [0.0, 0.0, 0.2], remnant=False) + + wf = make_sxs_stub(tmp_path, base_metadata()) + wf.load_horizon() + d = wf.dyn + + assert "chi1" in d # the binary is still loaded + assert "m_remnant" not in d + assert "chi_remnant" not in d + + +def test_download_simulation_requires_a_path(): + """ + path is both the download destination and the sxs cache directory. It + defaults to None, which used to die further down with a cryptic + "'NoneType' object has no attribute 'endswith'". + """ + wf = sxs.Waveform_SXS.__new__(sxs.Waveform_SXS) + wf.src = "BBH" + wf.level = 4 + wf.order = 2 + + with pytest.raises(ValueError, match="needs a path"): + wf.download_simulation(ID="0001", path=None) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..c9cca9c4 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,172 @@ +""" +Test the numerical helpers in PyART.utils.utils: +- D1 finite differencing +- upoly_fits polynomial extrapolation +- print_dict_comparison reporting +""" + +import logging + +import numpy as np +import pytest + +from PyART.utils.utils import D1, print_dict_comparison, upoly_fits + +############################## +# D1 +############################## + + +@pytest.mark.parametrize("order", [1, 2, 4]) +def test_D1_does_not_mutate_input(order): + """ + D1 must treat its input as read-only: callers pass waveform arrays that are + reused after differentiation. + """ + x = np.linspace(0.0, 1.0, 50) + f = np.sin(x) + f_ref = f.copy() + + D1(f, x, order) + + assert np.array_equal(f, f_ref), f"D1(order={order}) modified its input array" + + +@pytest.mark.parametrize("order", [1, 2, 4]) +def test_D1_exact_on_linear_function(order): + """ + Every differencing order is exact for a linear function, endpoints included. + """ + slope = 2.0 + x = np.linspace(0.0, 5.0, 40) + f = slope * x + 3.0 + + df = D1(f, x, order) + + assert np.allclose( + df, slope + ), f"D1(order={order}) is not exact on a linear function" + + +def test_D1_order1_fills_last_point(): + """ + Forward differencing has no stencil at the last point, so it is copied from + the previous one. It must not be left at zero. + """ + x = np.linspace(0.0, 1.0, 20) + f = np.exp(x) + + df = D1(f, x, order=1) + + assert df[-1] == df[-2] + assert df[-1] != 0.0 + + +@pytest.mark.parametrize("order, tol", [(2, 1e-4), (4, 1e-8)]) +def test_D1_converges_on_sine(order, tol): + """Interior accuracy against a known derivative.""" + x = np.linspace(0.0, 2.0 * np.pi, 400) + f = np.sin(x) + + df = D1(f, x, order) + + # skip the boundary stencils, which are lower order + assert np.allclose(df[5:-5], np.cos(x)[5:-5], atol=tol) + + +def test_D1_rejects_nonuniform_grid(): + x = np.array([0.0, 1.0, 3.0, 6.0]) + f = x**2 + with pytest.raises(RuntimeError, match="not uniformly spaced"): + D1(f, x, order=2) + + +############################## +# upoly_fits +############################## + + +def _upoly_data(): + """y = 1 + 2u + 3u^2 in u=1/r, so the r->inf extrapolation is exactly 1.""" + r = np.linspace(100.0, 10.0, 60) # inward-moving + u = 1.0 / r + y = 1.0 + 2.0 * u + 3.0 * u**2 + return r, y + + +def test_upoly_fits_extrapolates_to_known_value(): + r, y = _upoly_data() + out = upoly_fits(r, y, nmin=2, nmax=4, direction="in") + assert np.isclose(out["extrap"], 1.0, atol=1e-8) + + +def test_upoly_fits_n_extract_selects_fit_order(): + """'extrap' must be the entry of 'extrap_vec' belonging to n_extract.""" + r, y = _upoly_data() + nmin, nmax = 2, 5 + for n_extract in range(nmin, nmax + 1): + out = upoly_fits( + r, y, nmin=nmin, nmax=nmax, n_extract=n_extract, direction="in" + ) + idx = list(out["fit_orders"]).index(n_extract) + assert out["extrap"] == out["extrap_vec"][idx] + + +def test_upoly_fits_n_extract_defaults_to_nmax(): + r, y = _upoly_data() + nmin, nmax = 2, 5 + default = upoly_fits(r, y, nmin=nmin, nmax=nmax, direction="in") + explicit = upoly_fits(r, y, nmin=nmin, nmax=nmax, n_extract=nmax, direction="in") + assert default["extrap"] == explicit["extrap"] + + +@pytest.mark.parametrize("n_extract", [0, 1, 6, 99]) +def test_upoly_fits_rejects_n_extract_out_of_range(n_extract): + """ + An out-of-range n_extract used to fall through the loop and raise + UnboundLocalError; it must be reported as a ValueError instead. + """ + r, y = _upoly_data() + with pytest.raises(ValueError, match="n_extract"): + upoly_fits(r, y, nmin=2, nmax=5, n_extract=n_extract, direction="in") + + +def test_upoly_fits_rejects_nmin_above_nmax(): + r, y = _upoly_data() + with pytest.raises(ValueError, match="nmin>nmax"): + upoly_fits(r, y, nmin=5, nmax=2, direction="in") + + +############################## +# print_dict_comparison +############################## + + +def test_print_dict_comparison_reports_issues_only_when_nested_dicts_differ(caplog): + """ + The 'issues with ...' line must appear when the nested dicts differ, and + only then. + """ + base = {"sub": {"a": 1, "b": 2}} + same = {"sub": {"a": 1, "b": 2}} + diff = {"sub": {"a": 1, "b": 999}} + + with caplog.at_level(logging.INFO): + print_dict_comparison(base, same) + assert "issues with" not in caplog.text, "reported issues for identical dicts" + + caplog.clear() + with caplog.at_level(logging.INFO): + print_dict_comparison(base, diff) + assert "issues with" in caplog.text, "failed to report issues for differing dicts" + + +def test_print_dict_comparison_excluded_keys(caplog): + d1 = {"sub": {"a": 1}, "skipme": {"x": 1}} + d2 = {"sub": {"a": 1}, "skipme": {"x": 2}} + + with caplog.at_level(logging.INFO): + print_dict_comparison(d1, d2, excluded_keys=["skipme"]) + + assert "issues with" not in caplog.text + assert "skipme" not in caplog.text diff --git a/tests/test_waveform.py b/tests/test_waveform.py index cf586dff..9317be51 100644 --- a/tests/test_waveform.py +++ b/tests/test_waveform.py @@ -5,6 +5,7 @@ from PyART import waveform from PyART.utils import wf_utils import copy +import inspect import numpy as np import pytest @@ -113,6 +114,36 @@ def test_find_max_variants_and_errors(): assert t_mrg_psi4 == pytest.approx(3.0) +def test_find_max_first_max_after_t_without_peak_after_umin(): + """ + If no peak lies after umin the loop used to fall through and silently + return the last peak, i.e. a merger time *before* the requested umin. + """ + wf = waveform.Waveform() + wf._u = np.arange(6, dtype=float) + amp = np.array([0.0, 1.0, 0.0, 2.0, 0.0, 0.5]) + phase = np.linspace(0.0, 1.0, len(amp)) + wf._hlm[(2, 2)] = wf_utils.get_multipole_dict(amp * np.exp(-1j * phase)) + + with pytest.raises(ValueError, match="No peak found after"): + wf.find_max(kind="first-max-after-t", umin=4.0) + + +@pytest.mark.parametrize("kind", ["first-max-after-t", "last-peak", "global"]) +def test_find_max_without_any_peak(kind): + """ + A flat amplitude has no peaks; every kind must say so rather than raise + NameError on an unbound index. + """ + wf = waveform.Waveform() + wf._u = np.arange(6, dtype=float) + amp = np.ones(6) + wf._hlm[(2, 2)] = wf_utils.get_multipole_dict(amp * np.exp(-1j * wf.u)) + + with pytest.raises(ValueError, match="No peaks found"): + wf.find_max(kind=kind, umin=0.0) + + def test_compute_dothlm_and_psi4lm(): # mock waveform @@ -285,3 +316,182 @@ def test_waveform_plots(): assert len(ax) == len(wf.hlm) # should have one subplot per mode ax = wf.plot_modes(modes=[(2, 2), (3, 3)], show=False) assert len(ax) == 2 # should have two subplots + + +############################## +# to_frequency +############################## + + +def make_td_waveform(f0=2.0, dt=1.0 / 64.0, tmax=6.0): + """A time-domain waveform with hp/hc a quarter cycle out of phase.""" + wf = waveform.Waveform() + u = np.arange(0.0, tmax, dt) + wf._u = u + wf._t = u.copy() + wf._hp = np.sin(2 * np.pi * f0 * u) + wf._hc = np.cos(2 * np.pi * f0 * u) + wf._domain = "Time" + wf._units = "geom" + return wf, f0, dt + + +@pytest.mark.parametrize("taper", [False, True]) +def test_to_frequency_returns_arrays(taper): + """ + windowing() returns an (array, wfact) tuple; assigning it straight to + _hp/_hc left a tuple where an array was expected. + """ + wf, f0, dt = make_td_waveform() + + wf.to_frequency(taper=taper, pad=False) + + assert wf.domain == "Freq" + for arr in (wf.f, wf.hp, wf.hc): + assert isinstance(arr, np.ndarray), "to_frequency did not produce an array" + assert len(wf.f) == len(wf.hp) == len(wf.hc) + + +def test_to_frequency_pad_lengths_and_spectrum(): + """Padding must reach the next power of two and keep u/hp/hc consistent.""" + wf, f0, dt = make_td_waveform() + n_in = len(wf.u) + + wf.to_frequency(taper=False, pad=True) + + # seglen = nextpow2(6.0) = 8 -> 8/dt = 512 samples + n_pad = int(8.0 / dt) + assert len(wf.u) == n_pad + assert len(wf.t) == n_pad + assert n_pad > n_in + assert len(wf.f) == n_pad // 2 + 1 # rfft length + assert len(wf.hp) == len(wf.f) + assert len(wf.hc) == len(wf.f) + + # the spectrum still peaks at the injected frequency + df = wf.f[1] - wf.f[0] + assert wf.f[np.argmax(np.abs(wf.hp))] == pytest.approx(f0, abs=df) + assert wf.f[np.argmax(np.abs(wf.hc))] == pytest.approx(f0, abs=df) + + +def test_to_frequency_hc_is_not_a_copy_of_hp(): + """The pad branch built hc out of self.hp, silently duplicating hp.""" + wf, f0, dt = make_td_waveform() + + wf.to_frequency(taper=False, pad=True) + + assert not np.array_equal(wf.hp, wf.hc), "hc is a copy of hp" + # sin and cos differ by a 90 degree phase at the peak + ipk = np.argmax(np.abs(wf.hp)) + dphase = np.angle(wf.hc[ipk]) - np.angle(wf.hp[ipk]) + dphase = (dphase + np.pi) % (2 * np.pi) - np.pi + assert abs(dphase) == pytest.approx(np.pi / 2, abs=0.2) + + +############################## +# to_geom / to_SI +############################## + + +def make_unit_waveform(): + wf = waveform.Waveform() + u = np.linspace(1.0, 10.0, 20) + z = np.exp(-1j * 0.2 * u) + wf._u = u.copy() + wf._t = u.copy() + wf._t_psi4 = u.copy() + wf._hlm[(2, 2)] = wf_utils.get_multipole_dict(z) + wf._domain = "Time" + wf._units = "geom" + return wf + + +def test_to_SI_and_to_geom_roundtrip_with_t_psi4(): + """ + The attribute lists named the read-only 't_psi4' property instead of the + '_t_psi4' attribute, so setattr raised AttributeError whenever psi4 times + were set. + """ + wf = make_unit_waveform() + u0 = wf.u.copy() + t_psi40 = wf.t_psi4.copy() + z0 = wf.hlm[(2, 2)]["z"].copy() + M, distance = 50.0, 100.0 + + wf.to_SI(M, distance) + assert wf.units == "SI" + assert wf.t_psi4 is not None + assert not np.allclose(wf.t_psi4, t_psi40), "t_psi4 was not converted" + + wf.to_geom(M, distance) + assert wf.units == "geom" + assert np.allclose(wf.u, u0) + assert np.allclose(wf.t_psi4, t_psi40) + assert np.allclose(wf.hlm[(2, 2)]["z"], z0) + + +def test_to_SI_scales_t_psi4_like_u(): + wf = make_unit_waveform() + ratio0 = wf.t_psi4 / wf.u + + wf.to_SI(50.0, 100.0) + + assert np.allclose(wf.t_psi4 / wf.u, ratio0) + + +def test_unit_conversion_rejects_wrong_units(): + wf = make_unit_waveform() + with pytest.raises(RuntimeError, match="Already using geom"): + wf.to_geom(50.0, 100.0) + wf.to_SI(50.0, 100.0) + with pytest.raises(RuntimeError, match="Already using SI"): + wf.to_SI(50.0, 100.0) + + +############################## +# mutable default arguments +############################## + + +@pytest.mark.parametrize( + "func", + [waveform.Waveform.integrate_data, waveform.WaveIntegrated.__init__], + ids=["integrate_data", "WaveIntegrated.__init__"], +) +def test_no_mutable_dict_defaults(func): + """ + Both fill missing keys into integr_opts, so a dict default would be shared + state across every call that relies on it. + """ + for name, par in inspect.signature(func).parameters.items(): + assert not isinstance( + par.default, dict + ), f"{func.__qualname__} has a mutable dict default for '{name}'" + + +def test_integrate_data_does_not_mutate_caller_opts(): + wf = waveform.Waveform() + t = np.linspace(0.0, 40.0, 512) + z = np.exp(-1j * 0.3 * t) * np.exp(-(((t - 20.0) / 8.0) ** 2)) + wf._psi4lm[(2, 2)] = wf_utils.get_multipole_dict(z) + + integr_opts = {"method": "FFI", "f0": 0.01} + before = copy.deepcopy(integr_opts) + + wf.integrate_data(t_psi4=t, radius=100.0, integr_opts=integr_opts, M=1.0) + + assert integr_opts == before, "integrate_data mutated the caller's integr_opts" + assert (2, 2) in wf.hlm + + +def test_integrate_data_default_opts_are_not_shared(): + """Two successive default calls must behave identically.""" + + def run(): + wf = waveform.Waveform() + t = np.linspace(0.0, 40.0, 512) + z = np.exp(-1j * 0.3 * t) * np.exp(-(((t - 20.0) / 8.0) ** 2)) + wf._psi4lm[(2, 2)] = wf_utils.get_multipole_dict(z) + return wf.integrate_data(t_psi4=t, radius=100.0, M=1.0) + + assert run() == run()