Merge pull request #42 from ThibeauWouters/relative-binning-fix

Relative binning fix
kazewong · Dec 3, 2023 · f3f21e0 · f3f21e0
2 parents 1e16222 + 205cba6
commit f3f21e0
Show file tree

Hide file tree

Showing 8 changed files with 550 additions and 188 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -4,17 +4,17 @@ repos:
     hooks:
     - id: black
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.290'
+    rev: 'v0.1.6'
     hooks:
       - id: ruff
         args: ["--fix"]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.327
+    rev: v1.1.338
     hooks:
     - id: pyright
       additional_dependencies: [beartype, einops, jax, jaxtyping, pytest, tensorflow, tf2onnx, typing_extensions]
   - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.7.0
+    rev: 1.7.1
     hooks:
     - id: nbqa-black
       additional_dependencies: [ipython==8.12, black]

diff --git a/example/GW150914.py b/example/GW150914.py
@@ -17,15 +17,17 @@
 
 # first, fetch a 4s segment centered on GW150914
 gps = 1126259462.4
-start = gps - 2
-end = gps + 2
+duration = 4
+post_trigger_duration = 2
+start_pad = duration - post_trigger_duration
+end_pad = post_trigger_duration
 fmin = 20.0
 fmax = 1024.0
 
 ifos = ["H1", "L1"]
 
-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)
+H1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
+L1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
 
 Mc_prior = Unconstrained_Uniform(10.0, 80.0, naming=["M_c"])
 q_prior = Unconstrained_Uniform(
@@ -91,6 +93,7 @@
     post_trigger_duration=2,
 )
 
+likelihood = TransientLikelihoodFD([H1, L1], waveform=RippleIMRPhenomD(), trigger_time=gps, duration=4, post_trigger_duration=2)
 
 mass_matrix = jnp.eye(11)
 mass_matrix = mass_matrix.at[1, 1].set(1e-3)
@@ -100,7 +103,7 @@
 jim = Jim(
     likelihood,
     prior,
-    n_loop_training=200,
+    n_loop_training=100,
     n_loop_production=10,
     n_local_steps=150,
     n_global_steps=150,
@@ -117,5 +120,4 @@
     local_sampler_arg=local_sampler_arg,
 )
 
-# jim.maximize_likelihood([prior.xmin, prior.xmax])
 jim.sample(jax.random.PRNGKey(42))
diff --git a/example/GW150914_heterodyne.py b/example/GW150914_heterodyne.py
@@ -0,0 +1,89 @@
+import time
+from jimgw.jim import Jim
+from jimgw.detector import H1, L1
+from jimgw.likelihood import HeterodynedTransientLikelihoodFD, TransientLikelihoodFD
+from jimgw.waveform import RippleIMRPhenomD
+from jimgw.prior import Uniform
+import jax.numpy as jnp
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+###########################################
+########## 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
+fmin = 20.0
+fmax = 1024.0
+
+ifos = ["H1", "L1"]
+
+H1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
+L1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
+
+prior = Uniform(
+    xmin=[10, 0.125, -1.0, -1.0, 0.0, -0.05, 0.0, -1, 0.0, 0.0, -1.0],
+    xmax=[80.0, 1.0, 1.0, 1.0, 2000.0, 0.05, 2 * jnp.pi, 1.0, jnp.pi, 2 * jnp.pi, 1.0],
+    naming=[
+        "M_c",
+        "q",
+        "s1_z",
+        "s2_z",
+        "d_L",
+        "t_c",
+        "phase_c",
+        "cos_iota",
+        "psi",
+        "ra",
+        "sin_dec",
+    ],
+    transforms = {"q": ("eta", lambda params: params['q']/(1+params['q'])**2),
+                 "cos_iota": ("iota",lambda params: jnp.arccos(jnp.arcsin(jnp.sin(params['cos_iota']/2*jnp.pi))*2/jnp.pi)),
+                 "sin_dec": ("dec",lambda params: jnp.arcsin(jnp.arcsin(jnp.sin(params['sin_dec']/2*jnp.pi))*2/jnp.pi))} # sin and arcsin are periodize cos_iota and sin_dec
+)
+
+likelihood = HeterodynedTransientLikelihoodFD(
+    [H1, L1],
+    prior=prior,
+    bounds=[prior.xmin, prior.xmax],
+    waveform=RippleIMRPhenomD(),
+    trigger_time=gps,
+    duration=duration,
+    post_trigger_duration=post_trigger_duration,
+    n_loops=300
+)
+
+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}
+
+jim = Jim(
+    likelihood,
+    prior,
+    n_loop_training=100,
+    n_loop_production=10,
+    n_local_steps=150,
+    n_global_steps=150,
+    n_chains=500,
+    n_epochs=50,
+    learning_rate=0.001,
+    max_samples=45000,
+    momentum=0.9,
+    batch_size=50000,
+    use_global=True,
+    keep_quantile=0.0,
+    train_thinning=1,
+    output_thinning=10,
+    local_sampler_arg=local_sampler_arg,
+)
+
+jim.sample(jax.random.PRNGKey(42))
diff --git a/example/GW170817.py b/example/GW170817.py
@@ -0,0 +1,102 @@
+import time
+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
+from gwosc.datasets import event_gps
+import jax.numpy as jnp
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+###########################################
+########## First we grab data #############
+###########################################
+
+total_time_start = time.time()
+
+gps = event_gps("GW170817")
+duration = 128
+post_trigger_duration = 32
+start_pad = duration - post_trigger_duration
+end_pad = post_trigger_duration
+fmin = 20.0
+fmax = 1024.0
+
+ifos = ["H1", "L1"]#, "V1"]
+
+H1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=4*duration, tukey_alpha=0.05, gwpy_kwargs={"version": 2, "cache": False})
+L1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=4*duration, tukey_alpha=0.05, gwpy_kwargs={"version": 2, "cache": False})
+# V1.load_data(gps, start_pad, end_pad, fmin, fmax, psd_pad=16, tukey_alpha=0.05)
+
+prior = Uniform(
+    xmin=[1.18, 0.125, -0.3, -0.3, 1., -0.1, 0.0, -1, 0.0, 0.0, -1.0],
+    xmax=[1.21, 1.0, 0.3, 0.3, 75., 0.1, 2 * jnp.pi, 1.0, jnp.pi, 2 * jnp.pi, 1.0],
+    naming=[
+        "M_c",
+        "q",
+        "s1_z",
+        "s2_z",
+        "d_L",
+        "t_c",
+        "phase_c",
+        "cos_iota",
+        "psi",
+        "ra",
+        "sin_dec",
+    ],
+    transforms={
+        "q": ("eta", lambda params: params["q"] / (1 + params["q"]) ** 2),
+        "cos_iota": (
+            "iota",
+            lambda params: jnp.arccos(
+                jnp.arcsin(jnp.sin(params["cos_iota"] / 2 * jnp.pi)) * 2 / jnp.pi
+            ),
+        ),
+        "sin_dec": (
+            "dec",
+            lambda params: jnp.arcsin(
+                jnp.arcsin(jnp.sin(params["sin_dec"] / 2 * jnp.pi)) * 2 / jnp.pi
+            ),
+        ),
+    },  # sin and arcsin are periodize cos_iota and sin_dec
+)
+
+likelihood = HeterodynedTransientLikelihoodFD(
+    [H1],
+    prior=prior,
+    bounds=[prior.xmin, prior.xmax],
+    waveform=RippleIMRPhenomD(),
+    trigger_time=gps,
+    duration=duration,
+    post_trigger_duration=post_trigger_duration,
+    n_loops=1000
+)
+
+# 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}
+
+# jim = Jim(
+#     likelihood,
+#     prior,
+#     n_loop_training=100,
+#     n_loop_production=10,
+#     n_local_steps=150,
+#     n_global_steps=150,
+#     n_chains=500,
+#     n_epochs=50,
+#     learning_rate=0.001,
+#     max_samples=45000,
+#     momentum=0.9,
+#     batch_size=50000,
+#     use_global=True,
+#     keep_quantile=0.0,
+#     train_thinning=1,
+#     output_thinning=10,
+#     local_sampler_arg=local_sampler_arg,
+# )
+
+# jim.sample(jax.random.PRNGKey(42))