multivariate_gaussian_prior.py

#!/usr/bin/env python
"""
An example of how to use bilby with a (multi-modal) multivariate
Gaussian prior distribution.
"""

from __future__ import division
import bilby
import numpy as np
from scipy import linalg, stats
import matplotlib as mpl

from bilby.core.likelihood import GaussianLikelihood

# A few simple setup steps
label = 'multivariate_gaussian_prior'
outdir = 'outdir'
bilby.utils.check_directory_exists_and_if_not_mkdir(outdir)


# First, we define our "signal model", in this case a simple linear function
def model(time, m, c):
    return time * m + c


# For this example assume a very broad Gaussian likelihood, so that the
# posterior is completely dominated by the prior (and we can check the
# prior sampling looks correct!)
sigma = 300.

# These lines of code generate the fake (noise-free) data
sampling_frequency = 1
time_duration = 10
time = np.arange(0, time_duration, 1 / sampling_frequency)
N = len(time)
data = model(time, 0., 0.)  # noiseless data


# instantiate the GaussianLikelihood
likelihood = GaussianLikelihood(time, data, model, sigma=sigma)

# Create a Multivariate Gaussian prior distribution with two modes
names = ['m', 'c']
mus = [[-5., -5.], [5., 5.]]  # means of the two modes
corrcoefs = [[[1., -0.7], [-0.7, 1.]], [[1., 0.7], [0.7, 1.]]]  # correlation coefficients of the two modes
sigmas = [[1.5, 1.5], [2.1, 2.1]]  # standard deviations of the two modes
weights = [1., 3.]  # relative weights of each mode
nmodes = 2
mvg = bilby.core.prior.MultivariateGaussianDist(names, nmodes=2, mus=mus,
                                                corrcoefs=corrcoefs,
                                                sigmas=sigmas, weights=weights)
priors = dict()
priors['m'] = bilby.core.prior.MultivariateGaussian(mvg, 'm')
priors['c'] = bilby.core.prior.MultivariateGaussian(mvg, 'c')

result = bilby.run_sampler(
    likelihood=likelihood, priors=priors, sampler='dynesty', nlive=4000,
    outdir=outdir, label=label)

fig = result.plot_corner(save=False)

# plot the priors (to show that they look correct)
axs = fig.get_axes()

# plot the 1d marginal distributions
x = np.linspace(-12, 12, 5000)
aidx = [0, 3]
for j in range(2):  # loop over parameters
    gp = np.zeros(len(x))
    for i in range(nmodes):  # loop over modes
        gp += weights[i]*stats.norm.pdf(x, loc=mus[i][j], scale=mvg.sigmas[i][j])
    gp = gp/np.trapz(gp, x)  # renormalise
    
    axs[aidx[j]].plot(x, gp, 'k--', lw=2)

# plot the 2d distribution
for i in range(nmodes):
    v, w = linalg.eigh(mvg.covs[i])
    v = 2. * np.sqrt(2.) * np.sqrt(v)
    u = w[0] / linalg.norm(w[0])
    angle = np.arctan(u[1] / u[0])
    angle = 180. * angle / np.pi  # convert to degrees
    ell = mpl.patches.Ellipse(mus[i], v[0], v[1], 180. + angle, edgecolor='black', facecolor='none', lw=2, ls='--')
    axs[2].add_artist(ell)

fig.savefig('{}/{}_corner.png'.format(outdir, label), dpi=300)