Skip to content

Commit

Permalink
Merge pull request #166 from tsunhopang/example_and_test_fixing
Browse files Browse the repository at this point in the history 
Update examples and fixing the Pv2 test
  • Loading branch information
@kazewong
kazewong authored Oct 17, 2024
2 parents d0554fa + 38de2ac commit 6feea3d
Show file tree
Hide file tree
Showing 3 changed files with 343 additions and 154 deletions.
186 changes: 83 additions & 103 deletions example/GW150914_IMRPhenomD.py
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,29 @@
import time

import jax
import jax.numpy as jnp

from jimgw.jim import Jim
from jimgw.jim import Jim
from jimgw.prior import (
CombinePrior,
UniformPrior,
CosinePrior,
SinePrior,
PowerLawPrior,
UniformSpherePrior,
)
from jimgw.single_event.detector import H1, L1
from jimgw.single_event.likelihood import TransientLikelihoodFD
from jimgw.single_event.waveform import RippleIMRPhenomD
from jimgw.transforms import BoundToUnbound
from jimgw.single_event.transforms import (
ComponentMassesToChirpMassSymmetricMassRatioTransform,
SkyFrameToDetectorFrameSkyPositionTransform,
ComponentMassesToChirpMassMassRatioTransform,
SphereSpinToCartesianSpinTransform,
MassRatioToSymmetricMassRatioTransform,
DistanceToSNRWeightedDistanceTransform,
GeocentricArrivalTimeToDetectorArrivalTimeTransform,
GeocentricArrivalPhaseToDetectorArrivalPhaseTransform,
)
from jimgw.single_event.utils import Mc_q_to_m1_m2
from flowMC.strategy.optimization import optimization_Adam
Expand All @@ -27,136 +34,110 @@
########## First we grab data #############
###########################################

total_time_start = time.time()

# first, fetch a 4s segment centered on GW150914
gps = 1126259462.4
duration = 4
post_trigger_duration = 2
start_pad = duration - post_trigger_duration
end_pad = post_trigger_duration
start = gps - 2
end = gps + 2
fmin = 20.0
fmax = 1024.0

ifos = [H1, L1]

