Merge pull request #12 from TuringLang/csp/turing-connect

Make AdvancedMH better

Author: cpfiffer

.gitignore

@@ -1 +1,2 @@
 Manifest.toml
+.vscode/settings.json

Project.toml

@@ -1,21 +1,28 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.2.0"
+version = "0.3.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
-[extras]
-Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-
 [compat]
-AbstractMCMC = "0.1"
-Distributions = "0.21"
-MCMCChains = "1, 0.4"
+AbstractMCMC = "0.3"
+Distributions = "0.20, 0.21, 0.22, 0.23"
 Reexport = "0.2"
+Requires = "1.0"
+StructArrays = "^0"
 julia = "1"
+
+[extras]
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
+
+[targets]
+test = ["StructArrays", "MCMCChains"]

README.md

@@ -4,17 +4,17 @@ AdvancedMH.jl currently provides a robust implementation of random walk Metropol
 
 Further development aims to provide a suite of adaptive Metropolis-Hastings implementations.
 
-Currently there are two sampler types. The first is `RWMH`, which represents random-walk MH sampling, and the second is `StaticMH`, which draws proposals
-from only a prior distribution without incrementing the previous sample.
+AdvancedMH works by allowing users to define composable `Proposal` structs in different formats.
 
 ## Usage
 
-AdvancedMH works by accepting some log density function which is used to construct a `DensityModel`. The `DensityModel` is then used in a `sample` call.
+First, construct a `DensityModel`, which is a wrapper around the log density function for your inference problem. The `DensityModel` is then used in a `sample` call.
 
 ```julia
 # Import the package.
 using AdvancedMH
 using Distributions
+using MCMCChains
 
 # Generate a set of data from the posterior we want to estimate.
 data = rand(Normal(0, 1), 30)
@@ -27,11 +27,11 @@ density(θ) = insupport(θ) ? sum(logpdf.(dist(θ), data)) : -Inf
 # Construct a DensityModel.
 model = DensityModel(density)
 
-# Set up our sampler with initial parameters.
-spl = RWMH([0.0, 0.0])
+# Set up our sampler with a joint multivariate Normal proposal.
+spl = RWMH(MvNormal(2,1))
 
 # Sample from the posterior.
-chain = sample(model, spl, 100000; param_names=["μ", "σ"])
+chain = sample(model, spl, 100000; param_names=["μ", "σ"], chain_type=Chains)
 ```
 
 Output:
@@ -46,41 +46,74 @@ Samples per chain = 100000
 internals         = lp
 parameters        = μ, σ
 
-2-element Array{MCMCChains.ChainDataFrame,1}
+2-element Array{ChainDataFrame,1}
 
 Summary Statistics
 
-│ Row │ parameters │ mean      │ std      │ naive_se   │ mcse       │ ess     │ r_hat   │
-│     │ Symbol     │ Float64   │ Float64  │ Float64    │ Float64    │ Any     │ Any     │
-├─────┼────────────┼───────────┼──────────┼────────────┼────────────┼─────────┼─────────┤
-│ 1   │ μ          │ 0.0834188 │ 0.241418 │ 0.00076343 │ 0.00341067 │ 4693.04 │ 1.00008 │
-│ 2   │ σ          │ 1.33116   │ 0.184111 │ 0.00058221 │ 0.00258778 │ 4965.83 │ 1.00001 │
+│ Row │ parameters │ mean     │ std      │ naive_se    │ mcse       │ ess     │ r_hat   │
+│     │ Symbol     │ Float64  │ Float64  │ Float64     │ Float64    │ Any     │ Any     │
+├─────┼────────────┼──────────┼──────────┼─────────────┼────────────┼─────────┼─────────┤
+│ 1   │ μ          │ 0.156152 │ 0.19963  │ 0.000631285 │ 0.00323033 │ 3911.73 │ 1.00009 │
+│ 2   │ σ          │ 1.07493  │ 0.150111 │ 0.000474693 │ 0.00240317 │ 3707.73 │ 1.00027 │
 
 Quantiles
 
