Adaptive Metropolis algorithm

Project.toml

@@ -5,6 +5,8 @@ version = "0.6.8"
 [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,8 @@ export
     SymmetricRandomWalkProposal,
     Ensemble,
     StretchProposal,
-    MALA
+    MALA,
+    AMProposal
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed, MCMCSerial
@@ -120,5 +123,6 @@ end
 include("proposal.jl")
 include("mh-core.jl")
 include("emcee.jl")
+include("adaptivemetropolis.jl")
 
 end # module AdvancedMH

src/adaptivemetropolis.jl

@@ -0,0 +1,57 @@
+# 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
+    μ::Vector{FT}
+    M::Matrix{FT}
+    δ::Vector{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, zeros(size(sym,1)), zeros(size(sym)...), 
+               zeros(size(sym,1)), 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.μ .= params
+    proposal.M .= 0
+    proposal.proposal = MvNormal(zeros(size(proposal.μ)), proposal.epsilon)
+    proposal.N = 1
+end
+
+# Recompute the empirical posterior covariance matrix
+function trackstep!(proposal::AMProposal, params, ::Union{Val{true}, Val{false}})
+    proposal.N += 1
+    proposal.δ .= params .- proposal.μ
+    proposal.μ .+= proposal.δ ./ proposal.N
+
+    mul!(proposal.M, params .- proposal.μ, proposal.δ', true, true)
+
+    prop_cov = proposal.M .* proposal.scalefactor ./ proposal.N .+ proposal.epsilon
+    proposal.proposal = MvNormal(mean(proposal.proposal), Symmetric(prop_cov))
+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

test/runtests.jl

@@ -58,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 => [0.01, 0.01, 0.01, 0.01])))
+        spl2 = MetropolisHastings([AMProposal(diagm(0 => [0.01, 0.01])), 
+                                   AMProposal(diagm(0 => [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)])
 
@@ -159,7 +185,6 @@ 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)])