Adaptive Metropolis algorithm

Project.toml

@@ -5,6 +5,8 @@ version = "0.6.0"
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 

src/AdvancedMH.jl

@@ -4,6 +4,8 @@ module AdvancedMH
 using AbstractMCMC
 using Distributions
 using Requires
+using LinearAlgebra
+using PDMats
 
 import Random
 
@@ -19,7 +21,9 @@ export
     SymmetricRandomWalkProposal,
     Ensemble,
     StretchProposal,
-    MALA
+    MALA,
+    AMProposal,
+    AMSampler
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed
@@ -120,5 +124,6 @@ end
 include("proposal.jl")
 include("mh-core.jl")
 include("emcee.jl")
+include("adaptivemetropolis.jl")
 
 end # module AdvancedMH

src/adaptivemetropolis.jl

@@ -0,0 +1,56 @@
+# Simple Adaptive Metropolis Proposal
+# The proposal at each step is equal to a scalar multiple
+# of the empirical posterior covariance plus a fixed, small covariance
+# matrix epsilon which is also used for initial exploration.
+# 
+# Reference:
+#   H. Haario, E. Saksman, and J. Tamminen, "An adaptive Metropolis algorithm",
+#   Bernoulli 7(2): 223-242 (2001)
+mutable struct AMProposal{FT <: Real, CT <: AbstractMatrix{FT}, 
+                          MNT <: AbstractMvNormal} <: Proposal{MNT}
+    epsilon::CT
+    scalefactor::FT
+    proposal::MNT
+    samplemean::Vector{FT}
+    samplesqmean::Matrix{FT}
+    N::Int
+end
+
+function AMProposal(epsilon::AbstractMatrix{FT}, 
+                    scalefactor=FT(2.38^2 / size(epsilon, 1))) where {FT <: Real}
+    sym = PDMat(Symmetric(epsilon))
+    proposal = MvNormal(zeros(size(sym, 1)), sym)
+    AMProposal(sym, scalefactor, proposal, mean(proposal), zeros(size(sym)...), 0)
+end
+
+logratio_proposal_density(::AMProposal, params_prev, params) = 0
+
+# When the proposal is initialised the empirical posterior covariance is zero
+function trackstep(proposal::AMProposal, params)
+    proposal.samplemean .= params
+    proposal.samplesqmean .= params * params'
+    proposal.proposal = MvNormal(zeros(size(proposal.samplemean)), proposal.epsilon)
+    proposal.N = 1
+end
+
+# Recompute the empirical posterior covariance matrix
+function trackstep(proposal::AMProposal, params, ::Union{Val{true}, Val{false}})
+    proposal.samplemean .= (proposal.samplemean .* proposal.N .+ params) ./ (proposal.N + 1)
+
+    proposal.samplesqmean .= (proposal.samplesqmean .* proposal.N + params * params') ./ 
+                             (proposal.N + 1)
+
+    samplecov = proposal.samplesqmean .- proposal.samplemean * proposal.samplemean'
+    prop_cov = proposal.scalefactor .* samplecov
+    proposal.proposal = MvNormal(mean(proposal.proposal), prop_cov .+ proposal.epsilon)
+    proposal.N += 1
+end
+
+function propose(rng::Random.AbstractRNG, p::AMProposal, ::DensityModel)
+    return rand(rng, p.proposal)
+end
+
+function propose(rng::Random.AbstractRNG, p::AMProposal, ::DensityModel, t)
+    return t + rand(rng, p.proposal)
+end
+

src/mh-core.jl

@@ -141,6 +141,22 @@ end
 transition(sampler, model, params) = transition(model, params)
 transition(model, params) = Transition(model, params)
 
+# Called to update proposal when the first sample is drawn
+trackstep(proposal::Proposal, params) = nothing
+trackstep(proposal::AbstractArray, params) = foreach(trackstep, proposal, params)
+trackstep(proposal::NamedTuple, params) = foreach(trackstep, proposal, params)
+
+# Called to update proposal when a new step is performed
+# The last argument determines if the step is an acceptance step
+trackstep(proposal::Proposal, params, 
+          ::Union{Val{true}, Val{false}}) = nothing
+
+trackstep(proposal::AbstractArray, params, accept::Union{Val{true},Val{false}}) = 
+          foreach((prop, par) -> trackstep(prop, par, accept), proposal, params)
+
+trackstep(proposal::NamedTuple, params, accept::Union{Val{true},Val{false}}) = 
+          foreach((prop, par) -> trackstep(prop, par, accept), proposal, params)
+
 # Define the first sampling step.
 # Return a 2-tuple consisting of the initial sample and the initial state.
 # In this case they are identical.
@@ -157,6 +173,7 @@ function AbstractMCMC.step(
         transition = AdvancedMH.transition(spl, model, init_params)
     end
 
+    trackstep(spl.proposal, transition.params)
     return transition, transition
 end
 
@@ -180,8 +197,10 @@ function AbstractMCMC.step(
 
     # Decide whether to return the previous params or the new one.
     if -Random.randexp(rng) < logα
+        trackstep(spl.proposal, params.params, Val(true))
         return params, params
     else
+        trackstep(spl.proposal, params_prev.params, Val(false))
         return params_prev, params_prev
     end
 end

test/runtests.jl

@@ -2,6 +2,7 @@ using AdvancedMH
 using Distributions
 using StructArrays
 using MCMCChains
+using LinearAlgebra
 
 using Random
 using Test
@@ -57,6 +58,32 @@ include("util.jl")
         @test mean(chain2.σ) ≈ 1.0 atol=0.1
     end
 
+    @testset "AdaptiveMH" begin
+        # Set up our sampler with initial parameters.
+        doubledensity(θ::Vector{Float64}) = density(θ[1:2]) + density(θ[3:4])
+        doubledensity(θ::Vector{Vector{Float64}}) = density(θ[1]) + density(θ[2])
+        doublemodel = DensityModel(doubledensity)
+
+        spl1 = MetropolisHastings(AMProposal(diagm([0.01, 0.01, 0.01, 0.01])))
+        spl2 = MetropolisHastings([AMProposal(diagm([0.01, 0.01])), 
+                                   AMProposal(diagm([0.01, 0.01]))])
+
+        # Sample from the posterior.
+        chain1 = sample(doublemodel, spl1, 100000; chain_type=StructArray, 
+                        param_names=["μ1", "σ1", "μ2", "σ2"])
+        chain2 = sample(doublemodel, spl2, 100000; chain_type=StructArray, 
+                        param_names=["p1", "p2"])
+
+        # chn_mean ≈ dist_mean atol=atol_v
+        @test mean(chain1.μ1) ≈ 0.0 atol=0.1
+        @test mean(chain1.σ1) ≈ 1.0 atol=0.1
+        @test mean(chain1.μ2) ≈ 0.0 atol=0.1
+        @test mean(chain1.σ2) ≈ 1.0 atol=0.1
+
+        @test mean(chain2.p1) ≈ [ 0.0 1.0 ]' atol=0.1
+        @test mean(chain2.p2) ≈ [ 0.0 1.0 ]' atol=0.1
+    end
+
     @testset "parallel sampling" begin
         spl1 = StaticMH([Normal(0,1), Normal(0, 1)])
 
@@ -108,7 +135,7 @@ include("util.jl")
         @test keys(c3[1]) == (:a, :b, :lp)
         @test keys(c4[1]) == (:param_1, :lp)
     end
-
+    
     @testset "Initial parameters" begin
         # Set up our sampler with initial parameters.
         spl1 = StaticMH([Normal(0,1), Normal(0, 1)])
@@ -207,5 +234,4 @@ include("util.jl")
     end
 
     @testset "EMCEE" begin include("emcee.jl") end
-
 end