-│ Row │ parameters │ 2.5%      │ 25.0%      │ 50.0%     │ 75.0%    │ 97.5%    │
-│     │ Symbol     │ Float64   │ Float64    │ Float64   │ Float64  │ Float64  │
-├─────┼────────────┼───────────┼────────────┼───────────┼──────────┼──────────┤
-│ 1   │ μ          │ -0.393769 │ -0.0771134 │ 0.0801688 │ 0.241162 │ 0.564331 │
-│ 2   │ σ          │ 1.03685   │ 1.2044     │ 1.30992   │ 1.43609  │ 1.75745  │
+│ Row │ parameters │ 2.5%     │ 25.0%     │ 50.0%    │ 75.0%    │ 97.5%    │
+│     │ Symbol     │ Float64  │ Float64   │ Float64  │ Float64  │ Float64  │
+├─────┼────────────┼──────────┼───────────┼──────────┼──────────┼──────────┤
+│ 1   │ μ          │ -0.23361 │ 0.0297006 │ 0.159139 │ 0.283493 │ 0.558694 │
+│ 2   │ σ          │ 0.828288 │ 0.972682  │ 1.05804  │ 1.16155  │ 1.41349  │
+
 ```
 
-## Custom proposals
+## Proposals
 
-Custom proposal distributions can be specified by passing a distribution to `MetropolisHastings`:
+AdvancedMH offers various methods of defining your inference problem. Behind the scenes, a `MetropolisHastings` sampler simply holds
+some set of `Proposal` structs. AdvancedMH will return posterior samples in the "shape" of the proposal provided -- currently
+supported methods are `Array{Proposal}`, `Proposal`, and `NamedTuple{Proposal}`. For example, proposals can be created as:
 
 ```julia
-# Set up our sampler with initial parameters.
-spl1 = RWMH([0.0, 0.0], MvNormal(2, 0.5)) 
-spl2 = StaticMH([0.0, 0.0], MvNormal(2, 0.5)) 
+# Provide a univariate proposal.
+m1 = DensityModel(x -> logpdf(Normal(x,1), 1.0))
+p1 = Proposal(Static(), Normal(0,1))
+c1 = sample(m1, MetropolisHastings(p1), 100; chain_type=Vector{NamedTuple})
+
+# Draw from a vector of distributions.
+m2 = DensityModel(x -> logpdf(Normal(x[1], x[2]), 1.0))
+p2 = Proposal(Static(), [Normal(0,1), InverseGamma(2,3)])
+c2 = sample(m2, MetropolisHastings(p2), 100; chain_type=Vector{NamedTuple})
+
+# Draw from a `NamedTuple` of distributions.
+m3 = DensityModel(x -> logpdf(Normal(x.a, x.b), 1.0))
+p3 = (a=Proposal(Static(), Normal(0,1)), b=Proposal(Static(), InverseGamma(2,3)))
+c3 = sample(m3, MetropolisHastings(p3), 100; chain_type=Vector{NamedTuple})
+
+# Draw from a functional proposal.
+m4 = DensityModel(x -> logpdf(Normal(x,1), 1.0))
+p4 = Proposal(Static(), (x=1.0) -> Normal(x, 1))
+c4 = sample(m4, MetropolisHastings(p4), 100; chain_type=Vector{NamedTuple})
 ```
 
+## Static vs. Random Walk
+
+Currently there are only two methods of inference available. Static MH simply draws from the prior, with no
+conditioning on the previous sample. Random walk will add the proposal to the previously observed value.
+If you are constructing a `Proposal` by hand, you can determine whether the proposal is `Static` or `RandomWalk` using
+
+```julia
+static_prop = Proposal(Static(), Normal(0,1))
+rw_prop = Proposal(RandomWalk(), Normal(0,1))
+```
+
+Different methods are easily composeable. One parameter can be static and another can be a random walk,
+each of which may be drawn from separate distributions.
+
 ## Multithreaded sampling
 
 AdvancedMH.jl implements the interface of [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl/), which means you get multiple chain sampling
 in parallel for free:
 
 ```julia
 # Sample 4 chains from the posterior.
