add acceptance rate info to Transition struct

src/AdvancedMH.jl

@@ -27,13 +27,13 @@ export
 # Reexports
 export sample, MCMCThreads, MCMCDistributed, MCMCSerial
 
-# Abstract type for MH-style samplers. Needs better name? 
+# Abstract type for MH-style samplers. Needs better name?
 abstract type MHSampler <: AbstractMCMC.AbstractSampler end
 
 # Abstract type for MH-style transitions.
 abstract type AbstractTransition end
 
-# Define a model type. Stores the log density function and the data to 
+# Define a model type. Stores the log density function and the data to
 # evaluate the log density on.
 """
     DensityModel{F} <: AbstractModel
@@ -53,17 +53,19 @@ end
 
 const DensityModelOrLogDensityModel = Union{<:DensityModel,<:AbstractMCMC.LogDensityModel}
 
-# Create a very basic Transition type, only stores the
-# parameter draws and the log probability of the draw.
+# Create a very basic Transition type, stores the
+# parameter draws, the log probability of the draw,
+# and the draw information until this point
 struct Transition{T,L<:Real} <: AbstractTransition
     params :: T
     lp :: L
+    accepted :: Bool
 end
 
-# Store the new draw and its log density.
-Transition(model::DensityModelOrLogDensityModel, params) = Transition(params, logdensity(model, params))
-function Transition(model::AbstractMCMC.LogDensityModel, params)
-    return Transition(params, LogDensityProblems.logdensity(model.logdensity, params))
+# Store the new draw, its log density, and draw information
+Transition(model::DensityModelOrLogDensityModel, params, accepted) = Transition(params, logdensity(model, params), accepted)
+function Transition(model::AbstractMCMC.LogDensityModel, params, accepted)
+    return Transition(params, LogDensityProblems.logdensity(model.logdensity, params), accepted)
 end
 
 # Calculate the density of the model given some parameterization.
@@ -128,7 +130,7 @@ function __init__()
         if exc.f === logdensity_and_gradient && length(arg_types) == 2 && first(arg_types) <: DensityModel && isempty(kwargs)
             print(io, "\\nDid you forget to load ForwardDiff?")
         end
-    end    
+    end
     @static if !isdefined(Base, :get_extension)
         @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("../ext/AdvancedMHMCMCChainsExt.jl")
         @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("../ext/AdvancedMHStructArraysExt.jl")

src/MALA.jl

@@ -1,11 +1,13 @@
 struct MALA{D} <: MHSampler
     proposal::D
-
-    MALA{D}(proposal::D) where {D} = new{D}(proposal)
+    accepted_draws :: Int
+    total_draws :: Int
 end
 
+MALA(p,ad,td) = MALA(p,ad,td)
+
 # If we were given a RandomWalkProposal, just use that instead.
-MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)
+MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d,0,0)
 
 # Create a RandomWalkProposal if we weren't given one already.
 MALA(d) = MALA(RandomWalkProposal(d))
@@ -15,13 +17,14 @@ struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{
     params::T
     lp::L
     gradient::G
+    accepted :: Bool
 end
 
 logdensity(model::DensityModelOrLogDensityModel, t::GradientTransition) = t.lp
 
 propose(rng::Random.AbstractRNG, ::MALA, model) = error("please specify initial parameters")
-function transition(sampler::MALA, model::DensityModelOrLogDensityModel, params)
-    return GradientTransition(params, logdensity_and_gradient(model, params)...)
+function transition(sampler::MALA, model::DensityModelOrLogDensityModel, params, accepted)
+    return GradientTransition(params, logdensity_and_gradient(model, params)..., accepted)
 end
 
 check_capabilities(model::DensityModelOrLogDensityModel) = nothing
@@ -44,7 +47,7 @@ function AbstractMCMC.step(
     kwargs...
 )
     check_capabilities(model)
-    
+
     # Extract value and gradient of the log density of the current state.
     state = transition_prev.params
     logdensity_state = transition_prev.lp
@@ -69,10 +72,16 @@ function AbstractMCMC.step(
 
     # Decide whether to return the previous params or the new one.
     transition = if -Random.randexp(rng) < logα
-        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate)
+        accepted_draws = sampler.accepted_draws + 1
+        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true)
     else
-        transition_prev
+        accepted_draws = sampler.accepted_draws
+        candidate = transition_prev.params
+        lp = transition_prev.lp
+        gradient = transition_prev.gradient
+        GradientTransition(candidate, lp, gradient, false)
     end
+    sampler = MALA(sampler.proposal, accepted_draws, sampler.total_draws+1)
 
     return transition, transition
 end

src/emcee.jl

@@ -1,10 +1,14 @@
 struct Ensemble{D} <: MHSampler
     n_walkers::Int
     proposal::D
+    accepted_draws :: Int
+    total_draws :: Int
 end
 
-function transition(sampler::Ensemble, model::DensityModelOrLogDensityModel, params)
-    return [Transition(model, x) for x in params]
+Ensemble(n,p) = Ensemble{typeof(p)}(n,p,0,0)
+
+function transition(sampler::Ensemble, model::DensityModelOrLogDensityModel, params, accepted)
+    return [Transition(model, x, accepted) for x in params]
 end
 
 # Define the other sampling steps.
@@ -68,7 +72,7 @@ end
 StretchProposal(p) = StretchProposal(p, 2.0)
 
 function move(
-    rng::Random.AbstractRNG, 
+    rng::Random.AbstractRNG,
     spl::Ensemble{<:StretchProposal},
     model::DensityModelOrLogDensityModel,
     walker::Transition,
@@ -84,17 +88,25 @@ function move(
     # Make new parameters
     y = @. other_walker.params + z * (walker.params - other_walker.params)
 
-    # Construct a new walker
-    new_walker = Transition(model, y)
+    # Construct a temporary new walker
+    temp_walker = Transition(model, y, true)
 
     # Calculate accept/reject value.
-    alpha = alphamult + new_walker.lp - walker.lp
+    alpha = alphamult + temp_walker.lp - walker.lp
 
     if -Random.randexp(rng) <= alpha
+        accepted_draws = spl.accepted_draws + 1
+        new_walker = Transition(model, y, true)
         return new_walker
     else
-        return walker
+        accepted_draws = spl.accepted_draws
+        params = walker.params
+        lp = walker.lp
+        old_walker = Transition(params, lp, false)
+        return old_walker
     end
+
+    spl = Ensemble(spl.n_walkers, spl.proposal, accepted_draws, spl.total_draws+1)
 end
 
 #########################
@@ -124,7 +136,7 @@ end
 
 #     theta_min = theta - 2.0*π
 #     theta_max = theta
-    
+
 #     f = walker.params
 #     while true
 #         stheta, ctheta = sincos(theta)
@@ -136,15 +148,15 @@ end
 #         if new_walker.lp > y
 #             return new_walker
 #         else
-#             if theta < 0 
+#             if theta < 0
 #                 theta_min = theta
 #             else
 #                 theta_max = theta
 #             end
 
 #             theta = theta_min + (theta_max - theta_min) * rand()
 #         end
-#     end 
+#     end
 # end
 
 #####################
@@ -180,15 +192,15 @@ end
 
 #     theta_min = theta - 2.0*π
 #     theta_max = theta
-    
+
 #     f = walker.params
 
 #     i = 0
 #     while true
 #         i += 1
-        
+
 #         stheta, ctheta = sincos(theta)
-        
+
 #         f_prime = f .* ctheta + nu .* stheta
 
 #         new_walker = Transition(model, f_prime)
@@ -198,13 +210,13 @@ end
 #         if new_walker.lp > y
 #             return new_walker
 #         else
-#             if theta < 0 
+#             if theta < 0
 #                 theta_min = theta
 #             else
 #                 theta_max = theta
 #             end
 
 #             theta = theta_min + (theta_max - theta_min) * rand()
 #         end
-#     end 
+#     end
 # end

src/mh-core.jl

@@ -1,7 +1,7 @@
 """
     MetropolisHastings{D}
 
