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,113 @@
+struct AMSampler{D} <: MHSampler
+    proposal::D
+end
+
+propose(spl::AMSampler, args...) = propose(Random.GLOBAL_RNG, spl, args...)
+
+function propose(rng::Random.AbstractRNG, spl::AMSampler{<:Proposal}, model::DensityModel)
+    proposal = propose(rng, spl.proposal, model)
+    return Transition(model, proposal)
+end
+
+function propose(rng::Random.AbstractRNG, spl::AMSampler{<:Proposal}, model::DensityModel,
+                 params_prev::Transition)
+    proposal = propose(rng, spl.proposal, model, params_prev.params)
+    return Transition(model, proposal)
+end
+
+# Same as for MetropolisHastings, but with a `trackstep` callback
+function AbstractMCMC.step(rng::Random.AbstractRNG, model::DensityModel, spl::AMSampler;
+                           init_params=nothing, kwargs...)
+    if init_params === nothing
+        transition = propose(rng, spl, model)
+    else
+        transition = transition(spl, model, init_params)
+    end
+
+    trackstep(spl.proposal, transition)
+    return transition, transition
+end
+
+# Same as for MetropolisHastings but with a `trackstep` callback
+function AbstractMCMC.step(rng::Random.AbstractRNG, model::DensityModel, spl::AMSampler,
+                           params_prev::AbstractTransition; kwargs...)
+    # Generate a new proposal.
+    params = propose(rng, spl, model, params_prev)
+
+    # Calculate the log acceptance probability.
+    logα = logdensity(model, params) - logdensity(model, params_prev) +
+        logratio_proposal_density(spl, params_prev, params)
+
+    # Decide whether to return the previous params or the new one.
+    if -Random.randexp(rng) < logα
+        trackstep(spl.proposal, params, Val(true))
+        return params, params
+   else
+        trackstep(spl.proposal, params_prev, Val(false))
+        return params_prev, params_prev
+    end
+end
+
+# Called when the first sample is drawn
+trackstep(proposal::Proposal, params::Transition) = nothing
+
+# Called when a new step is performed, the last argument determines if the step is an acceptance step
+trackstep(proposal::Proposal, params::Transition, 
+          ::Union{Val{true}, Val{false}}) = nothing
+
+function logratio_proposal_density(sampler::AMSampler, params_prev::Transition, params::Transition)
+    return logratio_proposal_density(sampler.proposal, params_prev.params, params.params)
+end
+
+# 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.
+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(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, trans::Transition)
+    proposal.samplemean .= trans.params
+    proposal.samplesqmean .= trans.params * trans.params'
+    proposal.proposal = MvNormal(zeros(size(proposal.samplemean)), proposal.epsilon)
+    proposal.N = 1
+end
+
+# Recompute the empirical posterior covariance matrix
+function trackstep(proposal::AMProposal, trans::Transition, ::Union{Val{true}, Val{false}})
+    proposal.samplemean .= (proposal.samplemean .* proposal.N .+ trans.params) ./ (proposal.N + 1)
+
+    proposal.samplesqmean .= (proposal.samplesqmean .* proposal.N + trans.params * trans.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
+

test/runtests.jl

@@ -25,6 +25,18 @@ include("util.jl")
     # Construct a DensityModel.
     model = DensityModel(density)
 
+    @testset "AdaptiveMH" begin
+        # Set up our sampler with initial parameters.
+        spl = AMSampler(AMProposal([0.01 0. ; 0. 0.01]))
+
+        # Sample from the posterior.
+        chain = sample(model, spl, 100000; chain_type=StructArray, param_names=["μ", "σ"])
+
+        # chn_mean ≈ dist_mean atol=atol_v
+        @test mean(chain.μ) ≈ 0.0 atol=0.1
+        @test mean(chain.σ) ≈ 1.0 atol=0.1
+    end
+
     @testset "StaticMH" begin
         # Set up our sampler with initial parameters.
         spl1 = StaticMH([Normal(0,1), Normal(0, 1)])
@@ -207,5 +219,4 @@ include("util.jl")
     end
 
     @testset "EMCEE" begin include("emcee.jl") end
-
 end