Commit 8c8acd4a authored by Moritz Huebner's avatar Moritz Huebner

Added slab-and-spike prior example

parent a0ccb680
Pipeline #158310 failed with stage
in 18 minutes and 48 seconds
#!/usr/bin/env python
An example of how to use slab-and-spike priors in bilby.
In this example we look at a simple example with the sum
of two Gaussian distributions, and we try to fit with
up to three Gaussians.
import bilby
import numpy as np
import matplotlib.pyplot as plt
outdir = 'outdir'
label = 'slabspike'
# Here we define our model. We want to inject two Gaussians and recover with up to three.
def gaussian(xs, amplitude, mu, sigma):
return amplitude/np.sqrt(2*np.pi*sigma**2)*np.exp(-0.5 * (xs - mu)**2 / sigma**2)
def triple_gaussian(xs, amplitude_0, amplitude_1, amplitude_2, mu_0, mu_1, mu_2, sigma_0, sigma_1, sigma_2, **kwargs):
return gaussian(xs, amplitude_0, mu_0, sigma_0) + \
gaussian(xs, amplitude_1, mu_1, sigma_1) + \
gaussian(xs, amplitude_2, mu_2, sigma_2)
# Let's create our data set. We create 200 points on a grid.
xs = np.linspace(-5, 5, 200)
dx = xs[1] - xs[0]
# Note for our injection parameters we set the amplitude of the second component to 0.
injection_params = dict(amplitude_0=-3, mu_0=-4, sigma_0=4,
amplitude_1=0, mu_1=0, sigma_1=1,
amplitude_2=4, mu_2=3, sigma_2=3)
# We calculate the injected curve and add some Gaussian noise on the data points
sigma = 0.02
p = bilby.core.prior.Gaussian(mu=0, sigma=sigma)
ys = triple_gaussian(xs=xs, **injection_params) + p.sample(len(xs))
plt.errorbar(xs, ys, yerr=sigma, fmt=".k", capsize=0, label='Injected data')
plt.plot(xs, triple_gaussian(xs=xs, **injection_params), label='True signal')
# Now we want to set up our priors.
priors = bilby.core.prior.PriorDict()
# For the slab-and-spike prior, we first need to define the 'slab' part, which is just a regular bilby prior.
amplitude_slab_0 = bilby.core.prior.Uniform(minimum=-10, maximum=10, name='amplitude_0', latex_label='$A_0$')
amplitude_slab_1 = bilby.core.prior.Uniform(minimum=-10, maximum=10, name='amplitude_1', latex_label='$A_1$')
amplitude_slab_2 = bilby.core.prior.Uniform(minimum=-10, maximum=10, name='amplitude_2', latex_label='$A_2$')
# We do the following to create the slab-and-spike prior. The spike height is somewhat arbitrary and can
# be corrected in post-processing.
priors['amplitude_0'] = bilby.core.prior.SlabSpikePrior(slab=amplitude_slab_0, spike_location=0, spike_height=0.1)
priors['amplitude_1'] = bilby.core.prior.SlabSpikePrior(slab=amplitude_slab_1, spike_location=0, spike_height=0.1)
priors['amplitude_2'] = bilby.core.prior.SlabSpikePrior(slab=amplitude_slab_2, spike_location=0, spike_height=0.1)
# Our problem has a degeneracy in the ordering. In general, this problem is somewhat difficult to resolve properly.
# See e.g. for an implementation.
# We resolve this by not letting the priors overlap in this case.
priors['mu_0'] = bilby.core.prior.Uniform(minimum=-5, maximum=-2, name='mu_0', latex_label='$\mu_0$')
priors['mu_1'] = bilby.core.prior.Uniform(minimum=-2, maximum=2, name='mu_1', latex_label='$\mu_1$')
priors['mu_2'] = bilby.core.prior.Uniform(minimum=2, maximum=5, name='mu_2', latex_label='$\mu_2$')
priors['sigma_0'] = bilby.core.prior.LogUniform(minimum=0.01, maximum=10, name='sigma_0', latex_label='$\sigma_0$')
priors['sigma_1'] = bilby.core.prior.LogUniform(minimum=0.01, maximum=10, name='sigma_1', latex_label='$\sigma_1$')
priors['sigma_2'] = bilby.core.prior.LogUniform(minimum=0.01, maximum=10, name='sigma_2', latex_label='$\sigma_2$')
# Setting up the likelihood and running the samplers works the same as elsewhere.
likelihood = bilby.core.likelihood.GaussianLikelihood(x=xs, y=ys, func=triple_gaussian, sigma=sigma)
result = bilby.run_sampler(likelihood=likelihood, priors=priors, outdir=outdir, label=label,
sampler='dynesty', nlive=400)
# Let's also plot the maximum likelihood fit along with the data.
max_like_params = result.posterior.iloc[-1]
plt.errorbar(xs, ys, yerr=sigma, fmt=".k", capsize=0, label='Injected data')
plt.plot(xs, triple_gaussian(xs=xs, **injection_params), label='True signal')
plt.plot(xs, triple_gaussian(xs=xs, **max_like_params), label='Max likelihood fit')
# Finally, we can check what fraction of amplitude samples are exactly on the spike.
spike_samples_0 = len(np.where(result.posterior['amplitude_0'] == 0.0)[0])/len(result.posterior)
spike_samples_1 = len(np.where(result.posterior['amplitude_1'] == 0.0)[0])/len(result.posterior)
spike_samples_2 = len(np.where(result.posterior['amplitude_2'] == 0.0)[0])/len(result.posterior)
print(f"{spike_samples_0*100:.2f}% of amplitude_0 samples are exactly 0.0")
print(f"{spike_samples_1*100:.2f}% of amplitude_1 samples are exactly 0.0")
print(f"{spike_samples_2*100:.2f}% of amplitude_2 samples are exactly 0.0")
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment