move acceptance statistics to MHSampler

docs/src/api.md

@@ -6,7 +6,6 @@ Documentation for AdvancedMH.jl
 ## Structs
 ```@docs
 MetropolisHastings
-Transition
 ```
 
 ## Functions

src/AdvancedMH.jl

@@ -56,18 +56,16 @@ const DensityModelOrLogDensityModel = Union{<:DensityModel,<:AbstractMCMC.LogDen
 # 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,M<:Bool,D<:Int} <: AbstractTransition
+struct Transition{T,L<:Real} <: AbstractTransition
     params :: T
     lp :: L
-    accepted :: M
-    accepted_draws :: D
-    total_draws :: D
+    accepted :: Bool
 end
 
 # Store the new draw, its log density, and draw information
-Transition(model::DensityModelOrLogDensityModel, params, accepted, accepted_draws, total_draws) = Transition(params, logdensity(model, params), accepted, accepted_draws, total_draws)
-function Transition(model::AbstractMCMC.LogDensityModel, params, accepted, accepted_draws, total_draws)
-    return Transition(params, LogDensityProblems.logdensity(model.logdensity, params), accepted, accepted_draws, total_draws)
+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.

src/MALA.jl

@@ -1,30 +1,30 @@
 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))
 
 
-struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple},M<:Bool,D<:Int} <: AbstractTransition
+struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple}} <: AbstractTransition
     params::T
     lp::L
     gradient::G
-    accepted :: M
-    accepted_draws :: D
-    total_draws :: D
+    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, accepted, accepted_draws, total_draws)
-    return GradientTransition(params, logdensity_and_gradient(model, params)..., accepted, accepted_draws, total_draws)
+function transition(sampler::MALA, model::DensityModelOrLogDensityModel, params, accepted)
+    return GradientTransition(params, logdensity_and_gradient(model, params)..., accepted)
 end
 
 check_capabilities(model::DensityModelOrLogDensityModel) = nothing
@@ -71,17 +71,17 @@ function AbstractMCMC.step(
     logα = logdensity_candidate - logdensity_state + logratio_proposal_density
 
     # Decide whether to return the previous params or the new one.
-    total_draws = transition_prev.total_draws + 1
     transition = if -Random.randexp(rng) < logα
-        accepted_draws = transition_prev.accepted_draws + 1
-        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true, accepted_draws, total_draws)
+        accepted_draws = sampler.accepted_draws + 1
+        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true)
     else
-        accepted_draws = transition_prev.accepted_draws
+        accepted_draws = sampler.accepted_draws
         candidate = transition_prev.params
         lp = transition_prev.lp
         gradient = transition_prev.gradient
-        GradientTransition(candidate, lp, gradient, false, accepted_draws, total_draws)
+        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, accepted, accepted_draws, total_draws)
-    return [Transition(model, x, accepted, accepted_draws, total_draws) 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.
@@ -85,23 +89,24 @@ function move(
     y = @. other_walker.params + z * (walker.params - other_walker.params)
 
     # Construct a temporary new walker
-    temp_walker = Transition(model, y, true, walker.accepted_draws, walker.total_draws)
+    temp_walker = Transition(model, y, true)
 
     # Calculate accept/reject value.
     alpha = alphamult + temp_walker.lp - walker.lp
 
-    total_draws = walker.total_draws + 1
     if -Random.randexp(rng) <= alpha
-        accepted_draws = walker.accepted_draws + 1
-        new_walker = Transition(model, y, true, accepted_draws, total_draws)
+        accepted_draws = spl.accepted_draws + 1
+        new_walker = Transition(model, y, true)
         return new_walker
     else
-        accepted_draws = walker.accepted_draws
+        accepted_draws = spl.accepted_draws
         params = walker.params
         lp = walker.lp
-        old_walker = Transition(params, lp, false, accepted_draws, total_draws)
+        old_walker = Transition(params, lp, false)
         return old_walker
     end
+
+    spl = Ensemble(spl.n_walkers, spl.proposal, accepted_draws, spl.total_draws+1)
 end
 
 #########################

src/mh-core.jl

@@ -43,8 +43,12 @@ types are `chain_type=Chains` if `MCMCChains` is imported, or
 """
 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, accepted, accepted_draws, total_draws)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, accepted)
     logdensity = AdvancedMH.logdensity(model, params)
-    return transition(sampler, model, params, logdensity, accepted, accepted_draws, total_draws)
+    return transition(sampler, model, params, logdensity, accepted)
 end
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real, accepted, accepted_draws, total_draws)
-    return Transition(params, logdensity, accepted, accepted_draws, total_draws)
+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, false, 0, 0)
+    transition = AdvancedMH.transition(sampler, model, params, false)
     return transition, transition
 end
 
@@ -105,17 +109,18 @@ function AbstractMCMC.step(
         logratio_proposal_density(sampler, transition_prev, candidate)
 
     # Decide whether to return the previous params or the new one.
-    total_draws = transition_prev.total_draws + 1
     transition = if -Random.randexp(rng) < logα
-        accepted_draws = transition_prev.accepted_draws + 1
-        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true, accepted_draws, total_draws)
+        accepted_draws = sampler.accepted_draws + 1
+        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true)
     else
+        accepted_draws = sampler.accepted_draws
         params = transition_prev.params
         lp = transition_prev.lp
-        accepted_draws = transition_prev.accepted_draws
-        Transition(params, lp, false, accepted_draws, total_draws)
+        Transition(params, lp, false)
     end
 
+    sampler = typeof(sampler)(sampler.proposal, accepted_draws, sampler.total_draws+1)
+
     return transition, transition
 end