-chain = psample(model, RWMH(init_params), 100000, 4; param_names=["μ","σ"])
+chain = psample(model, RWMH(init_params), 100000, 4; param_names=["μ","σ"], chain_type=Chains)
 ```

src/AdvancedMH.jl

@@ -1,28 +1,31 @@
 module AdvancedMH
 
 # Import the relevant libraries.
-using Reexport
 using AbstractMCMC
-using Distributions
 using Random
+using Requires
+using Distributions
 
 # Import specific functions and types to use or overload.
-import MCMCChains: Chains
-import AbstractMCMC: step!, AbstractSampler, AbstractTransition, transition_type, bundle_samples
+import AbstractMCMC: step!, AbstractSampler, AbstractTransition, 
+    transition_type, bundle_samples
 
 # Exports
-export MetropolisHastings, DensityModel, sample, psample, RWMH, StaticMH
+export MetropolisHastings, DensityModel, sample, psample, RWMH, StaticMH, Proposal, Static, RandomWalk
 
 # Abstract type for MH-style samplers.
 abstract type Metropolis <: AbstractSampler end
 abstract type ProposalStyle end
 
+struct RandomWalk <: ProposalStyle end
+struct Static <: ProposalStyle end
+
 # Define a model type. Stores the log density function and the data to 
 # evaluate the log density on.
 """
     DensityModel{F<:Function} <: AbstractModel
 
-`DensityModel` wraps around a self-contained log-liklihood function `ℓπ`.
+`DensityModel` wraps around a self-contained log-liklihood function `logdensity`.
 
 Example:
 
@@ -32,57 +35,73 @@ DensityModel
 ```
 """
 struct DensityModel{F<:Function} <: AbstractModel
-    ℓπ :: F
+    logdensity :: F
 end
 
 # Create a very basic Transition type, only stores the 
 # parameter draws and the log probability of the draw.
-struct Transition{T<:Union{Vector{<:Real}, <:Real}, L<:Real} <: AbstractTransition
-    θ :: T
+struct Transition{T<:Union{Vector, Real, NamedTuple}, L<:Real} <: AbstractTransition
+    params :: T
     lp :: L
 end
 
 # Store the new draw and its log density.
-Transition(model::M, θ::T) where {M<:DensityModel, T} = Transition(θ, ℓπ(model, θ))
+Transition(model::M, params::T) where {M<:DensityModel, T} = Transition(params, logdensity(model, params))
 
 # Tell the interface what transition type we would like to use.
-transition_type(model::DensityModel, spl::Metropolis) = typeof(Transition(spl.init_θ, ℓπ(model, spl.init_θ)))
+transition_type(model::DensityModel, spl::Metropolis) = typeof(Transition(spl.init_params, logdensity(model, spl.init_params)))
 
 # Calculate the density of the model given some parameterization.
-ℓπ(model::DensityModel, θ::T) where T = model.ℓπ(θ)
-ℓπ(model::DensityModel, t::Transition) = t.lp
+logdensity(model::DensityModel, params) = model.logdensity(params)
+logdensity(model::DensityModel, t::Transition) = t.lp
 
 # A basic chains constructor that works with the Transition struct we defined.
 function bundle_samples(
     rng::AbstractRNG, 
-    ℓ::DensityModel, 
+    model::DensityModel, 
     s::Metropolis, 
     N::Integer, 
-    ts::Vector{T}; 
+    ts::Type{Any}; 
     param_names=missing,
     kwargs...
 ) where {ModelType<:AbstractModel, T<:AbstractTransition}
-    # Turn all the transitions into a vector-of-vectors.
-    vals = copy(reduce(hcat,[vcat(t.θ, t.lp) for t in ts])')
+    return ts
+end
 
+function bundle_samples(
+    rng::AbstractRNG, 
+    model::DensityModel, 
+    s::Metropolis, 
+    N::Integer, 
+    ts::Vector{T},
+    chain_type::Type{Vector{NamedTuple}}; 
+    param_names=missing,
+    kwargs...
+) where {ModelType<:AbstractModel, T<:AbstractTransition}
     # Check if we received any parameter names.
     if ismissing(param_names)
-        param_names = ["Parameter $i" for i in 1:length(s.init_θ)]
+        param_names = ["param_$i" for i in 1:length(keys(ts[1].params))]
     else
         # Deepcopy to be thread safe.
         param_names = deepcopy(param_names)
     end
 
-    # Add the log density field to the parameter names.
     push!(param_names, "lp")
 
-    # Bundle everything up and return a Chains struct.
-    return Chains(vals, param_names, (internals=["lp"],))
+    # Turn all the transitions into a vector-of-NamedTuple.
+    ks = tuple(Symbol.(param_names)...)
+    nts = [NamedTuple{ks}(tuple(t.params..., t.lp)) for t in ts]
+
+    return nts
+end
+
+function __init__()
+    @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("mcmcchains-connect.jl")
+    @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("structarray-connect.jl")
 end
 
 # Include inference methods.
+include("proposal.jl")
 include("mh-core.jl")
-include("rwmh.jl")
-include("staticmh.jl")
 
 end # module AdvancedMH

src/mcmcchains-connect.jl

@@ -0,0 +1,30 @@
+import .MCMCChains: Chains
+
+# A basic chains constructor that works with the Transition struct we defined.
+function bundle_samples(
+    rng::AbstractRNG, 
+    model::DensityModel, 
+    s::Metropolis, 
+    N::Integer, 
+    ts::Vector{T},
+    chain_type::Type{Chains}; 
+    param_names=missing,
+    kwargs...
+) where {ModelType<:AbstractModel, T<:AbstractTransition}
+    # Turn all the transitions into a vector-of-vectors.
+    vals = copy(reduce(hcat,[vcat(t.params, t.lp) for t in ts])')
+
+    # Check if we received any parameter names.
+    if ismissing(param_names)
+        param_names = ["Parameter $i" for i in 1:length(s.init_params)]
+    else
+        # Deepcopy to be thread safe.
+        param_names = deepcopy(param_names)
+    end
+
+    # Add the log density field to the parameter names.
+    push!(param_names, "lp")
+
+    # Bundle everything up and return a Chains struct.
+    return Chains(vals, param_names, (internals=["lp"],))
+end