From 71f10723e33a8942e46c1da0f0b5fb76a1682b41 Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Wed, 26 May 2021 11:30:46 +0100
Subject: [PATCH 1/9] Added Adaptive Metropolis algorithm

---
 Project.toml              |   1 +
 src/AdvancedMH.jl         |   6 ++-
 src/adaptivemetropolis.jl | 107 ++++++++++++++++++++++++++++++++++++++
 test/runtests.jl          |  13 ++++-
 4 files changed, 125 insertions(+), 2 deletions(-)
 create mode 100644 src/adaptivemetropolis.jl

diff --git a/Project.toml b/Project.toml
index 78f8cae..8085a08 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,6 +5,7 @@ version = "0.6.0"
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index d41c3e2..4aaa1f7 100644
--- a/src/AdvancedMH.jl
+++ b/src/AdvancedMH.jl
@@ -4,6 +4,7 @@ module AdvancedMH
 using AbstractMCMC
 using Distributions
 using Requires
+using LinearAlgebra
 
 import Random
 
@@ -19,7 +20,9 @@ export
     SymmetricRandomWalkProposal,
     Ensemble,
     StretchProposal,
-    MALA
+    MALA,
+    AMProposal,
+    AMSampler
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed
@@ -120,5 +123,6 @@ end
 include("proposal.jl")
 include("mh-core.jl")
 include("emcee.jl")
+include("adaptivemetropolis.jl")
 
 end # module AdvancedMH
diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
new file mode 100644
index 0000000..3e0f0fa
--- /dev/null
+++ b/src/adaptivemetropolis.jl
@@ -0,0 +1,107 @@
+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
+
+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
+
+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
+
+# Simple Adaptive Metropolis Proposal
+# The proposal at each step is equal to a scalar multiple
+# of the posterior covariance plus a fixed, small covariance
+# matrix epsilon which is also used for initial exploration.
+mutable struct AMProposal <: Proposal{MvNormal}
+    epsilon::Symmetric
+    scalefactor::Float64
+    proposal::MvNormal
+    samplemean::AbstractArray
+    samplesqmean::AbstractMatrix
+    N::Int
+end
+
+function AMProposal(epsilon::AbstractMatrix, scalefactor=2.38^2 / size(epsilon, 1))
+    sym = Symmetric(epsilon)
+    proposal = MvNormal(zeros(size(sym, 1)), sym)
+    AMProposal(sym, scalefactor, proposal, mean(proposal), zeros(size(sym)...), 0)
+end
+
+function logratio_proposal_density(sampler::AMSampler, params_prev::Transition, params::Transition)
+    return logratio_proposal_density(sampler.proposal, params_prev.params, params.params)
+end
+
+logratio_proposal_density(::AMProposal, params_prev, params) = 0
+
+function trackstep(proposal::AMProposal, trans::Transition)
+    proposal.samplemean = copy(trans.params)
+    proposal.samplesqmean = trans.params * trans.params'
+    proposal.proposal = MvNormal(zeros(size(proposal.samplemean)), proposal.epsilon)
+    proposal.N = 1
+end
+
+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
+
diff --git a/test/runtests.jl b/test/runtests.jl
index d437bdd..879af8d 100644
--- a/test/runtests.jl
+++ b/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

From 3a23129c338c1e4e6b3394d5b82865d464db5d34 Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Wed, 26 May 2021 11:47:41 +0100
Subject: [PATCH 2/9] More documentation

