From 21e940efd9c274703f330d0612ec02d2db70d6bd Mon Sep 17 00:00:00 2001
From: Connor Robertson <cnr.robertson@gmail.com>
Date: Wed, 13 Sep 2023 11:06:04 -0700
Subject: [PATCH 1/5] add acceptance rate info to Transition struct

---
 docs/src/api.md   |  1 +
 src/AdvancedMH.jl | 24 ++++++++++++++----------
 src/MALA.jl       | 21 +++++++++++++++------
 src/emcee.jl      | 37 ++++++++++++++++++++++---------------
 src/mh-core.jl    | 23 ++++++++++++++---------
 5 files changed, 66 insertions(+), 40 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 9c3100a..a4e0d2e 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -6,6 +6,7 @@ Documentation for AdvancedMH.jl
 ## Structs
 ```@docs
 MetropolisHastings
+Transition
 ```
 
 ## Functions
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index 830f63f..778fcdb 100644
--- a/src/AdvancedMH.jl
+++ b/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,21 @@ 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.
-struct Transition{T,L<:Real} <: AbstractTransition
+# 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
     params :: T
     lp :: L
+    accepted :: M
+    accepted_draws :: D
+    total_draws :: D
 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, 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)
 end
 
 # Calculate the density of the model given some parameterization.
@@ -128,7 +132,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")
diff --git a/src/MALA.jl b/src/MALA.jl
index 0e5703f..cd9847d 100644
--- a/src/MALA.jl
+++ b/src/MALA.jl
@@ -11,17 +11,20 @@ MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)
 MALA(d) = MALA(RandomWalkProposal(d))
 
 
-struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple}} <: AbstractTransition
+struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple},M<:Bool,D<:Int} <: AbstractTransition
     params::T
     lp::L
     gradient::G
+    accepted :: M
+    accepted_draws :: D
+    total_draws :: D
 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, accepted_draws, total_draws)
+    return GradientTransition(params, logdensity_and_gradient(model, params)..., accepted, accepted_draws, total_draws)
 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
@@ -68,10 +71,16 @@ 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α
-        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate)
+        accepted_draws = transition_prev.accepted_draws + 1
+        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true, accepted_draws, total_draws)
     else
-        transition_prev
+        accepted_draws = transition_prev.accepted_draws
+        candidate = transition_prev.params
+        lp = transition_prev.lp
+        gradient = transition_prev.gradient
+        GradientTransition(candidate, lp, gradient, false, accepted_draws, total_draws)
     end
 
     return transition, transition
diff --git a/src/emcee.jl b/src/emcee.jl
index 8ffa7e0..170c937 100644
--- a/src/emcee.jl
+++ b/src/emcee.jl
@@ -3,8 +3,8 @@ struct Ensemble{D} <: MHSampler
     proposal::D
 end
 
-function transition(sampler::Ensemble, model::DensityModelOrLogDensityModel, params)
-    return [Transition(model, x) for x in params]
+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]
 end
 
 # Define the other sampling steps.
@@ -68,7 +68,7 @@ end
 StretchProposal(p) = StretchProposal(p, 2.0)
 
 function move(
-    rng::Random.AbstractRNG, 
+    rng::Random.AbstractRNG,
     spl::Ensemble{<:StretchProposal},
     model::DensityModelOrLogDensityModel,
     walker::Transition,
@@ -84,16 +84,23 @@ 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, walker.accepted_draws, walker.total_draws)
 
     # Calculate accept/reject value.
-    alpha = alphamult + new_walker.lp - walker.lp
+    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)
         return new_walker
     else
-        return walker
+        accepted_draws = walker.accepted_draws
+        params = walker.params
+        lp = walker.lp
+        old_walker = Transition(params, lp, false, accepted_draws, total_draws)
+        return old_walker
     end
 end
 
@@ -124,7 +131,7 @@ end
 
 #     theta_min = theta - 2.0*π
 #     theta_max = theta
-    
+
 #     f = walker.params
 #     while true
 #         stheta, ctheta = sincos(theta)
@@ -136,7 +143,7 @@ end
 #         if new_walker.lp > y
 #             return new_walker
 #         else
-#             if theta < 0 
+#             if theta < 0
 #                 theta_min = theta
 #             else
 #                 theta_max = theta
@@ -144,7 +151,7 @@ end
 
 #             theta = theta_min + (theta_max - theta_min) * rand()
 #         end
