Merge pull request #10 from TuringLang/csp/ahm2

Add static MH and general code improvements

Author: cpfiffer

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

README.md

@@ -4,6 +4,9 @@ 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.
+
 ## 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.
@@ -25,7 +28,7 @@ density(θ) = insupport(θ) ? sum(logpdf.(dist(θ), data)) : -Inf
 model = DensityModel(density)
 
 # Set up our sampler with initial parameters.
-spl = MetropolisHastings([0.0, 0.0])
+spl = RWMH([0.0, 0.0])
 
 # Sample from the posterior.
 chain = sample(model, spl, 100000; param_names=["μ", "σ"])
@@ -68,5 +71,16 @@ Custom proposal distributions can be specified by passing a distribution to `Met
 
 ```julia
 # Set up our sampler with initial parameters.
-spl = MetropolisHastings([0.0, 0.0], MvNormal(2, 0.5)) 
+spl1 = RWMH([0.0, 0.0], MvNormal(2, 0.5)) 
+spl2 = StaticMH([0.0, 0.0], MvNormal(2, 0.5)) 
+```
+
+## 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=["μ","σ"])
 ```

src/AdvancedMH.jl

@@ -11,30 +11,11 @@ import MCMCChains: Chains
 import AbstractMCMC: step!, AbstractSampler, AbstractTransition, transition_type, bundle_samples
 
 # Exports
-export MetropolisHastings, DensityModel, sample
+export MetropolisHastings, DensityModel, sample, psample, RWMH, StaticMH
 
-"""
-    MetropolisHastings{T, F<:Function}
-
-Fields:
-
-- `init_θ` is the vector form of the parameters needed for the likelihood function.
-- `proposal` is a function that dynamically constructs a conditional distribution.
-
-Example:
-
-```julia
-MetropolisHastings([0.0, 0.0], x -> MvNormal(x, 1.0))
-````
-"""
-struct MetropolisHastings{T, D} <: AbstractSampler 
-    init_θ :: T
-    proposal :: D
-end
-
-# Default constructors.
-MetropolisHastings(init_θ::Real) = MetropolisHastings(init_θ, Normal(0,1))
-MetropolisHastings(init_θ::Vector{<:Real}) = MetropolisHastings(init_θ, MvNormal(length(init_θ),1))
+# Abstract type for MH-style samplers.
+abstract type Metropolis <: AbstractSampler end
+abstract type ProposalStyle end
 
 # Define a model type. Stores the log density function and the data to 
 # evaluate the log density on.
@@ -65,74 +46,17 @@ end
 Transition(model::M, θ::T) where {M<:DensityModel, T} = Transition(θ, ℓπ(model, θ))
 
 # Tell the interface what transition type we would like to use.
-transition_type(model::DensityModel, spl::MetropolisHastings) = typeof(Transition(spl.init_θ, ℓπ(model, spl.init_θ)))
-
-# Define a function that makes a basic proposal depending on a univariate
-# parameterization or a multivariate parameterization.
-propose(spl::MetropolisHastings, model::DensityModel, θ::Real) = 
-    Transition(model, θ + rand(spl.proposal))
-propose(spl::MetropolisHastings, model::DensityModel, θ::Vector{<:Real}) = 
-    Transition(model, θ + rand(spl.proposal))
-propose(spl::MetropolisHastings, model::DensityModel, t::Transition) = propose(spl, model, t.θ)
-
-"""
-    q(θ::Real, dist::Sampleable)
-    q(θ::Vector{<:Real}, dist::Sampleable)
-    q(t1::Transition, dist::Sampleable)
-
-Calculates the probability `q(θ | θcond)`, using the proposal distribution `spl.proposal`.
-"""
-q(spl::MetropolisHastings, θ::Real, θcond::Real) = logpdf(spl.proposal, θ - θcond)
-q(spl::MetropolisHastings, θ::Vector{<:Real}, θcond::Vector{<:Real}) = logpdf(spl.proposal, θ - θcond)
-q(spl::MetropolisHastings, t1::Transition, t2::Transition) = q(spl, t1.θ, t2.θ)
+transition_type(model::DensityModel, spl::Metropolis) = typeof(Transition(spl.init_θ, ℓπ(model, spl.init_θ)))
 
 # Calculate the density of the model given some parameterization.
 ℓπ(model::DensityModel, θ::T) where T = model.ℓπ(θ)
 ℓπ(model::DensityModel, t::Transition) = t.lp
 
