Make AdvancedMH way cooler

Author: cpfiffer

src/AdvancedMH.jl

@@ -7,15 +7,19 @@ using Requires
 using Distributions
 
 # Import specific functions and types to use or overload.
-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
 
 # 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.
 """
@@ -57,7 +61,7 @@ function bundle_samples(
     model::DensityModel, 
     s::Metropolis, 
     N::Integer, 
-    ts::Vector{T}; 
+    ts::Type{Any}; 
     param_names=missing,
     kwargs...
 ) where {ModelType<:AbstractModel, T<:AbstractTransition}
@@ -70,13 +74,13 @@ function bundle_samples(
     s::Metropolis, 
     N::Integer, 
     ts::Vector{T},
-    chain_type::Type{NamedTuple}; 
+    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 = ["param_$i" for i in 1:length(s.init_params)]
+        param_names = ["param_$i" for i in 1:length(keys(ts[1].params))]
     else
         # Deepcopy to be thread safe.
         param_names = deepcopy(param_names)
@@ -85,8 +89,8 @@ function bundle_samples(
     push!(param_names, "lp")
 
     # Turn all the transitions into a vector-of-NamedTuple.
-    keys = tuple(Symbol.(param_names)...)
-    nts = [NamedTuple{keys}(tuple(t.params..., t.lp)) for t in ts]
+    ks = tuple(Symbol.(param_names)...)
+    nts = [NamedTuple{ks}(tuple(t.params..., t.lp)) for t in ts]
 
     return nts
 end
@@ -97,8 +101,7 @@ function __init__()
 end
 
 # Include inference methods.
+include("proposal.jl")
 include("mh-core.jl")
-include("rwmh.jl")
-include("staticmh.jl")
 
 end # module AdvancedMH

src/mcmcchains-connect.jl

@@ -1,4 +1,4 @@
-import MCMCChains: Chains
+import .MCMCChains: Chains
 
 # A basic chains constructor that works with the Transition struct we defined.
 function bundle_samples(

src/mh-core.jl

@@ -1,21 +1,150 @@
 """
-    MetropolisHastings{P<:ProposalStyle, T, D}
+    MetropolisHastings{D}
 
-Fields:
+`MetropolisHastings` has one field, `proposal`. 
+`proposal` is a `Proposal`, `NamedTuple` of `Proposal`, or `Array{Proposal}` in the shape of your data.
+For example, if you wanted the sampler to return a `NamedTuple` with shape
 
-- `init_params` is the vector form of the parameters needed for the likelihood function.
-- `proposal` is a function that dynamically constructs a conditional distribution.
+```julia
+x = (a = 1.0, b=3.8)
+```
 
