Merge pull request #43 from kazewong/41-add-more-prior-classes-and-ad…

…d-composite-prior-example

41 add more prior classes and add composite prior example.

Unconstrained Uniform, Spherical Prior, and composite prior are now added to the prior pool, with composite prior being the default people should use. Additional class of prior will be added in a separate thread in favor of merging other PRs.

.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):