---
 src/adaptivemetropolis.jl | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index 3e0f0fa..1f24533 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -15,6 +15,7 @@ function propose(rng::Random.AbstractRNG, spl::AMSampler{<:Proposal}, model::Den
     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
@@ -27,6 +28,7 @@ function AbstractMCMC.step(rng::Random.AbstractRNG, model::DensityModel, spl::AM
     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.
@@ -53,9 +55,13 @@ trackstep(proposal::Proposal, params::Transition) = nothing
 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 posterior covariance plus a fixed, small covariance
+# of the empirical posterior covariance plus a fixed, small covariance
 # matrix epsilon which is also used for initial exploration.
 mutable struct AMProposal <: Proposal{MvNormal}
     epsilon::Symmetric
@@ -72,12 +78,9 @@ function AMProposal(epsilon::AbstractMatrix, scalefactor=2.38^2 / size(epsilon,
     AMProposal(sym, scalefactor, proposal, mean(proposal), zeros(size(sym)...), 0)
 end
 
-function logratio_proposal_density(sampler::AMSampler, params_prev::Transition, params::Transition)
-    return logratio_proposal_density(sampler.proposal, params_prev.params, params.params)
-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 = copy(trans.params)
     proposal.samplesqmean = trans.params * trans.params'
@@ -85,6 +88,7 @@ function trackstep(proposal::AMProposal, trans::Transition)
     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)
 

From c40d4dbf4aa0f6de31700f4b1ca03ef7730a8c4a Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Wed, 26 May 2021 13:07:06 +0100
Subject: [PATCH 3/9] Concrete types for AMProposal

---
 Project.toml              |  1 +
 src/AdvancedMH.jl         |  1 +
 src/adaptivemetropolis.jl | 26 ++++++++++++++------------
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8085a08..58c5d36 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.6.0"
 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"
 
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index 4aaa1f7..beb931c 100644
--- a/src/AdvancedMH.jl
+++ b/src/AdvancedMH.jl
@@ -5,6 +5,7 @@ using AbstractMCMC
 using Distributions
 using Requires
 using LinearAlgebra
+using PDMats
 
 import Random
 
diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index 1f24533..0abdb58 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -63,17 +63,19 @@ end
 # 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 <: Proposal{MvNormal}
-    epsilon::Symmetric
-    scalefactor::Float64
-    proposal::MvNormal
-    samplemean::AbstractArray
-    samplesqmean::AbstractMatrix
+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, scalefactor=2.38^2 / size(epsilon, 1))
-    sym = Symmetric(epsilon)
+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
@@ -82,17 +84,17 @@ 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 = copy(trans.params)
-    proposal.samplesqmean = trans.params * trans.params'
+    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.samplemean .= (proposal.samplemean .* proposal.N .+ trans.params) ./ (proposal.N + 1)
 
-    proposal.samplesqmean = (proposal.samplesqmean .* proposal.N + trans.params * trans.params') ./ 
+    proposal.samplesqmean .= (proposal.samplesqmean .* proposal.N + trans.params * trans.params') ./ 
                              (proposal.N + 1)
 
     samplecov = proposal.samplesqmean .- proposal.samplemean * proposal.samplemean'

From 5afb87e5bf02abd6536f2ee4dc4872ea4fec4cc8 Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Wed, 26 May 2021 13:10:37 +0100
Subject: [PATCH 4/9] Added reference for the AM algorithm

---
 src/adaptivemetropolis.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index 0abdb58..b61317d 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -63,6 +63,10 @@ end
 # 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

From 8cb275de5cf3f8afac03db639b75e3339e86b36d Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Wed, 26 May 2021 14:12:14 +0100
Subject: [PATCH 5/9] Merged AMSampler into MetropolisHastings

---
 src/adaptivemetropolis.jl | 75 ++++-----------------------------------
 src/mh-core.jl            | 19 ++++++++++
 test/runtests.jl          | 41 ++++++++++++++-------
 3 files changed, 54 insertions(+), 81 deletions(-)

diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index b61317d..b6f3fcf 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -1,64 +1,3 @@
-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
@@ -79,7 +18,7 @@ end
 
 function AMProposal(epsilon::AbstractMatrix{FT}, 
                     scalefactor=FT(2.38^2 / size(epsilon, 1))) where {FT <: Real}
-    sym = PDMat(epsilon)
+    sym = PDMat(Symmetric(epsilon))
     proposal = MvNormal(zeros(size(sym, 1)), sym)
     AMProposal(sym, scalefactor, proposal, mean(proposal), zeros(size(sym)...), 0)
 end
@@ -87,18 +26,18 @@ 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'
+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, trans::Transition, ::Union{Val{true}, Val{false}})
-    proposal.samplemean .= (proposal.samplemean .* proposal.N .+ trans.params) ./ (proposal.N + 1)
+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 + trans.params * trans.params') ./ 
+    proposal.samplesqmean .= (proposal.samplesqmean .* proposal.N + params * params') ./ 
                              (proposal.N + 1)
 
     samplecov = proposal.samplesqmean .- proposal.samplemean * proposal.samplemean'
diff --git a/src/mh-core.jl b/src/mh-core.jl
index cd195ce..05affe6 100644
--- a/src/mh-core.jl
+++ b/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
diff --git a/test/runtests.jl b/test/runtests.jl
index 879af8d..8a93fde 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,6 +2,7 @@ using AdvancedMH
 using Distributions
 using StructArrays
 using MCMCChains
+using LinearAlgebra
 
 using Random
 using Test
@@ -25,18 +26,6 @@ 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)])
@@ -69,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)])
 
@@ -120,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)])

From 46f77ca789c81d63b93b924a04d34866eedce55f Mon Sep 17 00:00:00 2001
From: kaandocal <>
Date: Wed, 26 May 2021 23:14:00 +0100
Subject: [PATCH 6/9] Added Welford's algorithm for adaptive covariance
 calculations

---
 src/adaptivemetropolis.jl | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index b6f3fcf..e76f2bc 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -11,8 +11,9 @@ mutable struct AMProposal{FT <: Real, CT <: AbstractMatrix{FT},
     epsilon::CT
     scalefactor::FT
     proposal::MNT
-    samplemean::Vector{FT}
-    samplesqmean::Matrix{FT}
+    μ::Vector{FT}
+    M::Matrix{FT}
+    δ::Vector{FT}
     N::Int
 end
 
@@ -20,30 +21,30 @@ 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)
+    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.samplemean .= params
-    proposal.samplesqmean .= params * params'
-    proposal.proposal = MvNormal(zeros(size(proposal.samplemean)), proposal.epsilon)
+    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.samplemean .= (proposal.samplemean .* proposal.N .+ params) ./ (proposal.N + 1)
+    proposal.N += 1
+    proposal.δ .= params .- proposal.μ
+    proposal.μ .+= proposal.δ ./ proposal.N
 