-`MetropolisHastings` has one field, `proposal`. 
+`MetropolisHastings` has one field, `proposal`.
 `proposal` is a `Proposal`, `NamedTuple` of `Proposal`, or `Array{Proposal}` in the shape of your data.
 For example, if you wanted the sampler to return a `NamedTuple` with shape
 
@@ -38,13 +38,17 @@ none is given, the initial parameters will be drawn from the sampler's proposals
 - `param_names` is a vector of strings to be assigned to parameters. This is only
 used if `chain_type=Chains`.
 - `chain_type` is the type of chain you would like returned to you. Supported
-types are `chain_type=Chains` if `MCMCChains` is imported, or 
+types are `chain_type=Chains` if `MCMCChains` is imported, or
 `chain_type=StructArray` if `StructArrays` is imported.
 """
 struct MetropolisHastings{D} <: MHSampler
     proposal::D
+    accepted_draws :: Int
+    total_draws :: Int
 end
 
+MetropolisHastings(p) = MetropolisHastings{typeof(p)}(p,0,0)
+
 StaticMH(d) = MetropolisHastings(StaticProposal(d))
 StaticMH(d::Int) = MetropolisHastings(StaticProposal(MvNormal(Zeros(d), I)))
 RWMH(d) = MetropolisHastings(RandomWalkProposal(d))
@@ -62,12 +66,12 @@ function propose(
     return propose(rng, sampler.proposal, model, transition_prev.params)
 end
 
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, accepted)
     logdensity = AdvancedMH.logdensity(model, params)
-    return transition(sampler, model, params, logdensity)
+    return transition(sampler, model, params, logdensity, accepted)
 end
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real)
-    return Transition(params, logdensity)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real, accepted)
+    return Transition(params, logdensity, accepted)
 end
 
 # Define the first sampling step.
@@ -81,7 +85,7 @@ function AbstractMCMC.step(
     kwargs...
 )
     params = init_params === nothing ? propose(rng, sampler, model) : init_params
-    transition = AdvancedMH.transition(sampler, model, params)
+    transition = AdvancedMH.transition(sampler, model, params, false)
     return transition, transition
 end
 
@@ -106,11 +110,17 @@ function AbstractMCMC.step(
 
     # Decide whether to return the previous params or the new one.
     transition = if -Random.randexp(rng) < logα
-        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate)
+        accepted_draws = sampler.accepted_draws + 1
+        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true)
     else
-        transition_prev
+        accepted_draws = sampler.accepted_draws
+        params = transition_prev.params
+        lp = transition_prev.lp
+        Transition(params, lp, false)
     end
 
+    sampler = typeof(sampler)(sampler.proposal, accepted_draws, sampler.total_draws+1)
+
     return transition, transition
 end