UltraNest.jl

This is a Julia wrapper for Python nested sampling package UltraNest.

Nested sampling allows Bayesian inference on arbitrary user-defined likelihoods. In particular, posterior probability distributions on model parameters are constructed, and the marginal likelihood ("evidence") Z is computed. The former can be used to describe the parameter constraints of the data, the latter can be used for model comparison (via Bayes factors) as a measure of the prediction parsimony of a model.

UltraNest provides novel, advanced techniques (see how it works). They are especially remarkable for being free of tuning parameters and theoretically justified. Beyond that, UltraNest has support for Big Data sets and high-performance computing applications.

This Julia package is currently just a very thin wrapper around the Python ultranest package, which it will (if not already present) try to automatically install via Conda.jl.

In the future, some Julianic wrapper functions may be added for ultranests main functionalities, but UltraNest.jl is primarily intended to be a lower-level package that can act as a sampling backend for high-level packages.

Currently UltraNest.jl exports a single global constant ultranest. PyCall.jl's magic makes using ultranest from Julia very similar to the equivalent Python code.

Python:

import ultranest
# ... define log-likelihood and prior_transform ...
paramnames = ["a", "b", "c"]
smplr = ultranest.ReactiveNestedSampler(paramnames, loglikelihood, transform = prior_transform, vectorized = True)
result = smplr.run(min_num_live_points = 400, cluster_num_live_points = 40)

Julia:

using UltraNest
# ... define log-likelihood and prior_transform ...
paramnames = ["a", "b", "c"]
smplr = ultranest.ReactiveNestedSampler(paramnames, loglikelihood, transform = prior_transform, vectorized = true)
result = smplr.run(min_num_live_points = 4000, cluster_num_live_points = 400)

See the the UltraNest Python documentation regarding usage.

Usage example

As an example, we'll sample a multi-modal distribution that is tricky to handle using other methods like MCMC.

In addition to UltraNest.jl, we'll be using the Julia packages Distributions.jl, Plots.jl and Measurements.jl, so they need to be added to your current Julia environment to run this example.

Let's define a multi-modal distribution with fully separated modes that will act as our likelihood:

using Distributions

dist = product_distribution([
    MixtureModel([truncated(Normal(-1, 0.1), -2, 0), truncated(Normal(1, 0.1), 0, 2)], [0.5, 0.5]),
    MixtureModel([truncated(Normal(-2, 0.25), -3, -1), truncated(Normal(2, 0.25), 1, 3)], [0.3, 0.7]),
    MixtureModel([truncated(Normal(-5, 0.25), -6, -4), truncated(Normal(5, 0.25), 4, 6)], [0.2, 0.8]),
])

To use UltraNest, we need to define express the prior of our model as a transformation from the unit hypercube to the model parameters. Here we will simply use a flat prior over the support dist:

prior_transform = let dist=dist
    (u::AbstractVector{<:Real}) -> begin
        x0 = minimum.(dist.v)
        Δx = maximum.(dist.v) .- x0
        u .* Δx .+ x0
    end
end

UltraNest supports vectorized operation, passing multiple parameter vectors to the transformation as a matrix. We need to take into account that in UltraNest.jl, parameters are stored as rows in the matrix:

prior_transform_vectorized = let trafo = prior_transform
    (U::AbstractMatrix{<:Real}) -> reduce(vcat, (u -> trafo(u)').(eachrow(U)))
end

Our log-likelihood will simply be the log-PDF of dist for a given set up parameters:

loglikelihood = let dist = dist
    function (x::AbstractVector{<:Real})
        ll = logpdf(dist, x)
        # lofpdf on MixtureModel returns NaN in gaps between distributions, and UltraNest
        # doesn't like -Inf, so return -1E10:
        T = promote_type(Float32, typeof(ll))
        isnan(ll) ? T(-1E10) : T(ll)
    end
end

To use UltraNest's vectorized mode, we need a vectorized version of the likelihood as well:

loglikelihood_vectorized = let loglikelihood = loglikelihood
    # UltraNest has variate in rows:
    (X::AbstractMatrix{<:Real}) -> loglikelihood.(eachrow(X))
end

For computationally expensive likelihood functions, it will often be beneficial to parallelize this using Julia's multi-threaded and distributed computing capabilities. Note that the Python UltraNest package comes with MPI support as well - it should also be possible to use this for distributed parallelization, in combination with Julia's muli-threading for host-level parallelization.

Now we're ready to sample our posterior. We'll use UltraNest's ReactiveNestedSampler:

using UltraNest

paramnames = ["a", "b", "c"]
smplr = ultranest.ReactiveNestedSampler(paramnames, loglikelihood_vectorized, transform = prior_transform_vectorized, vectorized = true)
result = smplr.run(min_num_live_points = 4000, cluster_num_live_points = 400)

Dict{Any,Any} with 18 entries:
  "maximum_likelihood"       => Dict{Any,Any}("point"=>[1.00344, 2.00677, 4.996…
  "Herr"                     => 0.0476054
  "paramnames"               => ["a", "b", "c"]
  "logzerr_tail"             => 0.00994509
  "H"                        => 4.67062
  "niter"                    => 14273
  "logz_single"              => -5.64119
  "weighted_samples"         => Dict{Any,Any}("bootstrapped_weights"=>[0.0 0.0 …
  "ncall"                    => 30499
  "logzerr"                  => 0.0969422
  "logzerr_single"           => 0.034171
  "ess"                      => 4819.82
  "logz_bs"                  => -5.64204
  "samples"                  => [-0.987307 1.85425 5.11725; -1.10802 2.12304 5.…
  "posterior"                => Dict{Any,Any}("median"=>[-0.781364, 1.8606, 4.9…
  "logz"                     => -5.64119
  "logzerr_bs"               => 0.0964307
  "insertion_order_MWW_test" => Dict{Any,Any}("independent_iterations"=>36.0,"c…

result is a Dict that contains the samples (in weighted and unweighted form), the logarithm of the posterior's integral and additional valuable information. Since our prior was flat, we can compare the samples generated by UltraNest with truly random (IID) samples drawn from dist:

X_ultranest = result["samples"]

X_iid = Array(rand(dist, 10^4)')

using Plots

function plot_marginal(d::NTuple{2,Integer})
    if d[1] == d[2]
        plt = stephist(X_iid[:,d[1]].+0.01, nbins = 200, normalize = true, label = "IID sampling", xlabel = paramnames[d[1]])
        stephist!(plt, X_ultranest[:,d[1]], nbins = 200, normalize = true, label = "UltraNest")
        plt
    else
        histogram2d(
            X_ultranest[:,d[1]], X_ultranest[:,d[2]],
            nbins = 200, xlabel = paramnames[d[1]], ylabel = paramnames[d[2]], legend = :none
        )
    end
end

plot(
    plot_marginal.([(i, j) for i in 1:3, j in 1:3])...,
)

The log-integral of the posterior computed by UltraNest is

using Measurements

logz = result["logz"]
logzerr = result["logzerr"]
logz ± logzerr

which matches our expectation based on the volume of our flat prior:

logz_expected = -log(prod(maximum.(dist.v) .- minimum.(dist.v)))

-5.662960480135946

In addition to processing UltraNest's output with Julia directly, you can also let ultranest.ReactiveNestedSampler create on-disk output via keyword argument log_dir = "somepath/".