-#     end 
+#     end
 # end
 
 #####################
@@ -180,15 +187,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,7 +205,7 @@ end
 #         if new_walker.lp > y
 #             return new_walker
 #         else
-#             if theta < 0 
+#             if theta < 0
 #                 theta_min = theta
 #             else
 #                 theta_max = theta
@@ -206,5 +213,5 @@ end
 
 #             theta = theta_min + (theta_max - theta_min) * rand()
 #         end
-#     end 
+#     end
 # end
diff --git a/src/mh-core.jl b/src/mh-core.jl
index bbf04e6..0e8d5d0 100644
--- a/src/mh-core.jl
+++ b/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,7 +38,7 @@ 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
@@ -62,12 +62,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, accepted_draws, total_draws)
     logdensity = AdvancedMH.logdensity(model, params)
-    return transition(sampler, model, params, logdensity)
+    return transition(sampler, model, params, logdensity, accepted, accepted_draws, total_draws)
 end
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real)
-    return Transition(params, logdensity)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real, accepted, accepted_draws, total_draws)
+    return Transition(params, logdensity, accepted, accepted_draws, total_draws)
 end
 
 # Define the first sampling step.
@@ -81,7 +81,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, 0, 0)
     return transition, transition
 end
 
@@ -105,10 +105,15 @@ 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α
-        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate)
+        accepted_draws = transition_prev.accepted_draws + 1
+        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true, accepted_draws, total_draws)
     else
-        transition_prev
+        params = transition_prev.params
+        lp = transition_prev.lp
+        accepted_draws = transition_prev.accepted_draws
+        Transition(params, lp, false, accepted_draws, total_draws)
     end
 
     return transition, transition

From fcfafaf383f223bbf8a0f578620e73286667ab9e Mon Sep 17 00:00:00 2001
From: Connor Robertson <cnr.robertson@gmail.com>
Date: Tue, 19 Sep 2023 16:30:21 -0700
Subject: [PATCH 2/5] move acceptance statistics to MHSampler

---
 docs/src/api.md   |  1 -
 src/AdvancedMH.jl | 12 +++++-------
 src/MALA.jl       | 28 ++++++++++++++--------------
 src/emcee.jl      | 21 +++++++++++++--------
 src/mh-core.jl    | 25 +++++++++++++++----------
 5 files changed, 47 insertions(+), 40 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index a4e0d2e..9c3100a 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -6,7 +6,6 @@ Documentation for AdvancedMH.jl
 ## Structs
 ```@docs
 MetropolisHastings
-Transition
 ```
 
 ## Functions
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index 778fcdb..0d87524 100644
--- a/src/AdvancedMH.jl
+++ b/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.
diff --git a/src/MALA.jl b/src/MALA.jl
index cd9847d..f5342e3 100644
--- a/src/MALA.jl
+++ b/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
diff --git a/src/emcee.jl b/src/emcee.jl
index 170c937..deb69b5 100644
--- a/src/emcee.jl
+++ b/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
 
 #########################
diff --git a/src/mh-core.jl b/src/mh-core.jl
index 0e8d5d0..00badf0 100644
--- a/src/mh-core.jl
+++ b/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
 

From f142046d184307f934780faadaef5b44831012d0 Mon Sep 17 00:00:00 2001
From: Connor Robertson <cnr.robertson@gmail.com>
Date: Tue, 10 Oct 2023 16:15:42 -0700
Subject: [PATCH 3/5] Revert "move acceptance statistics to MHSampler"

This reverts commit fcfafaf383f223bbf8a0f578620e73286667ab9e.
---
 docs/src/api.md   |  1 +
 src/AdvancedMH.jl | 12 +++++++-----
 src/MALA.jl       | 28 ++++++++++++++--------------
 src/emcee.jl      | 21 ++++++++-------------
 src/mh-core.jl    | 25 ++++++++++---------------
 5 files changed, 40 insertions(+), 47 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 9c3100a..a4e0d2e 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -6,6 +6,7 @@ Documentation for AdvancedMH.jl
 ## Structs
 ```@docs
 MetropolisHastings
+Transition
 ```
 
 ## Functions
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index 0d87524..778fcdb 100644
--- a/src/AdvancedMH.jl
+++ b/src/AdvancedMH.jl
@@ -56,16 +56,18 @@ 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} <: AbstractTransition
+struct Transition{T,L<:Real,M<:Bool,D<:Int} <: AbstractTransition
     params :: T
     lp :: L