for ifo in ifos:
ifo.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
H1.load_data(gps, 2, 2, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
L1.load_data(gps, 2, 2, fmin, fmax, psd_pad=16, tukey_alpha=0.2)

waveform = RippleIMRPhenomD(f_ref=20)

###########################################
########## Set up priors ##################
###########################################

prior = []

# Mass prior
M_c_min, M_c_max = 10.0, 80.0
eta_min, eta_max = 0.2, 0.25
# m_1_prior = UniformPrior(Mc_q_to_m1_m2(M_c_min, q_max)[0], Mc_q_to_m1_m2(M_c_max, q_min)[0], parameter_names=["m_1"])
# m_2_prior = UniformPrior(Mc_q_to_m1_m2(M_c_min, q_min)[1], Mc_q_to_m1_m2(M_c_max, q_max)[1], parameter_names=["m_2"])
q_min, q_max = 0.125, 1.0
Mc_prior = UniformPrior(M_c_min, M_c_max, parameter_names=["M_c"])
eta_prior = UniformPrior(eta_min, eta_max, parameter_names=["eta"])
s1z_prior = UniformPrior(-1.0, 1.0, parameter_names=["s1_z"])
s2z_prior = UniformPrior(-1.0, 1.0, parameter_names=["s2_z"])
q_prior = UniformPrior(q_min, q_max, parameter_names=["q"])

prior = prior + [Mc_prior, q_prior]

# Spin prior
s1_prior = UniformPrior(-1.0, 1.0, parameter_names=["s1_z"])
s2_prior = UniformPrior(-1.0, 1.0, parameter_names=["s2_z"])
iota_prior = SinePrior(parameter_names=["iota"])

prior = prior + [
s1_prior,
s2_prior,
iota_prior,
]

# Extrinsic prior
dL_prior = PowerLawPrior(1.0, 2000.0, 2.0, parameter_names=["d_L"])
t_c_prior = UniformPrior(-0.05, 0.05, parameter_names=["t_c"])
phase_c_prior = UniformPrior(0.0, 2 * jnp.pi, parameter_names=["phase_c"])
iota_prior = SinePrior(parameter_names=["iota"])
psi_prior = UniformPrior(0.0, jnp.pi, parameter_names=["psi"])
ra_prior = UniformPrior(0.0, 2 * jnp.pi, parameter_names=["ra"])
dec_prior = CosinePrior(parameter_names=["dec"])

prior = CombinePrior(
[
Mc_prior,
eta_prior,
s1z_prior,
s2z_prior,
dL_prior,
t_c_prior,
phase_c_prior,
iota_prior,
psi_prior,
ra_prior,
dec_prior,
]
)
prior = prior + [
dL_prior,
t_c_prior,
phase_c_prior,
psi_prior,
ra_prior,
dec_prior,
]

prior = CombinePrior(prior)

# Defining Transforms

sample_transforms = [
# ComponentMassesToChirpMassMassRatioTransform,
BoundToUnbound(
name_mapping=(["M_c"], ["M_c_unbounded"]),
original_lower_bound=M_c_min,
original_upper_bound=M_c_max,
),
BoundToUnbound(
name_mapping=(["eta"], ["eta_unbounded"]),
original_lower_bound=eta_min,
original_upper_bound=eta_max,
),
BoundToUnbound(
name_mapping=(["s1_z"], ["s1_z_unbounded"]),
original_lower_bound=-1.0,
original_upper_bound=1.0,
),
BoundToUnbound(
name_mapping=(["s2_z"], ["s2_z_unbounded"]),
original_lower_bound=-1.0,
original_upper_bound=1.0,
),
BoundToUnbound(
name_mapping=(["d_L"], ["d_L_unbounded"]),
original_lower_bound=1.0,
original_upper_bound=2000.0,
),
BoundToUnbound(
name_mapping=(["t_c"], ["t_c_unbounded"]),
original_lower_bound=-0.05,
original_upper_bound=0.05,
),
BoundToUnbound(
name_mapping=(["phase_c"], ["phase_c_unbounded"]),
original_lower_bound=0.0,
original_upper_bound=2 * jnp.pi,
),
BoundToUnbound(
name_mapping=(["iota"], ["iota_unbounded"]),
original_lower_bound=0.0,
original_upper_bound=jnp.pi,
),
BoundToUnbound(
name_mapping=(["psi"], ["psi_unbounded"]),
original_lower_bound=0.0,
original_upper_bound=jnp.pi,
),
DistanceToSNRWeightedDistanceTransform(gps_time=gps, ifos=ifos, dL_min=dL_prior.xmin, dL_max=dL_prior.xmax),
GeocentricArrivalPhaseToDetectorArrivalPhaseTransform(gps_time=gps, ifo=ifos[0]),
GeocentricArrivalTimeToDetectorArrivalTimeTransform(tc_min=t_c_prior.xmin, tc_max=t_c_prior.xmax, gps_time=gps, ifo=ifos[0]),
SkyFrameToDetectorFrameSkyPositionTransform(gps_time=gps, ifos=ifos),
BoundToUnbound(
name_mapping=(["zenith"], ["zenith_unbounded"]),
original_lower_bound=0.0,
original_upper_bound=jnp.pi,
),
BoundToUnbound(
name_mapping=(["azimuth"], ["azimuth_unbounded"]),
original_lower_bound=0.0,
original_upper_bound=2 * jnp.pi,
),
BoundToUnbound(name_mapping = (["M_c"], ["M_c_unbounded"]), original_lower_bound=M_c_min, original_upper_bound=M_c_max),
BoundToUnbound(name_mapping = (["q"], ["q_unbounded"]), original_lower_bound=q_min, original_upper_bound=q_max),
BoundToUnbound(name_mapping = (["s1_z"], ["s1_z_unbounded"]) , original_lower_bound=-1.0, original_upper_bound=1.0),
BoundToUnbound(name_mapping = (["s2_z"], ["s2_z_unbounded"]) , original_lower_bound=-1.0, original_upper_bound=1.0),
BoundToUnbound(name_mapping = (["iota"], ["iota_unbounded"]) , original_lower_bound=0.0, original_upper_bound=jnp.pi),
BoundToUnbound(name_mapping = (["phase_det"], ["phase_det_unbounded"]), original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
BoundToUnbound(name_mapping = (["psi"], ["psi_unbounded"]), original_lower_bound=0.0, original_upper_bound=jnp.pi),
BoundToUnbound(name_mapping = (["zenith"], ["zenith_unbounded"]), original_lower_bound=0.0, original_upper_bound=jnp.pi),
BoundToUnbound(name_mapping = (["azimuth"], ["azimuth_unbounded"]), original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
]

likelihood_transforms = [
# ComponentMassesToChirpMassSymmetricMassRatioTransform,
MassRatioToSymmetricMassRatioTransform,
]


likelihood = TransientLikelihoodFD(
ifos,
waveform=RippleIMRPhenomD(),
trigger_time=gps,
duration=4,
post_trigger_duration=2,
[H1, L1], waveform=waveform, trigger_time=gps, duration=4, post_trigger_duration=2
)


mass_matrix = jnp.eye(11)
mass_matrix = mass_matrix.at[1, 1].set(1e-3)
mass_matrix = mass_matrix.at[5, 5].set(1e-3)
local_sampler_arg = {"step_size": mass_matrix * 3e-3}
mass_matrix = jnp.eye(prior.n_dim)
# mass_matrix = mass_matrix.at[1, 1].set(1e-3)
# mass_matrix = mass_matrix.at[9, 9].set(1e-3)
local_sampler_arg = {"step_size": mass_matrix * 1e-3}

Adam_optimizer = optimization_Adam(n_steps=3000, learning_rate=0.01, noise_level=1)

n_epochs = 30
n_loop_training = 20
learning_rate = 1e-4
import optax

n_epochs = 20
n_loop_training = 100
total_epochs = n_epochs * n_loop_training
start = total_epochs // 10
learning_rate = optax.polynomial_schedule(
1e-3, 1e-4, 4.0, total_epochs - start, transition_begin=start
)

jim = Jim(
likelihood,
Expand All @@ -175,13 +156,12 @@
momentum=0.9,
batch_size=30000,
use_global=True,
keep_quantile=0.0,
train_thinning=1,
output_thinning=10,
local_sampler_arg=local_sampler_arg,
strategies=[Adam_optimizer, "default"],
verbose=True,
# strategies=[Adam_optimizer,"default"],
)


jim.sample(jax.random.PRNGKey(42))
# jim.get_samples()
# jim.print_summary()
Loading

0 comments on commit 6feea3d

Please sign in to comment.