From f92b3ea417098df331442c024f56a493338f910d Mon Sep 17 00:00:00 2001 From: Aasim Z Jan <58118417+AasimZahoor@users.noreply.github.com> Date: Sun, 7 Jun 2026 10:40:58 -0500 Subject: [PATCH 1/5] plot_RIFT.py: made changes so that it handles subdags better. Earlier version was optimized for basic iteration structure. Also, the diagnostics file is now thorough. --- .../Code/bin/plot_RIFT.py | 1200 +++++++++++------ 1 file changed, 757 insertions(+), 443 deletions(-) diff --git a/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py b/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py index 9ea92f30..45cb64d7 100755 --- a/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py +++ b/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py @@ -1,7 +1,9 @@ #!/usr/bin/env python """This code is meant to check the health of a RIFT run as it progresses and after it has finished.""" # TO DO: plot_cip_max_lnL needs to be robust. -# Usage: In the rundirectory, run plot_RIFT.py --precessing or plot_RIFT.py --precessing --eccentric or plot_RIFT.py --eccentric --LISA +# JSD convergence for subdags +# Two all.nets when subdags being used +# Usage: In the rundirectory, run plot_RIFT.py --precessing or plot_RIFT.py --rundir-path `pwd`/rundir --save-path `pwd` --precessing ########################################################################################### # Import ########################################################################################### @@ -42,18 +44,22 @@ # Path and Configuration Setup ########################################################################################### parser = ArgumentParser() -parser.add_argument("--path", default = os.getcwd(), help = "path to run directory") +parser.add_argument("--rundir-path", default = os.getcwd(), type = str, help = "path to run directory. Default is current directory.") +parser.add_argument("--save-path", default = os.getcwd(), type = str, help = "Where to create the plots folder. Default is current directory.") parser.add_argument("--LISA", action = "store_true", help = "Use this argument if analyzing a LISA run") parser.add_argument("--eccentric", action = "store_true", help = "Use this argument if the run has eccentricity and meanPerAno") parser.add_argument("--precessing", action = "store_true", help = "Use this argument if the run is precessing") parser.add_argument("--non-spinning", action = "store_true", help = "Use this argument if the run is non-spinning") parser.add_argument("--truth-file", default = None, type = str, help = "Path to the truth file. If not provided, the code will search for a truth file in the parent directory or the frames directory. The frames directory is expected to be located in the parent directory.") opts = parser.parse_args() -path = opts.path + +path = opts.rundir_path +save_path = opts.save_path LISA = opts.LISA eccentricity = opts.eccentric precessing = opts.precessing non_spinning = opts.non_spinning + # Print kind of analysis messages = [] if eccentricity: @@ -99,24 +105,38 @@ "JSD": {}, "JSD_3": {}, } - +JSD_threshold = 0.005 ########################################################################################### # Functions ########################################################################################### def get_lnL_cut_points(all_net_path, lnL_cut=15, error_threshold=0.4, composite=False, provide_max_lnL_point=False): """ - Analyzes the lnL values from an all.net file to find high likelihood points - and assess their Monte Carlo error. + Computes high-likelihood statistics from an all.net or composite file. Args: - all_net_path (str): Path to the all.net file. - lnL_cut (float): Cutoff for determining high likelihood points. - error_threshold (float): Maximum allowed error for high likelihood points. - composite (bool): is the file a composite file or an all.net file. + all_net_path (str): Path to the all.net or composite file. + lnL_cut (float): Threshold used to define high-likelihood points. Default is 15. + error_threshold (float): Maximum allowed Monte Carlo error. Default is 0.4 + composite (bool): Whether the input file is a composite file. + provide_max_lnL_point (bool): If True, return the maximum-likelihood + sample and its lnL value. Returns: - tuple: Maximum lnL value (rounded) and the number of high likelihood points with low error. + tuple: + If provide_max_lnL_point is True: + - Maximum-likelihood sample. + - Maximum lnL value. + + If composite is True: + - Global maximum lnL. + - Number of high-likelihood points in the composite. + - Maximum lnL in the composite file. + - Total number of points. + + Otherwise: + - Maximum lnL value. + - Number of high-likelihood points. """ # Load data from all.net file data = np.loadtxt(all_net_path) @@ -175,11 +195,11 @@ def get_lnL_cut_points(all_net_path, lnL_cut=15, error_threshold=0.4, composite= if composite: max_lnL_composite = max_lnL max_lnL = run_diagnostics["max_lnL"] - return np.round(max_lnL, 3), no_points, np.round(max_lnL_composite, 2), total_points + return np.round(max_lnL, 2), no_points, np.round(max_lnL_composite, 2), total_points if not(composite): run_diagnostics.update({ "total_lnL_evaluations":total_points, - "max_lnL": np.round(max_lnL, 3), + "max_lnL": np.round(max_lnL, 2), "high_lnL_points": no_points, "high_lnL_points_with_large_error": high_lnL_points - no_points, "total_high_lnL_points": high_lnL_points, @@ -195,28 +215,34 @@ def get_lnL_cut_points(all_net_path, lnL_cut=15, error_threshold=0.4, composite= def lnL_at_credible_interval(ci=0.99, run_diagnostics=run_diagnostics, return_ndim=False): """ Computes the log-likelihood corresponding to the boundary of a - specified credible interval, assuming a multivariate Gaussian posterior. + credible interval, assuming a multivariate Gaussian posterior. Args: - ci (float): Credible interval enclosed by the likelihood contour - (e.g., 0.90, 0.95, 0.99). - run_diagnostics (dict): Dictionary containing run diagnostics, - including the maximum log-likelihood stored under the key - ``"max_lnL"``. + ci (float): Credible interval enclosed by the likelihood contour. + Default is 0.99. + run_diagnostics (dict): Dictionary containing run diagnostics. + return_ndim (bool): If True, return the dimensionality of the + inferred parameter space instead of the log-likelihood. Returns: - - Log-likelihood at the credible interval boundary (float) - """ + float: + If return_ndim is True: + - Number of dimensions used in the calculation. + + Otherwise: + - Log-likelihood at the credible interval boundary. +""" if not (0 < ci < 1): raise ValueError("ci must be between 0 and 1, e.g. 0.99") ndim = 4 if opts.non_spinning: ndim = 2 - elif opts.precessing: - ndim += 4 - elif opts.eccentric: - ndim += 2 + else: + if opts.precessing: + ndim += 4 + if opts.eccentric: + ndim += 2 if return_ndim: return ndim max_lnL = run_diagnostics["max_lnL"] @@ -227,40 +253,44 @@ def lnL_at_credible_interval(ci=0.99, run_diagnostics=run_diagnostics, return_nd return lnL -def create_plots_folder(base_dir_path): +def create_plots_folder(save_path): """ - Creates a 'plots' folder in the specified base directory if it does not already exist. + Creates the plots directory structure if it does not already exist. Args: - base_dir_path (str): Path to the base directory where the 'plots' folder will be created. + save_path (str): Path where the plots directory will be created. + + Returns: + None """ - if not(os.path.exists(base_dir_path + "/plots")): - print(f"--> plots folder does not exist. Creating one in {base_dir_path}") - os.mkdir(base_dir_path + "/plots") - os.mkdir(base_dir_path + "/plots/histograms") - os.mkdir(base_dir_path + "/plots/corner_plots") - os.mkdir(base_dir_path + "/plots/1_D_plots") + if not(os.path.exists(save_path + "/plots")): + print(f"--> plots folder does not exist. Creating one in {save_path}") + os.mkdir(save_path + "/plots") + os.mkdir(save_path + "/plots/histograms") + os.mkdir(save_path + "/plots/corner_plots") + os.mkdir(save_path + "/plots/1_D_plots") else: - print(f"--> plots folder exists, saving plots in directory {base_dir_path}/plots") + print(f"--> plots folder exists, saving plots in directory {save_path}/plots") def get_chirpmass_massratio_eta_totalmass_from_componentmasses(m1, m2): """ - Computes chirp mass, mass ratio, symmetric mass ratio, and total mass from component masses. + Computes chirp mass, mass ratio, symmetric mass ratio, and total mass + from component masses. Args: - m1 (array): Array of primary masses. - m2 (array): Array of secondary masses. + m1 (array): Primary masses. + m2 (array): Secondary masses. Returns: - tuple: A tuple containing: - - Chirp mass (array) - - Mass ratio (array) - - Symmetric mass ratio (array) - - Total mass (array) + tuple: + - Chirp masses. + - Mass ratios (q = m2 / m1). + - Symmetric mass ratios. + - Total masses. """ return np.array((m1*m2)**(3/5) / (m1+m2)**(1/5)).reshape(-1,1), np.array(m2/m1).reshape(-1,1), np.array((m1*m2) / (m1+m2)**(2)).reshape(-1,1), np.array(m1+m2).reshape(-1,1) -def get_index_for_parameter(parameter): +def get_index_for_parameter(parameter, extrinsic = False): """ Retrieves the index corresponding to a given parameter name. @@ -270,57 +300,43 @@ def get_index_for_parameter(parameter): Returns: int or None: The index of the parameter if found, otherwise None. """ - # m1 m2 a1x a1y a1z a2x a2y a2z mc eta indx Npts ra dec tref phiorb incl psi dist p ps lnL mtotal q eccentricity meanPerAno - parameter_indices = { - "mc": 8, - "mtot": 22, - "a1x":2, - "s1x":2, - "a1y":3, - "s1y":3, - "a1z": 4, - "s1z": 4, - "a2x":5, - "s2x":5, - "a2y":6, - "s2y":6, - "a2z": 7, - "s2z": 7, - "eta": 9, - "m1": 0, - "m2": 1, - "q": 23, - "dec": 13, - "ra": 12, - "eccentricity":24, - "meanPerAno":25, - } - - return parameter_indices.get(parameter, None) # Return None if parameter is not found + rift_parameters = ["m1", "m2", "s1x", "s1y", "s1z", "s2x", "s2y", "s2z", "mc", "eta", "indx", + "Npts", "ra", "dec", "tref", "phiorb", "incl", "psi", "dist", "p", "ps", + "lnL", "mtot", "q", "eccentricity", "meanPerAno"] + if extrinsic: + rift_parameters = ["m1", "m2", "s1x", "s1y", "s1z", "s2x", "s2y", "s2z", "mc", "eta", "ra", "dec", "time", "phiorb", "incl", "psi", "distance", "Npts", "lnL", "p", "ps", "neff", + "mtot", "q", "chi_eff", "chi_p", "m1_source", "m2_source", "mc_source", "mtotal_source", "redshift", "eccentricity", "meanPerAno"] + return rift_parameters.index(parameter) if parameter in rift_parameters else None def get_sample_statistics(samples): """ - Computes statistical percentiles (5th, 50th, and 95th) for a given sample set. + Computes the 5th, 50th, and 95th percentiles of a sample set. + Args: - samples (array-like): A list or NumPy array of samples. + samples (array-like): Input samples. Returns: - sample statistics (numpy.ndarray): An array containing the 16th, 50th, and 84th percentiles - of the input samples. + numpy.ndarray: + - 5th percentile. + - 50th percentile (median). + - 95th percentile. """ return np.percentile(samples, [5,50,95]) def get_combination_from_mass_and_spin(posterior, parameter): """ - Computes spin combinations from posterior samples. + Computes spin combinations and derived spin parameters from posterior + samples. Args: posterior (numpy.ndarray): Posterior samples. - parameter (str): One of - 'chi1', 'chi2', 'chi_eff', 'chi_minus', 'chi_p' + parameter (str): Parameter to compute. Supported values are + 'chi1', 'chi2', 'chi_eff', 'chiMinus', 'chi_p', + 'chi1_perp', and 'chi2_perp'. Returns: - numpy.ndarray: Requested spin combination. + numpy.ndarray: + Requested spin combination evaluated for all samples. """ parameter_m1 = get_index_for_parameter("m1") @@ -394,12 +410,15 @@ def find_posteriors_in_main(path_to_main_folder, limit_iterations=None): Finds and sorts posterior sample files in the main folder. Args: - path_to_main_folder (str): Path to the main folder containing posterior sample files. - limit_iterations (int, optional): Number of files to limit the results to. + path_to_main_folder (str): Path to the main folder containing + posterior sample files. + limit_iterations (int, optional): Number of posterior files to + return. Returns: - posteriors (list of str): Sorted list of paths to posterior sample files. - indices (numpy array): Indices of the selected files. + tuple: + - Sorted posterior file paths. + - Corresponding iteration indices. """ posteriors_in_main = glob.glob(path_to_main_folder + "/posterior_samples*") posteriors_in_main.sort(key = os.path.getctime) # sort them according to creation time @@ -409,25 +428,47 @@ def find_posteriors_in_main(path_to_main_folder, limit_iterations=None): return np.array(posteriors_in_main, dtype = str)[index], index + 1 return posteriors_in_main, np.arange(len(posteriors_in_main)) + 1 -def find_posteriors_in_sub(path_to_main_folder, limit_iterations = None): +def find_posteriors_in_sub(path_to_main_folder, limit_iterations = None, return_subdag_folder_path=False): """ - Finds posterior sample files in the sub-directory specified by the path. + Finds posterior sample files in the subdag directory. Args: - path_to_main_folder (str): Path to the main folder containing sub-directory with posterior files. - limit_iterations (int, optional): Number of files to limit the results to. + path_to_main_folder (str): Path to the main folder containing + iteration_*_cip directories. + limit_iterations (int, optional): Number of posterior files to + return. + return_subdag_folder_path (bool): If True, return the path to the + subdag directory containing posterior samples. Returns: - posteriors (list of str): List of paths to posterior sample files in the sub-directory. - indices (numpy array): Indices of the selected files. + If return_subdag_folder_path is True: + str or None: + Path to the subdag directory, or None if no posterior + samples are found. + + Otherwise: + tuple: + - Posterior file paths. + - Iteration indices. """ + if return_subdag_folder_path: + subdag_folder_path = None + subdag_dirs = glob.glob(path_to_main_folder + "/iteration_*_cip") + for d in subdag_dirs: + posteriors = glob.glob(f"{d}/posterior_samples*.dat") + if len(posteriors) > 0: + subdag_folder_path = d + break + return subdag_folder_path posteriors_in_subdag, iterations = find_posteriors_in_main(path_to_main_folder + "/iteration*cip*") + + if limit_iterations: index = np.linspace(0, len(posteriors_in_subdag)-1, limit_iterations) index = np.array(index, dtype=int) return np.array(posteriors_in_subdag, dtype = str)[index], index + 1 else: - return posteriors_in_subdag, np.arange(len(posteriors_in_subdag)) + return posteriors_in_subdag, np.arange(len(posteriors_in_subdag)) + 1 def calculate_JS_divergence(data1, data2): """ @@ -465,63 +506,27 @@ def calc_median_error(jsvalues, quantiles=(0.16, 0.84)): return calculate_js(data1, data2) -def plot_high_likelihood_expoloration(path_to_main_folder): - """ - Plots high likelihood points over iterations. - - Args: - path_to_main_folder (str): Path to the main folder containing the composite files. - """ - print("\n--> Plotting likelihood exploration.") - run_diagnostics["composite_information"] = {} - fig, ax = plt.subplots() - ax.set_xlabel("iteration") - ax.set_ylabel("high lnL points") - ax.set_title(f"Total high lnL points = {run_diagnostics['high_lnL_points']}, max_lnL = {run_diagnostics['max_lnL']}") - collect_data = [] - collect_iter = [] - print("iteration, max lnL (global), high lnL points, max lnL(iteration), total lnL points(iteration)") - for iteration in np.arange(0, run_diagnostics["latest_iteration"]+1, 1): - run_diagnostics["composite_information"][iteration] = {} - try: - max_lnL, no_points, max_lnL_composite, total_points = get_lnL_cut_points(f"{path_to_main_folder}/consolidated_{iteration}.composite", composite=True) - except Exception as e: - print(f"Error loading file {path_to_main_folder}/consolidated_{iteration}.composite: {e}") - continue - print(iteration, max_lnL, no_points, max_lnL_composite, total_points) - percent_high_lnL_points = np.round(no_points/total_points*100, 2) - collect_data.append(no_points) - collect_iter.append(iteration) - ax.scatter(iteration, no_points, label = f"{max_lnL_composite} ({percent_high_lnL_points})", s=25) - run_diagnostics["composite_information"][iteration].update({ - "max_lnL":max_lnL_composite, - "high_lnL_points":no_points, - "percent_high_lnL_points": percent_high_lnL_points}) - ax.grid(alpha=0.4) - ax.plot(collect_iter, collect_data, color = "black", linestyle = "--", linewidth = 1.5, alpha = 0.5) - ax.set_xticks(np.arange(0, run_diagnostics["latest_iteration"]+1, 1)) - ax.legend(loc="upper left") - fig.savefig(path+f"/plots/Likelihood_exploration_plot.png", bbox_inches='tight') - plt.close(fig) - -def plot_neff_data(path_to_main_folder): +def plot_neff_data(path_to_main_folder, plot_title): """ Plot effective number of samples (neff) data from CIP iterations. Args: - path_to_main_folder (str): Path to the main folder containing CIP iteration subfolders. + path_to_main_folder (str): Path to the folder containing CIP iteration subfolders. + plot_title (str): suffix for the plots and filenames. """ - print("\n--> Plotting n-eff for CIP.") + print(f"\n--> Plotting n-eff for CIP for {plot_title} iterations.") # find CIP folders - cip_iteration_folders= glob.glob(path_to_main_folder + "/iteration*cip*") + cip_iteration_folders = glob.glob(path_to_main_folder + "/iteration*cip*") fig, ax = plt.subplots() ax.set_xlabel("iteration") ax.set_ylabel("neff") - iterations=np.arange(len(cip_iteration_folders)-1) # last folders don't usually have anything + iterations = np.arange(len(cip_iteration_folders)) + run_diagnostics[plot_title] = {} + run_diagnostics[plot_title]["CIP_neff"] = {} + # read requested neff from CIP sub files - run_diagnostics["CIP_neff"] = {} - for i in range(3): + for i in range(4): filename = f"{path}/CIP_worker{i}.sub" try: with open(filename, "r") as f: @@ -529,15 +534,8 @@ def plot_neff_data(path_to_main_folder): matches = re.findall(r"--n-eff\s+([+-]?\d+(?:\.\d+)?)", content) if matches: neff_value = float(matches[-1]) - ax.axhline( - y=neff_value, - linestyle="--", - color=default_colors[i], - alpha=1.0, - linewidth=1.0, - label=f"worker {i} neff" - ) - run_diagnostics["CIP_neff"][f"CIP_worker{i}"] = np.round(neff_value, 2) + ax.axhline(y=neff_value, linestyle="--", color=default_colors[i], alpha=1.0, linewidth=1.0, label=f"worker {i} neff") + run_diagnostics[plot_title]["CIP_neff"][f"CIP_worker{i}"] = np.round(neff_value, 2) else: continue except Exception as e: @@ -545,7 +543,7 @@ def plot_neff_data(path_to_main_folder): ax.legend(loc="upper left") # read neff achived for each iteration from each instance of CIP - run_diagnostics["CIP_neff_achieved"] = {} + run_diagnostics[plot_title]["CIP_neff_achieved"] = {} for n in iterations: i = path_to_main_folder + f"/iteration_{n}_cip" # remove existing data file because I append @@ -555,51 +553,120 @@ def plot_neff_data(path_to_main_folder): os.system(cmd) # calculate neff statistics try: + if os.path.getsize(f"{i}/neff_data.txt") == 0: + continue tmp_ESS_data=np.loadtxt(f"{i}/neff_data.txt", usecols=[2]) - low, avg, high = np.percentile(tmp_ESS_data, [2.5,50,97.5]) - low_1_std, avg, high_1_std = np.percentile(tmp_ESS_data, [16,50,84]) + low, avg, high = np.percentile(tmp_ESS_data, [2.5,50,97.5]) # 2 std + low_1_std, avg, high_1_std = np.percentile(tmp_ESS_data, [16,50,84]) # 1 std mini, maxi = np.min(tmp_ESS_data), np.max(tmp_ESS_data) ax.plot(iterations[n], mini, marker="x", color="black") ax.plot(iterations[n], maxi, marker="x", color="black") - print(f"neff detail iteration = {iterations[n]}: Average={avg:0.2f}, low std={low:0.2f}, high std={high:0.2f}") + print(f"n-eff summary | iteration {iterations[n]:>2} | median = {avg:8.2f} | 95% interval = [{low:8.2f}, {high:8.2f}]") ax.errorbar(iterations[n], avg, yerr=np.array([avg-low,high-avg]).reshape(-1, 1), color = "royalblue", ecolor = "red", fmt ='o') ax.errorbar(iterations[n], avg, yerr=np.array([avg-low_1_std,high_1_std-avg]).reshape(-1, 1), color = "royalblue", ecolor = "green", fmt ='.') - run_diagnostics["CIP_neff_achieved"][f"iteration_{n}_neff"] = np.round(avg, 2) + run_diagnostics[plot_title]["CIP_neff_achieved"][f"iteration_{n}_neff"] = np.round(avg, 2) iteration_prog = n except Exception as e: - print(f"Couldn't plot neff for iteration = {iterations[n]}") - break + #print(f"Couldn't plot neff for iteration = {iterations[n]}") + #break + continue # try plotting all possible cip's incase there is one folder where subdag exists + # read max lnL data from CIP output files lnL_files_last_iteration = glob.glob(path_to_main_folder + f"/iteration_{iterations[iteration_prog]}_cip/*lnL*") - run_diagnostics["latest_grid"] = f"overlap-grid-{iteration_prog+1}.xml.gz" - run_diagnostics["latest_iteration"] = int(iteration_prog) + run_diagnostics[plot_title]["latest_grid"] = f"overlap-grid-{iteration_prog+1}.xml.gz" + run_diagnostics[plot_title]["latest_iteration"] = int(iteration_prog) max_lnL, no_points = get_lnL_cut_points(all_net_path) - print(f"Total number of workers in final iteration = {len(lnL_files_last_iteration)-1}") - ax.set_title(f"Workers = {len(lnL_files_last_iteration)-1}") - ax.set_xticks(np.arange(0, run_diagnostics["latest_iteration"]+1, 1)) - fig.savefig(path+f"/plots/Effective_samples_per_CIPworker.png", bbox_inches='tight') + ax.set_title(f"Workers ({plot_title})= {len(lnL_files_last_iteration)-1}") + ax.set_xticks(np.arange(0, run_diagnostics[plot_title]["latest_iteration"]+1, 1)) + fig.savefig(save_path + f"/plots/Effective_samples_per_CIPworker_{plot_title}.png", bbox_inches='tight') plt.close(fig) +def plot_exploration_corner(all_net_path): + """ + Generates and saves a corner plot for all the points at which marginalized likelihood was evaluated, effectively acting as the exploration plot. -def plot_cip_max_lnL(path_to_main_folder): + Args: + all_net_path (str): File path to all.net """ - Plot the maximum log-likelihood (lnL) values sampled from different iterations. + print('\n--> Plotting exploration corner.') + use_cols = [1,2,5,8] + if use_truths: + P = lsu.xml_to_ChooseWaveformParams_array(truth_file_path)[0] + truths = [ P.extract_param('m1')/lsu.lsu_MSUN, P.extract_param('m2')/lsu.lsu_MSUN, P.extract_param('s1z'), P.extract_param('s2z')] + if LISA and not(eccentricity): + use_cols.append([9,10]) + labels=[r"$m_1$ $(\times 10^6 M_\odot)$", r"$m_2$ $(\times 10^6 M_\odot)$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\lambda$", r"$\beta$"] + if use_truths: + truths.append([P.extract_param('lambda'), P.extract_param('beta')]) + if LISA and eccentricity: + use_cols.append([9,10,11,12]) + labels=[r"$m_1$ $(\times 10^6 M_\odot)$", r"$m_2$ $(\times 10^6 M_\odot)$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\lambda$", r"$\beta$", r'$e$', '$\ell$'] + if use_truths: + truths.append([P.extract_param('lambda'), P.extract_param('beta'), P.extract_param('eccentricity'), P.extract_param('meanPerAno')]) + if not(LISA) and eccentricity: + use_cols.append([9,10]) + labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r'$e$', '$\ell$'] + if use_truths: + truths.append([P.extract_param('eccentricity'), P.extract_param('meanPerAno')]) + if precessing: + use_cols.append([3,4,6,7]) + labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r'$e$', '$\ell$', r"$\chi_{2z}$", r"$\chi_{1x}$", r"$\chi_{1y}$", r"$\chi_{2x}$", r"$\chi_{2y}$"] + if use_truths: + truths.append([P.extract_param('s1x'), P.extract_param('s1y'), P.extract_param('s2x'), P.extract_param('s2y')]) + if not(LISA) and precessing and not(eccentricity): + use_cols.append([3,4,6,7]) + labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\chi_{2z}$", r"$\chi_{1x}$", r"$\chi_{1y}$", r"$\chi_{2x}$", r"$\chi_{2y}$"] + if use_truths: + truths.append([P.extract_param('s1x'), P.extract_param('s1y'), P.extract_param('s2x'), P.extract_param('s2y')]) + # Load all.net + def flatten(arg): + if not isinstance(arg, list): # if not list + return [arg] + return [x for sub in arg for x in flatten(sub)] + + use_cols = flatten(use_cols) + if use_truths: + truths = flatten(truths) + data = np.loadtxt(all_net_path, usecols = use_cols) + # If else statement to check if truths are provided are not + if use_truths: + P = lsu.xml_to_ChooseWaveformParams_array(truth_file_path)[0] + fig = corner.corner(data, truth_color="black", truths=truths, color='cornflowerblue', smooth=None,smooth1d =None, linewidth = 1.0, plot_datapoints=True, plot_density=False, no_fill_contours=True, contours=False, levels=[0.0], contour_kwargs={"linewidths":1.0},hist_kwargs={"linewidth":1.0, "density": True},labels=labels) + else: + fig = corner.corner(data, color='cornflowerblue', smooth=None,smooth1d =None, linewidth = 1.0, plot_datapoints=True, plot_density=False, no_fill_contours=True, contours=False, levels=[0.0], contour_kwargs={"linewidths":1.0},hist_kwargs={"linewidth":1.0, "density": True},labels=labels) + # Save this figure + fig.savefig(save_path + f'/plots/Exploration_corner_plot.png', bbox_inches='tight') - This function iterates over all available iterations, collects maximum lnL values from files in each iteration's directory, - calculates the mean and standard deviation (using percentiles) of these values, and plots them with error bars. It also - adds a horizontal line indicating the maximum lnL value in all.net. +def plot_cip_max_lnL(path_to_main_folder, plot_title): + """ + Plot the maximum log-likelihood (lnL) values sampled from different iterations. Args: - path_to_main_folder (str): The path to the main folder containing iteration subfolders with lnL data files. + path_to_main_folder (str): The path to main folder containing iteration subfolders with lnL data files. + plot_title (str): suffix for the plots and filenames. """ - print("\n--> Plotting sampled lnL by CIP") - iterations = np.arange(0, run_diagnostics["latest_iteration"]+1, 1) - run_diagnostics['cip_sampled_lnL'] = {} + print(f"\n--> Plotting sampled lnL by CIP for {plot_title} iterations.") + + iterations = np.arange(0, run_diagnostics[plot_title]["latest_iteration"]+1, 1) + + # get bins for histograms + all_lnL_hist = [] + for iteration in iterations: + try: + files_iteration = glob.glob(path_to_main_folder + f"/iteration_{iteration}_cip/*lnL*") + except: + continue + for file_name in files_iteration: + all_lnL_hist.append(np.loadtxt(file_name)) + all_lnL_hist = np.concatenate(all_lnL_hist) + bins = np.linspace(np.min(all_lnL_hist), np.max(all_lnL_hist), 40) + fig, ax = plt.subplots() fig_hist, ax_hist = plt.subplots() + run_diagnostics[plot_title]['cip_sampled_lnL'] = {} for iteration in iterations: - run_diagnostics['cip_sampled_lnL'][iteration] = {} + run_diagnostics[plot_title]['cip_sampled_lnL'][iteration] = {} try: files_iteration = glob.glob(path_to_main_folder + f"/iteration_{iteration}_cip/*lnL*") except: @@ -615,33 +682,77 @@ def plot_cip_max_lnL(path_to_main_folder): samples_total_per_worker = samples_total/len(files_iteration) collec_lnL_hist = np.delete(collec_lnL_hist, 0, 0) collect_lnL = np.array(collect_lnL) - low_1_std, max_lnL_avg_this_iteration, high_1_std = np.percentile(collect_lnL, [16,50,84]) - low_2_std, max_lnL_avg_this_iteration, high_2_std = np.percentile(collect_lnL, [2.5,50,97.5]) - run_diagnostics['cip_sampled_lnL'][iteration].update({ - 'avg':max_lnL_avg_this_iteration, - '+':high_1_std, - '-':low_1_std}) + low_1_std, max_lnL_avg_this_iteration, high_1_std = np.percentile(collect_lnL, [16,50,84]) # 1 std + low_2_std, max_lnL_avg_this_iteration, high_2_std = np.percentile(collect_lnL, [2.5,50,97.5]) # 2 std + run_diagnostics[plot_title]['cip_sampled_lnL'][iteration].update({ + 'avg':np.round(max_lnL_avg_this_iteration, 2), + '+':np.round(high_2_std, 2), + '-':np.round(low_2_std, 2)}) ax.errorbar(iteration, max_lnL_avg_this_iteration, yerr = np.array([max_lnL_avg_this_iteration-low_2_std, high_2_std-max_lnL_avg_this_iteration]).reshape(-1,1), color = "royalblue", ecolor = "red", fmt ='.') ax.errorbar(iteration, max_lnL_avg_this_iteration, yerr = np.array([max_lnL_avg_this_iteration-low_1_std, high_1_std-max_lnL_avg_this_iteration]).reshape(-1,1), color = "royalblue", ecolor = "green", fmt ='o') - ax_hist.hist(collec_lnL_hist, label=iteration, histtype='step', linewidth = 1.0, bins=30, density=True) - run_diagnostics["cip_average_max_lnL_sampled"] = np.round(np.mean(collect_lnL), 2) - run_diagnostics["cip_std_max_lnL_sampled"] = np.round(np.std(collect_lnL), 3) - ndim = lnL_at_credible_interval(return_ndim=True) + ax_hist.hist(collec_lnL_hist, label=iteration, histtype='step', linewidth = 1.0, bins=bins, density=True) + + run_diagnostics[plot_title]["cip_average_max_lnL_sampled"] = run_diagnostics[plot_title]['cip_sampled_lnL'][iterations[-1]]['avg']#np.round(np.mean(collect_lnL), 2) + run_diagnostics[plot_title]["cip_std_max_lnL_sampled"] = [run_diagnostics[plot_title]['cip_sampled_lnL'][iterations[-1]]['+']-run_diagnostics[plot_title]["cip_average_max_lnL_sampled"], run_diagnostics[plot_title]["cip_average_max_lnL_sampled"]-run_diagnostics[plot_title]['cip_sampled_lnL'][iterations[-1]]['-']]#np.round(np.std(collect_lnL), 2) + ax.set_xlabel('iteration') ax.set_ylabel('lnL') ax.axhline(y = run_diagnostics['max_lnL'], linestyle = "--", color="black") expected_median = lnL_at_credible_interval(ci=0.5) ax.fill_between(iterations, expected_median-1.5, run_diagnostics['max_lnL'], color="green", alpha=0.5) ax.set_xticks(iterations) - fig.savefig(path+f"/plots/Maximum_sampled_lnL_per_CIPworker.png", bbox_inches="tight") + fig.savefig(save_path + f"/plots/Maximum_sampled_lnL_per_CIPworker_{plot_title}.png", bbox_inches="tight") ax_hist.set_xlabel('lnL') ax_hist.legend(loc='upper left', frameon=False) ax_hist.axvline(x = run_diagnostics['max_lnL'], linestyle = "--", color="black") - fig_hist.savefig(path+f"/plots/CIP_lnL_distribution_per_iteration", bbox_inches="tight") + fig_hist.savefig(save_path + f"/plots/CIP_lnL_distribution_per_iteration_{plot_title}.png", bbox_inches="tight") plt.close() +def plot_high_likelihood_expoloration(path_to_main_folder, plot_title): + """ + Plots high likelihood points over iterations. + + Args: + path_to_main_folder (str): Path to the folder containing the composite files. + plot_title (str): suffix for the plots and filenames. + """ + print(f"\n--> Plotting likelihood exploration for {plot_title} iterations.") + + run_diagnostics[plot_title]["composite_information"] = {} + + fig, ax = plt.subplots() + ax.set_xlabel("iteration") + ax.set_ylabel("high lnL points") + ax.set_title(f"Total high lnL points = {run_diagnostics['high_lnL_points']}, max_lnL = {run_diagnostics['max_lnL']}") + collect_data = [] + collect_iter = [] + print("\nIteration | Global max lnL | High-lnL points | Iteration max lnL | Total lnL points") + for iteration in np.arange(0, run_diagnostics[plot_title]["latest_iteration"]+1, 1): + run_diagnostics[plot_title]["composite_information"][iteration] = {} + try: + max_lnL, no_points, max_lnL_composite, total_points = get_lnL_cut_points(f"{path_to_main_folder}/consolidated_{iteration}.composite", composite=True) + except Exception as e: + print(f"Error loading file {path_to_main_folder}/consolidated_{iteration}.composite: {e}") + continue + print(f"{iteration:9d} | {max_lnL:14.2f} | {no_points:15d} | {max_lnL_composite:17.2f} | {total_points:16d}") + percent_high_lnL_points = np.round(no_points/total_points*100, 2) + collect_data.append(no_points) + collect_iter.append(iteration) + ax.scatter(iteration, no_points, label = f"{max_lnL_composite} ({percent_high_lnL_points})", s=25) + run_diagnostics[plot_title]["composite_information"][iteration].update({ + "max_lnL":max_lnL_composite, + "high_lnL_points":no_points, + "percent_high_lnL_points": percent_high_lnL_points}) + + ax.grid(alpha=0.4) + ax.plot(collect_iter, collect_data, color = "black", linestyle = "--", linewidth = 1.5, alpha = 0.5) + ax.set_xticks(np.arange(0, run_diagnostics[plot_title]["latest_iteration"]+1, 1)) + ax.legend(loc="upper left") + fig.savefig(save_path + f"/plots/Likelihood_exploration_plot_{plot_title}.png", bbox_inches='tight') + plt.close(fig) + def plot_histograms(sorted_posterior_file_paths, plot_title, iterations = None, plot_legend = True, JSD = True): """ Plots histograms for specified parameters across different posterior samples. @@ -653,7 +764,7 @@ def plot_histograms(sorted_posterior_file_paths, plot_title, iterations = None, plot_legend (bool): Whether to include a legend in the histograms. Defaults to True. JSD (bool): Whether to calculate and display Jensen-Shannon Divergence between iterations. Defaults to True. """ - print("\n--> Plotting histograms") + print(f"\n--> Plotting histograms for {plot_title} iterations.") # when you just want to plot final iterations histograms if iterations is None: iterations = [-1] @@ -663,7 +774,7 @@ def plot_histograms(sorted_posterior_file_paths, plot_title, iterations = None, # all_net_data = convert_all_net_to_posterior_format(all_net_path) # not_nan_lnL = np.argwhere(all_net_data[:,-3]>=np.max(all_net_data[:,-3]) - 15).flatten()#np.argwhere(~np.isnan(all_net_data[:,-3])).flatten() # all_net_data = np.array(all_net_data[not_nan_lnL]) - parameters = ["mc", "q", "eta", "m1", "m2", "mtot","s1z", "s2z", "chi_eff","chiMinus","chi1", "chi2"] + parameters = ["mc", "q", "eta", "m1", "m2", "mtot", "s1z", "s2z", "chi_eff", "chiMinus", "chi1", "chi2"] # for LISA include skylocation if LISA: parameters.append("dec") @@ -679,69 +790,49 @@ def plot_histograms(sorted_posterior_file_paths, plot_title, iterations = None, parameters.append("chi_p") parameters.append("chi1_perp") parameters.append("chi2_perp") + for parameter in parameters: print(f"Plotting histogram for {parameter}") + fig, ax = plt.subplots() ax.set_title(plot_title) ax.set_xlabel(parameter) ax.set_yticks([]) - for i in np.arange(len(sorted_posterior_file_paths)): - line_label = str(iterations[i]) + + # get bins + all_data = [] + for file_path in sorted_posterior_file_paths: if parameter in ["chi_eff", "chiMinus", "chi_p", "chi1", "chi2", "chi1_perp", "chi2_perp"]: - posterior_data = np.loadtxt(sorted_posterior_file_paths[i]) - #data = get_chi_eff_from_mass_and_spins(posterior_data) - data = get_combination_from_mass_and_spin(posterior_data, parameter) + data = get_combination_from_mass_and_spin(np.loadtxt(file_path), parameter) else: - parameter_index = get_index_for_parameter(parameter) - data = np.loadtxt(sorted_posterior_file_paths[i])[:,parameter_index] + data = np.loadtxt(file_path)[:, get_index_for_parameter(parameter)] + + all_data.append(data) + combined_data = np.concatenate(all_data) + bins = np.linspace(np.min(combined_data), np.max(combined_data), 40) + + for i, data in enumerate(all_data): + line_label = str(iterations[i]) + if i > 0 and JSD: JS_test = calculate_JS_divergence(data, data_previous) - line_label +=f" ({calculate_JS_divergence(data, data_previous).median:0.3f})" - ax.hist(data, label = line_label, histtype="step", bins = 25, density=True, linewidth=1.0, color=default_colors[i]) + line_label += f" ({JS_test.median:0.3f})" + + ax.hist(data, bins=bins, density=True, histtype="step", linewidth=1.0, color=default_colors[i], label=line_label) + if use_truths: - factor = 1 - parameter_extract = parameter - if parameter in ["mc", "m1", "m2", "mtot"]: - factor = lsu.lsu_MSUN - if parameter == "chi_eff": - parameter_extract = "xi" - if parameter == "ra": - parameter_extract = "phi" - if parameter == "dec": - parameter_extract = "theta" - ax.axvline(x = P.extract_param(parameter_extract)/factor, linestyle="--", linewidth=1.0, color="black") + factor = lsu.lsu_MSUN if parameter in ["mc", "m1", "m2", "mtot"] else 1 + parameter_extract = {"chi_eff": "xi", "ra": "phi", "dec": "theta"}.get(parameter, parameter) + ax.axvline(x=P.extract_param(parameter_extract)/factor, linestyle="--", linewidth=1.0, color="black") + data_previous = data - #try: (this isn't really helpful, so commenting it out) - # likelihood = np.exp(np.array(all_net_data[:,-3])) - # reweighted_all_net = np.random.choice(all_net_data[:,parameter_index], p = likelihood /np.sum(likelihood), size = 1000, replace = True) - # ax.hist(reweighted_all_net, label = "Likelihood", histtype="step", bins = 50, density=True, alpha = 0.7, linewidth=1.0, color = "grey") - #except: - # print("Couldn't plot likelihood distribution") + # don't create legend when only plotting finals iteration's histograms if plot_legend: ax.legend(loc = "upper right") - fig.savefig(path+f"/plots/histograms/histogram_{plot_title}_{parameter}.png", bbox_inches='tight') + fig.savefig(save_path + f"/plots/histograms/histogram_{plot_title}_{parameter}.png", bbox_inches='tight') plt.close() -def plot_log_likelihood(extrinsic_path): - """ - Plots a histogram of the log-likelihood (lnL) distribution from extrinsic samples file. - - Args: - extrinsic_path (str): File path to the extrinsic data file containing sampled parameters. - The log-likelihood values are assumed to be in column index 18. - """ - extrinsic_data = np.loadtxt(extrinsic_path+"/extrinsic_posterior_samples.dat") - index_lnL = 18 - plt.cla() - plt.title('lnL distribution') - plt.xlabel('lnL') - plt.ylabel('Points') - plt.hist(extrinsic_data[:, index_lnL], histtype='step', color='black', bins=40) - plt.savefig(path+f"/plots/lnL_distribution_from_extrinsic_step.png", bbox_inches='tight') - plt.cla() - plt.close() - def plot_corner(sorted_posterior_file_paths, plot_title, iterations = None, parameters = ["mc", "eta", "chi_eff"], use_truths = False): """ Generates corner plots for posterior samples using a specified plotting executable. @@ -759,12 +850,24 @@ def plot_corner(sorted_posterior_file_paths, plot_title, iterations = None, para parameters.append("dec") print(f"\n--> Plotting corner plot for params ({plot_title}) {parameters}") - max_lnL, no_points = run_diagnostics["max_lnL"], run_diagnostics["high_lnL_points"] + + # title + if "Subdag" in plot_title: + subdag_path = find_posteriors_in_sub(path, return_subdag_folder_path=True) + _, no_points, max_lnL, _ = get_lnL_cut_points(f"{subdag_path}/all.net", composite=True) + else: + max_lnL, no_points = run_diagnostics["max_lnL"], run_diagnostics["high_lnL_points"] title = f"max_lnL={max_lnL:0.2f},points_cut={no_points}" + + # Plotting command begins plotting_command = f"python {corner_plot_exe} --plot-1d-extra --quantiles None --ci-list [0.9] --use-title {title} " - if plot_title != "extrinsic": + + if plot_title != "extrinsic" and "Main" in plot_title: plotting_command += f"--composite-file {all_net_path} --lnL-cut 15 --sigma-cut 0.4 " - # Append iteration-related options to the command + elif plot_title != "extrinsic" and "Subdag" in plot_title: + plotting_command += f"--composite-file {subdag_path}/all.net --lnL-cut 15 --sigma-cut 0.4 " + + # Append iteration-related options to the command if iterations is not None: plotting_command += "--use-legend " else: @@ -775,7 +878,7 @@ def plot_corner(sorted_posterior_file_paths, plot_title, iterations = None, para plotting_command += f"--truth-file {truth_file_path} " # plot grey points (low lnL) when showing multiple iterations - if plot_title != "Final": + if "Final" not in plot_title: plotting_command += "--use-all-composite-but-grayscale " # Add parameter options to the command @@ -806,22 +909,31 @@ def plot_corner(sorted_posterior_file_paths, plot_title, iterations = None, para # Move and rename output files corner_plot_filename = f"corner_{'_'.join(parameters)}.png" new_corner_plot_path = f"plots/corner_plots/corner_{'_'.join(parameters)}_{plot_title}.png" - os.system(f"mv {corner_plot_filename} {new_corner_plot_path}") + os.system(f"mv {save_path}/{corner_plot_filename} {save_path}/{new_corner_plot_path}") # Move and rename individual parameter plots for parameter in parameters: - os.system(f"mv {parameter}.png plots/1_D_plots/{parameter}_{plot_title}.png") - os.system(f"mv {parameter}_cum.png plots/1_D_plots/{parameter}_cum_{plot_title}.png") + os.system(f"mv {save_path}/{parameter}.png {save_path}/plots/1_D_plots/{parameter}_{plot_title}.png") + os.system(f"mv {save_path}/{parameter}_cum.png {save_path}/plots/1_D_plots/{parameter}_cum_{plot_title}.png") -def plot_JS_divergence(posterior_1_path, posterior_2_path, posterior_3_path=None, plot_title=None, threshold=0.005, parameters = ["mc","eta", "m1", "m2"]): +def plot_JS_divergence(posterior_1_path, posterior_2_path, posterior_3_path=None, plot_title=None, threshold=JSD_threshold, parameters = ["mc","eta", "m1", "m2"]): """ - Plots Jensen-Shannon Divergence (JSD) between two posterior datasets for specified parameters. + Plots the Jensen-Shannon divergence between posterior samples for + selected parameters. Args: - posterior_1_path (str): File path to the first posterior dataset. - posterior_2_path (str): File path to the second posterior dataset. - plot_title (str): Title for the plot and filename. - parameters (list of str): List of parameters to calculate and plot JSD for. + posterior_1_path (str): Path to the first posterior sample file. + posterior_2_path (str): Path to the second posterior sample file. + posterior_3_path (str or None): Path to an optional third posterior + sample file. Default is None. + plot_title (str or None): Title for the plot and output filename. + Default is None. + threshold (float): JSD threshold shown as a horizontal line. + parameters (list of str): Parameters for which the JSD is computed. + Default is ["mc", "eta", "m1", "m2"]. + + Returns: + None """ if not(non_spinning): parameters.append("s1z") @@ -842,17 +954,23 @@ def plot_JS_divergence(posterior_1_path, posterior_2_path, posterior_3_path=None parameters.append("chi_p") parameters.append("chi1_perp") parameters.append("chi2_perp") + print(f"\n--> Plotting Jensen Shannon Divergence for {parameters} with threshold {threshold}\n") + posterior_data1 = np.loadtxt(posterior_1_path) posterior_data2 = np.loadtxt(posterior_2_path) if not(posterior_3_path is None): posterior_data3 = np.loadtxt(posterior_3_path) + JSD_array = [] # collect for last and second-to-last JSD_error = [] + JSD_array_third = [] # collect for last and third-to-last JSD_error_third = [] + run_diagnostics["JSD"][plot_title] = {} run_diagnostics["JSD_3"][plot_title] = {} + for parameter in parameters: if parameter in ["chi_eff", "chiMinus", "chi_p", "chi1", "chi2", "chi1_perp", "chi2_perp"]: data1, data2 = get_combination_from_mass_and_spin(posterior_data1, parameter), get_combination_from_mass_and_spin(posterior_data2, parameter) @@ -866,33 +984,50 @@ def plot_JS_divergence(posterior_1_path, posterior_2_path, posterior_3_path=None if not(posterior_3_path is None): parameter_n = get_index_for_parameter(parameter) JSD_3 = calculate_JS_divergence(posterior_data1[:, parameter_n], posterior_data3[:, parameter_n]) + JSD_array.append(JSD.median) JSD_error.append([JSD.minus, JSD.plus]) + run_diagnostics["JSD"][plot_title][parameter] = np.round(JSD.median, 3) + if not(posterior_3_path is None): JSD_array_third.append(JSD_3.median) JSD_error_third.append([JSD_3.minus, JSD_3.plus]) run_diagnostics["JSD_3"][plot_title][parameter] = np.round(JSD_3.median, 3) + fig, ax = plt.subplots() ax.set_title(plot_title) ax.set_ylabel("JSD") ax.axhline( y = threshold, linewidth = 1.0, linestyle = "--", color = "red") - ax.errorbar(parameters, JSD_array, np.array(JSD_error).T, color = "royalblue", ecolor = "red", fmt ='o', markersize = 5, label='latest-secondlatest') - if not(posterior_3_path is None): - ax.errorbar(parameters, JSD_array_third, np.array(JSD_error_third).T, color = "green", ecolor = "black", fmt ='o', markersize = 5, label='latest-thirdlatest') + + if "Main_subdag" in plot_title: + main_iteration = run_diagnostics["Main"]["latest_iteration"] + subdag_iteration = run_diagnostics["Subdag"]["latest_iteration"] + label_second = f"Main {main_iteration} vs Subdag {subdag_iteration}" + else: + latest_iteration = run_diagnostics[plot_title.split("_")[0]]["latest_iteration"] + label_second = f"Iter {latest_iteration} vs Iter {latest_iteration-1}" + + ax.errorbar(parameters, JSD_array, np.array(JSD_error).T, color="royalblue", ecolor="red", fmt='o', markersize=5, label=label_second) + if posterior_3_path is not None: + label_third = f"Iter {latest_iteration} vs Iter {latest_iteration-2}" + ax.errorbar(parameters, JSD_array_third, np.array(JSD_error_third).T, color="green", ecolor="black", fmt='o', markersize=5, label=label_third) + #ax.errorbar(parameters, JSD_array, np.array(JSD_error).T, color = "royalblue", ecolor = "red", fmt ='o', markersize = 5, label='latest-secondlatest') + #if not(posterior_3_path is None): + # ax.errorbar(parameters, JSD_array_third, np.array(JSD_error_third).T, color = "green", ecolor = "black", fmt ='o', markersize = 5, label='latest-thirdlatest') ax.legend(loc='upper right') ax.tick_params(axis='x', labelrotation=60) - fig.savefig(path+f"/plots/JSD_per_parameter_{plot_title}.png", bbox_inches='tight') + fig.savefig(save_path+f"/plots/JSD_per_parameter_{plot_title}.png", bbox_inches='tight') plt.close(fig) -def write_sample_statistics(posterior, parameters=["mc","eta", "m1", "m2"]): +def write_sample_statistics(posterior, parameters=["mc", "eta", "m1", "m2"], extrinsic = False): """ Computes and writes sample statistics for specified parameters to a file. Args: posterior (str): Path to the file containing posterior samples. parameters (list, optional): List of parameter names for which statistics will be computed. Defaults to - ["mc", "eta", "m1", "m2", "s1z", "s2z", "chi_eff"]. + ["mc", "eta", "m1", "m2"]. """ if not(non_spinning): parameters.append("s1z") @@ -917,20 +1052,25 @@ def write_sample_statistics(posterior, parameters=["mc","eta", "m1", "m2"]): parameters.append("chi2") if use_truths: P = lsu.xml_to_ChooseWaveformParams_array(truth_file_path)[0] + print(f"\n--> Writing sample statistics for parameters: {parameters}") + posterior = np.loadtxt(posterior) - f = open(path+f"/plots/sample_statistics.txt", "w") + + f = open(save_path + f"/plots/sample_statistics.txt", "w") f.write("Note: limits are equal-tailed 90th percentile\n") run_diagnostics["sample_statistics"] = {} for parameter in parameters: if parameter in ["chi_eff", "chiMinus", "chi_p", "chi1", "chi2", "chi1_perp", "chi2_perp"]: samples_here = get_combination_from_mass_and_spin(posterior, parameter) else: - parameter_n = get_index_for_parameter(parameter) + parameter_n = get_index_for_parameter(parameter, extrinsic = extrinsic) samples_here = posterior[:,parameter_n] statistics = get_sample_statistics(samples_here) - run_diagnostics["sample_statistics"][parameter] = np.round(statistics,3) + run_diagnostics["sample_statistics"][parameter] = np.round(statistics, 2) line = f"{parameter}: median = {statistics[1]:0.3f}, upper limit = {statistics[2]:0.3f}, lower limit = {statistics[0]:0.3f}" + line = (f"{parameter:12s} = "f"{statistics[1]:8.3f} "f"[{statistics[0]:8.3f}, {statistics[2]:8.3f}]") + if use_truths: factor = 1 parameter_extract = parameter @@ -942,171 +1082,318 @@ def write_sample_statistics(posterior, parameters=["mc","eta", "m1", "m2"]): parameter_extract = "phi" if parameter == "dec": parameter_extract = "theta" - line += f", truth here = {P.extract_param(parameter_extract)/factor:0.3f}" + line += f" truth: {P.extract_param(parameter_extract)/factor:8.3f}" f.write(line + "\n") max_sample, lnL = get_lnL_cut_points(all_net_path, lnL_cut=15, error_threshold=0.4, composite=False, provide_max_lnL_point=True) f.close() -def plot_exploration_corner(all_net_path): +def plot_log_likelihood(extrinsic_path): """ - Generates and saves a corner plot for all the points at which marginalized likelihood was evaluated, effectively acting as the exploration plot. + Plots a histogram of the log-likelihood (lnL) distribution from extrinsic samples file. Args: - all_net_path (str): File path to all.net + extrinsic_path (str): File path to the extrinsic data file containing sampled parameters. + The log-likelihood values are assumed to be in column index 18. """ - print('\n--> Plotting exploration corner') - use_cols = [1,2,5,8] - if use_truths: - P = lsu.xml_to_ChooseWaveformParams_array(truth_file_path)[0] - truths = [ P.extract_param('m1')/lsu.lsu_MSUN, P.extract_param('m2')/lsu.lsu_MSUN, P.extract_param('s1z'), P.extract_param('s2z')] - if LISA and not(eccentricity): - use_cols.append([9,10]) - labels=[r"$m_1$ $(\times 10^6 M_\odot)$", r"$m_2$ $(\times 10^6 M_\odot)$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\lambda$", r"$\beta$"] - if use_truths: - truths.append([P.extract_param('lambda'), P.extract_param('beta')]) - if LISA and eccentricity: - use_cols.append([9,10,11,12]) - labels=[r"$m_1$ $(\times 10^6 M_\odot)$", r"$m_2$ $(\times 10^6 M_\odot)$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\lambda$", r"$\beta$", r'$e$', '$\ell$'] - if use_truths: - truths.append([P.extract_param('lambda'), P.extract_param('beta'), P.extract_param('eccentricity'), P.extract_param('meanPerAno')]) - if not(LISA) and eccentricity: - use_cols.append([9,10]) - labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r'$e$', '$\ell$'] - if use_truths: - truths.append([P.extract_param('eccentricity'), P.extract_param('meanPerAno')]) - if precessing: - use_cols.append([3,4,6,7]) - labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r'$e$', '$\ell$', r"$\chi_{2z}$", r"$\chi_{1x}$", r"$\chi_{1y}$", r"$\chi_{2x}$", r"$\chi_{2y}$"] - if use_truths: - truths.append([P.extract_param('s1x'), P.extract_param('s1y'), P.extract_param('s2x'), P.extract_param('s2y')]) - if not(LISA) and precessing and not(eccentricity): - use_cols.append([3,4,6,7]) - labels=[r"$m_1$", r"$m_2$", r"$\chi_{1z}$", r"$\chi_{2z}$", r"$\chi_{2z}$", r"$\chi_{1x}$", r"$\chi_{1y}$", r"$\chi_{2x}$", r"$\chi_{2y}$"] - if use_truths: - truths.append([P.extract_param('s1x'), P.extract_param('s1y'), P.extract_param('s2x'), P.extract_param('s2y')]) - # Load all.net - def flatten(arg): - if not isinstance(arg, list): # if not list - return [arg] - return [x for sub in arg for x in flatten(sub)] - - use_cols = flatten(use_cols) - if use_truths: - truths = flatten(truths) - data = np.loadtxt(all_net_path, usecols = use_cols) - # If else statement to check if truths are provided are not - if use_truths: - P = lsu.xml_to_ChooseWaveformParams_array(truth_file_path)[0] - fig = corner.corner(data, truth_color="black", truths=truths, color='cornflowerblue', smooth=None,smooth1d =None, linewidth = 1.0, plot_datapoints=True, plot_density=False, no_fill_contours=True, contours=False, levels=[0.0], contour_kwargs={"linewidths":1.0},hist_kwargs={"linewidth":1.0, "density": True},labels=labels) - else: - fig = corner.corner(data, color='cornflowerblue', smooth=None,smooth1d =None, linewidth = 1.0, plot_datapoints=True, plot_density=False, no_fill_contours=True, contours=False, levels=[0.0], contour_kwargs={"linewidths":1.0},hist_kwargs={"linewidth":1.0, "density": True},labels=labels) - # Save this figure - fig.savefig(f'plots/Exploration_corner_plot.png') - + extrinsic_data = np.loadtxt(extrinsic_path+"/extrinsic_posterior_samples.dat") + index_lnL = 18 + plt.cla() + plt.title('lnL distribution') + plt.xlabel('lnL') + plt.ylabel('Points') + plt.hist(extrinsic_data[:, index_lnL], histtype='step', color='black', bins=40) + plt.savefig(path+f"/plots/lnL_distribution_from_extrinsic_step.png", bbox_inches='tight') + plt.cla() + plt.close() def evaluate_run(run_diagnostics): """ - Evaluates and writes diagnostics information to a file. + Evaluates the run and writes diagnostics information to a file. Args: run_diagnostics (dict): Dictionary containing diagnostics data. + + Returns: + None """ - f = open(path+f"/plots/Diagnostics.txt", "w") - # ILE - f.write("###########################################################################################\n") - f.write("# ILE diagnostics\n") - f.write("###########################################################################################\n") - # Monte carlo error and number of high lnL points - f.write(f"Total number of marginalized lnL evaluations = {run_diagnostics['total_lnL_evaluations']}\n") - f.write(f"Total number of high marginalized lnL points = {run_diagnostics['total_high_lnL_points']}\n") - f.write(f"Total number of high marginalized lnL points used = {run_diagnostics['high_lnL_points']}\n") - f.write(f"Total number of high marginalized lnL points not used due to large error = {run_diagnostics['high_lnL_points_with_large_error']}\n") - f.write(f"Approximate SNR captured = {np.sqrt(2*run_diagnostics['max_lnL']):0.2f}\n") - f.write(f"Number of high lnL points with lnL cut [12, 10, 5, 2] = [{run_diagnostics['high_lnL_points_with_lnLcut_12'], run_diagnostics['high_lnL_points_with_lnLcut_10'], run_diagnostics['high_lnL_points_with_lnLcut_5'], run_diagnostics['high_lnL_points_with_lnLcut_2']}]\n") - print(run_diagnostics['composite_information']) - f.write(f"Likelihood exploration data per iteration: \n{run_diagnostics['composite_information']}\n") - - ILE_is_good = True - if run_diagnostics['high_lnL_points_with_large_error']/run_diagnostics['total_high_lnL_points'] > 0.5: - f.write(f"\t--> Large number of points have a high Monte Carlo error (sigma = 0.4). Consider reducing d-max, increasing d-min, increasing n-max and/or changing the sampler.\n") - ILE_is_good = False - if run_diagnostics['high_lnL_points'] <= 500: - f.write(f"\t--> Number of high likelihood points is less than 500, which could be caused due to initial grid not having sufficient resolution. Considering reducing the parameter space and/or increasing the number of points on the grid.\n") - ILE_is_good = False - if 500 < run_diagnostics['high_lnL_points'] < 5000: - f.write(f"\t--> Number of high likelihood points is less than 5000, consider rerunning with {run_diagnostics['latest_grid']} as your starting grid and copying this run's all.net as bonus.composite in your new run directory.\n") - ILE_is_good = False - f.write("\n") - if ILE_is_good: - f.write("\t--> ILE status: GOOD! <--\n") - else: - f.write("\t--> ILE status: BAD! <--\n") - f.write("\n\n") - # CIP - f.write("###########################################################################################\n") - f.write("# CIP diagnostics\n") - f.write("###########################################################################################\n") - # CIP neff - f.write(f"CIP neff requested = {run_diagnostics['CIP_neff']}]\n") - f.write(f"CIP neff achieved = {run_diagnostics['CIP_neff_achieved']}\n") - CIP_is_good = True - first_iter_neff = run_diagnostics['CIP_neff'][list(run_diagnostics['CIP_neff'].keys())[0]] - last_iter_neff = run_diagnostics['CIP_neff'][list(run_diagnostics['CIP_neff'].keys())[-1]] - last_iter_neff_achieved = run_diagnostics['CIP_neff_achieved'][list(run_diagnostics['CIP_neff_achieved'].keys())[-1]] - if first_iter_neff <= last_iter_neff_achieved <= last_iter_neff and first_iter_neff!=last_iter_neff: - f.write(f"\t--> neff has not been reached, the posterior distribution may be wider and/or less smooth. To address this, try narrowing the parameter space or switching to a different sampler. Alternatively, you can reduce the neff for each CIP job (>10) and increase the number of CIP jobs submitted per iteration. However, it seems like it has reached neff as set by the CIP for earlier iterations (CIP_worker0.sub), so if the run is ongoing let it continue.\n") - elif last_iter_neff > last_iter_neff_achieved: - f.write(f"\t--> neff has not been reached, the posterior distribution may be wider and/or less smooth. To address this, try narrowing the parameter space or switching to a different sampler. Alternatively, you can reduce the neff for each CIP job (>10) and increase the number of CIP jobs submitted per iteration.\n") - CIP_is_good = False - # CIP JSD - f.write(f"\nCIP Jensen-Shannon divergence:\n{run_diagnostics['JSD']}\n") - JSD_not_good = {} - JSD_is_good = True - for iteration_type in run_diagnostics['JSD']: - JSD_not_good[iteration_type] = {} - for param in run_diagnostics['JSD'][iteration_type]: - if run_diagnostics['JSD'][iteration_type][param] > 0.005: - JSD_not_good[iteration_type][param] = run_diagnostics['JSD'][iteration_type][param] - for iteration_type in run_diagnostics['JSD']: - if len(JSD_not_good[iteration_type]) > 0: - JSD_is_good = False - if JSD_is_good is False: - f.write(f"\n\t If the Jensen-Shannon Divergence for any parameter between the last and second-to-last iterations is greater than 0.005, it means the run has not yet converged. In this case, you should rerun the analysis using {run_diagnostics['latest_grid']} as the starting grid. Additionally, copy this run's all.net file to a new file named bonus.composite in your new run directory. \n") - CIP_is_good = False - # CIP sampling - f.write(f"\nAverage max lnL sampled by CIP in iteration {run_diagnostics['latest_iteration']} is: {run_diagnostics['cip_average_max_lnL_sampled']} +- {run_diagnostics['cip_std_max_lnL_sampled']}. Max lnL in all.net is {run_diagnostics['max_lnL']}.\n") - expected_median = lnL_at_credible_interval(ci=0.5) - if expected_median - run_diagnostics['cip_average_max_lnL_sampled'] > 1.5: - f.write(f"\t--> The difference between the expected max lnL {expected_median} and the average maximum lnL value sampled by CIP {run_diagnostics['cip_average_max_lnL_sampled']} is more than 1.5, which could cause the peak to be slightly shifted. This discrepancy might be because CIP hasn't sampled the peak well enough or due to interpolation errors. If the issue is inadequate sampling (which is more likely in high signal-to-noise ratio cases), you should increase the neff parameter in the CIP sub file, reduce the number of samples requested, and run the CIP script more times (cip-explode-jobs option) (this can be done without setting up a new run). This will help ensure that the peak is accurately sampled and that the number of samples is close to neff. If the problem is due to interpolation errors, consider running additional iterations to have sufficient lnL evaluations around the peak. \n") - CIP_is_good = False - f.write("\n") - if CIP_is_good: - f.write("\t--> CIP status: GOOD! <--\n") - else: - f.write("\t--> CIP status: BAD! <--\n") - f.write("\n###########################################################################################\n") - f.write("# Visual diagnostics\n") - f.write("###########################################################################################\n") - f.write("\t 1) Is the 90% credible interval mostly around the red points? If not, it could be that the run needs more iterations. If the SNR < 30, then the prior might impact it and the shift is expected.") - f.write(f"\n\t 2) Has the parameter space been sufficiently explored? Are there blue points around the red points? Continuing the run will help if this is true with {run_diagnostics['latest_grid']} as your starting grid and copying this run's all.net as bonus.composite in your new run directory") - f.write("\n\t 3) Is the approximate SNR captured close to True SNR? A significant difference implies the inference got stuck at a local lnL maxima. Happens rarely") + def write_header(f, title): + f.write("\n") + f.write("###########################################################################################\n") + f.write(f"# {title}\n") + f.write("###########################################################################################\n") - f.close() - print("###########################################################################################") - print("# Run diagnositcs") + def write_dict_information(f, title, data): + f.write(f"\n\t{title}:\n") + if not data: + f.write("\tNo data available.\n") + return + + for key, value in data.items(): + f.write(f"\t\t{key}: {value}\n") + + def write_composite_information(f, plot_title): + composite_information = run_diagnostics.get(plot_title, {}).get("composite_information", {}) + + f.write(f"\n{plot_title} likelihood exploration:\n") + if not composite_information: + f.write("\tNo composite information available.\n") + return + + f.write("\titeration | max lnL | high lnL points | percent high lnL points\n") + for iteration, values in composite_information.items(): + f.write(f"\t{iteration:>9} | {values.get('max_lnL', 'NA'):>7} | {values.get('high_lnL_points', 'NA'):>15} | {values.get('percent_high_lnL_points', 'NA')}%\n") + + def write_neff_information(f, plot_title): + requested_neff = run_diagnostics.get(plot_title, {}).get("CIP_neff", {}) + achieved_neff = run_diagnostics.get(plot_title, {}).get("CIP_neff_achieved", {}) + + f.write(f"\n{plot_title} CIP n-eff:") + write_dict_information(f, "Requested n-eff", requested_neff) + write_dict_information(f, "Achieved n-eff", achieved_neff) + + if not requested_neff or not achieved_neff: + return True + + first_requested_neff = requested_neff[list(requested_neff.keys())[0]] + last_requested_neff = requested_neff[list(requested_neff.keys())[-1]] + last_achieved_neff = achieved_neff[list(achieved_neff.keys())[-1]] + + neff_is_good = True + + if last_achieved_neff < last_requested_neff: + f.write(f"\t--> Latest achieved n-eff ({last_achieved_neff}) is below the requested value ({last_requested_neff}). To address this, try narrowing the parameter space or switching to a different sampler. Alternatively, you can reduce the neff for each CIP job (>10) and increase the number of CIP jobs submitted per iteration.\n") + neff_is_good = False + if first_requested_neff <= last_achieved_neff < last_requested_neff and first_requested_neff != last_requested_neff: + if run_diagnostics.get("run_is_complete") == False: + f.write(f"\t--> Latest achieved n-eff ({last_achieved_neff}) is below the final target ({last_requested_neff}), although it exceeds the initial target ({first_requested_neff}). Since the run is incomplete, this warning can be ignored.\n") + + return neff_is_good + + def write_cip_sampled_lnL_information(f, plot_title): + diagnostics = run_diagnostics.get(plot_title, {}) + + latest_iteration = diagnostics.get("latest_iteration", None) + average_lnL = diagnostics.get("cip_average_max_lnL_sampled", None) + std_lnL = diagnostics.get("cip_std_max_lnL_sampled", None) + sampled_lnL = diagnostics.get("cip_sampled_lnL", {}) + + #f.write(f"\n{plot_title} CIP sampled lnL:\n") + + #if sampled_lnL: + # f.write("\titeration | median max lnL | low | high\n") + # for iteration, values in sampled_lnL.items(): + # f.write(f"\t{iteration:>9} | {values.get('avg', 'NA'):>14} | {values.get('-', 'NA'):>3} | {values.get('+', 'NA'):>4}\n") + + if average_lnL is None or std_lnL is None: + f.write("\tNo sampled lnL summary available.\n") + return True + + sampling_is_good = True + + try: + expected_median = lnL_at_credible_interval(ci=0.5) + if expected_median - average_lnL > 1.5: + f.write(f"\t--> The expected median max lnL ({expected_median:0.2f}) is more than 1.5 above the average sampled max lnL ({average_lnL}). CIP may not be sampling the peak well enough, or there may be interpolation error near the peak.\n") + sampling_is_good = False + if run_diagnostics.get("run_is_complete") == False: + f.write(f"\t--> Since the run is incomplete, this warning can be ignored.\n") + except Exception: + f.write("\tCould not compute the expected median lnL.\n") + + return sampling_is_good + + def write_jsd_information(f, jsd_labels, title, threshold=JSD_threshold): + f.write(f"\n{title}:\n") + + found_jsd = False + jsd_is_good = True + + for jsd_label in jsd_labels: + jsd_data = run_diagnostics.get("JSD", {}).get(jsd_label, {}) + jsd_data_3 = run_diagnostics.get("JSD_3", {}).get(jsd_label, {}) + + if not jsd_data and not jsd_data_3: + continue + + found_jsd = True + #f.write(f"\n\t{jsd_label}:\n") + + if jsd_data: + if "Main_subdag" in jsd_label: + main_iteration = run_diagnostics.get("Main", {}).get("latest_iteration", "latest") + subdag_iteration = run_diagnostics.get("Subdag", {}).get("latest_iteration", "latest") + f.write(f"\t Main iteration {main_iteration} vs Subdag iteration {subdag_iteration}:\n") + else: + latest_iteration = run_diagnostics.get(jsd_label.split("_")[0], {}).get("latest_iteration", "latest") + f.write(f"\t iteration {latest_iteration} vs iteration {latest_iteration - 1}:\n") + + for parameter, value in jsd_data.items(): + marker = " <-- above threshold" if value > threshold else "" + f.write(f"\t\t{parameter}: {value}{marker}\n") + if value > threshold: + jsd_is_good = False + + if jsd_data_3: + latest_iteration = run_diagnostics.get(jsd_label.split("_")[0], {}).get("latest_iteration", "latest") + f.write(f"\t iteration {latest_iteration} vs iteration {latest_iteration - 2}:\n") + + for parameter, value in jsd_data_3.items(): + marker = " <-- above threshold" if value > threshold else "" + f.write(f"\t\t{parameter}: {value}{marker}\n") + if value > threshold: + jsd_is_good = False + if not found_jsd: + f.write("\tNo JSD information available.\n") + + return jsd_is_good + + diagnostics_path = save_path + "/plots/Diagnostics.txt" + + with open(diagnostics_path, "w") as f: + + write_header(f, "Run summary") + f.write(f"Run directory = {path}\n") + f.write(f"Save directory = {save_path}/plots\n") + if run_diagnostics.get("Subdag"): + f.write(f"Subdag directory = {find_posteriors_in_sub(path, return_subdag_folder_path=True)}\n") + f.write(f"LISA analysis = {LISA}\n") + f.write(f"Eccentric analysis = {eccentricity}\n") + f.write(f"Precessing analysis = {precessing}\n") + f.write(f"Non-spinning analysis = {non_spinning}\n") + f.write(f"Truth file used = {truth_file_path if use_truths else None}\n") + f.write(f"Run complete = {run_diagnostics.get('run_is_complete', False)}\n") + + write_header(f, "ILE diagnostics") + + ILE_is_good = True + + f.write(f"Total number of marginalized lnL evaluations = {run_diagnostics.get('total_lnL_evaluations', 'NA')}\n") + f.write(f"Total number of high marginalized lnL points = {run_diagnostics.get('total_high_lnL_points', 'NA')}\n") + f.write(f"Total number of high marginalized lnL points used = {run_diagnostics.get('high_lnL_points', 'NA')}\n") + f.write(f"Total number of high marginalized lnL points not used due to large error = {run_diagnostics.get('high_lnL_points_with_large_error', 'NA')}\n") + + if "max_lnL" in run_diagnostics: + f.write(f"Maximum marginalized lnL = {run_diagnostics['max_lnL']}\n") + f.write(f"Approximate SNR captured = {np.sqrt(2 * run_diagnostics['max_lnL']):0.2f}\n") + + f.write(f"Number of high lnL points with lnL cuts [12, 10, 5, 2] = [{run_diagnostics.get('high_lnL_points_with_lnLcut_12', 'NA')}, {run_diagnostics.get('high_lnL_points_with_lnLcut_10', 'NA')}, {run_diagnostics.get('high_lnL_points_with_lnLcut_5', 'NA')}, {run_diagnostics.get('high_lnL_points_with_lnLcut_2', 'NA')}]\n") + + total_high_points = run_diagnostics.get("total_high_lnL_points", 0) + large_error_points = run_diagnostics.get("high_lnL_points_with_large_error", 0) + high_points_used = run_diagnostics.get("high_lnL_points", None) + + if total_high_points > 0 and large_error_points / total_high_points > 0.5: + f.write("\t--> More than half of the high-likelihood points have large Monte Carlo error. Consider reducing d-max, increasing d-min, increasing n-max, or changing the sampler.\n") + ILE_is_good = False + + if high_points_used is not None: + if high_points_used <= 500: + f.write("\t--> Number of high-likelihood points is less than 500. The initial grid may not have sufficient resolution. Consider reducing the parameter space or increasing the number of grid points.\n") + ILE_is_good = False + + elif high_points_used < 5000: + latest_grid = run_diagnostics.get("Main", {}).get("latest_grid", "latest grid") + f.write(f"\t--> Number of high-likelihood points is less than 5000. Consider rerunning with {latest_grid} as the starting grid and copying this run's all.net as bonus.composite.\n") + ILE_is_good = False + if ILE_is_good ==False and run_diagnostics.get("run_is_complete") == False: + f.write(f"\t--> Since the run is incomplete, this warning can be ignored.\n") + write_composite_information(f, "Main") + + if run_diagnostics.get("Subdag"): + write_composite_information(f, "Subdag") + + f.write("\n") + if ILE_is_good: + f.write("\t--> ILE status: GOOD! <--\n") + else: + f.write("\t--> ILE status: BAD! <--\n") + + write_header(f, "CIP diagnostics") + + CIP_is_good = True + + for plot_title in ["Main", "Subdag"]: + if plot_title not in run_diagnostics: + continue + + if not run_diagnostics[plot_title]: + continue + + diagnostics = run_diagnostics[plot_title] + + f.write(f"\n{plot_title} CIP summary:\n") + f.write(f"\tLatest iteration = {diagnostics.get('latest_iteration', 'NA')}\n") + f.write(f"\tLatest grid = {diagnostics.get('latest_grid', 'NA')}\n") + + latest_iteration = diagnostics.get("latest_iteration", None) + average_lnL = diagnostics.get("cip_average_max_lnL_sampled", None) + std_lnL = diagnostics.get("cip_std_max_lnL_sampled", None) + sampled_lnL = diagnostics.get("cip_sampled_lnL", {}) + + if latest_iteration is None: + f.write(f"\tAverage max lnL sampled by CIP = {average_lnL:0.2f} (+{std_lnL[0]:0.2f}, -{std_lnL[1]:0.2f})\n") + else: + f.write(f"\tAverage max lnL sampled by CIP in iteration {latest_iteration} = {average_lnL:0.2f} (+{std_lnL[0]:0.2f}, -{std_lnL[1]:0.2f})\n") + if "max_lnL" in run_diagnostics: + f.write(f"\tMaximum lnL in all.net = {run_diagnostics['max_lnL']}\n") + + sampling_is_good = write_cip_sampled_lnL_information(f, plot_title) + neff_is_good = write_neff_information(f, plot_title) + + if not neff_is_good or not sampling_is_good: + CIP_is_good = False + + jsd_is_good_main = write_jsd_information(f, ["Main_iteration"], "Main iteration JSD") + jsd_is_good_subdag = write_jsd_information(f, ["Subdag_iteration"], "Subdag iteration JSD") + jsd_is_good_main_subdag = write_jsd_information(f, ["Main_subdag_iteration"], "Main-subdag JSD") + + if not jsd_is_good_main or not jsd_is_good_subdag or not jsd_is_good_main_subdag: + latest_grid = run_diagnostics.get("Main", {}).get("latest_grid", "latest grid") + f.write(f"\n\t--> At least one JSD value is above {JSD_threshold}. This suggests that the run may not have converged. Consider rerunning with {latest_grid} as the starting grid and copying this run's all.net as bonus.composite.\n") + if run_diagnostics.get("run_is_complete") is False: + f.write(f"\t--> Since the run is not complete, this warning can be ignored.\n") + CIP_is_good = False + + f.write("\n") + if CIP_is_good: + f.write("\t--> CIP status: GOOD! <--\n") + else: + f.write("\t--> CIP status: BAD! <--\n") + + #if "sample_statistics" in run_diagnostics: + # write_header(f, "Posterior sample statistics") + # write_dict_information(f, "Equal-tailed 90% intervals", run_diagnostics["sample_statistics"]) + + write_header(f, "Visual diagnostics") + f.write("\t 1) Is the 90% credible interval mostly around the red points? If not, it could be that the run needs more iterations. If the SNR < 30, then the prior might impact it and the shift is expected.") + f.write(f"\n\t 2) Has the parameter space been sufficiently explored? Are there blue points around the red points? Continuing the run will help if this is true with {run_diagnostics['Main']['latest_grid']} as your starting grid and copying this run's all.net as bonus.composite in your new run directory") + f.write("\n\t 3) Is the approximate SNR captured close to True SNR? A significant difference implies the inference got stuck at a local lnL maxima. Happens rarely") + print("\n###########################################################################################") + print("# Run diagnostics") print("###########################################################################################") + print(f"Diagnostics written to {diagnostics_path}") - for key in run_diagnostics: - print(f"{key}: {run_diagnostics[key]}") + for key, value in run_diagnostics.items(): + print(f"{key}: {value}") def check_extrinsic_present(path): + """ + Checks whether an extrinsic posterior sample file exists. + + Args: + path (str): Path to the run directory. + + Returns: + bool: + True if extrinsic_posterior_samples.dat exists, otherwise False. + """ return os.path.exists(f"{path}/extrinsic_posterior_samples.dat") ########################################################################################### # Generate plots ########################################################################################### # create plots folder -create_plots_folder(path) +create_plots_folder(save_path) # finding posterior files main_posterior_files, main_iterations = find_posteriors_in_main(path) @@ -1117,11 +1404,12 @@ def check_extrinsic_present(path): # plot neff try: - plot_neff_data(path) -except: + plot_neff_data(path, plot_title='Main') +except Exception as e: + print(e) # run this function so some information in run_diagnostics dict gets populated. get_lnL_cut_points(all_net_path, lnL_cut=15, error_threshold=0.4, composite=False) - print("Couldn't plot CIP neff per worker for each iteration.") + print("Couldn't plot CIP neff per worker for Main iterations.") # plot exploration corner try: @@ -1132,22 +1420,21 @@ def check_extrinsic_present(path): # plot sampled max lnL try: - plot_cip_max_lnL(path) + plot_cip_max_lnL(path, plot_title='Main') except Exception as e: print(e) - print("Couldn't plot max lnL sampled by CIP per iteration.") + print("Couldn't plot max lnL sampled by CIP for Main iterations.") # plot likelihood exploration try: - plot_high_likelihood_expoloration(path) -except: - print("Couldn't plot high likelihod exploration plot.") - -# write sample statistics -write_sample_statistics(main_posterior_files[-1]) + plot_high_likelihood_expoloration(path, plot_title='Main') +except Exception as e: + print(e) + print("Couldn't plot high likelihod exploration plot for Main iterations.") # plot histograms plot_histograms(main_posterior_files, plot_title="Main", iterations=main_iterations, JSD = False) +plot_histograms([main_posterior_files[-1]], plot_title="Main_Final", iterations=None, JSD = False) # plot corner plots if LISA: @@ -1155,22 +1442,22 @@ def check_extrinsic_present(path): if eccentricity: plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths = use_truths) else: plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta", "chi_eff", "dec", "ra"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "a1z", "a2z", "dec", "ra"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta", "chi_eff", "dec", "ra"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "a1z", "a2z", "dec", "ra"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "a1z", "a2z", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta", "chi_eff", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "a1z", "a2z", "dec", "ra"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "a1z", "a2z", "dec", "ra"], use_truths = use_truths) else: # block with no spin and no eccentricity + precession if non_spinning: plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2"], use_truths = use_truths) else: plot_corner(main_posterior_files, "Main", iterations = main_iterations, use_truths = use_truths) @@ -1178,26 +1465,26 @@ def check_extrinsic_present(path): if eccentricity and not(precessing): plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths = use_truths) elif precessing and not(eccentricity): plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta", "chi_eff", "chi_p"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta", "chi_eff", "chi_p"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta", "chi_eff", "chi_p"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths = use_truths) elif precessing and eccentricity: plot_corner(main_posterior_files, "Main", parameters = ["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], iterations = main_iterations, use_truths = use_truths) plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths = use_truths) elif not(precessing) and not(eccentricity) and not(non_spinning): plot_corner(main_posterior_files, "Main", parameters = ["m1", "m2", "s1z", "s2z"], iterations = main_iterations, use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["m1", "m2", "s1z", "s2z"], use_truths = use_truths) - plot_corner([main_posterior_files[-1]], "Final", parameters = ["mtot", "q", "s1z", "s2z"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["m1", "m2", "s1z", "s2z"], use_truths = use_truths) + plot_corner([main_posterior_files[-1]], "Main_Final", parameters = ["mtot", "q", "s1z", "s2z"], use_truths = use_truths) # plot JS test try: @@ -1220,54 +1507,78 @@ def check_extrinsic_present(path): subdag_posterior_files, subdag_iterations = find_posteriors_in_sub(path, limit_iterations=limit_subdag_iterations) if analyse_subdag: + subdag_path = find_posteriors_in_sub(path, return_subdag_folder_path=True) + print(f'\t-->Analyzing Subdag iterations found in {subdag_path}<--') + try: + plot_neff_data(subdag_path, plot_title='Subdag') + except Exception as e: + print(e) + print("Couldn't plot CIP neff per worker for Subdag iterations.") + + # plot sampled max lnL + try: + plot_cip_max_lnL(subdag_path, plot_title='Subdag') + except Exception as e: + print(e) + print("Couldn't plot max lnL sampled by CIP for Subdag iterations.") + + # plot likelihood exploration + try: + plot_high_likelihood_expoloration(subdag_path, plot_title='Subdag') + except Exception as e: + print(e) + print("Couldn't plot high likelihod exploration plot for Subdag iterations.") + + plot_histograms(subdag_posterior_files, plot_title="Subdag", iterations=subdag_iterations, JSD = False) + # plot corner plots if LISA: plot_corner(subdag_posterior_files, "Subdag", iterations=subdag_iterations, use_truths=use_truths) if eccentricity: plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "a1z", "a2z", "eccentricity", "meanPerAno", "dec", "ra"], use_truths=use_truths) else: plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta", "chi_eff", "dec", "ra"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "a1z", "a2z", "dec", "ra"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta", "chi_eff", "dec", "ra"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "a1z", "a2z", "dec", "ra"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "a1z", "a2z", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta", "chi_eff", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "a1z", "a2z", "dec", "ra"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "a1z", "a2z", "dec", "ra"], use_truths=use_truths) else: if non_spinning: plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2"], use_truths=use_truths) else: plot_corner(subdag_posterior_files, "Subdag", iterations=subdag_iterations, use_truths=use_truths) - if eccentricity and not(precessing): + if eccentricity and not(precessing) and not(non_spinning): plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths=use_truths) - elif precessing and not(eccentricity): + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta", "chi_eff", "eccentricity", "meanPerAno"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "s1z", "s2z", "eccentricity", "meanPerAno"], use_truths=use_truths) + elif precessing and not(eccentricity) and not(non_spinning): plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta", "chi_eff", "chi_p"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta", "chi_eff", "chi_p"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths=use_truths) - elif precessing and eccentricity: + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta", "chi_eff", "chi_p"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y"], use_truths=use_truths) + elif precessing and eccentricity and not(non_spinning): plot_corner(subdag_posterior_files, "Subdag", parameters=["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], iterations=subdag_iterations, use_truths=use_truths) plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mc", "eta", "chi_eff", "chi_p", "eccentricity", "meanPerAno"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "s1z", "s2z", "s1x", "s1y", "s2x", "s2y", "eccentricity", "meanPerAno"], use_truths=use_truths) elif not(precessing) and not(eccentricity) and not(non_spinning): plot_corner(subdag_posterior_files, "Subdag", parameters=["m1", "m2", "s1z", "s2z"], iterations=subdag_iterations, use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["m1", "m2", "s1z", "s2z"], use_truths=use_truths) - plot_corner([subdag_posterior_files[-1]], "Subdag", parameters=["mtot", "q", "s1z", "s2z"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["m1", "m2", "s1z", "s2z"], use_truths=use_truths) + plot_corner([subdag_posterior_files[-1]], "Subdag_Final", parameters=["mtot", "q", "s1z", "s2z"], use_truths=use_truths) try: plot_JS_divergence(subdag_posterior_files[-1], subdag_posterior_files[-2], None, "Subdag_iteration") # the last two subdag iterations @@ -1278,10 +1589,13 @@ def check_extrinsic_present(path): print("Couldn't plot Jensen Shannon Divergence plot.") plot_JS_divergence(main_posterior_files[-1], subdag_posterior_files[-1], None, "Main_subdag_iteration") # the last main and subdag iteration +run_diagnostics['run_is_complete'] = False if check_extrinsic_present(path): plot_corner([f"{path}/extrinsic_posterior_samples.dat"], "extrinsic", parameters = ["distance", "incl", "phiorb", "psi", "time"], use_truths = use_truths) plot_corner([f"{path}/extrinsic_posterior_samples.dat"], "extrinsic_source_mass", parameters = ["m1_source", "m2_source", "mtotal_source"], use_truths = False) + write_sample_statistics(f"{path}/extrinsic_posterior_samples.dat", extrinsic = True) plot_log_likelihood(path) + run_diagnostics['run_is_complete'] = True # run diagnostics evaluate_run(run_diagnostics) From 1e27d31acdebfa278db814ce77fc9482dee26b35 Mon Sep 17 00:00:00 2001 From: Aasim Z Jan <58118417+AasimZahoor@users.noreply.github.com> Date: Sun, 7 Jun 2026 10:58:08 -0500 Subject: [PATCH 2/5] util_LALNRWriteFrame.py: NR frame writer. This version allows the user to bypass the RIFT NR infrastructure and allows the user to use the new SXS interface, where they do not provide the NR files in h5 format. While writing the SXS function, I ran into a number of bugs that I reported to the SXS people (here:https://github.com/sxs-collaboration/sxs/issues/188) and as such I have hardcoded a few settings to avoid those bugs. As SXS developers fix those bugs, we should remove these hardcoded settings. I have compared the LVK format function with NR infrastructure frame writer for multiple NR files, and they match within a time sample. --- .../Code/bin/util_LALNRWriteFrame.py | 393 ++++++++++++++++++ 1 file changed, 393 insertions(+) create mode 100755 MonteCarloMarginalizeCode/Code/bin/util_LALNRWriteFrame.py diff --git a/MonteCarloMarginalizeCode/Code/bin/util_LALNRWriteFrame.py b/MonteCarloMarginalizeCode/Code/bin/util_LALNRWriteFrame.py new file mode 100755 index 00000000..d36aeba3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/bin/util_LALNRWriteFrame.py @@ -0,0 +1,393 @@ +#! /usr/bin/env python +# +# To generate NR injections using LALSimulations's function. This bypasses RIFT's NR infrastructure, the base code is utiL_LALWriteFrame.py with added function to take in NRhdf5 file and generate polarizations. + +import argparse +import numpy as np +import RIFT.lalsimutils as lalsimutils +import lalsimulation as lalsim +import lal +import h5py +from astropy.time import Time +import romspline +try: + import sxs +except: + print('SXS package not installed.') + +parser = argparse.ArgumentParser() +parser.add_argument("--fname", default=None, help = "Base name for output frame file. Otherwise auto-generated ") +parser.add_argument("--instrument", default="H1",help="Use H1, L1,V1") +parser.add_argument("--inj", dest='inj', default=None,help="inspiral XML file containing injection information.") +parser.add_argument("--event",type=int, dest="event_id", default=None,help="event ID of injection XML to use.") +parser.add_argument("--approx",type=str,default=None) +parser.add_argument("--srate",type=int,default=16384,help="Sampling rate") +parser.add_argument("--seglen", type=float,default=16., help="Default window size for processing.") +parser.add_argument("--start", type=int,default=None) +parser.add_argument("--stop", type=int,default=None) +parser.add_argument("--incl",default=None,help="Set the inclination of L (at fref). Particularly helpful for aligned spin tests") +parser.add_argument("--mass1",default=10,type=float,help='Mass 1 (solar masses)') +parser.add_argument("--mass2",default=1.4,type=float,help='Mass 2 (solar masses)') +parser.add_argument("--l-max",default=None,type=int,help='Inclusion of modes in injection') +parser.add_argument("--path-to-hdf5", help='Path to NRhdf5 file. This needs to be in the LVK format') +parser.add_argument("--modes-list", type=str, default=None, help="List of specific modes you want to use. Set l-max to None if you want to use this option.") +parser.add_argument("--sxs-simulation-name", type=str, default=None, help="SXS simulation name you want to use to generate the injection") +parser.add_argument("--taper-percent", default=None, type=float, help="Amount of waveform to taper, if None tapering will not be performed. Should be between 0 and 1.") +parser.add_argument("--verbose", action="store_true",default=False) +opts= parser.parse_args() + +def get_lvk_modes_from_NRhdf5(P, path_to_hdf5, modes_list=opts.modes_list, l_max=opts.l_max): + + print(f"Reading waveform from {path_to_hdf5}") + + # Converstion factors using lal + kgs_to_sec = lal.G_SI/lal.C_SI**3 + code_units_to_sec = lal.MTSUN_SI + meters_to_sec = 1/lal.C_SI + + # get mtotal based on user input. This is in kgs. + mtotal= P.m1 + P.m2 + mtot_in_sec = mtotal * kgs_to_sec + dist_in_sec = P.dist * meters_to_sec # distance in m + + # load in hdf5 file to get masses and fmin + data_1 = h5py.File(path_to_hdf5,"r") + m1, m2 = data_1.attrs["mass1"] * mtotal, data_1.attrs["mass2"] * mtotal + fmin = data_1.attrs["f_lower_at_1MSUN"] * lal.MSUN_SI/mtotal + fref = 0.0 # set to zero to avoid errors + print(f"Smallest possible frequency for this waveform {fmin} Hz. Frequency at 1 solar mass is {data_1.attrs['f_lower_at_1MSUN']}.\nReference frequency is set to {fref} Hz since the lalsimulation function does not take non-zero reference frequency.") + + # get spins, useful for precessing case + s1x, s1y, s1z, s2x, s2y, s2z = lalsim.SimInspiralNRWaveformGetSpinsFromHDF5File(fref, mtotal/lal.MSUN_SI, path_to_hdf5) + + # Collect modes based on input + if modes_list == None and l_max is not None: + modes = [(l, m) for l in range(2, l_max + 1) for m in range(-l, l + 1) if m != 0] + elif modes_list is not None and l_max is None: + modes = list(eval(modes_list)) + elif modes_list is not None and l_max is not None: + raise ValueError("Use either l_max or modes_list, not both.") + else: + raise ValueError("One of l_max or modes_list must be provided.") + + print(f"Generating waveform with m1 = {m1/lal.MSUN_SI:0.4f} MSUN, m2 = {m2/lal.MSUN_SI:0.4f} MSUN \n s1 = {[s1x, s1y, s1z]}, s2 = {[s2x, s2y, s2z]}, eccentricity = {data_1.attrs['eccentricity']}, meanPerAno = {data_1.attrs['mean_anomaly']}, \n fmin = {fmin} Hz, fref= {fref}") + print(f"Modes requested = {modes}") + print(f"WARNING: The provided fmin has no effect; the waveform starts at the lowest available frequency of the NR simulation, which is {fmin} Hz.") + + #interpolating using romspline + hlm = {} + dt_in_code_units = P.deltaT / code_units_to_sec * lal.MSUN_SI/mtotal + for i in range(len(modes)): + amp22_time_0=np.array(data_1[f"phase_l{modes[i][0]}_m{modes[i][1]}"]["X"]) + + amp = romspline.readSpline(path_to_hdf5, f"amp_l{modes[i][0]}_m{modes[i][1]}") + phase = romspline.readSpline(path_to_hdf5, f"phase_l{modes[i][0]}_m{modes[i][1]}") + + amp22_time_0 = np.arange(np.min(amp22_time_0), np.max(amp22_time_0), dt_in_code_units) + generated_amp = amp(amp22_time_0) + generated_phase = phase(amp22_time_0) + generated_phase = np.unwrap(generated_phase) + + mode_content_at_distance = mtot_in_sec/dist_in_sec * generated_amp * np.exp(1j*generated_phase) + + # Save as a lal object + wf = lal.CreateCOMPLEX16TimeSeries("hlm", 0, 0, P.deltaT, lal.DimensionlessUnit, len(mode_content_at_distance)) + wf.data.data *= 0 + wf.data.data = mode_content_at_distance + + # resize + if P.deltaF: + TDlen = int(1./P.deltaF * 1./P.deltaT) + if TDlen < wf.data.length: # Truncate the series to the desired length, removing data at the *start* (left) + wf = lal.ResizeCOMPLEX16TimeSeries(wf, wf.data.length-TDlen, TDlen) + elif TDlen > wf.data.length: # Zero pad, extend at end + wf = lal.ResizeCOMPLEX16TimeSeries(wf, 0, TDlen) + + # tapering + taper = opts.taper_percent is not None + if taper and P.deltaF is not None: + #TDlen = int(1./P.deltaF * 1./P.deltaT) + TDlen = wf.data.length + ntaper = int(opts.taper_percent*TDlen) + vectaper= 0.5 - 0.5*np.cos(np.pi*np.arange(ntaper)/(1.*ntaper)) + # Taper at the start of the segment + wf.data.data[:ntaper]*=vectaper + + hlm[modes[i][0],modes[i][1]] = wf + + # set epoch based on GWsignal approach + rhosq = np.zeros(TDlen) + for mode in hlm: + rhosq += np.abs(hlm[mode].data.data)**2 + indx_max = np.argmax(rhosq) + new_epoch = - indx_max*P.deltaT + for mode in hlm: + hlm[mode].epoch = new_epoch + + return hlm + + +def get_lvk_modes_from_SXS(P, simulation_name, modes_list=opts.modes_list, l_max=opts.l_max, set_fref_equal_to_fmin=True): + print(f"Loading SXS waveform {simulation_name}") + # Load metadata + sim = sxs.load(simulation_name) + + # Converstion factors using lal + kgs_to_sec = lal.G_SI/lal.C_SI**3 + code_units_to_sec = lal.MTSUN_SI + meters_to_sec = 1/lal.C_SI + + # Conversion of mass and distance to seconds + mtotal = P.m1 + P.m2 # in kgs + mtot_in_sec = mtotal * kgs_to_sec + dist_in_sec = P.dist * meters_to_sec # distance in m + + # Get mass in real units based on simulation parameters + m1, m2 = sim.metadata['reference_mass1'] * mtotal, sim.metadata['reference_mass2'] * mtotal + + # Get minimum frequency. Note: initial_orbital_frequency gives frequency that is almost of f_low + # This did not match with NRHybSur3dq8 comparisons + # _, _, dyn_test = sim.to_lvk(t_ref=0.0, ell_max=8) + # flow_at_1MSUN = dyn_test['f_low'] / code_units_to_sec + flow_at_1MSUN = sim.metadata['initial_orbital_frequency']/ (np.pi*code_units_to_sec) + flow = flow_at_1MSUN * lal.MSUN_SI/mtotal + + # Get fref. f_22 = |omega|/pi + Omega_vec = sim.metadata["reference_orbital_frequency"] + Omega = np.linalg.norm(Omega_vec) + fref_at_1MSUN = np.array(Omega) / (np.pi*code_units_to_sec) + fref = fref_at_1MSUN * lal.MSUN_SI/mtotal + print(f"Smallest possible frequency for this waveform {flow} Hz. Frequency at 1 solar mass is {flow_at_1MSUN} Hz.\nReference frequency for this waveform is {fref} Hz. Reference frequency at 1 solar mass in {fref_at_1MSUN} Hz.") + + # Sanity checks for fmin. Don't go below what the waveform can provide. + if P.fmin <= flow and P.fmin != 0.0: + fmin = flow #+ 0.5*10**(-2)*flow + print(f"WARNING: Can't have fmin less than that of the NR waveform. Provided fmin is {P.fmin} Hz, defaulting to fmin={fmin} Hz.") + else: + fmin = P.fmin + print(f"Generating waveform with fmin is {P.fmin} Hz.") + print(f"Generating waveform with m1 = {m1/lal.MSUN_SI:0.4f} MSUN, m2 = {m2/lal.MSUN_SI:0.4f} MSUN \n s1 = {sim.metadata['reference_dimensionless_spin1']}, s2 = {sim.metadata['reference_dimensionless_spin2']}, eccentricity = {sim.metadata['reference_eccentricity']}, meanAnomaly = {sim.metadata['reference_mean_anomaly']}\n fmin = {fmin} Hz") + + # Collect modes based on input + if modes_list == None and l_max is not None: + modes = [(l, m) for l in range(2, l_max + 1) for m in range(-l, l + 1) if m != 0] + elif modes_list is not None and l_max is None: + modes = list(eval(modes_list)) + l_max = np.max(np.array(modes)[:, 0]) + elif modes_list is not None and l_max is not None: + raise ValueError("Use either l_max or modes_list, not both.") + else: + raise ValueError("One of l_max or modes_list must be provided.") + print(f"Modes requested = {modes}") + + # Now we use dt and fmin to generate modes + dt_in_code_units = P.deltaT / code_units_to_sec * lal.MSUN_SI/mtotal + print(f"WARNING: The provided fmin has no effect; the waveform starts at the lowest available frequency of the NR simulation, which is {flow} Hz.") + fmin_in_code_units = fmin * (mtotal/lal.MSUN_SI) * code_units_to_sec + fref_in_code_units = fref * (mtotal/lal.MSUN_SI) * code_units_to_sec + # This function gives weird results for eccentric waveform if f_ref is different than f_low + if set_fref_equal_to_fmin: + fref_in_code_units = fmin_in_code_units + times, hlms_dict, dyn = sim.to_lvk(f_ref=fref_in_code_units, ell_max=l_max, dt=dt_in_code_units, f_low=fmin_in_code_units) + if modes_list is not None: + missing_modes = [mode for mode in modes_list if mode not in hlms_dict] + if missing_modes: + print(f"WARNING: The following modes are not present and will be ignored: {missing_modes}") + modes = [mode for mode in modes_list if mode in hlms_dict] + + # collect modes + hlm = {} + # remove junk radiation, which I have found to be usually within 150M of the start. + tvals = np.arange(len(times))*dt_in_code_units + index = np.argwhere(tvals >= 150).flatten()[0] + print(f"Removing content from 0 M upto {tvals[index]} M") + for i, mode in enumerate(modes): + # extract modes from the SXS object + mode_content_here = hlms_dict[mode][index:] + + # Scale based on distance and mtotal + mode_content_at_distance = mtot_in_sec/dist_in_sec * mode_content_here + + # Save as a lal object + wf = lal.CreateCOMPLEX16TimeSeries("hlm", 0, 0, P.deltaT, lal.DimensionlessUnit, len(mode_content_at_distance)) + wf.data.data *= 0 + wf.data.data = mode_content_at_distance + + # resize + if P.deltaF: + TDlen = int(1./P.deltaF * 1./P.deltaT) + if TDlen < wf.data.length: # Truncate the series to the desired length, removing data at the *start* (left) + wf = lal.ResizeCOMPLEX16TimeSeries(wf, wf.data.length-TDlen, TDlen) + elif TDlen > wf.data.length: # Zero pad, extend at end + wf = lal.ResizeCOMPLEX16TimeSeries(wf, 0, TDlen) + + # tapering + taper = opts.taper_percent is not None + if taper and P.deltaF is not None: + TDlen = wf.data.length + ntaper = int(opts.taper_percent*TDlen) + vectaper= 0.5 - 0.5*np.cos(np.pi*np.arange(ntaper)/(1.*ntaper)) + # Taper at the start of the segment + wf.data.data[:ntaper]*=vectaper + + hlm[modes[i][0],modes[i][1]] = wf + + # set epoch based on GWsignal approach + rhosq = np.zeros(TDlen) + for mode in hlm: + rhosq += np.abs(hlm[mode].data.data)**2 + indx_max = np.argmax(rhosq) + new_epoch = - indx_max*P.deltaT + for mode in hlm: + hlm[mode].epoch = new_epoch + return hlm + +def get_polarizations_from_modes(P, hlms): + hp = lal.CreateREAL8TimeSeries("hp", lal.LIGOTimeGPS(0.), 0., hlms[2,2].deltaT, lal.DimensionlessUnit, hlms[2,2].data.length) + hc = lal.CreateREAL8TimeSeries("hc", lal.LIGOTimeGPS(0.), 0., hlms[2,2].deltaT, lal.DimensionlessUnit, hlms[2,2].data.length) + hp.epoch = hlms[(2,2)].epoch + hc.epoch = hlms[(2,2)].epoch + hp.data.data *= 0 + hc.data.data *= 0 + + wfmTS = lal.CreateCOMPLEX16TimeSeries("wfmTS", lal.LIGOTimeGPS(0.), 0., hlms[2,2].deltaT, lal.DimensionlessUnit, hlms[2,2].data.length) + wfmTS.epoch = hlms[(2,2)].epoch + wfmTS.data.data *= 0 + for mode in list(hlms.keys()): + wfmTS.data.data += hlms[mode].data.data*lal.SpinWeightedSphericalHarmonic(P.incl, -P.phiref, -2, int(mode[0]), int(mode[1])) + + hp.data.data = np.real(wfmTS.data.data) + hc.data.data = -1*np.imag(wfmTS.data.data) + + return hp, hc + +def generate_hoft(P, hp, hc, Fp=None, Fc=None): + + # Apply detector response + if Fp!=None and Fc!=None: + hp.data.data *= Fp + hc.data.data *= Fc + hp = lal.AddREAL8TimeSeries(hp, hc) + ht = hp + elif P.radec==False: + fp = Fplus(P.theta, P.phi, P.psi) + fc = Fcross(P.theta, P.phi, P.psi) + hp.data.data *= fp + hc.data.data *= fc + hp = lal.AddREAL8TimeSeries(hp, hc) + ht = hp + else: + # If astropy Time function, overwrite with GPS time, otherwise use normal addition + if isinstance(hp.epoch, Time): + dT = hp.epoch.to_value('gps','long') # pull out the time + hp.epoch = P.tref + dT + hc.epoch = P.tref +dT + else: + hp.epoch = hp.epoch + P.tref + hc.epoch = hc.epoch + P.tref + ht = lalsim.SimDetectorStrainREAL8TimeSeries(hp, hc, + P.phi, P.theta, P.psi, + lalsim.DetectorPrefixToLALDetector(str(P.detector))) + + return ht + +# Generate signal +P = lalsimutils.ChooseWaveformParams() +P.deltaT = 1./opts.srate +P.radec = True # use a real source with a real instrument +if not opts.inj: + P.randomize(aligned_spin_Q=True,default_inclination=opts.incl) + P.m1 = opts.mass1*lalsimutils.lsu_MSUN + P.m2 = opts.mass2*lalsimutils.lsu_MSUN + P.taper = lalsimutils.lsu_TAPER_START + P.tref =1000000000 # default + if opts.approx: + P.approx = lalsim.GetApproximantFromString(str(opts.approx)) + else: + P.approx = lalsim.GetApproximantFromString("SpinTaylorT2") +else: + from igwn_ligolw import lsctables, table, utils # check all are needed + #from ligo.lw import lsctables, table, utils + filename = opts.inj + event = opts.event_id + xmldoc = utils.load_filename(filename, verbose = True, contenthandler =lalsimutils.cthdler) + sim_inspiral_table = lsctables.SimInspiralTable.get_table(xmldoc) + P.copy_sim_inspiral(sim_inspiral_table[int(event)]) + P.taper = lalsimutils.lsu_TAPER_START + #if opts.approx: + # P.approx = lalsim.GetApproximantFromString(str(opts.approx)) + +P.taper = lalsimutils.lsu_TAPER_START # force taper +P.detector = opts.instrument +if opts.approx == "EccentricTD": + P.phaseO = 3 +P.print_params() + +T_est = lalsimutils.estimateWaveformDuration(P) +T_est = P.deltaT*lalsimutils.nextPow2(T_est/P.deltaT) +if T_est > opts.seglen: + print(" WARNING: THE SIGNAL WILL LIKELY BE TRUNCATED when writing the frame, which is VERY BAD ") +T_est = opts.seglen +P.deltaF = 1./T_est +print(" Duration ", T_est) +if T_est < opts.seglen: + print(" Buffer length too short, automating retuning forced ") + +# Generate signal +if opts.sxs_simulation_name is not None: + hlms = get_lvk_modes_from_SXS(P, opts.sxs_simulation_name) +else: + hlms = get_lvk_modes_from_NRhdf5(P, opts.path_to_hdf5) +hp, hc = get_polarizations_from_modes(P, hlms) +hoft = generate_hoft(P, hp, hc) +epoch_orig = hoft.epoch +# zero pad to be opts.seglen long, if necessary +if opts.seglen/hoft.deltaT > hoft.data.length: + TDlenGoal = int(opts.seglen/hoft.deltaT) + hoft = lal.ResizeREAL8TimeSeries(hoft, 0, TDlenGoal) + +# zero pad some more on either side, to make sure the segment covers start to stop +if opts.start and hoft.epoch > opts.start: + nToAddBefore = int((float(hoft.epoch)-opts.start)/hoft.deltaT) + # hoft.epoch - nToAddBefore*hoft.deltaT # this is close to the epoch, but not quite ... we are adjusting it to be within 1 time sample + print(nToAddBefore, hoft.data.length) + ht = lal.CreateREAL8TimeSeries("Template h(t)", + opts.start , 0, hoft.deltaT, lalsimutils.lsu_DimensionlessUnit, + hoft.data.length+nToAddBefore) + ht.data.data = np.zeros(ht.data.length) # clear + ht.data.data[nToAddBefore:nToAddBefore+hoft.data.length] = hoft.data.data + hoft = ht + +if opts.stop and hoft.epoch+hoft.data.length*hoft.deltaT < opts.stop: + nToAddAtEnd = int( (-(hoft.epoch+hoft.data.length*hoft.deltaT)+opts.stop)/hoft.deltaT) + print("Padding end ", nToAddAtEnd, hoft.data.length) + hoft = lal.ResizeREAL8TimeSeries(hoft,0, int(hoft.data.length+nToAddAtEnd)) +channel = opts.instrument+":FAKE-STRAIN" + +tstart = int(hoft.epoch) +duration = int(round(hoft.data.length*hoft.deltaT)) +if not opts.fname: + fname = opts.instrument.replace("1","")+"-fake_strain-"+str(tstart)+"-"+str(duration)+".gwf" + +print("Writing signal with ", hoft.data.length*hoft.deltaT, " to file ", fname) +lalsimutils.hoft_to_frame_data(fname,channel,hoft) + +# TEST: Confirm it works by reading the frame +if opts.verbose: + print(" ----- Plotting data ------ ") + import os + from matplotlib import pyplot as plt + # First must create corresponding cache file + os.system("echo "+ fname+ " | lal_path2cache > test.cache") + # Now I can read it + # Beware that the results are OFFSET FROM ONE ANOTHER due to PADDING, + # but that the time associations are correct + hoft2 = lalsimutils.frame_data_to_hoft("test.cache", channel) + tvals2 = (float(hoft2.epoch) - float(P.tref)) + np.arange(hoft2.data.length)*hoft2.deltaT + tvals = (float(hoft.epoch) - float(P.tref)) + np.arange(hoft.data.length)*hoft.deltaT + plt.plot(tvals2,hoft2.data.data,label='Fr') + plt.plot(tvals,hoft.data.data,label='orig') + plt.xlim(float(epoch_orig)- float(P.tref), 0.2) + plt.xlabel('t - tref') + plt.legend(); #plt.show() + plt.savefig("injected-data_"+opts.instrument +".png") From c9e9c7b2a83484cdc6e292dc4cb6b98fa127e953 Mon Sep 17 00:00:00 2001 From: Aasim Z Jan <58118417+AasimZahoor@users.noreply.github.com> Date: Sun, 7 Jun 2026 11:03:59 -0500 Subject: [PATCH 3/5] EFPE.py: implemented the EFPE (eccentric+precessing PN) model. Since they do not provide modes, this calls the waveform at 5 different inclinations and phases to compute modes. This approach is copied from the hlmoft_IMRPv2_dict function in lalsimutils. Has been tested on 3 cases with satisfactory performance, however it needs thorough testing. Also, we need to implement changes in factored_likelihood and util_RIFT_pseudopipe so this can be used. --- .../Code/RIFT/physics/EFPE.py | 324 ++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py diff --git a/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py b/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py new file mode 100644 index 00000000..d300ffa6 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py @@ -0,0 +1,324 @@ +# Wrapper for EFPE models, for now just pyEFPE + +# What is the plan? +# 1) Write a function that can output polarizations in time domain. +# 2) Use the PhenomPv2 approach to get the modes +# Need to be careful about epoch. Also the input frequency array should mimic that of time series. + +# Functions 1) get_ISCO_fval 2) get_frequency_array 3) get_polarizations 4) complex_hoft 5) hoft 6) hlmoft 7) std_and_conj_hlmoff +# What is this code doing? First, get_polarization gets the waveform polarizations from pyEFPE in FD, FFTs it to TD (while properly setting epoch) and returns it. I learned the hardway to have a cutoff on fmax so the fmax cutoff is ISCO. This function is utlized in complex_hoft which return hp-1jhc *exp(2j*psi). complex_hoft is used by hlmoft which is in turn used by std_and_conj_hlmoff (that is used in recovery). hoft uses get_polarizations to create injections. +# Potential sources of problems 1) How I construct hlms 2) Epoch 3) Injection creation. The lal detector series function adds more points to the time series. + +import lal +import RIFT.lalsimutils as lalsimutils +import lalsimulation as lalsim +import pyEFPE +import numpy as np +import RIFT.lalsimutils as lalsimutils +from astropy.time import Time + +debug=True +use_ISCO_cutoff=False # FFT routines sometimes fail if I set this to True. Bizarre! + +def get_ISCO_fval(P): + """ + Compute the ISCO (innermost stable circular orbit) frequency for the binary system. This frequency serves as an upper frequency cutoff for waveform generation. + + Args: + P: ChooseWaveformParams object. + + Returns: + isco_fval (float): ISCO frequency in Hz. + """ + # c^3/ (6^(3/2)*pi*G*M_tot) + isco_fval = 4331.648896/(P.m1 + P.m2) * lal.MSUN_SI + return isco_fval + +def get_frequency_array(P, use_ISCO_cutoff=use_ISCO_cutoff): + """ + Generate frequency arrays used for waveform generation + Args: + P: ChooseWaveformParams object + Returns: + fvals_wf (np.ndarray): Frequency array starting from the closest frequency to fmin, used for waveform generation. + fvals (np.ndarray): Full frequency array from 0 to Nyquist frequency (inclusive). + index_fmin (int): Index in `fvals` corresponding to the first frequency, which is closes to fmin. + index_fmax (int): Index in `fvals` corresponding to the first frequency, which is closes to f_isco. + """ + # for real series the fvals go from 0 -> fNyq + fvals = np.arange(0, 1/P.deltaT/2 + P.deltaF, P.deltaF) # 0 and fNyq included + + # find the fval closest to fmin + index_fmin = np.argmin(np.abs(fvals - P.fmin)) + + # fvals for waveform generation (NOTE: no upper cutoff here!) + fvals_wf = fvals[index_fmin:] + + # generate the waveform only upto ISCO + index_fmax = -1 + if use_ISCO_cutoff: + isco_fval = get_ISCO_fval(P) + index_fmax = np.argmin(np.abs(fvals - isco_fval)) + fvals_wf = fvals[index_fmin:index_fmax+1] + + if debug: + print(f'Waveform being generated from {fvals_wf[0]} to {fvals_wf[-1]} Hz with deltaF {P.deltaF} Hz') + + return fvals_wf, fvals, index_fmin, index_fmax + +def get_polarizations(P, return_in_FD=False, use_ISCO_cutoff=use_ISCO_cutoff): + """ + Generate gravitational wave polarizations (h₊, h×) for a binary system using pyEFPE. + Args: + P: ChooseWaveformParams object + return_in_FD (bool): If True, return frequency-domain polarizations; otherwise, return time-domain versions. + Returns: + hp, hc: + If return_in_FD is False: + hp (lal.REAL8TimeSeries): Time-domain h₊ polarization. + hc (lal.REAL8TimeSeries): Time-domain h× polarization. + If return_in_FD is True: + hp_f (lal.COMPLEX16FrequencySeries): Frequency-domain h₊ polarization. + hc_f (lal.COMPLEX16FrequencySeries): Frequency-domain h× polarization. + """ + params = { + 'mass1': P.m1/lal.MSUN_SI, # Mass of companion 1 (solar masses) + 'mass2': P.m2/lal.MSUN_SI, # Mass of companion 2 (solar masses) + 'e_start':P.eccentricity, # Initial eccentricity + 'mean_anomaly_start': P.meanPerAno, # initial mean anomaly of quasi keplerian parametrization + 'spin1x': P.s1x, # Spin components of companion 1 + 'spin1y': P.s1y, + 'spin1z': P.s1z, + 'spin2x': P.s2x, # Spin components of companion 2 + 'spin2y': P.s2y, + 'spin2z': P.s2z, + 'inclination': P.incl, # Initial binary inclination (radians) + 'phi_start': P.phiref, # Initial orbital phase (radians) + 'f22_start': P.fmin, # Starting (simulation) waveform frequency of GW 22 mode (Hz) + 'distance': P.dist/1e6/lal.PC_SI # Distance to the source, in Mpc + } + + if debug: + print(f'\nGenerating polarizations for params = {params}') + + # Initialize pyEFPE waveform model + wf = pyEFPE.pyEFPE(params) + + # Define frequency array for waveform generation + freqs, freqs_full, index_fmin, index_fmax = get_frequency_array(P, use_ISCO_cutoff=use_ISCO_cutoff) + + # Compute frequency-domain gravitational wave polarizations + hpf, hcf = wf.generate_waveform(freqs) + + # Create lal objects so we can compute FFT + # hp + hp_f = lal.CreateCOMPLEX16FrequencySeries("Template hp(f)", + 0.0, 0.0, P.deltaF, lal.HertzUnit, # epoch set to 0.0 and so is f0. Will set epoch when in TD + len(freqs_full)) + hp_f.data.data *= 0.0 + if use_ISCO_cutoff: + hp_f.data.data[index_fmin:index_fmax+1] = hpf + else: + hp_f.data.data[index_fmin:] = hpf + + # hc + hc_f = lal.CreateCOMPLEX16FrequencySeries("Template hc(f)", + 0.0, 0.0, P.deltaF, lal.HertzUnit, # epoch set to 0.0 and so is f0. Will set epoch when in TD + len(freqs_full)) + hc_f.data.data *= 0.0 + if use_ISCO_cutoff: + hc_f.data.data[index_fmin:index_fmax+1] = hcf + else: + hc_f.data.data[index_fmin:] = hcf + + # return FD polarizations + if return_in_FD: + return hp_f, hc_f + + # FFT + hp_t, hc_t = lalsimutils.DataInverseFourierREAL8(hp_f), lalsimutils.DataInverseFourierREAL8(hc_f) + + # roll them so the polarizations epoch is positioned at center + tchirp = lalsim.SimInspiralChirpTimeBound(P.fmin, P.m1,P.m2, P.s1z, P.s2z) + s = lalsim.SimInspiralFinalBlackHoleSpinBound(P.s1z,P.s2z) + t_merge = lalsim.SimInspiralMergeTimeBound(P.m1,P.m2) + lalsim.SimInspiralRingdownTimeBound(P.m1+P.m2,s) + factor_roll = 0.98 + assert t_merge < (1-factor_roll)*hp_t.data.length*hp_t.deltaT, f'The waveform is being rolled by factor {factor_roll} and this is causing wraparound. The predicted merge time is {t_merge}s and time left for it is {(1-factor_roll)*hp_t.data.length*hp_t.deltaT}, tchirp = {tchirp} {factor_roll*hp_t.data.length*hp_t.deltaT}' + + hp_t.data.data = np.roll(hp_t.data.data, int(factor_roll*hp_t.data.length)) + hc_t.data.data = np.roll(hc_t.data.data, int(factor_roll*hc_t.data.length)) + + # find epoch: based on the apporach in GWSignal.py + amplitude = np.sqrt(hp_t.data.data**2 + hc_t.data.data**2) + max_amp_index = np.argmax(amplitude) + hp_t.epoch, hc_t.epoch = -max_amp_index * hp_t.deltaT, -max_amp_index * hc_t.deltaT + if debug: + print(f'Length in FD = {hp_f.data.length}, length in TD = {hp_t.data.length, hp_t.data.length * hp_t.deltaT}, epoch set to {hp_t.epoch}') + + + return hp_t, hc_t + + +def complex_hoft(P, sgn=-1, use_ISCO_cutoff=use_ISCO_cutoff): + """ + Generate the complex strain h(t) = h₊(t) - i·h×(t) using the pyEFPE waveform model. + Args: + P: ChooseWaveformParams object + sgn (int, optional): Sign convention for constructing complex strain. + Default is -1, corresponding to h(t) = h₊(t) - i·h×(t). + Returns: + hT (lal.COMPLEX16TimeSeries): Time-domain complex gravitational wave strain. + """ + # Generate time-domain polarizations + hp_t, hc_t = get_polarizations(P, return_in_FD=False, use_ISCO_cutoff=use_ISCO_cutoff) + + # Create complex strain time series h(t) = h₊(t) - i·h×(t) + hT = lal.CreateCOMPLEX16TimeSeries("Complex h(t)", hp_t.epoch, hp_t.f0, + hp_t.deltaT, lal.DimensionlessUnit, hp_t.data.length) + hT.data.data *= 0.0 + hT.data.data = np.real(hp_t.data.data) + 1j * sgn * np.real(hc_t.data.data) + + # Include polarization + hT.data.data *= np.exp(2j*sgn*P.psi) + + return hT + +def hoft(P, Fp=None, Fc=None, use_ISCO_cutoff=use_ISCO_cutoff, **kwargs): + """ + Generate the detector strain time series h(t) from gravitational wave polarizations. + Args: + P: ChooseWaveformParams object + Fp (float, optional): Antenna pattern projection factor for the plus polarization. + Fc (float, optional): Antenna pattern projection factor for the cross polarization. + **kwargs: Additional keyword arguments (currently unused). + Returns: + ht (lal.REAL8TimeSeries): Detector strain time series h(t). + """ + # Call polarizations in TD + P_copy = P.manual_copy() + P_copy.deltaF = 1.01 * P.deltaF # Why? 1/DeltaF is T, but SimDetectorStrainREAL8TimeSeries adds more points so by increasing it, I am decreasing the length, allowing the addition of points to not break the code. This should be checked for signals where the signal length is almost the same as data length. + hp, hc = get_polarizations(P_copy, use_ISCO_cutoff=use_ISCO_cutoff) + + # Apply detector response + if Fp!=None and Fc!=None: + hp.data.data *= Fp + hc.data.data *= Fc + hp = lal.AddREAL8TimeSeries(hp, hc) + ht = hp + elif P.radec==False: + fp = Fplus(P.theta, P.phi, P.psi) + fc = Fcross(P.theta, P.phi, P.psi) + hp.data.data *= fp + hc.data.data *= fc + hp = lal.AddREAL8TimeSeries(hp, hc) + ht = hp + else: + # If astropy Time function, overwrite with GPS time, otherwise use normal addition + if isinstance(hp.epoch, Time): + dT = hp.epoch.to_value('gps','long') # pull out the time + hp.epoch = P.tref + dT + hc.epoch = P.tref +dT + else: + hp.epoch = hp.epoch + P.tref + hc.epoch = hc.epoch + P.tref + ht = lalsim.SimDetectorStrainREAL8TimeSeries(hp, hc, + P.phi, P.theta, P.psi, + lalsim.DetectorPrefixToLALDetector(str(P.detector))) + + # lalsimulation tapering: not sufficient for high SNRs + if P.taper != lalsimutils.lsu_TAPER_NONE: + lalsim.SimInspiralREAL8WaveTaper(ht.data, P.taper) + + # Resize such that TDlen = 1/deltaF + if P.deltaF is not None: + TDlen = int(1./P.deltaF * 1./P.deltaT) + # hp and hc were already at length such that TDlen = 1/deltaF, but lalsim.SimDetectorStrainREAL8TimeSeries adds a few points. Removing this assert so it resizes to desired size + assert TDlen >= ht.data.length, f"TDlen = {TDlen}, data_length = {ht.data.length}, 1/deltaT = {1/P.deltaT}, 1/deltaF = {1/P.deltaF}" + if TDlen < ht.data.length: + print(f'Data removed from {TDlen}:{ht.data.length}. Values = {ht.data.data[TDlen:ht.data.length]}') + ht = lal.ResizeREAL8TimeSeries(ht, 0, TDlen) + + # Match lalsimutils tapering: Do we need to taper given it is a FD model? + try: + taper=False # Don't taper, it is already an FD model. + if taper : + ntaper = int(0.01*TDlen) + if P.fmin > 0: # avoid failure if waveform start frequency 0 is nominally specified + ntaper = np.max([ntaper, int(1./(P.fmin*P.deltaT))]) + vectaper= 0.5 - 0.5*np.cos(np.pi*np.arange(ntaper)/(1.*ntaper)) + # Taper at the start of the segment + ht.data.data[:ntaper]*=vectaper + except Exception as e: + print("Couldn't apply tapering", e) + return ht + +def hlmoft(P, Lmax=2, use_ISCO_cutoff=use_ISCO_cutoff, **kwargs): + mtxAngularForward = np.zeros((5,5),dtype=np.complex64) + # Organize points so the (2,-2) and (2,2) mode clearly have contributions only from one point + # Problem with points in equatorial plane: pure real signal in aligned-spin limit! Degeneracy! + paramsForward =[ [np.pi,0], [np.pi/3,0], [np.pi/2, 0], [2*np.pi/3, 0], [0,0]]; + mvals = [-2,-1,0,1,2] + for indx in np.arange(len(paramsForward)): + for indx2 in np.arange(5): + m = mvals[indx2] + th,ph = paramsForward[indx] + mtxAngularForward[indx,indx2] = lal.SpinWeightedSphericalHarmonic(th,-ph,-2,2,m) # note phase sign + mtx = np.linalg.inv(mtxAngularForward) + + + # Now generate solutions at these values + P_copy = P.manual_copy() + P_copy.tref =0 # we do not need or want this offset when constructing hlm + P_copy.psi = 0 # we want to force a certain polarization + + hTC_list = [] + for indx in np.arange(len(paramsForward)): + th,ph = paramsForward[indx] +# P_copy.assign_param('thetaJN', th) # to be CONSISTENT with h(t) produced by ChooseTD, need to use incl here (!!) + P_copy.incl = th # to be CONSISTENT with h(t) produced by ChooseTD, need to use incl here (!!) + P_copy.phiref = ph + # Note argument change! Hack to recover reasonable hlm modes for (2,+/-1), (2,0) + # Something is funny about how SimInspiralFD/etc applies all these coordinate transforms + if P_copy.fref ==0: + P_copy.fref = P_copy.fmin + + hTC = complex_hoft(P_copy, use_ISCO_cutoff=use_ISCO_cutoff) + hTC_list.append(hTC) + + # Now construct the hlm from this sequence + hlmT = {} + for indx2 in np.arange(5): + m = mvals[indx2] + hlmT[(2,m)] = lal.CreateCOMPLEX16TimeSeries("Complex h(t)", hTC.epoch, hTC.f0, + hTC.deltaT, lal.DimensionlessUnit, hTC.data.length) + hlmT[(2,m)].epoch = float(hTC_list[0].epoch) + hlmT[(2,m)].data.data *=0 # this is needed, memory is not reliably cleaned + for indx in np.arange(len(paramsForward)): + hlmT[(2,m)].data.data += mtx[indx2,indx] * hTC_list[indx].data.data + return hlmT + +def std_and_conj_hlmoff(P, Lmax=2, use_ISCO_cutoff=use_ISCO_cutoff, **kwargs): + """ + Generate frequency-domain spherical harmonic modes (hlm) and their complex conjugates for RIFT precomputation. + + Args: + P : ChooseWaveformParams + Lmax : int, optional + Maximum spherical harmonic multipole order to include (default is 2). + **kwargs : dict + Additional keyword arguments passed to `hlmoft`. + + Returns: + hlmsF : dict + Dictionary mapping (l, m) mode tuples to the frequency-domain h_{lm}(f). + hlms_conj_F : dict + Dictionary mapping (l, m) mode tuples to the frequency-domain of the conjugated h_{lm}^*(f). + """ + hlms = hlmoft(P, Lmax, use_ISCO_cutoff=use_ISCO_cutoff, **kwargs) + hlmsF = {} + hlms_conj_F = {} + for mode in hlms: + hlmsF[mode] = lalsimutils.DataFourier(hlms[mode]) + hlms[mode].data.data = np.conj(hlms[mode].data.data) + hlms_conj_F[mode] = lalsimutils.DataFourier(hlms[mode]) + return hlmsF, hlms_conj_F From 7b34111a56f673b1535ed088ba67334b66c0f884 Mon Sep 17 00:00:00 2001 From: Aasim Z Jan <58118417+AasimZahoor@users.noreply.github.com> Date: Mon, 8 Jun 2026 10:22:56 -0500 Subject: [PATCH 4/5] plot_RIFT.py: fixed a minor bug, subdag CIP have a slightly different name than main iterations that this was not taking into account --- MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py b/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py index 45cb64d7..2e3bed61 100755 --- a/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py +++ b/MonteCarloMarginalizeCode/Code/bin/plot_RIFT.py @@ -526,16 +526,18 @@ def plot_neff_data(path_to_main_folder, plot_title): run_diagnostics[plot_title]["CIP_neff"] = {} # read requested neff from CIP sub files - for i in range(4): - filename = f"{path}/CIP_worker{i}.sub" + CIP_file_names = ["CIP_worker.sub", "CIP_worker0.sub", "CIP_worker1.sub", "CIP_worker2.sub", "CIP_worker3.sub"] + for i, name in enumerate(CIP_file_names): + filename = f"{path_to_main_folder}/{name}" try: with open(filename, "r") as f: content = f.read() matches = re.findall(r"--n-eff\s+([+-]?\d+(?:\.\d+)?)", content) if matches: neff_value = float(matches[-1]) - ax.axhline(y=neff_value, linestyle="--", color=default_colors[i], alpha=1.0, linewidth=1.0, label=f"worker {i} neff") - run_diagnostics[plot_title]["CIP_neff"][f"CIP_worker{i}"] = np.round(neff_value, 2) + worker_label = name.replace(".sub", "") + ax.axhline(y=neff_value, linestyle="--", color=default_colors[i], alpha=1.0, linewidth=1.0, label=f"{worker_label} neff") + run_diagnostics[plot_title]["CIP_neff"][worker_label] = np.round(neff_value, 2) else: continue except Exception as e: From 4a57e32e33c694ee6773e4482b6a352c69f5a2f7 Mon Sep 17 00:00:00 2001 From: Aasim Z Jan <58118417+AasimZahoor@users.noreply.github.com> Date: Mon, 8 Jun 2026 10:23:18 -0500 Subject: [PATCH 5/5] EFPE.py: conditional pyEFPE import --- MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py b/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py index d300ffa6..2ebc254a 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/physics/EFPE.py @@ -12,7 +12,11 @@ import lal import RIFT.lalsimutils as lalsimutils import lalsimulation as lalsim -import pyEFPE +try: + import pyEPFE + HAVE_PYEPFE = True +except ImportError: + HAVE_PYEPFE = False import numpy as np import RIFT.lalsimutils as lalsimutils from astropy.time import Time