diff --git a/Project.toml b/Project.toml index 78f8cae..6f17776 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedMH" uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -version = "0.6.0" +version = "0.6.1" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" diff --git a/src/MALA.jl b/src/MALA.jl index abb789d..f7a2a03 100644 --- a/src/MALA.jl +++ b/src/MALA.jl @@ -19,48 +19,60 @@ struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{ gradient::G end -transition(::MALA, model, params) = GradientTransition(model, params) - -# Store the new draw, its log density and its gradient -GradientTransition(model::DensityModel, params) = GradientTransition(params, logdensity_and_gradient(model, params)...) +logdensity(model::DensityModel, t::GradientTransition) = t.lp propose(rng::Random.AbstractRNG, ::MALA, model) = error("please specify initial parameters") +function transition(sampler::MALA, model::DensityModel, params) + return GradientTransition(params, logdensity_and_gradient(model, params)...) +end -function propose( +function AbstractMCMC.step( rng::Random.AbstractRNG, - spl::MALA{<:Proposal}, model::DensityModel, - params_prev::GradientTransition -) - proposal = propose(rng, spl.proposal(params_prev.gradient), model, params_prev.params) - return GradientTransition(model, proposal) -end - - -function q( - spl::MALA{<:Proposal}, - t::GradientTransition, - t_cond::GradientTransition -) - return q(spl.proposal(-t_cond.gradient), t.params, t_cond.params) -end - -function logratio_proposal_density( - sampler::MALA{<:Proposal}, state::GradientTransition, candidate::GradientTransition + sampler::MALA, + transition_prev::GradientTransition; + kwargs... ) - return q(sampler, state, candidate) - q(sampler, candidate, state) + # Extract value and gradient of the log density of the current state. + state = transition_prev.params + logdensity_state = transition_prev.lp + gradient_logdensity_state = transition_prev.gradient + + # Generate a new proposal. + proposal = sampler.proposal + candidate = propose(rng, proposal(gradient_logdensity_state), model, state) + + # Compute both the value of the log density and its gradient + logdensity_candidate, gradient_logdensity_candidate = logdensity_and_gradient( + model, candidate + ) + + # Compute the log ratio of proposal densities. + logratio_proposal_density = q( + proposal(-gradient_logdensity_candidate), state, candidate + ) - q(proposal(-gradient_logdensity_state), candidate, state) + + # Compute the log acceptance probability. + logα = logdensity_candidate - logdensity_state + logratio_proposal_density + + # Decide whether to return the previous params or the new one. + transition = if -Random.randexp(rng) < logα + GradientTransition(candidate, logdensity_candidate, gradient_logdensity_candidate) + else + transition_prev + end + + return transition, transition end """ logdensity_and_gradient(model::DensityModel, params) -Efficiently returns the value and gradient of the model +Return the value and gradient of the log density of the parameters `params` for the `model`. """ function logdensity_and_gradient(model::DensityModel, params) res = GradientResult(params) gradient!(res, model.logdensity, params) - return (value(res), gradient(res)) + return value(res), gradient(res) end - -logdensity(model::DensityModel, t::GradientTransition) = t.lp diff --git a/src/emcee.jl b/src/emcee.jl index 28fb81d..2504122 100644 --- a/src/emcee.jl +++ b/src/emcee.jl @@ -3,23 +3,8 @@ struct Ensemble{D} <: MHSampler proposal::D end -# Define the first sampling step. -# Return a 2-tuple consisting of the initial sample and the initial state. -# In this case they are identical. -function AbstractMCMC.step( - rng::Random.AbstractRNG, - model::DensityModel, - spl::Ensemble; - init_params = nothing, - kwargs..., -) - if init_params === nothing - transitions = propose(rng, spl, model) - else - transitions = [Transition(model, x) for x in init_params] - end - - return transitions, transitions +function transition(sampler::Ensemble, model::DensityModel, params) + return [Transition(model, x) for x in params] end # Define the other sampling steps. diff --git a/src/mh-core.jl b/src/mh-core.jl index cd195ce..def5553 100644 --- a/src/mh-core.jl +++ b/src/mh-core.jl @@ -48,115 +48,38 @@ end StaticMH(d) = MetropolisHastings(StaticProposal(d)) RWMH(d) = MetropolisHastings(RandomWalkProposal(d)) -# default function without RNG -propose(spl::MetropolisHastings, args...) = propose(Random.GLOBAL_RNG, spl, args...) - -# Propose from a vector of proposals -function propose( - rng::Random.AbstractRNG, - spl::MetropolisHastings{<:AbstractArray}, - model::DensityModel -) - proposal = map(p -> propose(rng, p, model), spl.proposal) - return Transition(model, proposal) -end - -function propose( - rng::Random.AbstractRNG, - spl::MetropolisHastings{<:AbstractArray}, - model::DensityModel, - params_prev::Transition -) - proposal = map(spl.proposal, params_prev.params) do p, params - propose(rng, p, model, params) - end - return Transition(model, proposal) -end - -# Make a proposal from one Proposal struct. -function propose( - rng::Random.AbstractRNG, - spl::MetropolisHastings{<:Proposal}, - model::DensityModel -) - proposal = propose(rng, spl.proposal, model) - return Transition(model, proposal) -end - -function propose( - rng::Random.AbstractRNG, - spl::MetropolisHastings{<:Proposal}, - model::DensityModel, - params_prev::Transition -) - proposal = propose(rng, spl.proposal, model, params_prev.params) - return Transition(model, proposal) -end - -# Make a proposal from a NamedTuple of Proposal. -function propose( - rng::Random.AbstractRNG, - spl::MetropolisHastings{<:NamedTuple}, - model::DensityModel -) - proposal = _propose(rng, spl.proposal, model) - return Transition(model, proposal) +function propose(rng::Random.AbstractRNG, sampler::MHSampler, model::DensityModel) + return propose(rng, sampler.proposal, model) end - function propose( rng::Random.AbstractRNG, - spl::MetropolisHastings{<:NamedTuple}, + sampler::MHSampler, model::DensityModel, - params_prev::Transition + transition_prev::Transition, ) - proposal = _propose(rng, spl.proposal, model, params_prev.params) - return Transition(model, proposal) + return propose(rng, sampler.proposal, model, transition_prev.params) end -@generated function _propose( - rng::Random.AbstractRNG, - proposal::NamedTuple{names}, - model::DensityModel -) where {names} - isempty(names) && return :(NamedTuple()) - expr = Expr(:tuple) - expr.args = Any[:($name = propose(rng, proposal.$name, model)) for name in names] - return expr +function transition(sampler::MHSampler, model::DensityModel, params) + logdensity = AdvancedMH.logdensity(model, params) + return transition(sampler, model, params, logdensity) end - -@generated function _propose( - rng::Random.AbstractRNG, - proposal::NamedTuple{names}, - model::DensityModel, - params_prev::NamedTuple -) where {names} - isempty(names) && return :(NamedTuple()) - expr = Expr(:tuple) - expr.args = Any[ - :($name = propose(rng, proposal.$name, model, params_prev.$name)) for name in names - ] - return expr +function transition(sampler::MHSampler, model::DensityModel, params, logdensity::Real) + return Transition(params, logdensity) end -transition(sampler, model, params) = transition(model, params) -transition(model, params) = Transition(model, params) - # Define the first sampling step. # Return a 2-tuple consisting of the initial sample and the initial state. # In this case they are identical. function AbstractMCMC.step( rng::Random.AbstractRNG, model::DensityModel, - spl::MHSampler; + sampler::MHSampler; init_params=nothing, kwargs... ) - if init_params === nothing - transition = propose(rng, spl, model) - else - transition = AdvancedMH.transition(spl, model, init_params) - end - + params = init_params === nothing ? propose(rng, sampler, model) : init_params + transition = AdvancedMH.transition(sampler, model, params) return transition, transition end @@ -167,27 +90,30 @@ end function AbstractMCMC.step( rng::Random.AbstractRNG, model::DensityModel, - spl::MHSampler, - params_prev::AbstractTransition; + sampler::MHSampler, + transition_prev::AbstractTransition; kwargs... ) # Generate a new proposal. - params = propose(rng, spl, model, params_prev) + candidate = propose(rng, sampler, model, transition_prev) - # Calculate the log acceptance probability. - logα = logdensity(model, params) - logdensity(model, params_prev) + - logratio_proposal_density(spl, params_prev, params) + # Calculate the log acceptance probability and the log density of the candidate. + logdensity_candidate = logdensity(model, candidate) + logα = logdensity_candidate - logdensity(model, transition_prev) + + logratio_proposal_density(sampler, transition_prev, candidate) # Decide whether to return the previous params or the new one. - if -Random.randexp(rng) < logα - return params, params + transition = if -Random.randexp(rng) < logα + AdvancedMH.transition(sampler, model, candidate, logdensity_candidate) else - return params_prev, params_prev + transition_prev end + + return transition, transition end function logratio_proposal_density( - sampler::MetropolisHastings, params_prev::Transition, params::Transition + sampler::MetropolisHastings, transition_prev::AbstractTransition, candidate ) - return logratio_proposal_density(sampler.proposal, params_prev.params, params.params) + return logratio_proposal_density(sampler.proposal, transition_prev.params, candidate) end diff --git a/src/proposal.jl b/src/proposal.jl index 13f3bf0..c02e5f7 100644 --- a/src/proposal.jl +++ b/src/proposal.jl @@ -125,6 +125,55 @@ function q( return q(proposal(t_cond), t, t_cond) end +#################### +# Multiple proposals +#################### + +function propose( + rng::Random.AbstractRNG, + proposals::AbstractArray{<:Proposal}, + model::DensityModel, +) + return map(proposals) do proposal + return propose(rng, proposal, model) + end +end +function propose( + rng::Random.AbstractRNG, + proposals::AbstractArray{<:Proposal}, + model::DensityModel, + ts, +) + return map(proposals, ts) do proposal, t + return propose(rng, proposal, model, t) + end +end + +@generated function propose( + rng::Random.AbstractRNG, + proposals::NamedTuple{names}, + model::DensityModel, +) where {names} + isempty(names) && return :(NamedTuple()) + expr = Expr(:tuple) + expr.args = Any[:($name = propose(rng, proposals.$name, model)) for name in names] + return expr +end + +@generated function propose( + rng::Random.AbstractRNG, + proposals::NamedTuple{names}, + model::DensityModel, + ts, +) where {names} + isempty(names) && return :(NamedTuple()) + expr = Expr(:tuple) + expr.args = Any[ + :($name = propose(rng, proposals.$name, model, ts.$name)) for name in names + ] + return expr +end + """ logratio_proposal_density(proposal, state, candidate)