attempt to fix Pv2 test script

test/integration/test_GW150914_Pv2.py

@@ -4,12 +4,28 @@
 import jax.numpy as jnp
 
 from jimgw.jim import Jim
-from jimgw.prior import CombinePrior, UniformPrior, CosinePrior, SinePrior, PowerLawPrior
+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.single_event.waveform import RippleIMRPhenomPv2
 from jimgw.transforms import BoundToUnbound
-from jimgw.single_event.transforms import MassRatioToSymmetricMassRatioTransform, SpinToCartesianSpinTransform
+from jimgw.single_event.transforms import (
+    SkyFrameToDetectorFrameSkyPositionTransform,
+    SphereSpinToCartesianSpinTransform,
+    MassRatioToSymmetricMassRatioTransform,
+    DistanceToSNRWeightedDistanceTransform,
+    GeocentricArrivalTimeToDetectorArrivalTimeTransform,
+    GeocentricArrivalPhaseToDetectorArrivalPhaseTransform,
+)
+from jimgw.single_event.utils import Mc_q_to_m1_m2
 from flowMC.strategy.optimization import optimization_Adam
 
 jax.config.update("jax_enable_x64", True)
@@ -34,68 +50,77 @@
 for ifo in ifos:
     ifo.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
 
-Mc_prior = UniformPrior(10.0, 80.0, parameter_names=["M_c"])
-q_prior = UniformPrior(0.125, 1., parameter_names=["q"])
-theta_jn_prior = SinePrior(parameter_names=["theta_jn"])
-phi_jl_prior = UniformPrior(0.0, 2 * jnp.pi, parameter_names=["phi_jl"])
-theta_1_prior = SinePrior(parameter_names=["theta_1"])
-theta_2_prior = SinePrior(parameter_names=["theta_2"])
-phi_12_prior = UniformPrior(0.0, 2 * jnp.pi, parameter_names=["phi_12"])
-a_1_prior = UniformPrior(0.0, 1.0, parameter_names=["a_1"])
-a_2_prior = UniformPrior(0.0, 1.0, parameter_names=["a_2"])
-dL_prior = PowerLawPrior(10.0, 2000.0, 2.0, parameter_names=["d_L"])
+prior = []
+
+# Mass prior
+M_c_min, M_c_max = 10.0, 80.0
+q_min, q_max = 0.125, 1.0
+Mc_prior = UniformPrior(M_c_min, M_c_max, parameter_names=["M_c"])
+q_prior = UniformPrior(q_min, q_max, parameter_names=["q"])
+
+prior = prior + [Mc_prior, q_prior]
+
+# Spin prior
+s1_prior = UniformSpherePrior(parameter_names=["s1"])
+s2_prior = UniformSpherePrior(parameter_names=["s2"])
+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"])
 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,
-        q_prior,
-        theta_jn_prior,
-        phi_jl_prior,
-        theta_1_prior,
-        theta_2_prior,
-        phi_12_prior,
-        a_1_prior,
-        a_2_prior,
-        dL_prior,
-        t_c_prior,
-        phase_c_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 = [
-    BoundToUnbound(name_mapping = [["M_c"], ["M_c_unbounded"]], original_lower_bound=10.0, original_upper_bound=80.0),
-    BoundToUnbound(name_mapping = [["q"], ["q_unbounded"]], original_lower_bound=0.125, original_upper_bound=1.),
-    BoundToUnbound(name_mapping = [["theta_jn"], ["theta_jn_unbounded"]] , original_lower_bound=0.0, original_upper_bound=jnp.pi),
-    BoundToUnbound(name_mapping = [["phi_jl"], ["phi_jl_unbounded"]] , original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
-    BoundToUnbound(name_mapping = [["theta_1"], ["theta_1_unbounded"]] , original_lower_bound=0.0, original_upper_bound=jnp.pi),
-    BoundToUnbound(name_mapping = [["theta_2"], ["theta_2_unbounded"]] , original_lower_bound=0.0, original_upper_bound=jnp.pi),
-    BoundToUnbound(name_mapping = [["phi_12"], ["phi_12_unbounded"]] , original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
-    BoundToUnbound(name_mapping = [["a_1"], ["a_1_unbounded"]] , original_lower_bound=0.0, original_upper_bound=1.0),
-    BoundToUnbound(name_mapping = [["a_2"], ["a_2_unbounded"]] , original_lower_bound=0.0, original_upper_bound=1.0),
-    BoundToUnbound(name_mapping = [["d_L"], ["d_L_unbounded"]] , original_lower_bound=0.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 = [["psi"], ["psi_unbounded"]], original_lower_bound=0.0, original_upper_bound=jnp.pi),
-    BoundToUnbound(name_mapping = [["ra"], ["ra_unbounded"]], original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
-    BoundToUnbound(name_mapping = [["dec"], ["dec_unbounded"]],original_lower_bound=-jnp.pi / 2, original_upper_bound=jnp.pi / 2)
+    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 = (["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_phi"], ["s1_phi_unbounded"]) , original_lower_bound=0.0, original_upper_bound=2 * jnp.pi),
+    BoundToUnbound(name_mapping = (["s2_phi"], ["s2_phi_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 = (["s1_theta"], ["s1_theta_unbounded"]) , original_lower_bound=0.0, original_upper_bound=jnp.pi),
+    BoundToUnbound(name_mapping = (["s2_theta"], ["s2_theta_unbounded"]) , original_lower_bound=0.0, original_upper_bound=jnp.pi),
+    BoundToUnbound(name_mapping = (["s1_mag"], ["s1_mag_unbounded"]) , original_lower_bound=0.0, original_upper_bound=0.99),
+    BoundToUnbound(name_mapping = (["s2_mag"], ["s2_mag_unbounded"]) , original_lower_bound=0.0, original_upper_bound=0.99),
+    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 = [
-    SpinToCartesianSpinTransform(freq_ref=20.0),
     MassRatioToSymmetricMassRatioTransform,
+    SphereSpinToCartesianSpinTransform("s1"),
+    SphereSpinToCartesianSpinTransform("s2"),
 ]
 
 likelihood = TransientLikelihoodFD(
     ifos,
-    waveform=RippleIMRPhenomD(),
+    waveform=RippleIMRPhenomPv2(),
     trigger_time=gps,
     duration=4,
     post_trigger_duration=2,
@@ -139,4 +164,4 @@
 
 jim.sample(jax.random.PRNGKey(42))
 jim.get_samples()
-jim.print_summary()
+jim.print_summary()