Merge pull request #56 from TuringLang/ml/models

docs/Project.toml

@@ -4,6 +4,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+NestedSamplers = "41ceaf6f-1696-4a54-9b49-2e7a9ec3782e"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"

docs/make.jl

@@ -12,6 +12,10 @@ makedocs(
     sitename = "NestedSamplers.jl",
     pages = [
         "Home" => "index.md",
+        "Examples" => [
+            "Gaussian Shells" => "examples/shells.md",
+            "Correlated Gaussian" => "examples/correlated.md"
+        ],
         "API/Reference" => "api.md"
     ],
     format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),

docs/src/api.md

@@ -9,6 +9,7 @@
 NestedModel
 Nested
 ```
+
 ### Convergence
 
  There are a few convergence criteria available, by default the `dlogz` criterion will be used.
@@ -35,4 +36,12 @@ Proposals.RWalk
 Proposals.RStagger
 Proposals.Slice
 Proposals.RSlice
-```
+```
+
+## Models
+
+```@docs
+Models
+Models.GaussianShells
+Models.CorrelatedGaussian
+```

docs/src/examples/correlated.md

@@ -0,0 +1,81 @@
+# Correlated Gaussian
+
+This example will explore a highly-correlated Gaussian using [`Models.CorrelatedGaussian`](@ref). This model uses a conjuage Gaussian prior, see the docstring for the mathematical definition.
+
+## Setup
+
+For this example, you'll need to add the following packages
+```julia
+julia>]add Distributions MCMCChains Measurements NestedSamplers StatsBase StatsPlots
+```
+
+```@setup correlated
+using AbstractMCMC
+using Random
+AbstractMCMC.setprogress!(false)
+Random.seed!(8452)
+```
+
+## Define model
+
+```@example correlated
+using NestedSamplers
+
+# set up a 4-dimensional Gaussian
+D = 4
+model, logz = Models.CorrelatedGaussian(D)
+nothing; # hide
+```
+
+let's take a look at a couple of parameters to see what the likelihood surface looks like
+
+```@example correlated
+using StatsPlots
+
+θ1 = range(-1, 1, length=1000)
+θ2 = range(-1, 1, length=1000)
+logf = [model.loglike([t1, t2, 0, 0]) for t2 in θ2, t1 in θ1]
+heatmap(
+    θ1, θ2, exp.(logf),
+    aspect_ratio=1,
+    xlims=extrema(θ1),
+    ylims=extrema(θ2),
+    xlabel="θ1",
+    ylabel="θ2"
+)
+```
+
+## Sample
+
+```@example correlated
+using MCMCChains
+using StatsBase
+# using single Ellipsoid for bounds
+# using Gibbs-style slicing for proposing new points
+sampler = Nested(D, 50 * (D + 1); 
+    bounds=Bounds.Ellipsoid,
+    proposal=Proposals.Slice()
+)
+names = ["θ_$i" for i in 1:D]
+chain, state = sample(model, sampler; dlogz=0.01, param_names=names)
+# resample chain using statistical weights
+chain_resampled = sample(chain, Weights(vec(chain[:weights])), length(chain));
+nothing # hide
+```
+
+## Results
+
+```@example correlated
+chain_resampled
+```
+
+```@example correlated
+corner(chain_resampled)
+```
+
+```@example correlated
+using Measurements
+logz_est = state.logz ± state.logzerr
+diff = logz_est - logz
+print("logz: ", logz, "\nestimate: ", logz_est, "\ndiff: ", diff)
+```

docs/src/examples/shells.md

@@ -0,0 +1,82 @@
+# Gaussian Shells
+
+This example will explore the classic Gaussian shells model using [`Models.GaussianShells`](@ref).
+
+## Setup
+
+For this example, you'll need to add the following packages
+```julia
+julia>]add Distributions MCMCChains Measurements NestedSamplers StatsBase StatsPlots
+```
+
+```@setup shells
+using AbstractMCMC
+using Random
+AbstractMCMC.setprogress!(false)
+Random.seed!(8452)
+```
+
+## Define model
+
+```@example shells
+using NestedSamplers
+
+model, logz = Models.GaussianShells()
+nothing; # hide
+```
+
+let's take a look at a couple of parameters to see what the likelihood surface looks like
+
+```@example shells
+using StatsPlots
+
+x = range(-6, 6, length=1000)
+y = range(-6, 6, length=1000)
+logf = [model.loglike([xi, yi]) for yi in y, xi in x]
+heatmap(
+    x, y, exp.(logf),
+    aspect_ratio=1,
+    xlims=extrema(x),
+    ylims=extrema(y),
+    xlabel="x",
+    ylabel="y",
+    size=(400, 400)
+)
+```
+
+## Sample
+
+```@example shells
+using MCMCChains
+using StatsBase
+# using multi-ellipsoid for bounds
+# using default rejection sampler for proposals
+sampler = Nested(2, 1000)
+chain, state = sample(model, sampler; dlogz=0.05, param_names=["x", "y"])
+# resample chain using statistical weights
+chain_resampled = sample(chain, Weights(vec(chain[:weights])), length(chain));
+nothing # hide
+```
+
+## Results
+
+```@example shells
+chain_resampled
+```
+
+```@example shells
+marginalkde(chain[:x], chain[:y])
+```
+
+```@example shells
+density(chain_resampled)
+vline!([-5.5, -1.5, 1.5, 5.5], c=:black, ls=:dash, sp=1)
+vline!([-2, 2], c=:black, ls=:dash, sp=2)
+```
+
+```@example shells
+using Measurements
+logz_est = state.logz ± state.logzerr
+diff = logz_est - logz
+print("logz: ", logz, "\nestimate: ", logz_est, "\ndiff: ", diff)
+```

src/NestedSamplers.jl

@@ -6,7 +6,6 @@ using .Bounds
 include("proposals/Proposals.jl")
 using .Proposals
 
-
 using LinearAlgebra
 using Random
 using Random: AbstractRNG, GLOBAL_RNG
@@ -29,6 +28,7 @@ using StatsFuns: logaddexp,
 
 export Bounds,
        Proposals,
+       Models,
        NestedModel,
        Nested
 
@@ -37,4 +37,7 @@ include("staticsampler.jl") # The static nested sampler
 include("step.jl")          # The stepping mechanics (extends AbstractMCMC)
 include("sample.jl")        # Custom sampling (extends AbstractMCMC)
 
+include("models/Models.jl")
+using .Models
+
 end

src/models/Models.jl

@@ -0,0 +1,18 @@
+"""
+This module contains various statistical models in the form of [`NestedModel`](@ref)s. These models can be used for examples and for testing.
+
+* [`Models.GaussianShells`](@ref)
+* [`Models.CorrelatedGaussian`](@ref)
+"""
+module Models
+
+using ..NestedSamplers
+
+using Distributions
+using LinearAlgebra
+using StatsFuns
+
+include("shells.jl")
+include("correlated.jl")
+
+end # module

src/models/correlated.jl

@@ -0,0 +1,41 @@
+
+@doc raw"""
+Models.CorrelatedGaussian(ndims)
+
+Creates a highly-correlated Gaussian with the given dimensionality.
+
+```math
+\mathbf\theta \sim \mathcal{N}\left(2\mathbf{1}, \mathbf{I}\right)
+```
+```math
+\Sigma_{ij} = \begin{cases} 1 &\quad i=j \\ 0.95 &\quad i\neq j \end{cases}
+```
+```math
+\mathcal{L}(\mathbf\theta) = \mathcal{N}\left(\mathbf\theta | \mathbf{0}, \mathbf\Sigma \right)
+```
+
+the analytical evidence of the model is
+
+```math
+Z = \mathcal{N}\left(2\mathbf{1} | \mathbf{0}, \mathbf\Sigma + \mathbf{I} \right)
+```
+
+## Examples
+```jldoctest
+julia> model, lnZ = Models.CorrelatedGaussian(10);
+
+julia> lnZ
+-12.482738597926607
+```
+"""
+function CorrelatedGaussian(ndims)
+    priors = fill(Normal(2, 1), ndims)
+    Σ = fill(0.95, ndims, ndims)
+    Σ[diagind(Σ)] .= 1
+    cent_dist = MvNormal(Σ)
+    loglike(X) = logpdf(cent_dist, X)
+
+    model = NestedModel(loglike, priors)
+    true_lnZ = logpdf(MvNormal(fill(2, ndims), Σ + I), zeros(ndims))
+    return model, true_lnZ
+end

src/models/shells.jl

@@ -0,0 +1,31 @@
+using StatsFuns
+
+"""
+    Models.GaussianShells()
+
+2-D Gaussian shells centered at `[-3.5, 0]` and `[3.5, 0]` with a radius of 2 and a shell width of 0.1
+
+# Examples
+```jldoctest
+julia> model, lnZ = Models.GaussianShells();
+
+julia> lnZ
+-1.75
+```
+"""
+function GaussianShells()
+    μ1 = [-3.5, 0]
+    μ2 = [3.5, 0]
+
+    prior(X) = 12 .* X .- 6
+    loglike(X) = logaddexp(logshell(X, μ1), logshell(X, μ2))
+
+    lnZ = -1.75
+    return NestedModel(loglike, prior), lnZ
+end
+
+function logshell(X, μ, radius=2, width=0.1)
+    d = LinearAlgebra.norm(X - μ)
+    norm = -log(sqrt(2 * π * width^2))
+    return norm - (d - radius)^2 / (2 * width^2)
+end

src/step.jl

@@ -29,6 +29,7 @@ function step(rng, model, sampler::Nested; kwargs...)
     logz = logwt
     h = logl_dead - logz
     logzerr = sqrt(h / sampler.nactive)
+    logvol -= 1 / sampler.nactive
 
     sample = (u = u_dead, v = v_dead, logwt = logwt, logl = logl_dead)
     state = (it = 1, ncall = ncall, us = us, vs = vs, logl = logl, logl_dead = logl_dead,
@@ -85,14 +86,14 @@ function step(rng, model, sampler, state; kwargs...)
     since_update += nc
 
     # update weight
-    logvol = state.logvol - 1 / sampler.nactive
     logwt = state.logvol + logl_dead
 
     # update evidence and information
     logz = logaddexp(state.logz, logwt)
     h = (exp(logwt - logz) * logl_dead +
          exp(state.logz - logz) * (state.h + state.logz) - logz)
     logzerr = sqrt(h / sampler.nactive)
+    logvol = state.logvol - 1 / sampler.nactive
 
     ## prepare returns
     sample = (u = u_dead, v = v_dead, logwt = logwt, logl = logl_dead)
@@ -219,8 +220,6 @@ function add_live_points(samples, model, sampler, state)
 
         sample = (u = u, v = v, logwt = logwt, logl = logl)
         save!!(samples, sample, length(samples) + i, model, sampler)
-
-        i += 1
     end
 
     state = (it = state.it + sampler.nactive, us = state.us, vs = state.vs, logl = logl,