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

@@ -15,13 +15,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 +45,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,9 +70,12 @@ 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)
+ GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate, true)
  else
- transition_prev
+ candidate = transition_prev.params
+ lp = transition_prev.lp
+ gradient = transition_prev.gradient
+ GradientTransition(candidate, lp, gradient, false)
  end
 
  return transition, transition

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)
+ return [Transition(model, x, accepted) 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,20 @@ 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
+ new_walker = Transition(model, y, true)
  return new_walker
  else
- return walker
+ params = walker.params
+ lp = walker.lp
+ old_walker = Transition(params, lp, false)
+ return old_walker
  end
 end
 
@@ -124,7 +128,7 @@ end
 
 # theta_min = theta - 2.0*π
 # theta_max = theta
- 
+
 # f = walker.params
 # while true
 # stheta, ctheta = sincos(theta)
@@ -136,15 +140,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 +184,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 +202,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,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)
  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 +81,7 @@ function AbstractMCMC.step(
  kwargs...
 )
  params = initial_params === nothing ? propose(rng, sampler, model) : initial_params
- transition = AdvancedMH.transition(sampler, model, params)
+ transition = AdvancedMH.transition(sampler, model, params, false)
  return transition, transition
 end
 
@@ -106,9 +106,11 @@ 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)
+ AdvancedMH.transition(sampler, model, candidate, logdensity_candidate, true)
  else
- transition_prev
+ params = transition_prev.params
+ lp = transition_prev.lp
+ Transition(params, lp, false)
  end
 
  return transition, transition