Integration with Turing and NestedSamplers

[yebai]

NestedSamplers is an independent sampler library at the moment and lacks an interface for models specified by the domain-specific language (DSL) in Turing. However, since NestedSampler supports the AbstractMCMC interface, it should be a relatively lightweight task to add such an interface. A good starting point is the elliptical slice sampler (ESS) interface, which can be found at

• Interface: https://github.com/TuringLang/Turing.jl/blob/master/src/inference/ess.jl
• Elliptical slice sampler implementation: https://github.com/TuringLang/EllipticalSliceSampling.jl

[SaranjeetKaur]

"""
Nested sampling algorithm.

Example

#Note-to-self: set a seed here
julia> using NestedSamplers
julia> using Distributions

# eggbox likelihood function
tmax = 3π
julia> function logl(x)
       t = @. 2*tmax*x - tmax
       return 2 + cos(t[1]/2)*cos(t[2]/2)^5
       end
logl (generic function with 1 method)

# prior constrined to line in the range (0, 20)
julia> prior = [
           Uniform(0, 20),
           Uniform(0, 20)
       ]
2-element Array{Uniform{Float64},1}:
 Uniform{Float64}(a=0.0, b=20.0)
 Uniform{Float64}(a=0.0, b=20.0)

# creating the model
julia> model = NestedModel(logl, prior)
NestedModel{typeof(logl),Uniform{Float64}}(logl, Uniform{Float64}[Uniform{Float64}(a=0.0, b=20.0), Uniform{Float64}(a=0.0, b=20.0)])

julia> using StatsBase: sample, Weights
julia> using MCMCChains: Chains

# creating the sampler
# 2 parameters, 100 active points, multi-ellipsoid

julia> spl = Nested(2, 100; bounds = Bounds.MultiEllipsoid)
Nested(ndims=2, nactive=100, enlarge=1.25, update_interval=150)
  bounds=MultiEllipsoid{Float64}(ndims=2)
  proposal=NestedSamplers.Proposals.Uniform
  logz=-1.0e300
  log_vol=-4.610166019324897
  H=0.0

julia> spl = Nested(2, 100; bounds = Bounds.MultiEllipsoid)
Nested(ndims=2, nactive=100, enlarge=1.25, update_interval=150)
  bounds=MultiEllipsoid{Float64}(ndims=2)
  proposal=NestedSamplers.Proposals.Uniform
  logz=-1.0e300
  log_vol=-4.610166019324897
  H=0.0

julia> chain = sample(model, spl;
                      dlogz = 0.2,
                      param_names = ["x", "y"],
                      chain_type = Chains)
Object of type Chains, with data of type 358×3×1 Array{Float64,3}

Log evidence      = 2.086139406693275
Iterations        = 1:358
Thinning interval = 1
Chains            = 1
Samples per chain = 358
internals         = weights
parameters        = x, y

2-element Array{MCMCChains.ChainDataFrame,1}

Summary Statistics
  parameters     mean     std  naive_se    mcse       ess   r_hat
  ──────────  ───────  ──────  ────────  ──────  ────────  ──────
           x   9.9468  5.7734    0.3051  0.5474  315.7439  1.0017
           y  10.0991  5.5946    0.2957  0.2016  396.0076  0.9995

Quantiles
  parameters    2.5%   25.0%   50.0%    75.0%    97.5%
  ──────────  ──────  ──────  ──────  ───────  ───────
           x  0.4682  5.2563  9.5599  14.9525  19.3789
           y  0.5608  5.4179  9.8601  14.8477  19.2020

"""

[SaranjeetKaur]

Getting the following error while using bounds = Bounds.MultiEllipsoid argument in the function Nested
ERROR: UndefVarError: Bounds not defined

[mileslucas]

Are you on version 0.3?

[SaranjeetKaur]

Not yet.. i used NestedSamplers version 0.1 with Julia version 1.2 .. updating it all now

…

On Wed, Jun 3, 2020 at 12:08 AM Miles Lucas ***@***.***> wrote: Are you on version 0.3? — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#29 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AGZ32SB72SZV3FCUCUPN2JDRUVBLXANCNFSM4M3NTL4Q> .

[mvsoom]

Any news on this? Integration between TuringLang and NestedSamplers would be a dream.

[yebai]

The integration with nested sampling should become trivial once TuringLang/DynamicPPL.jl#309 is merged.

cc @ParadaCarleton

[mvsoom]

Great! Does the PR include an example of this integration by any chance?