-    accepted :: Bool
+    accepted :: M
+    accepted_draws :: D
+    total_draws :: D
 end
 
 # 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)
+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)
 end
 
 # Calculate the density of the model given some parameterization.
diff --git a/src/MALA.jl b/src/MALA.jl
index f5342e3..cd9847d 100644
--- a/src/MALA.jl
+++ b/src/MALA.jl
@@ -1,30 +1,30 @@
 struct MALA{D} <: MHSampler
     proposal::D
-    accepted_draws :: Int
-    total_draws :: Int
-end
 
-MALA(p,ad,td) = MALA(p,ad,td)
+    MALA{D}(proposal::D) where {D} = new{D}(proposal)
+end
 
 # If we were given a RandomWalkProposal, just use that instead.
-MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d,0,0)
+MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)
 
 # 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}} <: AbstractTransition
+struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple},M<:Bool,D<:Int} <: AbstractTransition
     params::T
     lp::L
     gradient::G
-    accepted :: Bool
+    accepted :: M
+    accepted_draws :: D
+    total_draws :: D
 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)
-    return GradientTransition(params, logdensity_and_gradient(model, params)..., accepted)
+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)
 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 = sampler.accepted_draws + 1
-        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true)
+        accepted_draws = transition_prev.accepted_draws + 1
+        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true, accepted_draws, total_draws)
     else
-        accepted_draws = sampler.accepted_draws
+        accepted_draws = transition_prev.accepted_draws
         candidate = transition_prev.params
         lp = transition_prev.lp
         gradient = transition_prev.gradient
-        GradientTransition(candidate, lp, gradient, false)
+        GradientTransition(candidate, lp, gradient, false, accepted_draws, total_draws)
     end
-    sampler = MALA(sampler.proposal, accepted_draws, sampler.total_draws+1)
 
     return transition, transition
 end
diff --git a/src/emcee.jl b/src/emcee.jl
index deb69b5..170c937 100644
--- a/src/emcee.jl
+++ b/src/emcee.jl
@@ -1,14 +1,10 @@
 struct Ensemble{D} <: MHSampler
     n_walkers::Int
     proposal::D
-    accepted_draws :: Int
-    total_draws :: Int
 end
 
-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]
+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]
 end
 
 # Define the other sampling steps.
@@ -89,24 +85,23 @@ function move(
     y = @. other_walker.params + z * (walker.params - other_walker.params)
 
     # Construct a temporary new walker
-    temp_walker = Transition(model, y, true)
+    temp_walker = Transition(model, y, true, walker.accepted_draws, walker.total_draws)
 
     # Calculate accept/reject value.
     alpha = alphamult + temp_walker.lp - walker.lp
 
+    total_draws = walker.total_draws + 1
     if -Random.randexp(rng) <= alpha
-        accepted_draws = spl.accepted_draws + 1
-        new_walker = Transition(model, y, true)
+        accepted_draws = walker.accepted_draws + 1
+        new_walker = Transition(model, y, true, accepted_draws, total_draws)
         return new_walker
     else
-        accepted_draws = spl.accepted_draws
+        accepted_draws = walker.accepted_draws
         params = walker.params
         lp = walker.lp
-        old_walker = Transition(params, lp, false)
+        old_walker = Transition(params, lp, false, accepted_draws, total_draws)
         return old_walker
     end
-
-    spl = Ensemble(spl.n_walkers, spl.proposal, accepted_draws, spl.total_draws+1)
 end
 
 #########################
diff --git a/src/mh-core.jl b/src/mh-core.jl
index 00badf0..0e8d5d0 100644
--- a/src/mh-core.jl
+++ b/src/mh-core.jl
@@ -43,12 +43,8 @@ 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))
@@ -66,12 +62,12 @@ function propose(
     return propose(rng, sampler.proposal, model, transition_prev.params)
 end
 
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, accepted)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, accepted, accepted_draws, total_draws)
     logdensity = AdvancedMH.logdensity(model, params)
-    return transition(sampler, model, params, logdensity, accepted)
+    return transition(sampler, model, params, logdensity, accepted, accepted_draws, total_draws)
 end
