slabspike_example.py

#!/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.

We will use the `PyMultiNest` sampler which is fast but can be unreliable when
significant correlations exist in the likelihood.

To install `PyMultiNest` call

$ conda install -c conda-forge pymultinest
"""

import matplotlib.pyplot as plt
import numpy as np

import bilby

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=r"$\mu_0$"
)
priors["mu_1"] = bilby.core.prior.Uniform(
    minimum=-2, maximum=2, name="mu_1", latex_label=r"$\mu_1$"
)
priors["mu_2"] = bilby.core.prior.Uniform(
    minimum=2, maximum=5, name="mu_2", latex_label=r"$\mu_2$"
)
priors["sigma_0"] = bilby.core.prior.LogUniform(
    minimum=0.01, maximum=10, name="sigma_0", latex_label=r"$\sigma_0$"
)
priors["sigma_1"] = bilby.core.prior.LogUniform(
    minimum=0.01, maximum=10, name="sigma_1", latex_label=r"$\sigma_1$"
)
priors["sigma_2"] = bilby.core.prior.LogUniform(
    minimum=0.01, maximum=10, name="sigma_2", latex_label=r"$\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="pymultinest",
    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")