Skip to content

Resolve "Add implementation of the UltraNest sampler"

Matthew Pitkin requested to merge matthew-pitkin/bilby:ultranest into master

This adds a wrapper to the Ultranest sampler into bilby.

The sampler has two integrator methods:

  • NestedSampler: a "standard" simple nested sampler. This is used if the number of live points is specified (with num_live_points or any of the other aliases).
  • ReactiveNestedSampler: a nested sampler with a reactive exploration strategy. This is used if the number of live points is not specified.

This MR includes test codes. This has been checked to work using a version of the linear_regression_example:

#!/usr/bin/env python
"""
An example of how to use bilby to perform paramater estimation for
non-gravitational wave data. In this case, fitting a linear function to
data with background Gaussian noise

"""
from __future__ import division
import bilby
import numpy as np
import matplotlib.pyplot as plt

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


np.random.seed(235641)

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


# Now we define the injection parameters which we make simulated data with
injection_parameters = dict(m=0.5, c=0.2)

# For this example, we'll use standard Gaussian noise

# These lines of code generate the fake data. Note the ** just unpacks the
# contents of the injection_parameters when calling the model function.
sampling_frequency = 10
time_duration = 10
time = np.arange(0, time_duration, 1 / sampling_frequency)
N = len(time)
sigma = np.random.normal(1, 0.01, N)
data = model(time, **injection_parameters) + np.random.normal(0, sigma, N)

# We quickly plot the data to check it looks sensible
fig, ax = plt.subplots()
ax.plot(time, data, 'o', label='data')
ax.plot(time, model(time, **injection_parameters), '--r', label='signal')
ax.set_xlabel('time')
ax.set_ylabel('y')
ax.legend()
fig.savefig('{}/{}_data.png'.format(outdir, label))

# Now lets instantiate a version of our GaussianLikelihood, giving it
# the time, data and signal model
likelihood = bilby.likelihood.GaussianLikelihood(time, data, model, sigma)

# From hereon, the syntax is exactly equivalent to other bilby examples
# We make a prior
priors = dict()
priors['m'] = bilby.core.prior.Uniform(0, 5, 'm')
priors['c'] = bilby.core.prior.Uniform(-2, 2, 'c')

# And run sampler
result = bilby.run_sampler(
    likelihood=likelihood, priors=priors, sampler='ultranest', nlive=500,
    injection_parameters=injection_parameters, outdir=outdir, label=label)
result.plot_corner()

Resumption and checkpointing works if this script is exited during running and restarted. It has been checked to work for both the standard and reactive nested sampling by removing the nlive option from run_sampler().

Ultranest has a variety of built-in MCMC step samplers that can be used. In this implementation they can be used by creating the step sampler and passing it to run_sampler() with the step_sampler keyword argument, e.g., the above example could be changed to include

import ultranest.stepsampler

stepsampler = ultranest.stepsampler.RegionSliceSampler(nsteps=10)

# And run sampler
result = bilby.run_sampler(
    likelihood=likelihood, priors=priors, sampler='ultranest',
    injection_parameters=injection_parameters, outdir=outdir, label=label,
    step_sampler=stepsampler)

Resolves #456 (closed).

Edited by Matthew Pitkin

Merge request reports