Thank you for your hard work.

[mileslucas]

The few times I've tried figuring out how to do this integration, I've given up because I didn't understand how Turing worked internally. It was never clear to me how to compute the prior transform or loglikelihood methods using the DynamicPPL API. For example, in the current NestedSamplers framework, I need to know the number of parameters, and have a loglike(X::AbstractVector) and prior_transform(X::AbstractVector) methods. It was never clear to me how I would construct, e.g., the prior_transform from a bunch of assume calls, or the likelihood function from observe calls. As far as I know, nested sampling is not build to operate on single parameters at a time, but rather samples from the entire high-dimensional parameter space, since we have to generate new samples contingent on the current iso-likelihood contour.

@yebai how will TuringLang/DynamicPPL.jl#309 make this trivial? It is not clear to me from the PR discussion.

[torfjelde]

That PR will essentially just make it easier to do something like

loglikelihood(model, namedtuple)

or

loglikelihood(model, vector)

etc. with very high-performance, i.e. make it super-easy to evaluate the density of (a large subset of but not all definable) models.

This would mean that you could make NestedSamplers.jl work with Model without implenting stuff like assume, observe, dot_assume, etc. for the particular samplers.

[ParadaCarleton]

That PR will essentially just make it easier to do something like

loglikelihood(model, namedtuple)

or

loglikelihood(model, vector)

etc. with very high-performance, i.e. make it super-easy to evaluate the density of (a large subset of but not all definable) models.

This would mean that you could make NestedSamplers.jl work with Model without implenting stuff like assume, observe, dot_assume, etc. for the particular samplers.

What kinds of models might not work? And is there a timeline on the PR?

[mvsoom]

Next to evaluating the likelihood, how would one efficiently implement the prior_transform?

[torfjelde]

What kinds of models might not work?

Well, all models would technically work but depending on what the underlying storage is, you might need to do some manual labour:

1. Using NamedTuple, if you have access patterns that looks like x[i] ~ ..., etc. i.e. you're sampling parts of some container, then you need to pre-allocate the container when constructing the "trace"/SimpleVarInfo, e.g. SimpleVarInfo((x = zeros(n), )). If you're not sampling parts of some container, i.e. all your statements are of the form x ~ ..., then it will just work and be blazingly fast.
2. Using Dict we can support statements such as x[i] ~ ... without any input from the user, but it'll be much slower than using NamedTuple.

You can check out the docs for the PR here: https://turinglang.github.io/DynamicPPL.jl/previews/PR309/#DynamicPPL.SimpleVarInfo

And docstring for logjoint, etc.: https://turinglang.github.io/DynamicPPL.jl/previews/PR309/#DynamicPPL.logjoint-Tuple{Model,%20Union{AbstractDict,%20NamedTuple}}

And is there a timeline on the PR?

This is open-source bby, there are no timelines! 🎸

Nah I kid:) That PR is my highest priority in my open-source work, so it shouldn't be long now. I'm just ironing out some quirks and making sure that all test-cases are covered. This is a pretty significant change, so we need to make sure that it's not a bad idea.

Next to evaluating the likelihood, how would one efficiently implement the prior_transform?

So this is actually something I'm somewhat confused about: why do we need to work on the unit-cube? I'm new to nested sampling, so please forgive my ignorance 🙃

But in general, we deal with transformations between spaces in Turing.jl too; in fact we have an entire package to deal with this (Bijectors.jl). So IIUC, this shouldn't be an issue since we can just provide the transform using a composition of the bijectors. But we could also provide a function which extracts the distributions used, if need be.

[yebai]

So this is actually something I'm somewhat confused about: why do we need to work on the unit-cube? I'm new to nested sampling, so please forgive my ignorance 🙃

This is not strictly necessary from a technical point of view. In order to perform nested sampling, one only need to sample from a constrained prior, or sample from the unconstrained prior, then perform a rejection sampling step. However, it seems very popular in astrophysics to sample from a unit cube to simplify implementation. @mileslucas Would it be possible to directly work with the prior?

[yebai]

It seems the current ‘prior_transform’ is coupled with proposal implementations. That makes working directly with priors slightly more involved - we need to implement new proposals that operate in the parameter space instead of the unit-cube space. The current implementation also limits the applicable model of nested sampling to cases whose prior can be factorised into univariate distributions with known quantile functions.

I think a near-term integration goal should be supporting the same family of models that nested sampling can handle, i.e. models with priors that can be written as a product of independent univariate distributions.