-function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real, accepted)
-    return Transition(params, logdensity, accepted)
+function transition(sampler::MHSampler, model::DensityModelOrLogDensityModel, params, logdensity::Real, accepted, accepted_draws, total_draws)
+    return Transition(params, logdensity, accepted, accepted_draws, total_draws)
 end
 
 # Define the first sampling step.
@@ -85,7 +81,7 @@ function AbstractMCMC.step(
     kwargs...
 )
     params = init_params === nothing ? propose(rng, sampler, model) : init_params
-    transition = AdvancedMH.transition(sampler, model, params, false)
+    transition = AdvancedMH.transition(sampler, model, params, false, 0, 0)
     return transition, transition
 end
 
@@ -109,18 +105,17 @@ 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 = sampler.accepted_draws + 1
-        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true)
+        accepted_draws = transition_prev.accepted_draws + 1
+        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true, accepted_draws, total_draws)
     else
-        accepted_draws = sampler.accepted_draws
         params = transition_prev.params
         lp = transition_prev.lp
-        Transition(params, lp, false)
+        accepted_draws = transition_prev.accepted_draws
+        Transition(params, lp, false, accepted_draws, total_draws)
     end
 
-    sampler = typeof(sampler)(sampler.proposal, accepted_draws, sampler.total_draws+1)
-
     return transition, transition
 end
 

From 06dec95dd1997b34f5ecb1899a92541f148afe11 Mon Sep 17 00:00:00 2001
From: Connor Robertson <cnr.robertson@gmail.com>
Date: Tue, 10 Oct 2023 16:29:03 -0700
Subject: [PATCH 4/5] incorporate feedback - use only accepted boolean in
 Transition

---
 docs/src/api.md   |  1 -
 src/AdvancedMH.jl | 12 +++++-------
 src/MALA.jl       | 17 ++++++-----------
 src/emcee.jl      | 13 +++++--------
 src/mh-core.jl    | 17 +++++++----------
 5 files changed, 23 insertions(+), 37 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index a4e0d2e..9c3100a 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -6,7 +6,6 @@ Documentation for AdvancedMH.jl
 ## Structs
 ```@docs
 MetropolisHastings
-Transition
 ```
 
 ## Functions
diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl
index 778fcdb..0d87524 100644
--- a/src/AdvancedMH.jl
+++ b/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.
diff --git a/src/MALA.jl b/src/MALA.jl
index cd9847d..4fa75eb 100644
--- a/src/MALA.jl
+++ b/src/MALA.jl
@@ -11,20 +11,18 @@ MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)
 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,16 +69,13 @@ 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)
+        GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true)
     else
-        accepted_draws = transition_prev.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
 
     return transition, transition
diff --git a/src/emcee.jl b/src/emcee.jl
index 170c937..566867f 100644
--- a/src/emcee.jl
+++ b/src/emcee.jl
@@ -3,8 +3,8 @@ struct Ensemble{D} <: MHSampler
     proposal::D
 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]
+function transition(sampler::Ensemble, model::DensityModelOrLogDensityModel, params, accepted)
+    return [Transition(model, x, accepted) for x in params]
 end
 
 # Define the other sampling steps.
@@ -85,21 +85,18 @@ 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)
+        new_walker = Transition(model, y, true)
         return new_walker
     else
-        accepted_draws = walker.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
 end
diff --git a/src/mh-core.jl b/src/mh-core.jl
index 0e8d5d0..3faf46e 100644
--- a/src/mh-core.jl
+++ b/src/mh-core.jl
@@ -62,12 +62,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 +81,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,15 +105,12 @@ 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)
+        AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true)
     else
         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
 
     return transition, transition

From b9dc1c54b0656288ad77a111d952592ad05df0c2 Mon Sep 17 00:00:00 2001
From: Cameron Pfiffer <cameron@pfiffer.org>
Date: Thu, 15 Feb 2024 11:22:49 -0800
Subject: [PATCH 5/5] Update src/mh-core.jl

---
 src/mh-core.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mh-core.jl b/src/mh-core.jl
index 2cc8980..147e53f 100644
--- a/src/mh-core.jl
+++ b/src/mh-core.jl
@@ -80,7 +80,7 @@ function AbstractMCMC.step(
     initial_params=nothing,
     kwargs...
 )
-    params = init_params === nothing ? propose(rng, sampler, model) : init_params
+    params = initial_params === nothing ? propose(rng, sampler, model) : initial_params
     transition = AdvancedMH.transition(sampler, model, params, false)
     return transition, transition
 end