-# Define the first step! function, which is called at the 
-# beginning of sampling. Return the initial parameter used
-# to define the sampler.
-function step!(
-    rng::AbstractRNG,
-    model::DensityModel,
-    spl::MetropolisHastings,
-    N::Integer;
-    kwargs...
-)
-    return Transition(model, spl.init_θ)
-end
-
-# Define the other step functions. Returns a Transition containing
-# either a new proposal (if accepted) or the previous proposal 
-# (if not accepted).
-function step!(
-    rng::AbstractRNG,
-    model::DensityModel,
-    spl::MetropolisHastings,
-    ::Integer,
-    θ_prev::Transition;
-    kwargs...
-)
-    # Generate a new proposal.
-    θ = propose(spl, model, θ_prev)
-
-    # Calculate the log acceptance probability.
-    α = ℓπ(model, θ) - ℓπ(model, θ_prev) + q(spl, θ_prev, θ) - q(spl, θ, θ_prev)
-
-    # Decide whether to return the previous θ or the new one.
-    if log(rand(rng)) < min(α, 0.0)
-        return θ
-    else
-        return θ_prev
-    end
-end
-
 # A basic chains constructor that works with the Transition struct we defined.
 function bundle_samples(
     rng::AbstractRNG, 
     ℓ::DensityModel, 
-    s::MetropolisHastings, 
+    s::Metropolis, 
     N::Integer, 
     ts::Vector{T}; 
     param_names=missing,
@@ -143,7 +67,10 @@ function bundle_samples(
 
     # Check if we received any parameter names.
     if ismissing(param_names)
-        param_names = ["Parameter $i" for i in 1:(length(first(vals))-1)]
+        param_names = ["Parameter $i" for i in 1:length(s.init_θ)]
+    else
+        # Deepcopy to be thread safe.
+        param_names = deepcopy(param_names)
     end
 
     # Add the log density field to the parameter names.
@@ -153,4 +80,9 @@ function bundle_samples(
     return Chains(vals, param_names, (internals=["lp"],))
 end
 
-end # module AdvancedMH
+# Include inference methods.
+include("mh-core.jl")
+include("rwmh.jl")
+include("staticmh.jl")
+
+end # module AdvancedMH

src/mh-core.jl

@@ -0,0 +1,76 @@
+"""
+    MetropolisHastings{P<:ProposalStyle, T, D}
+
+Fields:
+
+- `init_θ` is the vector form of the parameters needed for the likelihood function.
+- `proposal` is a function that dynamically constructs a conditional distribution.
+
+Example:
+
+```julia
+MetropolisHastings([0.0, 0.0], x -> MvNormal(x, 1.0))
+````
+"""
+struct MetropolisHastings{P<:ProposalStyle, T, D} <: Metropolis
+    proposal_type :: P
+    init_θ :: T
+    proposal :: D
+end
+
+# Default constructors.
+MetropolisHastings(init_θ::Real) = MetropolisHastings(init_θ, Normal(0,1))
+MetropolisHastings(init_θ::Vector{<:Real}) = MetropolisHastings(init_θ, MvNormal(length(init_θ),1))
+
+"""
+    propose(spl::MetropolisHastings, model::DensityModel, t::Transition)
+
+Generates a new parameter proposal conditional on the model, the sampler, and the previous
+sample.
+"""
+@inline propose(spl::MetropolisHastings, model::DensityModel, t::Transition) = propose(spl, model, t.θ)
+
+"""
+    q(spl::MetropolisHastings, t1::Transition, t2::Transition)
+
+Calculates the probability `q(θ | θcond)`, using the proposal distribution `spl.proposal`.
+"""
+@inline q(spl::MetropolisHastings, t1::Transition, t2::Transition) = q(spl, t1.θ, t2.θ)
+
+# Define the first step! function, which is called at the 
+# beginning of sampling. Return the initial parameter used
+# to define the sampler.
+function step!(
+    rng::AbstractRNG,
+    model::DensityModel,
+    spl::MetropolisHastings,
+    N::Integer;
+    kwargs...
+)
+    return Transition(model, spl.init_θ)
+end
+
+# Define the other step functions. Returns a Transition containing
+# either a new proposal (if accepted) or the previous proposal 
+# (if not accepted).
+function step!(
+    rng::AbstractRNG,
+    model::DensityModel,
+    spl::MetropolisHastings,
+    ::Integer,
+    θ_prev::Transition;
+    kwargs...
+)
+    # Generate a new proposal.
+    θ = propose(spl, model, θ_prev)
+
+    # Calculate the log acceptance probability.
+    α = ℓπ(model, θ) - ℓπ(model, θ_prev) + q(spl, θ_prev, θ) - q(spl, θ, θ_prev)
+
+    # Decide whether to return the previous θ or the new one.
+    if log(rand(rng)) < min(α, 0.0)
+        return θ
+    else
+        return θ_prev
+    end
+end

src/rwmh.jl

@@ -0,0 +1,36 @@
+struct RandomWalk <: ProposalStyle end
+
+"""
+    RWMH(init_theta::Real, proposal = Normal(init_theta, 1))
+    RWMH(init_theta::Vector{Real}, proposal = MvNormal(init_theta, 1))
+
+Random walk Metropolis-Hastings.
+
+Fields:
+
+- `init_θ` is the vector form of the parameters needed for the likelihood function.
+- `proposal` is a function that dynamically constructs a conditional distribution.
+
+Example:
+
+```julia
+RWMH([0.0, 0.0], x -> MvNormal(x, 1.0))
+````
+"""
+RWMH(init_theta::Real, proposal = Normal(init_theta, 1)) = MetropolisHastings(RandomWalk(), init_theta, proposal)
+RWMH(init_theta::Vector{<:Real}, proposal = MvNormal(init_theta, 1)) = MetropolisHastings(RandomWalk(), init_theta, proposal)
+
+# Define a function that makes a basic proposal depending on a univariate
+# parameterization or a multivariate parameterization.
+propose(spl::MetropolisHastings{RandomWalk}, model::DensityModel, θ::Real) = Transition(model, θ + rand(spl.proposal))
+propose(spl::MetropolisHastings{RandomWalk}, model::DensityModel, θ::Vector{<:Real}) = Transition(model, θ + rand(spl.proposal))
+
+"""
+    q(θ::Real, dist::Sampleable)
+    q(θ::Vector{<:Real}, dist::Sampleable)
+    q(t1::Transition, dist::Sampleable)
+
+Calculates the probability `q(θ | θcond)`, using the proposal distribution `spl.proposal`.
+"""
+@inline q(spl::MetropolisHastings{RandomWalk}, θ::Real, θcond::Real) = logpdf(spl.proposal, θ - θcond)
+@inline q(spl::MetropolisHastings{RandomWalk}, θ::Vector{<:Real}, θcond::Vector{<:Real}) = logpdf(spl.proposal, θ - θcond)

src/staticmh.jl

@@ -0,0 +1,37 @@
+struct Static <: ProposalStyle end
+
+"""
+    StaticMH(init_theta::Real, proposal = Normal(init_theta, 1))
+    StaticMH(init_theta::Vector{Real}, proposal = MvNormal(init_theta, 1))
+
+Static Metropolis-Hastings. Proposes only from the prior distribution.
+
+Fields:
+
+- `init_θ` is the vector form of the parameters needed for the likelihood function.
+- `proposal` is a distribution.
+
+Example:
+
+```julia
+RWMH([0.0, 0.0], MvNormal(x, 1.0))
+````
+"""
+StaticMH(init_theta::Real, proposal = Normal(init_theta, 1)) = MetropolisHastings(Static(), init_theta, proposal)
+StaticMH(init_theta::Vector{<:Real}, proposal = MvNormal(init_theta, 1)) = MetropolisHastings(Static(), init_theta, proposal)
+
+# Define a function that makes a basic proposal depending on a univariate
+# parameterization or a multivariate parameterization.
+propose(spl::MetropolisHastings{Static}, model::DensityModel, θ::Real) = Transition(model, rand(spl.proposal))
+propose(spl::MetropolisHastings{Static}, model::DensityModel, θ::Vector{<:Real}) = Transition(model, rand(spl.proposal))
+
+"""
+    q(θ::Real, dist::Sampleable)
+    q(θ::Vector{<:Real}, dist::Sampleable)
+    q(t1::Transition, dist::Sampleable)
+
+Calculates the probability `q(θ | θcond)`, using the proposal distribution `spl.proposal`.
+"""
+q(spl::MetropolisHastings{Static}, θ::Real, θcond::Real) = logpdf(spl.proposal, θ)
+q(spl::MetropolisHastings{Static}, θ::Vector{<:Real}, θcond::Vector{<:Real}) = logpdf(spl.proposal, θ)
+

test/runtests.jl

@@ -20,12 +20,26 @@ using Random
     model = DensityModel(density)
 
     # Set up our sampler with initial parameters.
-    spl = MetropolisHastings([0.0, 0.0])
+    spl1 = RWMH([0.0, 0.0])
+    spl2 = StaticMH([0.0, 0.0], MvNormal([0.0, 0.0], 1))
 
-    # Sample from the posterior.
-    chain = sample(model, spl, 100000; param_names=["μ", "σ"])
+    @testset "Inference" begin
 
-    # chn_mean ≈ dist_mean atol=atol_v
-    @test mean(chain["μ"].value) ≈ 0.0 atol=0.1
-    @test mean(chain["σ"].value) ≈ 1.0 atol=0.1
+        # Sample from the posterior.
+        chain1 = sample(model, spl1, 100000; param_names=["μ", "σ"])
+        chain2 = sample(model, spl2, 100000; param_names=["μ", "σ"])
+
+        # chn_mean ≈ dist_mean atol=atol_v
+        @test mean(chain1["μ"].value) ≈ 0.0 atol=0.1
+        @test mean(chain1["σ"].value) ≈ 1.0 atol=0.1
+        @test mean(chain2["μ"].value) ≈ 0.0 atol=0.1
+        @test mean(chain2["σ"].value) ≈ 1.0 atol=0.1
+    end
+
+    @testset "psample" begin
+        chain1 = psample(model, spl1, 10000, 4; param_names=["μ", "σ"])
+        @test mean(chain1["μ"].value) ≈ 0.0 atol=0.1
+        @test mean(chain1["σ"].value) ≈ 1.0 atol=0.1
+    end
 end
+