41 add more prior classes and add composite prior example

.github/workflows/python-package.yml

@@ -29,6 +29,7 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install pytest
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+        python -m pip install .
     - name: Test with pytest
       run: |
         pytest

example/GW150914.py

@@ -3,7 +3,7 @@
 from jimgw.detector import H1, L1
 from jimgw.likelihood import HeterodynedTransientLikelihoodFD, TransientLikelihoodFD
 from jimgw.waveform import RippleIMRPhenomD
-from jimgw.prior import Uniform
+from jimgw.prior import Unconstrained_Uniform, Composite
 import jax.numpy as jnp
 import jax
 
@@ -27,28 +27,69 @@
 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)
 
-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
+Mc_prior = Unconstrained_Uniform(10.0, 80.0, naming=["M_c"])
+q_prior = Unconstrained_Uniform(
+    0.125,
+    1.0,
+    naming=["q"],
+    transforms={"q": ("eta", lambda params: params["q"] / (1 + params["q"]) ** 2)},
+)
+s1z_prior = Unconstrained_Uniform(-1.0, 1.0, naming=["s1_z"])
+s2z_prior = Unconstrained_Uniform(-1.0, 1.0, naming=["s2_z"])
+dL_prior = Unconstrained_Uniform(0.0, 2000.0, naming=["d_L"])
+t_c_prior = Unconstrained_Uniform(-0.05, 0.05, naming=["t_c"])
+phase_c_prior = Unconstrained_Uniform(0.0, 2 * jnp.pi, naming=["phase_c"])
+cos_iota_prior = Unconstrained_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 = Unconstrained_Uniform(0.0, jnp.pi, naming=["psi"])
+ra_prior = Unconstrained_Uniform(0.0, 2 * jnp.pi, naming=["ra"])
+sin_dec_prior = Unconstrained_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,
+    ]
+)
+likelihood = TransientLikelihoodFD(
+    [H1, L1],
+    waveform=RippleIMRPhenomD(),
+    trigger_time=gps,
+    duration=4,
+    post_trigger_duration=2,
 )
-likelihood = TransientLikelihoodFD([H1, L1], waveform=RippleIMRPhenomD(), trigger_time=gps, duration=4, post_trigger_duration=2)
-# likelihood = HeterodynedTransientLikelihoodFD([H1, L1], prior=prior, bounds=[prior.xmin, prior.xmax], waveform=RippleIMRPhenomD(), trigger_time=gps, duration=4, post_trigger_duration=2)
 
 
 mass_matrix = jnp.eye(11)
@@ -76,5 +117,5 @@
     local_sampler_arg=local_sampler_arg,
 )
 
-jim.maximize_likelihood([prior.xmin, prior.xmax])
+# jim.maximize_likelihood([prior.xmin, prior.xmax])
 jim.sample(jax.random.PRNGKey(42))

example/GW150914_PV2.py

@@ -2,8 +2,8 @@
 from jimgw.jim import Jim
 from jimgw.detector import H1, L1
 from jimgw.likelihood import HeterodynedTransientLikelihoodFD, TransientLikelihoodFD
-from jimgw.waveform import RippleIMRPhenomD, RippleIMRPhenomPv2
-from jimgw.prior import Uniform
+from jimgw.waveform import RippleIMRPhenomPv2
+from jimgw.prior import Uniform, Composite, Sphere
 import jax.numpy as jnp
 import jax
 
@@ -28,19 +28,66 @@
 L1.load_data(gps, 2, 2, fmin, fmax, psd_pad=16, tukey_alpha=0.2)
 
 waveform = RippleIMRPhenomPv2(f_ref=20)