-Example:
+The proposal would be
 
 ```julia
-MetropolisHastings([0.0, 0.0], x -> MvNormal(x, 1.0))
+proposal = (a=Proposal(Static(), Normal(0,1)), b=Proposal(Static(), Normal(0,1)))
 ````
+
+Other allowed proposal styles are
+
+```
+p1 = Proposal(Static(), Normal(0,1))
+p2 = Proposal(Static(), [Normal(0,1), InverseGamma(2,3)])
+p3 = Proposal(Static(), (a=Normal(0,1), b=InverseGamma(2,3)))
+p4 = Proposal(Static(), (x=1.0) -> Normal(x, 1))
+```
+
+The sampler is constructed using
+
+```julia
+spl = MetropolisHastings(proposal)
+```
 """
-mutable struct MetropolisHastings{P<:ProposalStyle, D, T} <: Metropolis
-    proposal_type :: P
+mutable struct MetropolisHastings{D} <: Metropolis
     proposal :: D
-    init_params :: T
+end
+
+StaticMH(d) = MetropolisHastings(Proposal(Static(), d))
+RWMH(d) = MetropolisHastings(Proposal(RandomWalk(), d))
+
+# Propose from a vector of proposals
+function propose(
+    spl::MetropolisHastings{<:AbstractArray},
+    model::DensityModel
+)
+    proposal = map(i -> propose(spl.proposal[i], model), 1:length(spl.proposal))
+    return Transition(model, proposal)
+end
+
+function propose(
+    spl::MetropolisHastings{<:AbstractArray},
+    model::DensityModel,
+    params_prev::Transition
+)
+    proposal = map(i -> propose(spl.proposal[i], model, params_prev.params[i]), 1:length(spl.proposal))
+    return Transition(model, proposal)
+end
+
+# Make a proposal from one Proposal struct.
+function propose(
+    spl::MetropolisHastings{<:Proposal},
+    model::DensityModel
+)
+    proposal = propose(spl.proposal, model)
+    return Transition(model, proposal)
+end
+
+function propose(
+    spl::MetropolisHastings{<:Proposal},
+    model::DensityModel,
+    params_prev::Transition
+)
+    proposal = propose(spl.proposal, model, params_prev.params)
+    return Transition(model, proposal)
+end
+
+# Make a proposal from a NamedTuple of Proposal.
+function propose(
+    spl::MetropolisHastings{<:NamedTuple},
+    model::DensityModel
+)
+    proposal = _propose(spl.proposal, model)
+    return Transition(model, proposal)
+end
+
+@generated function _propose(
+    proposals::NamedTuple{names},
+    model::DensityModel
+) where {names}
+    expr = Expr(:tuple)
+    map(names) do f
+        push!(expr.args, Expr(:(=), f, :(propose(proposals.$f, model)) ))
+    end
+    return expr
+end
+
+function propose(
+    spl::MetropolisHastings{<:NamedTuple},
+    model::DensityModel,
+    params_prev::Transition
+)
+    proposal = _propose(spl.proposal, model, params_prev.params)
+    return Transition(model, proposal)
+end
+
+@generated function _propose(
+    proposals::NamedTuple{names},
+    model::DensityModel, 
+    params_prev::NamedTuple
+) where {names}
+    expr = Expr(:tuple)
+    map(names) do f
+        push!(expr.args, Expr(:(=), f, :(propose(proposals.$f, model, params_prev.$f)) ))
+    end
+    return expr
+end
+
+
+# Evaluate the likelihood of t conditional on t_cond.
+function q(
+    spl::MetropolisHastings{<:AbstractArray},
+    t::Transition,
+    t_cond::Transition
+)
+    return sum(map(i -> q(spl.proposal[i], t.params[i], t_cond.params[i]), 1:length(spl.proposal)))
+end
+
+function q(
+    spl::MetropolisHastings{<:Proposal},
+    t::Transition,
+    t_cond::Transition
+)
+    return q(spl.proposal, t.params, t_cond.params)
+end
+
+function q(
+    spl::MetropolisHastings{<:NamedTuple},
+    t::Transition,
+    t_cond::Transition
+)
+    ks = keys(t.params)
+    total = 0.0
+
+    for k in ks
+        total += q(spl.proposal[k], t.params[k], t_cond.params[k])
+    end
+
+    return total
 end
 
 # Define the first step! function, which is called at the 
@@ -28,7 +157,7 @@ function step!(
     N::Integer;
     kwargs...
 )
-    return Transition(model, spl.init_params)
+    return propose(spl, model)
 end
 
 # Define the other step functions. Returns a Transition containing
@@ -50,7 +179,7 @@ function step!(
         q(spl, params_prev, params) - q(spl, params, params_prev)
 
     # Decide whether to return the previous params or the new one.
-    if log(rand(rng)) < min(α, 0.0)
+    if log(rand(rng)) < min(0.0, α)
         return params
     else
         return params_prev

src/proposal.jl

@@ -0,0 +1,131 @@
+struct Proposal{T<:ProposalStyle, P}
+    type :: T
+    proposal :: P
+end
+
+# Random draws
+Base.rand(p::Proposal{<:ProposalStyle, <:Distribution}) = rand(p.proposal)
+function Base.rand(p::Proposal{<:ProposalStyle, <:AbstractArray})
+    return map(rand, p.proposal)
+end
+
+# Densities
+function Distributions.logpdf(p::Proposal{<:ProposalStyle, <:UnivariateDistribution}, v) 
+    return sum(logpdf(p.proposal, v))
+end
+function Distributions.logpdf(p::Proposal{<:ProposalStyle, <:MultivariateDistribution}, v) 
+    return sum(logpdf(p.proposal, v))
+end
+function Distributions.logpdf(p::Proposal{<:ProposalStyle, <:MatrixDistribution}, v) 
+    return sum(logpdf(p.proposal, v))
+end
+function Distributions.logpdf(p::Proposal{<:ProposalStyle, <:AbstractArray}, v)
+    return sum(map(x -> logpdf(x[1], x[2]), zip(p.proposal, v)))
+end
+function Distributions.logpdf(p::Proposal{<:ProposalStyle, <:Function}, v)
+    return logpdf(p.proposal(v), v)
+end
+
+###############
+# Random Walk #
+###############
+
+function propose(
+    proposal::Proposal{RandomWalk, <:Distribution}, 
+    model::DensityModel, 
+    t
+)
+    return t + rand(proposal)
+end
+
+function propose(
+    proposal::Proposal{RandomWalk, <:AbstractArray}, 
+    model::DensityModel, 
+    t
+)
+    return t + rand(proposal)
+end
+
+function q(
+    proposal::Proposal{RandomWalk, <:Distribution}, 
+    t,
+    t_cond
+)
+    return sum(logpdf(proposal, t - t_cond))
+end
+
+function q(
+    proposal::Proposal{RandomWalk, <:AbstractArray}, 
+    t,
+    t_cond
+)
+    return sum(logpdf(proposal, t - t_cond))
+end
+
+##########
+# Static #
+##########
+
+propose(p::Proposal{RandomWalk}, m::DensityModel) = propose(Proposal(Static(), p.proposal), m)
+function propose(
+    proposal::Proposal{Static, <:Distribution}, 
+    model::DensityModel, 
+    t=nothing
+)
+    return rand(proposal)
+end
+
+function propose(
+    p::Proposal{Static, <:AbstractArray}, 
+    model::DensityModel, 
+    t=nothing
+)
+    props = map(x -> rand(x), p.proposal)
+    return props
+end
+
+function q(
+    proposal::Proposal{Static, <:Distribution}, 
+    t,
+    t_cond
+)
+    return sum(logpdf(proposal, t))
+end
+
+function q(
+    proposal::Proposal{Static, <:AbstractArray}, 
+    t,
+    t_cond
+)
+    return sum(logpdf(proposal, t))
+end
+
+############
+# Function #
+############
+
+function propose(
+    proposal::Proposal{<:ProposalStyle, <:Function}, 
+    model::DensityModel
+)
+    p = proposal.proposal 
+    return rand(proposal.proposal())
+end
+
+function propose(
+    proposal::Proposal{<:ProposalStyle, <:Function}, 
+    model::DensityModel,
+    t
+)
+    p = proposal.proposal 
+    return rand(proposal.proposal(t))
+end
+
+function q(
+    proposal::Proposal{<:ProposalStyle, <:Function}, 
+    t,
+    t_cond
+)
+    p = proposal.proposal
+    return sum(logpdf.(p(t_cond), t))
+end