From b4e338d986fa7d8a202aaf10c3a18ecc904b8494 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 16:32:16 +0200 Subject: [PATCH 1/7] Generally support container of proposals --- src/mh-core.jl | 109 ++++++++---------------------------------------- src/proposal.jl | 49 ++++++++++++++++++++++ 2 files changed, 66 insertions(+), 92 deletions(-) diff --git a/src/mh-core.jl b/src/mh-core.jl index cd195ce..52d8938 100644 --- a/src/mh-core.jl +++ b/src/mh-core.jl @@ -48,94 +48,16 @@ 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 + params_prev::Transition, ) - proposal = _propose(rng, spl.proposal, model, params_prev.params) - return Transition(model, proposal) -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 -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 + return propose(rng, sampler.proposal, model, params_prev.params) end transition(sampler, model, params) = transition(model, params) @@ -168,26 +90,29 @@ function AbstractMCMC.step( rng::Random.AbstractRNG, model::DensityModel, spl::MHSampler, - params_prev::AbstractTransition; + params_prev::Transition; kwargs... ) # Generate a new proposal. - params = propose(rng, spl, model, params_prev) + candidate = propose(rng, spl, model, params_prev) # Calculate the log acceptance probability. - logα = logdensity(model, params) - logdensity(model, params_prev) + - logratio_proposal_density(spl, params_prev, params) + logdensity_candidate = logdensity(model, candidate) + logα = logdensity_candidate - logdensity(model, params_prev) + + logratio_proposal_density(spl, params_prev, candidate) # Decide whether to return the previous params or the new one. - if -Random.randexp(rng) < logα - return params, params + params = if -Random.randexp(rng) < logα + Transition(candidate, logdensity_candidate) else - return params_prev, params_prev + params_prev end + + return params, params end function logratio_proposal_density( - sampler::MetropolisHastings, params_prev::Transition, params::Transition + sampler::MetropolisHastings, params_prev::Transition, candidate ) - return logratio_proposal_density(sampler.proposal, params_prev.params, params.params) + return logratio_proposal_density(sampler.proposal, params_prev.params, candidate) end diff --git a/src/proposal.jl b/src/proposal.jl index 13f3bf0..0809d27 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, + proposal::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) From 83480e0ed5dd6f772e5ad411b79c2c927a1f154a Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 16:32:39 +0200 Subject: [PATCH 2/7] Bump version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 39406a71c70d2df861c982d24dd25fbf812e4c9a Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 17:43:36 +0200 Subject: [PATCH 3/7] Fix some typos --- src/proposal.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/proposal.jl b/src/proposal.jl index 0809d27..c02e5f7 100644 --- a/src/proposal.jl +++ b/src/proposal.jl @@ -160,9 +160,9 @@ end return expr end -@generated function _propose( +@generated function propose( rng::Random.AbstractRNG, - proposal::NamedTuple{names}, + proposals::NamedTuple{names}, model::DensityModel, ts, ) where {names} From c7b73f5d82548eea37ac7d8799ce0f0ceffb8c0c Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 17:44:41 +0200 Subject: [PATCH 4/7] Wrap provided initial parameters correctly --- src/mh-core.jl | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/mh-core.jl b/src/mh-core.jl index 52d8938..ecd4bb6 100644 --- a/src/mh-core.jl +++ b/src/mh-core.jl @@ -69,16 +69,12 @@ transition(model, params) = Transition(model, params) 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 From 1224981640f45068446e3f96eb91be82e4e26c49 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 17:45:42 +0200 Subject: [PATCH 5/7] Fix MALA by extracting computation of acceptance probability --- src/MALA.jl | 54 ++++++++++++++++++++++++++++++-------------------- src/mh-core.jl | 51 +++++++++++++++++++++++++++++++---------------- 2 files changed, 67 insertions(+), 38 deletions(-) diff --git a/src/MALA.jl b/src/MALA.jl index abb789d..745361f 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 propose( 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) + return propose(rng, spl.proposal(params_prev.gradient), model, params_prev.params) end - -function q( - spl::MALA{<:Proposal}, - t::GradientTransition, - t_cond::GradientTransition -) - return q(spl.proposal(-t_cond.gradient), t.params, t_cond.params) +function transition(sampler::MALA, model::DensityModel, params) + logdensity_gradient = logdensity_and_gradient(model, params) + return transition(sampler, model, params, logdensity_gradient) +end +function transition(::MALA, model::DensityModel, params, (logdensity, gradient)) + return GradientTransition(params, logdensity, gradient) end -function logratio_proposal_density( - sampler::MALA{<:Proposal}, state::GradientTransition, candidate::GradientTransition +function logacceptance_logdensity( + sampler::MALA, + model::DensityModel, + transition_prev::GradientTransition, + candidate, ) - return q(sampler, state, candidate) - q(sampler, candidate, state) + # compute both the value of the log density and its gradient + logdensity_candidate, gradient_logdensity_candidate = logdensity_and_gradient( + model, candidate + ) + + # compute log ratio of proposal densities + proposal = sampler.proposal + state = transition_prev.params + gradient_logdensity_state = transition_prev.gradient + logratio_proposal_density = q( + proposal(-gradient_logdensity_candidate), state, candidate + ) - q(proposal(-gradient_logdensity_state), candidate, state) + + # compute log acceptance probability + logα = logdensity_candidate - logdensity(model, transition_prev) + + logratio_proposal_density + + return logα, (logdensity_candidate, gradient_logdensity_candidate) 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/mh-core.jl b/src/mh-core.jl index ecd4bb6..a35eed2 100644 --- a/src/mh-core.jl +++ b/src/mh-core.jl @@ -55,13 +55,18 @@ function propose( rng::Random.AbstractRNG, sampler::MHSampler, model::DensityModel, - params_prev::Transition, + transition_prev::Transition, ) - return propose(rng, sampler.proposal, model, params_prev.params) + return propose(rng, sampler.proposal, model, transition_prev.params) end -transition(sampler, model, params) = transition(model, params) -transition(model, params) = Transition(model, params) +function transition(sampler::MHSampler, model::DensityModel, params) + logdensity = AdvancedMH.logdensity(model, params) + return transition(sampler, model, params, logdensity) +end +function transition(sampler::MHSampler, model::DensityModel, params, logdensity::Real) + return Transition(params, logdensity) +end # Define the first sampling step. # Return a 2-tuple consisting of the initial sample and the initial state. @@ -85,30 +90,42 @@ end function AbstractMCMC.step( rng::Random.AbstractRNG, model::DensityModel, - spl::MHSampler, - params_prev::Transition; + sampler::MHSampler, + transition_prev::AbstractTransition; kwargs... ) # Generate a new proposal. - candidate = propose(rng, spl, model, params_prev) + candidate = propose(rng, sampler, model, transition_prev) - # Calculate the log acceptance probability. - logdensity_candidate = logdensity(model, candidate) - logα = logdensity_candidate - logdensity(model, params_prev) + - logratio_proposal_density(spl, params_prev, candidate) + # Calculate the log acceptance probability and the log density of the candidate. + logα, logdensity_candidate = logacceptance_logdensity( + sampler, model, transition_prev, candidate + ) # Decide whether to return the previous params or the new one. - params = if -Random.randexp(rng) < logα - Transition(candidate, logdensity_candidate) + transition = if -Random.randexp(rng) < logα + AdvancedMH.transition(sampler, model, candidate, logdensity_candidate) else - params_prev + transition_prev end - return params, params + return transition, transition +end + +function logacceptance_logdensity( + sampler::MHSampler, + model::DensityModel, + transition_prev::AbstractTransition, + candidate, +) + logdensity_candidate = logdensity(model, candidate) + logα = logdensity_candidate - logdensity(model, transition_prev) + + logratio_proposal_density(sampler, transition_prev, candidate) + return logα, logdensity_candidate end function logratio_proposal_density( - sampler::MetropolisHastings, params_prev::Transition, candidate + sampler::MetropolisHastings, transition_prev::AbstractTransition, candidate ) - return logratio_proposal_density(sampler.proposal, params_prev.params, candidate) + return logratio_proposal_density(sampler.proposal, transition_prev.params, candidate) end From 3040cafc691d505c528e0d1a375a1278b29231b3 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 17:46:31 +0200 Subject: [PATCH 6/7] Fix EMCEE --- src/emcee.jl | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) 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. From 8ce01d42e3375a4b7bc120c3e0abfd2cf98da1c0 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 26 May 2021 18:54:43 +0200 Subject: [PATCH 7/7] Implement `AbstractMCMC.step` for MALA --- src/MALA.jl | 54 +++++++++++++++++++++++++------------------------- src/mh-core.jl | 18 +++-------------- 2 files changed, 30 insertions(+), 42 deletions(-) diff --git a/src/MALA.jl b/src/MALA.jl index 745361f..f7a2a03 100644 --- a/src/MALA.jl +++ b/src/MALA.jl @@ -22,47 +22,47 @@ end logdensity(model::DensityModel, t::GradientTransition) = t.lp propose(rng::Random.AbstractRNG, ::MALA, model) = error("please specify initial parameters") -function propose( - rng::Random.AbstractRNG, - spl::MALA{<:Proposal}, - model::DensityModel, - params_prev::GradientTransition -) - return propose(rng, spl.proposal(params_prev.gradient), model, params_prev.params) -end - function transition(sampler::MALA, model::DensityModel, params) - logdensity_gradient = logdensity_and_gradient(model, params) - return transition(sampler, model, params, logdensity_gradient) -end -function transition(::MALA, model::DensityModel, params, (logdensity, gradient)) - return GradientTransition(params, logdensity, gradient) + return GradientTransition(params, logdensity_and_gradient(model, params)...) end -function logacceptance_logdensity( - sampler::MALA, +function AbstractMCMC.step( + rng::Random.AbstractRNG, model::DensityModel, - transition_prev::GradientTransition, - candidate, + sampler::MALA, + transition_prev::GradientTransition; + kwargs... ) - # compute both the value of the log density and its gradient + # 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 log ratio of proposal densities - proposal = sampler.proposal - state = transition_prev.params - gradient_logdensity_state = transition_prev.gradient + # 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 log acceptance probability - logα = logdensity_candidate - logdensity(model, transition_prev) + - logratio_proposal_density + # 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 logα, (logdensity_candidate, gradient_logdensity_candidate) + return transition, transition end """ diff --git a/src/mh-core.jl b/src/mh-core.jl index a35eed2..def5553 100644 --- a/src/mh-core.jl +++ b/src/mh-core.jl @@ -98,9 +98,9 @@ function AbstractMCMC.step( candidate = propose(rng, sampler, model, transition_prev) # Calculate the log acceptance probability and the log density of the candidate. - logα, logdensity_candidate = logacceptance_logdensity( - sampler, model, transition_prev, 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. transition = if -Random.randexp(rng) < logα @@ -112,18 +112,6 @@ function AbstractMCMC.step( return transition, transition end -function logacceptance_logdensity( - sampler::MHSampler, - model::DensityModel, - transition_prev::AbstractTransition, - candidate, -) - logdensity_candidate = logdensity(model, candidate) - logα = logdensity_candidate - logdensity(model, transition_prev) + - logratio_proposal_density(sampler, transition_prev, candidate) - return logα, logdensity_candidate -end - function logratio_proposal_density( sampler::MetropolisHastings, transition_prev::AbstractTransition, candidate )