-prior = Uniform(
-    xmin = [10, 0.125, 0, 0, 0, 0, 0, 0, 0., -0.05, 0., -1, 0., 0.,-1.],
-    xmax = [80., 1., jnp.pi, 2*jnp.pi, 1., jnp.pi, 2*jnp.pi, 1., 2000., 0.05, 2*jnp.pi, 1., jnp.pi, 2*jnp.pi, 1.],
-    naming = ["M_c", "q", "s1_theta", "s1_phi", "s1_mag", "s2_theta", "s2_phi", "s2_mag", "d_L", "t_c", "phase_c", "cos_iota", "psi", "ra", "sin_dec"],
-    transforms = {"q": ("eta", lambda params: params['q']/(1+params['q'])**2),
-                 "s1_theta": ("s1_x", lambda params: jnp.sin(params['s1_theta'])*jnp.cos(params['s1_phi'])*params['s1_mag']),
-                 "s1_phi": ("s1_y", lambda params: jnp.sin(params['s1_theta'])*jnp.sin(params['s1_phi'])*params['s1_mag']),
-                 "s1_mag": ("s1_z", lambda params: jnp.cos(params['s1_theta'])*params['s1_mag']),
-                 "s2_theta": ("s2_x", lambda params: jnp.sin(params['s2_theta'])*jnp.cos(params['s2_phi'])*params['s2_mag']),
-                 "s2_phi": ("s2_y", lambda params: jnp.sin(params['s2_theta'])*jnp.sin(params['s2_phi'])*params['s2_mag']),
-                 "s2_mag": ("s2_z", lambda params: jnp.cos(params['s2_theta'])*params['s2_mag']),
-                 "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
+
+###########################################
+########## Set up priors ##################
+###########################################
+
+Mc_prior = Uniform(10.0, 80.0, naming=["M_c"])
+q_prior = Uniform(
+    0.125,
+    1.0,
+    naming=["q"],
+    transforms={"q": ("eta", lambda params: params["q"] / (1 + params["q"]) ** 2)},
+)
+s1_prior = Sphere(naming="s1")
+s2_prior = Sphere(naming="s2")
+dL_prior = Uniform(0.0, 2000.0, naming=["d_L"])
+t_c_prior = Uniform(-0.05, 0.05, 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,
+        s1_prior,
+        s2_prior,
+        dL_prior,
+        t_c_prior,
+        phase_c_prior,
+        cos_iota_prior,
+        psi_prior,
+        ra_prior,
+        sin_dec_prior,
+    ],
 )
 likelihood = TransientLikelihoodFD([H1, L1], waveform=waveform, trigger_time=gps, duration=4, post_trigger_duration=2)
 # likelihood = HeterodynedTransientLikelihoodFD([H1, L1], prior=prior, bounds=[prior.xmin, prior.xmax], waveform=RippleIMRPhenomD(), trigger_time=gps, duration=4, post_trigger_duration=2)
@@ -54,7 +101,7 @@
 jim = Jim(
     likelihood,
     prior,
-    n_loop_training=200,
+    n_loop_training=100,
     n_loop_production=10,
     n_local_steps=300,
     n_global_steps=300,
@@ -63,17 +110,12 @@
     learning_rate=0.001,
     max_samples = 60000,
     momentum=0.9,
-    batch_size=30000,
+    batch_size=60000,
     use_global=True,
     keep_quantile=0.,
     train_thinning=1,
     output_thinning=30,
     local_sampler_arg=local_sampler_arg,
-    num_layers = 6,
-    hidden_size = [32,32],
-    num_bins = 8
 )
 
-jim.maximize_likelihood([prior.xmin, prior.xmax])
-# initial_guess = jnp.array(jnp.load('initial.npz')['chain'])
 jim.sample(jax.random.PRNGKey(42))

example/GW150914_PV2_newglobal.py

@@ -0,0 +1,81 @@
+import time
+from jimgw.jim import Jim
+from jimgw.detector import H1, L1
+from jimgw.likelihood import HeterodynedTransientLikelihoodFD, TransientLikelihoodFD
+from jimgw.waveform import RippleIMRPhenomD, RippleIMRPhenomPv2
+from jimgw.prior import Uniform, Unconstrained_Uniform, Composite, Sphere
+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
+start = gps - 2
+end = gps + 2
+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)
+
+waveform = RippleIMRPhenomPv2(f_ref=20)
+
+Mc_prior = Unconstrained_Uniform(10., 80., naming=["M_c"])
+q_prior = Unconstrained_Uniform(0.125, 1., naming=["q"], transforms={"q": ("eta", lambda params: params['q']/(1+params['q'])**2)})
+s1_prior = Sphere("s1")
+s2_prior = Sphere("s2")
+dL_prior = Unconstrained_Uniform(0., 2000., naming=["d_L"])
+t_c_prior = Unconstrained_Uniform(-0.05, 0.05, naming=["t_c"])
+phase_c_prior = Unconstrained_Uniform(0., 2*jnp.pi, naming=["phase_c"])
+cos_iota_prior = Unconstrained_Uniform(-1., 1., 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 = Unconstrained_Uniform(0., jnp.pi, naming=["psi"])
+ra_prior = Unconstrained_Uniform(0., 2*jnp.pi, naming=["ra"])
+sin_dec_prior = Unconstrained_Uniform(-1., 1., 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, s1_prior, s2_prior, dL_prior, t_c_prior, phase_c_prior, cos_iota_prior, psi_prior, ra_prior, sin_dec_prior])
+
+likelihood = TransientLikelihoodFD([H1, L1], waveform=waveform, trigger_time=gps, duration=4, post_trigger_duration=2)
+
+
+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 * 3e-3}
+
+
+jim = Jim(
+    likelihood,
+    prior,
+    n_loop_training=50,
+    n_loop_production=10,
+    n_local_steps=300,
+    n_global_steps=300,
+    n_chains=500,
+    n_epochs=300,
+    learning_rate=0.001,
+    max_samples = 60000,
+    momentum=0.9,
+    batch_size=30000,
+    use_global=True,
+    keep_quantile=0.,
+    train_thinning=1,
+    output_thinning=30,
+    local_sampler_arg=local_sampler_arg,
+    num_layers = 6,
+    hidden_size = [32,32],
+    num_bins = 8
+)
+
+# jim.maximize_likelihood([prior.xmin, prior.xmax])
+# initial_guess = jnp.array(jnp.load('initial.npz')['chain'])
+jim.sample(jax.random.PRNGKey(42))

