From f98fa17d80ca8dcd6449e3d162e6b17b5cef293a Mon Sep 17 00:00:00 2001 From: Arianna Renzini Date: Wed, 24 Jan 2024 09:44:21 -0800 Subject: [PATCH] adding script to generate training sets --- training_omegas.py | 283 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100755 training_omegas.py diff --git a/training_omegas.py b/training_omegas.py new file mode 100755 index 0000000..a246d7f --- /dev/null +++ b/training_omegas.py @@ -0,0 +1,283 @@ +#!/bin/env python +# -*- coding: utf-8 -*- +# Copyright (C) Arianna I. Renzini 2024 +# +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. +# +# This file is part of the stochmon package + +import os +import sys + +sys.path.append(os.getcwd()) + +from popstock.PopulationOmegaGW import PopulationOmegaGW +from gwpopulation.models.mass import SinglePeakSmoothedMassDistribution +from gwpopulation.models.redshift import MadauDickinsonRedshift +from gwpopulation.utils import xp + +import argparse +import numpy as np +import bilby +import tqdm +import json + +import os +from pathlib import Path + +from bilby.core.prior import Interped +from bilby.core.utils import infer_args_from_function_except_n_args + +""" +*** +Use popstock to generate a training set of Omegas from a set of +hyper-parameters describing the redshift and mass distribution +of a population of black hole binaries. + +popstock will calculate omega_GW once, and reweight the samples from +the fiducial distribution to the provided distributions. +*** +""" + +parser = argparse.ArgumentParser() +parser.add_argument( + "-ps", + "--parameter_set", + help="Parameter sets to generate training data for.", + action="store", + type=str, + default=10000, +) +parser.add_argument( + "-rd", + "--run_directory", + help="Run directory.", + action="store", + type=str, + default="./", +) +parser.add_argument( + "-sm", "--samples", help="Samples to use.", action="store", type=str, default=None +) +parser.add_argument( + "-hsm", "--hyper_samples", help="Values of the hyper-parameters to use.", action="store", type=str, default=None +) +parser.add_argument( + "-ns", + "--number_samples", + help="number of samples.", + action="store", + type=int, + default=None, +) +parser.add_argument( + "-t", + "--tag", + help="Tag to label the run.", + action="store", + type=str, + default="test", +) +parser.add_argument( + "-wf", + "--waveform_approximant", + help="Wavefrom approximant. Default is IMRPhenomD.", + action="store", + type=str, + default="IMRPhenomD", +) +args = parser.parse_args() + +""" +Unpack arguments +--- + +Loading in parameters from parser; frequencies are fixed here +(these can be easily customised later on by passing in a range and spacing). +""" + +N_proposal_samples = args.number_samples +wave_approx = args.waveform_approximant +rundir = Path(args.run_directory) +tag = args.tag +freqs = np.arange(10, 2000, 5) + +""" +Create mass and redshift distributions +--- + +These are *fixed* in this script. To use custom distributions, +more development is required. Current distributions are imported +from the gwpopulation package (https://github.com/ColmTalbot/gwpopulation) + +Mass distribution: Power-Law Plus Peak (PLPP) +(see App. B.2 of https://iopscience.iop.org/article/10.3847/2041-8213/abe949) + +Redshift distribution: Madau-Dickinson Star Formation Rate (SFR) +(see https://arxiv.org/abs/1805.10270, +and https://arxiv.org/abs/2003.12152 for the normalisation) +""" + +mass_obj = SinglePeakSmoothedMassDistribution() +redshift_obj = MadauDickinsonRedshift(z_max=10) + +models = { + "mass_model": mass_obj, + "redshift_model": redshift_obj, +} + +""" +Create the popstock object +--- + +To create the PopulationOmegaGW object, you just need the mass and +redshift models, and frequencies to use in the spectra calculations. +We can then either sample over the models, or use user-defined samples. +As sampling takes time (and is currently not optimized), it is recommended +to do it once and use a (healthy) set of samples for future calculations. +""" + +newpop = PopulationOmegaGW(models=models, frequency_array=freqs) + +if args.samples is not None: + with open(args.samples) as samples_file: + samples_dict = json.load(samples_file) + Lambda_0 = samples_dict["Lambda_0"] + samples_dict.pop("Lambda_0") + if args.number_samples is None: + args.number_samples = len(samples_dict["redshift"]) + else: + for key in samples_dict.keys(): + samples_dict[key] = samples_dict[key][: args.number_samples] + newpop.set_proposal_samples(proposal_samples=samples_dict) + print(f"Using {args.number_samples} samples...") + +else: + # fiducial distribution Lambda_0, should be good to generate + # healthy sample sets. + Lambda_0 = { + "alpha": 2.5, + "beta": 1, + "delta_m": 3, + "lam": 0.04, + "mmax": 100, + "mmin": 4, + "mpp": 33, + "sigpp": 5, + "gamma": 2.7, + "kappa": 3, + "z_peak": 1.9, + "rate": 15, + } + newpop.draw_and_set_proposal_samples( + Lambda_0, N_proposal_samples=N_proposal_samples + ) + +""" +Calculate omega_GW once +--- + +Omega_GW is calculated once from the sample set provided. This +can take a while depending on the number of samples, but only needs +to be done once. using 10^6 samples takes about 40 minutes. +This first omega_GW is saved to file, together with the samples used +in case these weren't loaded from file. +""" + +newpop.calculate_omega_gw(waveform_approximant=wave_approx, Lambda=Lambda_0) + +if args.samples is not None: + np.savez( + f"{os.path.join(rundir, f'omegagw_0_{tag}.npz')}", + omega_gw=newpop.omega_gw, + freqs=newpop.frequency_array, + Lambda_0=Lambda_0, + ) + +else: + np.savez( + f"{os.path.join(rundir, f'omegagw_0_{tag}.npz')}", + omega_gw=newpop.omega_gw, + freqs=newpop.frequency_array, + fiducial_samples=newpop.proposal_samples, + Lambda_0=Lambda_0, + draw_dict=newpop.pdraws, + ) + +""" +Calculate omega_GW for a set of hyper-parameters (Lambda) +--- + +This is the main step to create the omega_GW training set. +The desired set of Lambdas should be passed in as a json file. +In case none is passed, this script will look for a posterior file +from the LVK O3 Populations paper (arXiv 2010.14533) lying around +on CIT. All samples provided will be used. + +Result omega_GWs with their respective Lambdas are saved to a json file. +""" + +new_omegas = {} + +new_omegas["Lambdas"] = [] +new_omegas["Neff"] = [] +new_omegas["omega_gw"] = [] + +if args.hyper_samples is not None: + with open(args.hyper_samples) as samples_file: + lambda_samples = json.load(samples_file) + +else: + try: + result = bilby.core.result.read_in_result(filename='/home/jacob.golomb/o3b-population-data/analyses/PowerLawPeak/o1o2o3_mass_c_iid_mag_iid_tilt_powerlaw_redshift_result.json') + lambda_samples = result.posterior.sample(10000).to_dict('list') + except IOError: + raise ValueError("No samples for Lambda hyper-parameters have been provided, nor a default found.") +N_trials = len(lambda_samples["alpha"]) + +print("Running trials...") +for idx in tqdm.tqdm(range(N_trials)): + Lambda_new = { + "alpha": lambda_samples["alpha"][idx], + "beta": lambda_samples["beta"][idx], + "delta_m": lambda_samples["delta_m"][idx], + "lam": lambda_samples["lam"][idx], + "mmax": lambda_samples["mmax"][idx], + "mmin": lambda_samples["mmin"][idx], + "mpp": lambda_samples["mpp"][idx], + "sigpp": lambda_samples["sigpp"][idx], + "rate": lambda_samples["rate"][idx], + "gamma": lambda_samples['lamb'][idx], + "kappa": 3.83, #make this madau-dickinson; right now these two fixed parameters are fixed but we can pass in other ones! + "z_peak": 2.0, + } + new_omegas["Lambdas"].append(Lambda_new) + + newpop.calculate_omega_gw(sampling_frequency=2048, Lambda=Lambda_new) + new_omegas["Neff"].append( + float((xp.sum(newpop.weights) ** 2) / (xp.sum(newpop.weights ** 2))) + ) + new_omegas["omega_gw"].append(newpop.omega_gw.tolist()) + +new_omegas["freqs"] = newpop.frequency_array.tolist() + +omegas_dict = json.dumps(new_omegas) +f = open(f"{os.path.join(rundir, f'new_omegas_{tag}.json')}", "w") +f.write(omegas_dict) +f.close() + +print("Done!") + +exit()