From ed223ef4b514dcb7d514f37d7e6fd91d73711eac Mon Sep 17 00:00:00 2001 From: ThibeauWouters Date: Wed, 6 Dec 2023 00:46:08 -0800 Subject: [PATCH 1/3] transpose bounds if necessary, avoid confusion --- src/jimgw/likelihood.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/jimgw/likelihood.py b/src/jimgw/likelihood.py index bd37dcf3..de274e66 100644 --- a/src/jimgw/likelihood.py +++ b/src/jimgw/likelihood.py @@ -179,6 +179,10 @@ def __init__( self.freq_grid_low = freq_grid[:-1] print("Finding reference parameters..") + + # Check if shape of bounds is correct, otherwise, transpose bounds + if jnp.shape(bounds)[0] != 2: + bounds = bounds.T self.ref_params = self.maximize_likelihood( bounds=bounds, prior=prior, popsize=popsize, n_loops=n_loops From f006ca696e7df90c0dc6a11c22eb0b7cda6a987d Mon Sep 17 00:00:00 2001 From: ThibeauWouters Date: Wed, 6 Dec 2023 07:40:14 -0800 Subject: [PATCH 2/3] added GW170817 relative binning example --- example/GW170817_heterodyne.py | 140 +++++++++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 example/GW170817_heterodyne.py diff --git a/example/GW170817_heterodyne.py b/example/GW170817_heterodyne.py new file mode 100644 index 00000000..2e4f8e40 --- /dev/null +++ b/example/GW170817_heterodyne.py @@ -0,0 +1,140 @@ +from jimgw.jim import Jim +from jimgw.detector import H1, L1, V1 +from jimgw.likelihood import HeterodynedTransientLikelihoodFD +from jimgw.waveform import RippleIMRPhenomD +from jimgw.prior import Uniform, Powerlaw, Composite +import jax.numpy as jnp +import jax +import time +jax.config.update("jax_enable_x64", True) + + +### Fetching the data + +total_time_start = time.time() +gps = 1187008882.43 +trigger_time = gps +fmin = 20 +fmax = 2048 +minimum_frequency = fmin +maximum_frequency = fmax +T = 128 +duration = T +post_trigger_duration = 2 +epoch = duration - post_trigger_duration +f_ref = fmin + +### Getting ifos and overwriting with above data + +tukey_alpha = 2 / (duration / 2) +H1.load_data(gps, duration, 2, fmin, fmax, psd_pad=duration+16, tukey_alpha=tukey_alpha) +L1.load_data(gps, duration, 2, fmin, fmax, psd_pad=duration+16, tukey_alpha=tukey_alpha) +V1.load_data(gps, duration, 2, fmin, fmax, psd_pad=duration+16, tukey_alpha=tukey_alpha) + +### Define priors + +# Internal parameters +Mc_prior = Uniform(1.18, 1.21, naming=["M_c"]) +q_prior = Uniform( + 0.125, + 1.0, + naming=["q"], + transforms={"q": ("eta", lambda params: params["q"] / (1 + params["q"]) ** 2)}, +) +s1z_prior = Uniform(-0.05, 0.05, naming=["s1_z"]) +s2z_prior = Uniform(-0.05, 0.05, naming=["s2_z"]) + +# External parameters +dL_prior = Powerlaw(1.0, 75.0, 2.0, naming=["d_L"]) +t_c_prior = Uniform(-0.1, 0.1, naming=["t_c"]) +phase_c_prior = Uniform(0.0, 2 * jnp.pi, naming=["phase_c"]) +cos_iota_prior = Uniform( + -1.0, + 1.0, + naming=["cos_iota"], + transforms={ + "cos_iota": ( + "iota", + lambda params: jnp.arccos( + jnp.arcsin(jnp.sin(params["cos_iota"] / 2 * jnp.pi)) * 2 / jnp.pi + ), + ) + }, +) +psi_prior = Uniform(0.0, jnp.pi, naming=["psi"]) +ra_prior = Uniform(0.0, 2 * jnp.pi, naming=["ra"]) +sin_dec_prior = Uniform( + -1.0, + 1.0, + naming=["sin_dec"], + transforms={ + "sin_dec": ( + "dec", + lambda params: jnp.arcsin( + jnp.arcsin(jnp.sin(params["sin_dec"] / 2 * jnp.pi)) * 2 / jnp.pi + ), + ) + }, +) + +prior = Composite([ + Mc_prior, + q_prior, + s1z_prior, + s2z_prior, + dL_prior, + t_c_prior, + phase_c_prior, + cos_iota_prior, + psi_prior, + ra_prior, + sin_dec_prior, + ] +) + +# The following only works if every prior has xmin and xmax property, which is OK for Uniform and Powerlaw +bounds = jnp.array([[p.xmin, p.xmax] for p in prior.priors]).T + +### Create likelihood object +likelihood = HeterodynedTransientLikelihoodFD([H1, L1, V1], prior=prior, bounds=bounds, waveform=RippleIMRPhenomD(), trigger_time=gps, duration=T, n_bins=500) + +### Create sampler and jim objects +eps = 3e-2 +n_dim = 11 +mass_matrix = jnp.eye(n_dim) +mass_matrix = mass_matrix.at[0,0].set(1e-5) +mass_matrix = mass_matrix.at[1,1].set(1e-4) +mass_matrix = mass_matrix.at[2,2].set(1e-3) +mass_matrix = mass_matrix.at[3,3].set(1e-3) +mass_matrix = mass_matrix.at[5,5].set(1e-5) +mass_matrix = mass_matrix.at[9,9].set(1e-2) +mass_matrix = mass_matrix.at[10,10].set(1e-2) +local_sampler_arg = {"step_size": mass_matrix * eps} + +outdir_name = "./outdir/" + +jim = Jim( + likelihood, + prior, + n_loop_pretraining=0, + n_loop_training=200, + n_loop_production=20, + n_local_steps=200, + n_global_steps=200, + n_chains=1000, + n_epochs=300, + learning_rate=0.001, + max_samples=50000, + momentum=0.9, + batch_size=50000, + use_global=True, + keep_quantile=0.0, + train_thinning=10, + output_thinning=30, + n_loops_maximize_likelihood = 2000, + local_sampler_arg=local_sampler_arg, + outdir_name=outdir_name +) + +jim.sample(jax.random.PRNGKey(42)) +jim.print_summary() \ No newline at end of file From 17d404425dc0278f003418227011db7c8583aea0 Mon Sep 17 00:00:00 2001 From: Kaze Wong Date: Wed, 6 Dec 2023 12:56:08 -0500 Subject: [PATCH 3/3] Update likelihood.py Regarding this changes. The new styling changes pushed recently should give more explicit hints to address the ambiguity in shape of array. Reverting this in favor of the new changes --- src/jimgw/likelihood.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/jimgw/likelihood.py b/src/jimgw/likelihood.py index de274e66..3a236936 100644 --- a/src/jimgw/likelihood.py +++ b/src/jimgw/likelihood.py @@ -180,10 +180,6 @@ def __init__( print("Finding reference parameters..") - # Check if shape of bounds is correct, otherwise, transpose bounds - if jnp.shape(bounds)[0] != 2: - bounds = bounds.T - self.ref_params = self.maximize_likelihood( bounds=bounds, prior=prior, popsize=popsize, n_loops=n_loops )