-    proposal.samplesqmean .= (proposal.samplesqmean .* proposal.N + params * params') ./ 
-                             (proposal.N + 1)
+    mul!(proposal.M, params .- proposal.μ, proposal.δ', 1.0, 1.0)
 
-    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
+    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)
@@ -53,4 +54,3 @@ end
 function propose(rng::Random.AbstractRNG, p::AMProposal, ::DensityModel, t)
     return t + rand(rng, p.proposal)
 end
-

From 970b0f5621a55aeda1580353b761eee7e3d65b1b Mon Sep 17 00:00:00 2001
From: kaandocal <>
Date: Wed, 26 May 2021 23:26:40 +0100
Subject: [PATCH 7/9] trackstep -> trackstep!

---
 src/adaptivemetropolis.jl |  4 ++--
 src/mh-core.jl            | 24 ++++++++++++------------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index e76f2bc..670f47d 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -28,7 +28,7 @@ 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)
+function trackstep!(proposal::AMProposal, params)
     proposal.μ .= params
     proposal.M .= 0
     proposal.proposal = MvNormal(zeros(size(proposal.μ)), proposal.epsilon)
@@ -36,7 +36,7 @@ function trackstep(proposal::AMProposal, params)
 end
 
 # Recompute the empirical posterior covariance matrix
-function trackstep(proposal::AMProposal, params, ::Union{Val{true}, Val{false}})
+function trackstep!(proposal::AMProposal, params, ::Union{Val{true}, Val{false}})
     proposal.N += 1
     proposal.δ .= params .- proposal.μ
     proposal.μ .+= proposal.δ ./ proposal.N
diff --git a/src/mh-core.jl b/src/mh-core.jl
index 05affe6..59cd05d 100644
--- a/src/mh-core.jl
+++ b/src/mh-core.jl
@@ -142,20 +142,20 @@ 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)
+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::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::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)
+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.
@@ -173,7 +173,7 @@ function AbstractMCMC.step(
         transition = AdvancedMH.transition(spl, model, init_params)
     end
 
-    trackstep(spl.proposal, transition.params)
+    trackstep!(spl.proposal, transition.params)
     return transition, transition
 end
 
@@ -197,10 +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))
+        trackstep!(spl.proposal, params.params, Val(true))
         return params, params
     else
-        trackstep(spl.proposal, params_prev.params, Val(false))
+        trackstep!(spl.proposal, params_prev.params, Val(false))
         return params_prev, params_prev
     end
 end

From c800aafe4975cd699a7b1434bf6339cf698bdd0d Mon Sep 17 00:00:00 2001
From: kaandocal <>
Date: Thu, 27 May 2021 09:36:14 +0100
Subject: [PATCH 8/9] Fixed diagm issue in AdaptiveMH test

---
 test/runtests.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 8a93fde..1d64731 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -64,9 +64,9 @@ include("util.jl")
         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]))])
+        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, 

From eb91b36cb89064a8f991e5b658994557b5694cc8 Mon Sep 17 00:00:00 2001
From: kaandocal <26488673+kaandocal@users.noreply.github.com>
Date: Thu, 27 May 2021 18:11:22 +0100
Subject: [PATCH 9/9] Apply suggestions from code review

Co-authored-by: Cameron Pfiffer <cpfiffer@gmail.com>
Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/AdvancedMH.jl         | 3 +--
 src/adaptivemetropolis.jl | 7 ++++---
 test/runtests.jl          | 1 -
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index beb931c..8e63105 100644
--- a/src/AdvancedMH.jl
+++ b/src/AdvancedMH.jl
@@ -22,8 +22,7 @@ export
     Ensemble,
     StretchProposal,
     MALA,
-    AMProposal,
-    AMSampler
+    AMProposal
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed
diff --git a/src/adaptivemetropolis.jl b/src/adaptivemetropolis.jl
index 670f47d..eaade87 100644
--- a/src/adaptivemetropolis.jl
+++ b/src/adaptivemetropolis.jl
@@ -17,8 +17,9 @@ mutable struct AMProposal{FT <: Real, CT <: AbstractMatrix{FT},
     N::Int
 end
 
-function AMProposal(epsilon::AbstractMatrix{FT}, 
-                    scalefactor=FT(2.38^2 / size(epsilon, 1))) where {FT <: Real}
+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)...), 
@@ -41,7 +42,7 @@ function trackstep!(proposal::AMProposal, params, ::Union{Val{true}, Val{false}}
     proposal.δ .= params .- proposal.μ
     proposal.μ .+= proposal.δ ./ proposal.N
 
-    mul!(proposal.M, params .- proposal.μ, proposal.δ', 1.0, 1.0)
+    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))
diff --git a/test/runtests.jl b/test/runtests.jl
index 1d64731..02cce67 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -135,7 +135,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)])