Support for the ultranest sampler

[JohannesBuchner]

Dear all,

It would be great if you could offer an option to select ultranest as a sampler. UltraNest is a nested sampling package, which provides several algorithms. By default, UltraNest is designed to be very robust and careful about integrating. Its default sampling algorithm, MLFriends, slows down when complex shapes are encountered to sample carefully and ensure correctness. This solves some biases that MultiNest has on some problems. MLFriends is very fast in low dimensions, and decent up to ~15d. There may be faster algorithms, but it is a good, trustworthy default comparison. However, UltraNest also allows turning on slice sampling within nested sampling (like polychord and dynesty), which allows effective usage in high dimensions and complex posteriors. UltraNest also provides a careful uncertainty estimation (using bootstrapping to emulate reruns) and run-time diagnostics (insertion rank U test).

More information on the documentation page: https://johannesbuchner.github.io/UltraNest/ There are extensive user tutorials and a detailed API description, but feel free to send me questions if something is unclear.

The interface of ultranest is very similar to pymultinest, where a prior transform function and the log-likelihood need to be provided. Additionally, a list with the parameter names should be passed.

Simple usage would be:

from ultranest import ReactiveNestedSampler

# param_names: list of parameter names
# loglike: log-likelihood function
# transform: prior unit-cube transform function
sampler = ReactiveNestedSampler(param_names, loglike, transform=transform)
sampler.run()
sampler.print_results()
# sampler.run returns the results, but they can also be accessed with sampler.results
# e.g. sampler.results['logz'], sampler.results['logzerr'], sampler.results['samples']

I would suggest something like the following, as a decent initialisation that should work for many problems:

from ultranest import ReactiveNestedSampler
from ultranest.mlfriends import SimpleRegion

# param_names: list of parameter names
# loglike: log-likelihood function
# transform: prior unit-cube transform function
# log_dir: folder where to store results (this should be a user parameter)
# resume: whether to allow continuing from previous run
# vectorized: whether loglike and transform can take many points at the same time

sampler = ReactiveNestedSampler(param_names, loglike, transform=transform,
    log_dir=log_dir, resume=True, vectorized=vectorized)

# use the very robust MLFriends algorithm
# this is very safe but potentially slow for high-dimensional / complex problems
# max_ncalls: we let it run only for some number of calls (this should be a user parameter)
# max_improvement_loops: setting this to >0 would enable dynamic nested sampling, 
#   but it usually is not necessary.
sampler.run(max_ncalls=100000, max_improvement_loops=0)

# if it is slow we should use a slice sampler:
#    there are several slice samplers available
#    generate_mixture_random_direction is the fastest according to https://arxiv.org/abs/2211.09426

slice_steps = min(10, 4 * len(param_names))  ## here is a good default, but this should be a user parameter

import ultranest.stepsampler
sampler.stepsampler = ultranest.stepsampler.SliceSampler(
    nsteps=slice_steps,  # see above
    region_filter=False, # this has some computational cost and is not that useful in high-d
    generate_direction=ultranest.stepsampler.generate_mixture_random_direction,
    # log=open(log_dir + '/stepsampler.log', 'w')  # can optionally log to a file
)

# frac_remain: fraction of posterior mass allowed to be in the live points (default is 0.01)
# region_class: since we use a step sampler, no need to construct a complicated region
sampler.run(region_class=SimpleRegion, frac_remain=0.1, max_improvement_loops=0)

# sampler.run returns the results, but they can also be accessed with sampler.results
# e.g. sampler.results['logz'], sampler.results['logzerr'], sampler.results['samples']

# make a nice printout of the results
sampler.print_results()

# optionally, stores corner plots and diagnostic plots in the directory
sampler.plot()

If you want to provide feedback to the user while running, setting a viz_callback callback allows you to display progress.