diff --git a/src/AdvancedMH.jl b/src/AdvancedMH.jl index 830f63f..0d87524 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,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") diff --git a/src/MALA.jl b/src/MALA.jl index 0e5703f..4fa75eb 100644 --- a/src/MALA.jl +++ b/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 diff --git a/src/emcee.jl b/src/emcee.jl index 8ffa7e0..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) - 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,7 +140,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 +148,7 @@ 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,7 +202,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 +210,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 afb50dd..147e53f 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) 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