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'
bilby.utils.check_directory_exists_and_if_not_mkdir(outdir)
# 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')
plt.legend()
plt.savefig(f'{outdir}/{label}_injected_data')
plt.clf()
# 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. https://github.com/GregoryAshton/kookaburra/blob/master/src/priors.py#L72 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)
result.plot_corner(truths=injection_params)
# 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')
plt.legend()
plt.savefig(f'{outdir}/{label}_max_likelihood_recovery')
plt.clf()
# 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