src/jimgw/jim.py

@@ -8,6 +8,7 @@
 from jaxtyping import Array
 import jax
 import jax.numpy as jnp
+from flowMC.sampler.flowHMC import flowHMC
 
 class Jim(object):
     """
@@ -30,13 +31,30 @@ def __init__(self, likelihood: LikelihoodBase, prior: Prior, **kwargs):
 
         local_sampler = MALA(self.posterior, True, local_sampler_arg) # Remember to add routine to find automated mass matrix
 
+        flowHMC_params = kwargs.get("flowHMC_params", {})
         model = MaskedCouplingRQSpline(self.Prior.n_dim, num_layers, hidden_size, num_bins, rng_key_set[-1])
+        if len(flowHMC_params) > 0:
+            global_sampler = flowHMC(
+                self.posterior,
+                True,
+                model,
+                params={
+                    "step_size": flowHMC_params["step_size"],
+                    "n_leapfrog": flowHMC_params["n_leapfrog"],
+                    "condition_matrix": flowHMC_params["condition_matrix"],
+                },
+            )
+        else:
+            global_sampler = None
+
+
         self.Sampler = Sampler(
             self.Prior.n_dim,
             rng_key_set,
             None,
             local_sampler,
             model,
+            global_sampler = global_sampler,
             **kwargs)
 
 
@@ -59,13 +77,15 @@ def maximize_likelihood(self, bounds: tuple[Array,Array], set_nwalkers: int = 10
         return best_fit
 
     def posterior(self, params: Array, data: dict):
-        named_params = self.Prior.add_name(params, transform_name=True, transform_value=True)
-        return self.Likelihood.evaluate(named_params, data) + self.Prior.log_prob(params)
+        prior_params = self.Prior.add_name(params.T)
+        prior  = self.Prior.log_prob(prior_params)
+        return self.Likelihood.evaluate(self.Prior.transform(prior_params), data) + prior
 
     def sample(self, key: jax.random.PRNGKey,
                initial_guess: Array = None):
         if initial_guess is None:
             initial_guess = self.Prior.sample(key, self.Sampler.n_chains)
+            initial_guess = jnp.stack([i for i in initial_guess.values()]).T
         self.Sampler.sample(initial_guess, None)
 
     def print_summary(self):