From f36a3ccc7e29eec42ab380dc845afe235234f42e Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 6 Apr 2021 16:05:28 +0200 Subject: [PATCH 01/36] initial work on adopting the AbstractMCMC interface --- Project.toml | 1 + src/AdvancedHMC.jl | 2 ++ src/sampler.jl | 67 +++++++++++++++++++++++++++++++++++++++++++++- src/trajectory.jl | 6 +++-- 4 files changed, 73 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index c7904a1f..57d4df7e 100644 --- a/Project.toml +++ b/Project.toml @@ -3,6 +3,7 @@ uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" version = "0.2.28" [deps] +AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" InplaceOps = "505f98c9-085e-5b2c-8e89-488be7bf1f34" diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 25947dcc..6c543797 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -12,6 +12,8 @@ using ArgCheck: @argcheck using DocStringExtensions +import AbstractMCMC + import StatsBase: sample import Parameters: reconstruct diff --git a/src/sampler.jl b/src/sampler.jl index c8e9302d..1f4d1d5a 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -32,7 +32,6 @@ end ## ## Interface functions ## - function sample_init( rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}}, h::Hamiltonian, @@ -210,3 +209,69 @@ function sample( end return θs, stats end + +################################# +### AbstractMCMC.jl interface ### +################################# +struct HamiltonianModel{H} <: AbstractMCMC.AbstractModel + hamiltonian :: H +end + +struct HMCState{ + TTrans<:Transition, + TAdapt<:Adaptation.AbstractAdaptor +} + i::Int + transition::TTrans + adaptor::TAdapt +end + +function AbstractMCMC.step( + rng::AbstractRNG, + model::HamiltonianModel, + spl::HMCKernel; + init_params, # TODO: implement this. Just do `rand`? Need dimensionality though. + adaptor::AbstractAdaptor=NoAdaptation(), + kwargs... +) + # Get an initial smaple + h, t = AdvancedHMC.sample_init(rng, model.hamiltonian, init_params) + + # Compute next transition and state. + state = HMCState(0, t, adaptor) + + # Take actual first step + return AbstractMCMC.step(rng, model, spl, state; kwargs...) +end + +function AbstractMCMC.step( + rng::AbstractRNG, + model::HamiltonianModel, + spl::HMCKernel, + state::HMCState; + nadapts::Int=0, + kwargs... +) + # Get step size + @debug "current ϵ" getstepsize(spl, state) + + # Compute transition. + h = model.hamiltonian + i = state.i + 1 + t_old = state.transition + adaptor = state.adaptor + + # Make new transition. + t = transition(rng, h, spl, t_old.z) + + # Adapt h and spl. + tstat = stat(t) + h, spl, isadapted = adapt!(h, spl, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) + tstat = merge(tstat, (is_adapt=isadapted,)) + + # Compute next transition and state. + newstate = HMCState(i, t, adaptor) + + # Return `Transition` with additional stats added. + return Transition(t.z, tstat), newstate +end diff --git a/src/trajectory.jl b/src/trajectory.jl index e189fb44..159d601d 100644 --- a/src/trajectory.jl +++ b/src/trajectory.jl @@ -25,7 +25,9 @@ end "Returns the statistics for transition `t`." stat(t::Transition) = t.stat -abstract type AbstractMCMCKernel end +# abstract type AbstractMCMCKernel end + +const AbstractMCMCKernel = AbstractMCMC.AbstractSampler abstract type AbstractTerminationCriterion end @@ -211,7 +213,7 @@ nsteps(τ::Trajectory{TS, I, TC}) where {TS, I, TC<:FixedIntegrationTime} = ## Kernel interface ## -struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel +struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel refreshment::R τ::T end From 16a233c47f5cb02f9fc2d00ccd9a2d18e607a11c Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 6 Apr 2021 20:34:57 +0200 Subject: [PATCH 02/36] removed old sample method since is now redundant --- src/sampler.jl | 123 ------------------------------------------------- 1 file changed, 123 deletions(-) diff --git a/src/sampler.jl b/src/sampler.jl index 1f4d1d5a..e717d53e 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -87,129 +87,6 @@ function Adaptation.adapt!( return h, κ, isadapted end -""" -Progress meter update with all trajectory stats, iteration number and metric shown. -""" -function pm_next!(pm, stat::NamedTuple) - ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stat)]) -end - -""" -Simple progress meter update without any show values. -""" -simple_pm_next!(pm, stat::NamedTuple) = ProgressMeter.next!(pm) - -## -## Sampling functions -## - -sample( - h::Hamiltonian, - κ::AbstractMCMCKernel, - θ::AbstractVecOrMat{<:AbstractFloat}, - n_samples::Int, - adaptor::AbstractAdaptor=NoAdaptation(), - n_adapts::Int=min(div(n_samples, 10), 1_000); - drop_warmup=false, - verbose::Bool=true, - progress::Bool=false, - (pm_next!)::Function=pm_next! -) = sample( - GLOBAL_RNG, - h, - κ, - θ, - n_samples, - adaptor, - n_adapts; - drop_warmup=drop_warmup, - verbose=verbose, - progress=progress, - (pm_next!)=pm_next!, -) - -""" - sample( - rng::AbstractRNG, - h::Hamiltonian, - κ::AbstractMCMCKernel, - θ::AbstractVecOrMat{T}, - n_samples::Int, - adaptor::AbstractAdaptor=NoAdaptation(), - n_adapts::Int=min(div(n_samples, 10), 1_000); - drop_warmup::Bool=false, - verbose::Bool=true, - progress::Bool=false - ) - -Sample `n_samples` samples using the proposal `κ` under Hamiltonian `h`. -- The randomness is controlled by `rng`. - - If `rng` is not provided, `GLOBAL_RNG` will be used. -- The initial point is given by `θ`. -- The adaptor is set by `adaptor`, for which the default is no adaptation. - - It will perform `n_adapts` steps of adaptation, for which the default is the minimum of `1_000` and 10% of `n_samples` -- `drop_warmup` controls to drop the samples during adaptation phase or not -- `verbose` controls the verbosity -- `progress` controls whether to show the progress meter or not -""" -function sample( - rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}}, - h::Hamiltonian, - κ::HMCKernel, - θ::T, - n_samples::Int, - adaptor::AbstractAdaptor=NoAdaptation(), - n_adapts::Int=min(div(n_samples, 10), 1_000); - drop_warmup=false, - verbose::Bool=true, - progress::Bool=false, - (pm_next!)::Function=pm_next! -) where {T<:AbstractVecOrMat{<:AbstractFloat}} - @assert !(drop_warmup && (adaptor isa Adaptation.NoAdaptation)) "Cannot drop warmup samples if there is no adaptation phase." - # Prepare containers to store sampling results - n_keep = n_samples - (drop_warmup ? n_adapts : 0) - θs, stats = Vector{T}(undef, n_keep), Vector{NamedTuple}(undef, n_keep) - # Initial sampling - h, t = sample_init(rng, h, θ) - # Progress meter - pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing - time = @elapsed for i = 1:n_samples - # Make a transition - t = transition(rng, h, κ, t.z) - # Adapt h and κ; what mutable is the adaptor - tstat = stat(t) - h, κ, isadapted = adapt!(h, κ, adaptor, i, n_adapts, t.z.θ, tstat.acceptance_rate) - tstat = merge(tstat, (is_adapt=isadapted,)) - # Update progress meter - if progress - # Do include current iteration and mass matrix - pm_next!(pm, (iterations=i, tstat..., mass_matrix=h.metric)) - # Report finish of adapation - elseif verbose && isadapted && i == n_adapts - @info "Finished $n_adapts adapation steps" adaptor κ.τ.integrator h.metric - end - # Store sample - if !drop_warmup || i > n_adapts - j = i - drop_warmup * n_adapts - θs[j], stats[j] = t.z.θ, tstat - end - end - # Report end of sampling - if verbose - EBFMI_est = EBFMI(map(s -> s.hamiltonian_energy, stats)) - average_acceptance_rate = mean(map(s -> s.acceptance_rate, stats)) - if θ isa AbstractVector - n_chains = 1 - else - n_chains = size(θ, 2) - EBFMI_est = "[" * join(EBFMI_est, ", ") * "]" - average_acceptance_rate = "[" * join(average_acceptance_rate, ", ") * "]" - end - @info "Finished $n_samples sampling steps for $n_chains chains in $time (s)" h κ EBFMI_est average_acceptance_rate - end - return θs, stats -end - ################################# ### AbstractMCMC.jl interface ### ################################# From a1494ac1562cc919ebdca40d72c5813ad2bd6ea1 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 7 Apr 2021 00:02:51 +0200 Subject: [PATCH 03/36] AbstractMCMCKernel is now a subtype of AbstractSampler --- src/sampler.jl | 23 ++++++++++++++--------- src/trajectory.jl | 6 ++---- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/src/sampler.jl b/src/sampler.jl index e717d53e..d242c012 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -17,7 +17,7 @@ function reconstruct( return reconstruct(τ, integrator=integrator) end -reconstruct(κ::AbstractMCMCKernel, adaptor::AbstractAdaptor) = +reconstruct(κ::AbstractHMCKernel, adaptor::AbstractAdaptor) = reconstruct(κ, τ=reconstruct(κ.τ, adaptor)) function resize(h::Hamiltonian, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat} @@ -58,7 +58,7 @@ end Adaptation.adapt!( h::Hamiltonian, - κ::AbstractMCMCKernel, + κ::AbstractHMCKernel, adaptor::Adaptation.NoAdaptation, i::Int, n_adapts::Int, @@ -68,7 +68,7 @@ Adaptation.adapt!( function Adaptation.adapt!( h::Hamiltonian, - κ::AbstractMCMCKernel, + κ::AbstractHMCKernel, adaptor::AbstractAdaptor, i::Int, n_adapts::Int, @@ -106,27 +106,32 @@ end function AbstractMCMC.step( rng::AbstractRNG, model::HamiltonianModel, - spl::HMCKernel; - init_params, # TODO: implement this. Just do `rand`? Need dimensionality though. + spl::AbstractHMCKernel; + init_params = nothing, adaptor::AbstractAdaptor=NoAdaptation(), kwargs... ) - # Get an initial smaple + if init_params === nothing + init_params = randn(size(model.hamiltonian.metric, 1)) + end + + # Get an initial sample. h, t = AdvancedHMC.sample_init(rng, model.hamiltonian, init_params) # Compute next transition and state. state = HMCState(0, t, adaptor) - # Take actual first step + # Take actual first step. return AbstractMCMC.step(rng, model, spl, state; kwargs...) end function AbstractMCMC.step( rng::AbstractRNG, model::HamiltonianModel, - spl::HMCKernel, + spl::AbstractHMCKernel, state::HMCState; - nadapts::Int=0, + nadapts::Int = 0, + verbose::Bool = true, kwargs... ) # Get step size diff --git a/src/trajectory.jl b/src/trajectory.jl index 159d601d..f7cba0dc 100644 --- a/src/trajectory.jl +++ b/src/trajectory.jl @@ -25,9 +25,7 @@ end "Returns the statistics for transition `t`." stat(t::Transition) = t.stat -# abstract type AbstractMCMCKernel end - -const AbstractMCMCKernel = AbstractMCMC.AbstractSampler +abstract type AbstractHMCKernel <: AbstractMCMC.AbstractSampler end abstract type AbstractTerminationCriterion end @@ -213,7 +211,7 @@ nsteps(τ::Trajectory{TS, I, TC}) where {TS, I, TC<:FixedIntegrationTime} = ## Kernel interface ## -struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel +struct HMCKernel{R, T<:Trajectory} <: AbstractHMCKernel refreshment::R τ::T end From 5f01c5db9e7b5880b8d81ff9d1c7ba84cf077056 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 7 Apr 2021 01:03:18 +0200 Subject: [PATCH 04/36] fixed impl and added back old one --- src/sampler.jl | 176 ++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 168 insertions(+), 8 deletions(-) diff --git a/src/sampler.jl b/src/sampler.jl index d242c012..f7fbf186 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -87,28 +87,158 @@ function Adaptation.adapt!( return h, κ, isadapted end +""" +Progress meter update with all trajectory stats, iteration number and metric shown. +""" +function pm_next!(pm, stat::NamedTuple) + ProgressMeter.next!(pm; showvalues=[tuple(s...) for s in pairs(stat)]) +end + +""" +Simple progress meter update without any show values. +""" +simple_pm_next!(pm, stat::NamedTuple) = ProgressMeter.next!(pm) + +## +## Sampling functions +## + +sample( + h::Hamiltonian, + κ::AbstractHMCKernel, + θ::AbstractVecOrMat{<:AbstractFloat}, + n_samples::Int, + adaptor::AbstractAdaptor=NoAdaptation(), + n_adapts::Int=min(div(n_samples, 10), 1_000); + drop_warmup=false, + verbose::Bool=true, + progress::Bool=false, + (pm_next!)::Function=pm_next! +) = sample( + GLOBAL_RNG, + h, + κ, + θ, + n_samples, + adaptor, + n_adapts; + drop_warmup=drop_warmup, + verbose=verbose, + progress=progress, + (pm_next!)=pm_next!, +) + +""" + sample( + rng::AbstractRNG, + h::Hamiltonian, + κ::AbstractHMCKernel, + θ::AbstractVecOrMat{T}, + n_samples::Int, + adaptor::AbstractAdaptor=NoAdaptation(), + n_adapts::Int=min(div(n_samples, 10), 1_000); + drop_warmup::Bool=false, + verbose::Bool=true, + progress::Bool=false + ) +Sample `n_samples` samples using the proposal `κ` under Hamiltonian `h`. +- The randomness is controlled by `rng`. + - If `rng` is not provided, `GLOBAL_RNG` will be used. +- The initial point is given by `θ`. +- The adaptor is set by `adaptor`, for which the default is no adaptation. + - It will perform `n_adapts` steps of adaptation, for which the default is the minimum of `1_000` and 10% of `n_samples` +- `drop_warmup` controls to drop the samples during adaptation phase or not +- `verbose` controls the verbosity +- `progress` controls whether to show the progress meter or not +""" +function sample( + rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}}, + h::Hamiltonian, + κ::HMCKernel, + θ::T, + n_samples::Int, + adaptor::AbstractAdaptor=NoAdaptation(), + n_adapts::Int=min(div(n_samples, 10), 1_000); + drop_warmup=false, + verbose::Bool=true, + progress::Bool=false, + (pm_next!)::Function=pm_next! +) where {T<:AbstractVecOrMat{<:AbstractFloat}} + @assert !(drop_warmup && (adaptor isa Adaptation.NoAdaptation)) "Cannot drop warmup samples if there is no adaptation phase." + # Prepare containers to store sampling results + n_keep = n_samples - (drop_warmup ? n_adapts : 0) + θs, stats = Vector{T}(undef, n_keep), Vector{NamedTuple}(undef, n_keep) + # Initial sampling + h, t = sample_init(rng, h, θ) + # Progress meter + pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing + time = @elapsed for i = 1:n_samples + # Make a transition + t = transition(rng, h, κ, t.z) + # Adapt h and κ; what mutable is the adaptor + tstat = stat(t) + h, κ, isadapted = adapt!(h, κ, adaptor, i, n_adapts, t.z.θ, tstat.acceptance_rate) + tstat = merge(tstat, (is_adapt=isadapted,)) + # Update progress meter + if progress + # Do include current iteration and mass matrix + pm_next!(pm, (iterations=i, tstat..., mass_matrix=h.metric)) + # Report finish of adapation + elseif verbose && isadapted && i == n_adapts + @info "Finished $n_adapts adapation steps" adaptor κ.τ.integrator h.metric + end + # Store sample + if !drop_warmup || i > n_adapts + j = i - drop_warmup * n_adapts + θs[j], stats[j] = t.z.θ, tstat + end + end + # Report end of sampling + if verbose + EBFMI_est = EBFMI(map(s -> s.hamiltonian_energy, stats)) + average_acceptance_rate = mean(map(s -> s.acceptance_rate, stats)) + if θ isa AbstractVector + n_chains = 1 + else + n_chains = size(θ, 2) + EBFMI_est = "[" * join(EBFMI_est, ", ") * "]" + average_acceptance_rate = "[" * join(average_acceptance_rate, ", ") * "]" + end + @info "Finished $n_samples sampling steps for $n_chains chains in $time (s)" h κ EBFMI_est average_acceptance_rate + end + return θs, stats +end + ################################# ### AbstractMCMC.jl interface ### ################################# +struct HMCSampler{K, A} <: AbstractMCMC.AbstractSampler + κ::K + adaptor::A +end + struct HamiltonianModel{H} <: AbstractMCMC.AbstractModel - hamiltonian :: H + hamiltonian::H end struct HMCState{ TTrans<:Transition, + THam<:Hamiltonian, + TKernel<:AbstractHMCKernel, TAdapt<:Adaptation.AbstractAdaptor } i::Int transition::TTrans + hamiltonian::THam + κ::TKernel adaptor::TAdapt end function AbstractMCMC.step( rng::AbstractRNG, model::HamiltonianModel, - spl::AbstractHMCKernel; + spl::HMCSampler; init_params = nothing, - adaptor::AbstractAdaptor=NoAdaptation(), kwargs... ) if init_params === nothing @@ -119,7 +249,7 @@ function AbstractMCMC.step( h, t = AdvancedHMC.sample_init(rng, model.hamiltonian, init_params) # Compute next transition and state. - state = HMCState(0, t, adaptor) + state = HMCState(0, t, h, spl.κ, spl.adaptor) # Take actual first step. return AbstractMCMC.step(rng, model, spl, state; kwargs...) @@ -128,7 +258,7 @@ end function AbstractMCMC.step( rng::AbstractRNG, model::HamiltonianModel, - spl::AbstractHMCKernel, + spl::HMCSampler, state::HMCState; nadapts::Int = 0, verbose::Bool = true, @@ -142,18 +272,48 @@ function AbstractMCMC.step( i = state.i + 1 t_old = state.transition adaptor = state.adaptor + κ = state.κ # Make new transition. - t = transition(rng, h, spl, t_old.z) + t = transition(rng, h, κ, t_old.z) # Adapt h and spl. tstat = stat(t) - h, spl, isadapted = adapt!(h, spl, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) + h, κ, isadapted = adapt!(h, κ, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) tstat = merge(tstat, (is_adapt=isadapted,)) # Compute next transition and state. - newstate = HMCState(i, t, adaptor) + newstate = HMCState(i, t, h, κ, adaptor) # Return `Transition` with additional stats added. return Transition(t.z, tstat), newstate end + + +################ +### Callback ### +################ + +struct HMCCallback end + +function (cb::HMCCallback)( + rng, model, t, spl, state, i; + progress, verbose, kwargs... +) + h = model.hamiltonian + adaptor = spl.adaptor + κ = spl.kernel + isadapted = t.stat.is_adapt + + # Update progress meter + if progress + # Do include current iteration and mass matrix + pm_next!( + pm, + (iterations=i, tstat..., mass_matrix=h.metric) + ) + # Report finish of adapation + elseif verbose && isadapted && i == n_adapts + @info "Finished $n_adapts adapation steps" adaptor κ.τ.integrator h.metric + end +end From fc9e7474522392e6210a83ae11b362279d79acae Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 7 Apr 2021 13:30:06 +0200 Subject: [PATCH 05/36] overload AbstractMCMC.sample to make interface a bit nicer --- src/AdvancedHMC.jl | 1 + src/sampler.jl | 96 +++++++++++++++++++++++++++++++++++++++------- src/trajectory.jl | 2 +- 3 files changed, 84 insertions(+), 15 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 6c543797..3207fb46 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -5,6 +5,7 @@ const DEBUG = convert(Bool, parse(Int, get(ENV, "DEBUG_AHMC", "0"))) using Statistics: mean, var, middle using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky, UniformScaling using StatsFuns: logaddexp, logsumexp +import Random using Random: GLOBAL_RNG, AbstractRNG using ProgressMeter: ProgressMeter using Parameters: @unpack, reconstruct diff --git a/src/sampler.jl b/src/sampler.jl index f7fbf186..b4606596 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -212,44 +212,108 @@ end ################################# ### AbstractMCMC.jl interface ### ################################# -struct HMCSampler{K, A} <: AbstractMCMC.AbstractSampler - κ::K - adaptor::A +struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler + initial_kernel::K + initial_metric::M + initial_adaptor::A end -struct HamiltonianModel{H} <: AbstractMCMC.AbstractModel - hamiltonian::H +struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel + ℓπ::Tlogπ + ∂ℓπ∂θ::T∂logπ∂θ end struct HMCState{ TTrans<:Transition, - THam<:Hamiltonian, + TMetric<:AbstractMetric, TKernel<:AbstractHMCKernel, TAdapt<:Adaptation.AbstractAdaptor } i::Int transition::TTrans - hamiltonian::THam + metric::TMetric κ::TKernel adaptor::TAdapt end +function AbstractMCMC.sample( + model::DifferentiableDensityModel, + kernel::AbstractHMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + N::Integer; + kwargs... +) + return AbstractMCMC.sample(Random.GLOBAL_RNG, model, kernel, metric, adaptor, N; kwargs...) +end + +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DifferentiableDensityModel, + kernel::AbstractHMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + N::Integer; + kwargs... +) + sampler = HMCSampler(kernel, metric, adaptor) + return AbstractMCMC.mcmcsample(rng, model, sampler, N; kwargs...) +end + +function AbstractMCMC.sample( + model::DifferentiableDensityModel, + kernel::AbstractHMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + parallel::AbstractMCMC.AbstractMCMCParallel, + N::Integer, + nchains::Integer; + kwargs... +) + return AbstractMCMC.sample( + Random.GLOBAL_RNG, model, kernel, metric, adaptor, N, nchains; + kwargs... + ) +end + +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DifferentiableDensityModel, + kernel::AbstractHMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + parallel::AbstractMCMC.AbstractMCMCParallel, + N::Integer, + nchains::Integer; + kwargs... +) + sampler = HMCSampler(kernel, metric, adaptor) + return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, nchains; kwargs...) +end + function AbstractMCMC.step( rng::AbstractRNG, - model::HamiltonianModel, + model::DifferentiableDensityModel, spl::HMCSampler; init_params = nothing, kwargs... ) + metric = spl.initial_metric + κ = spl.initial_kernel + adaptor = spl.initial_adaptor + if init_params === nothing - init_params = randn(size(model.hamiltonian.metric, 1)) + init_params = randn(size(metric, 1)) end + # Construct the hamiltonian using the initial metric + hamiltonian = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) + # Get an initial sample. - h, t = AdvancedHMC.sample_init(rng, model.hamiltonian, init_params) + h, t = AdvancedHMC.sample_init(rng, hamiltonian, init_params) # Compute next transition and state. - state = HMCState(0, t, h, spl.κ, spl.adaptor) + state = HMCState(0, t, h.metric, κ, adaptor) # Take actual first step. return AbstractMCMC.step(rng, model, spl, state; kwargs...) @@ -257,7 +321,7 @@ end function AbstractMCMC.step( rng::AbstractRNG, - model::HamiltonianModel, + model::DifferentiableDensityModel, spl::HMCSampler, state::HMCState; nadapts::Int = 0, @@ -268,11 +332,14 @@ function AbstractMCMC.step( @debug "current ϵ" getstepsize(spl, state) # Compute transition. - h = model.hamiltonian i = state.i + 1 t_old = state.transition adaptor = state.adaptor κ = state.κ + metric = state.metric + + # Reconstruct hamiltonian. + h = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) # Make new transition. t = transition(rng, h, κ, t_old.z) @@ -283,7 +350,7 @@ function AbstractMCMC.step( tstat = merge(tstat, (is_adapt=isadapted,)) # Compute next transition and state. - newstate = HMCState(i, t, h, κ, adaptor) + newstate = HMCState(i, t, h.metric, κ, adaptor) # Return `Transition` with additional stats added. return Transition(t.z, tstat), newstate @@ -294,6 +361,7 @@ end ### Callback ### ################ +# TODO: add as a default once AbstractMCMC.jl has been updated struct HMCCallback end function (cb::HMCCallback)( diff --git a/src/trajectory.jl b/src/trajectory.jl index f7cba0dc..a40f3110 100644 --- a/src/trajectory.jl +++ b/src/trajectory.jl @@ -25,7 +25,7 @@ end "Returns the statistics for transition `t`." stat(t::Transition) = t.stat -abstract type AbstractHMCKernel <: AbstractMCMC.AbstractSampler end +abstract type AbstractHMCKernel end abstract type AbstractTerminationCriterion end From ba0d6231fa994ec3890739a7067c316b9a0a868c Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 7 Apr 2021 14:04:08 +0200 Subject: [PATCH 06/36] reverted renaming of abstract kernel type --- src/sampler.jl | 9 +++++---- src/trajectory.jl | 4 ++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/sampler.jl b/src/sampler.jl index b4606596..1efd6929 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -217,6 +217,7 @@ struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler initial_metric::M initial_adaptor::A end +HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation()) struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel ℓπ::Tlogπ @@ -238,7 +239,7 @@ end function AbstractMCMC.sample( model::DifferentiableDensityModel, - kernel::AbstractHMCKernel, + kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, N::Integer; @@ -250,7 +251,7 @@ end function AbstractMCMC.sample( rng::Random.AbstractRNG, model::DifferentiableDensityModel, - kernel::AbstractHMCKernel, + kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, N::Integer; @@ -262,7 +263,7 @@ end function AbstractMCMC.sample( model::DifferentiableDensityModel, - kernel::AbstractHMCKernel, + kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, parallel::AbstractMCMC.AbstractMCMCParallel, @@ -279,7 +280,7 @@ end function AbstractMCMC.sample( rng::Random.AbstractRNG, model::DifferentiableDensityModel, - kernel::AbstractHMCKernel, + kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, parallel::AbstractMCMC.AbstractMCMCParallel, diff --git a/src/trajectory.jl b/src/trajectory.jl index a40f3110..bbb20a4a 100644 --- a/src/trajectory.jl +++ b/src/trajectory.jl @@ -25,7 +25,7 @@ end "Returns the statistics for transition `t`." stat(t::Transition) = t.stat -abstract type AbstractHMCKernel end +abstract type AbstractMCMCKernel end abstract type AbstractTerminationCriterion end @@ -211,7 +211,7 @@ nsteps(τ::Trajectory{TS, I, TC}) where {TS, I, TC<:FixedIntegrationTime} = ## Kernel interface ## -struct HMCKernel{R, T<:Trajectory} <: AbstractHMCKernel +struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel refreshment::R τ::T end From b6ad39cb0fc83f11523b0b042aa278b5b5d7a62a Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 8 Apr 2021 15:51:06 +0200 Subject: [PATCH 07/36] replicate logging behavior of current AHMC when using AbstractMCMC --- Project.toml | 1 + src/sampler.jl | 82 +++++++++++++++++++++++++++++++++++++------------- 2 files changed, 62 insertions(+), 21 deletions(-) diff --git a/Project.toml b/Project.toml index 57d4df7e..577345a0 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" [compat] +AbstractMCMC = "3" ArgCheck = "1, 2" DocStringExtensions = "0.8" InplaceOps = "0.3" diff --git a/src/sampler.jl b/src/sampler.jl index 1efd6929..59bcd707 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -17,7 +17,7 @@ function reconstruct( return reconstruct(τ, integrator=integrator) end -reconstruct(κ::AbstractHMCKernel, adaptor::AbstractAdaptor) = +reconstruct(κ::AbstractMCMCKernel, adaptor::AbstractAdaptor) = reconstruct(κ, τ=reconstruct(κ.τ, adaptor)) function resize(h::Hamiltonian, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat} @@ -58,7 +58,7 @@ end Adaptation.adapt!( h::Hamiltonian, - κ::AbstractHMCKernel, + κ::AbstractMCMCKernel, adaptor::Adaptation.NoAdaptation, i::Int, n_adapts::Int, @@ -68,7 +68,7 @@ Adaptation.adapt!( function Adaptation.adapt!( h::Hamiltonian, - κ::AbstractHMCKernel, + κ::AbstractMCMCKernel, adaptor::AbstractAdaptor, i::Int, n_adapts::Int, @@ -105,7 +105,7 @@ simple_pm_next!(pm, stat::NamedTuple) = ProgressMeter.next!(pm) sample( h::Hamiltonian, - κ::AbstractHMCKernel, + κ::AbstractMCMCKernel, θ::AbstractVecOrMat{<:AbstractFloat}, n_samples::Int, adaptor::AbstractAdaptor=NoAdaptation(), @@ -132,7 +132,7 @@ sample( sample( rng::AbstractRNG, h::Hamiltonian, - κ::AbstractHMCKernel, + κ::AbstractMCMCKernel, θ::AbstractVecOrMat{T}, n_samples::Int, adaptor::AbstractAdaptor=NoAdaptation(), @@ -227,7 +227,7 @@ end struct HMCState{ TTrans<:Transition, TMetric<:AbstractMetric, - TKernel<:AbstractHMCKernel, + TKernel<:AbstractMCMCKernel, TAdapt<:Adaptation.AbstractAdaptor } i::Int @@ -255,10 +255,24 @@ function AbstractMCMC.sample( metric::AbstractMetric, adaptor::AbstractAdaptor, N::Integer; + progress = true, + verbose = false, + callback = nothing, kwargs... ) sampler = HMCSampler(kernel, metric, adaptor) - return AbstractMCMC.mcmcsample(rng, model, sampler, N; kwargs...) + if callback === nothing + callback = HMCProgressCallback(N, progress = progress, verbose = verbose) + progress = false # don't use AMCMC's progress-funtionality + end + + return AbstractMCMC.mcmcsample( + rng, model, sampler, N; + progress = progress, + verbose = verbose, + callback = callback, + kwargs... + ) end function AbstractMCMC.sample( @@ -286,10 +300,24 @@ function AbstractMCMC.sample( parallel::AbstractMCMC.AbstractMCMCParallel, N::Integer, nchains::Integer; + progress = true, + verbose = false, + callback = nothing, kwargs... ) sampler = HMCSampler(kernel, metric, adaptor) - return AbstractMCMC.mcmcsample(rng, model, sampler, parallel, N, nchains; kwargs...) + if callback === nothing + callback = HMCProgressCallback(N, progress = progress, verbose = verbose) + progress = false # don't use AMCMC's progress-funtionality + end + + return AbstractMCMC.mcmcsample( + rng, model, sampler, parallel, N, nchains; + progress = progress, + verbose = verbose, + callback = callback, + kwargs... + ) end function AbstractMCMC.step( @@ -326,7 +354,6 @@ function AbstractMCMC.step( spl::HMCSampler, state::HMCState; nadapts::Int = 0, - verbose::Bool = true, kwargs... ) # Get step size @@ -361,28 +388,41 @@ end ################ ### Callback ### ################ +struct HMCProgressCallback{P} + pm::P + progress::Bool + verbose::Bool +end -# TODO: add as a default once AbstractMCMC.jl has been updated -struct HMCCallback end +function HMCProgressCallback(n_samples; progress=true, verbose=false) + pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing + HMCProgressCallback(pm, progress, verbose) +end -function (cb::HMCCallback)( - rng, model, t, spl, state, i; - progress, verbose, kwargs... +function (cb::HMCProgressCallback)( + rng, model, spl, t, state, i; + nadapts = 0, + kwargs... ) - h = model.hamiltonian - adaptor = spl.adaptor - κ = spl.kernel - isadapted = t.stat.is_adapt + progress = cb.progress + verbose = cb.verbose + pm = cb.pm + + metric = state.metric + adaptor = state.adaptor + κ = state.κ + tstat = t.stat + isadapted = tstat.is_adapt # Update progress meter if progress # Do include current iteration and mass matrix pm_next!( pm, - (iterations=i, tstat..., mass_matrix=h.metric) + (iterations=i, tstat..., mass_matrix=metric) ) # Report finish of adapation - elseif verbose && isadapted && i == n_adapts - @info "Finished $n_adapts adapation steps" adaptor κ.τ.integrator h.metric + elseif verbose && isadapted && i == nadapts + @info "Finished $nadapts adapation steps" adaptor κ.τ.integrator metric end end From 66b37b5fb8a70aed58962090a04655b181643d08 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 8 Apr 2021 17:18:34 +0200 Subject: [PATCH 08/36] initial work on tests, but require update to testing deps --- test/abstractmcmc.jl | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 test/abstractmcmc.jl diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl new file mode 100644 index 00000000..ac79ebed --- /dev/null +++ b/test/abstractmcmc.jl @@ -0,0 +1,29 @@ +using Test, Random, AdvancedHMC, ForwardDiff, AbstractMCMC +using Statistics: mean +include("common.jl") + +@testset "`gdemo`" begin + rng = MersenneTwister(0) + + n_samples = 5_000 + n_adapts = 1_000 + + θ_init = randn(rng, 2) + + metric = DiagEuclideanMetric(2) + h = Hamiltonian(metric, ℓπ_gdemo, ForwardDiff) + init_eps = Leapfrog(0.1) + κ = NUTS(init_eps) + adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) + + # samples, _ = sample(rng, h, κ, θ_init, n_samples, adaptor, n_adapts) + samples = AbstractMCMC.sample( + model, κ, metric, adaptor, n_samples; + nadapts = n_adapts, + init_params = θ_init + ) + + m_est = mean(map(invlink_gdemo, samples[1000:end])) + + @test m_est ≈ [49 / 24, 7 / 6] atol=RNDATOL +end From f041ff618546256880f504e521bd7a1f7c67a331 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Fri, 9 Apr 2021 14:32:01 +0200 Subject: [PATCH 09/36] moved abstractmcmc interface into separate file --- src/AdvancedHMC.jl | 2 + src/abstractmcmc.jl | 214 +++++++++++++++++++++++++++++++++++++++++++ src/sampler.jl | 218 -------------------------------------------- 3 files changed, 216 insertions(+), 218 deletions(-) create mode 100644 src/abstractmcmc.jl diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index f10f45bf..0e307520 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -131,6 +131,8 @@ include("diagnosis.jl") include("sampler.jl") export sample +include("abstractmcmc.jl") + include("contrib/ad.jl") ### Init diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl new file mode 100644 index 00000000..25a47203 --- /dev/null +++ b/src/abstractmcmc.jl @@ -0,0 +1,214 @@ +struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler + initial_kernel::K + initial_metric::M + initial_adaptor::A +end +HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation()) + +struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel + ℓπ::Tlogπ + ∂ℓπ∂θ::T∂logπ∂θ +end + +struct HMCState{ + TTrans<:Transition, + TMetric<:AbstractMetric, + TKernel<:AbstractMCMCKernel, + TAdapt<:Adaptation.AbstractAdaptor +} + i::Int + transition::TTrans + metric::TMetric + κ::TKernel + adaptor::TAdapt +end + +function AbstractMCMC.sample( + model::DifferentiableDensityModel, + kernel::AbstractMCMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + N::Integer; + kwargs... +) + return AbstractMCMC.sample(Random.GLOBAL_RNG, model, kernel, metric, adaptor, N; kwargs...) +end + +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DifferentiableDensityModel, + kernel::AbstractMCMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + N::Integer; + progress = true, + verbose = false, + callback = nothing, + kwargs... +) + sampler = HMCSampler(kernel, metric, adaptor) + if callback === nothing + callback = HMCProgressCallback(N, progress = progress, verbose = verbose) + progress = false # don't use AMCMC's progress-funtionality + end + + return AbstractMCMC.mcmcsample( + rng, model, sampler, N; + progress = progress, + verbose = verbose, + callback = callback, + kwargs... + ) +end + +function AbstractMCMC.sample( + model::DifferentiableDensityModel, + kernel::AbstractMCMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + parallel::AbstractMCMC.AbstractMCMCParallel, + N::Integer, + nchains::Integer; + kwargs... +) + return AbstractMCMC.sample( + Random.GLOBAL_RNG, model, kernel, metric, adaptor, N, nchains; + kwargs... + ) +end + +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DifferentiableDensityModel, + kernel::AbstractMCMCKernel, + metric::AbstractMetric, + adaptor::AbstractAdaptor, + parallel::AbstractMCMC.AbstractMCMCParallel, + N::Integer, + nchains::Integer; + progress = true, + verbose = false, + callback = nothing, + kwargs... +) + sampler = HMCSampler(kernel, metric, adaptor) + if callback === nothing + callback = HMCProgressCallback(N, progress = progress, verbose = verbose) + progress = false # don't use AMCMC's progress-funtionality + end + + return AbstractMCMC.mcmcsample( + rng, model, sampler, parallel, N, nchains; + progress = progress, + verbose = verbose, + callback = callback, + kwargs... + ) +end + +function AbstractMCMC.step( + rng::AbstractRNG, + model::DifferentiableDensityModel, + spl::HMCSampler; + init_params = nothing, + kwargs... +) + metric = spl.initial_metric + κ = spl.initial_kernel + adaptor = spl.initial_adaptor + + if init_params === nothing + init_params = randn(size(metric, 1)) + end + + # Construct the hamiltonian using the initial metric + hamiltonian = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) + + # Get an initial sample. + h, t = AdvancedHMC.sample_init(rng, hamiltonian, init_params) + + # Compute next transition and state. + state = HMCState(0, t, h.metric, κ, adaptor) + + # Take actual first step. + return AbstractMCMC.step(rng, model, spl, state; kwargs...) +end + +function AbstractMCMC.step( + rng::AbstractRNG, + model::DifferentiableDensityModel, + spl::HMCSampler, + state::HMCState; + nadapts::Int = 0, + kwargs... +) + # Get step size + @debug "current ϵ" getstepsize(spl, state) + + # Compute transition. + i = state.i + 1 + t_old = state.transition + adaptor = state.adaptor + κ = state.κ + metric = state.metric + + # Reconstruct hamiltonian. + h = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) + + # Make new transition. + t = transition(rng, h, κ, t_old.z) + + # Adapt h and spl. + tstat = stat(t) + h, κ, isadapted = adapt!(h, κ, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) + tstat = merge(tstat, (is_adapt=isadapted,)) + + # Compute next transition and state. + newstate = HMCState(i, t, h.metric, κ, adaptor) + + # Return `Transition` with additional stats added. + return Transition(t.z, tstat), newstate +end + + +################ +### Callback ### +################ +struct HMCProgressCallback{P} + pm::P + progress::Bool + verbose::Bool +end + +function HMCProgressCallback(n_samples; progress=true, verbose=false) + pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing + HMCProgressCallback(pm, progress, verbose) +end + +function (cb::HMCProgressCallback)( + rng, model, spl, t, state, i; + nadapts = 0, + kwargs... +) + progress = cb.progress + verbose = cb.verbose + pm = cb.pm + + metric = state.metric + adaptor = state.adaptor + κ = state.κ + tstat = t.stat + isadapted = tstat.is_adapt + + # Update progress meter + if progress + # Do include current iteration and mass matrix + pm_next!( + pm, + (iterations=i, tstat..., mass_matrix=metric) + ) + # Report finish of adapation + elseif verbose && isadapted && i == nadapts + @info "Finished $nadapts adapation steps" adaptor κ.τ.integrator metric + end +end diff --git a/src/sampler.jl b/src/sampler.jl index 67005436..558f3102 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -209,221 +209,3 @@ function sample( end return θs, stats end - -################################# -### AbstractMCMC.jl interface ### -################################# -struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler - initial_kernel::K - initial_metric::M - initial_adaptor::A -end -HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation()) - -struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel - ℓπ::Tlogπ - ∂ℓπ∂θ::T∂logπ∂θ -end - -struct HMCState{ - TTrans<:Transition, - TMetric<:AbstractMetric, - TKernel<:AbstractMCMCKernel, - TAdapt<:Adaptation.AbstractAdaptor -} - i::Int - transition::TTrans - metric::TMetric - κ::TKernel - adaptor::TAdapt -end - -function AbstractMCMC.sample( - model::DifferentiableDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, - N::Integer; - kwargs... -) - return AbstractMCMC.sample(Random.GLOBAL_RNG, model, kernel, metric, adaptor, N; kwargs...) -end - -function AbstractMCMC.sample( - rng::Random.AbstractRNG, - model::DifferentiableDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, - N::Integer; - progress = true, - verbose = false, - callback = nothing, - kwargs... -) - sampler = HMCSampler(kernel, metric, adaptor) - if callback === nothing - callback = HMCProgressCallback(N, progress = progress, verbose = verbose) - progress = false # don't use AMCMC's progress-funtionality - end - - return AbstractMCMC.mcmcsample( - rng, model, sampler, N; - progress = progress, - verbose = verbose, - callback = callback, - kwargs... - ) -end - -function AbstractMCMC.sample( - model::DifferentiableDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, - parallel::AbstractMCMC.AbstractMCMCParallel, - N::Integer, - nchains::Integer; - kwargs... -) - return AbstractMCMC.sample( - Random.GLOBAL_RNG, model, kernel, metric, adaptor, N, nchains; - kwargs... - ) -end - -function AbstractMCMC.sample( - rng::Random.AbstractRNG, - model::DifferentiableDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, - parallel::AbstractMCMC.AbstractMCMCParallel, - N::Integer, - nchains::Integer; - progress = true, - verbose = false, - callback = nothing, - kwargs... -) - sampler = HMCSampler(kernel, metric, adaptor) - if callback === nothing - callback = HMCProgressCallback(N, progress = progress, verbose = verbose) - progress = false # don't use AMCMC's progress-funtionality - end - - return AbstractMCMC.mcmcsample( - rng, model, sampler, parallel, N, nchains; - progress = progress, - verbose = verbose, - callback = callback, - kwargs... - ) -end - -function AbstractMCMC.step( - rng::AbstractRNG, - model::DifferentiableDensityModel, - spl::HMCSampler; - init_params = nothing, - kwargs... -) - metric = spl.initial_metric - κ = spl.initial_kernel - adaptor = spl.initial_adaptor - - if init_params === nothing - init_params = randn(size(metric, 1)) - end - - # Construct the hamiltonian using the initial metric - hamiltonian = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) - - # Get an initial sample. - h, t = AdvancedHMC.sample_init(rng, hamiltonian, init_params) - - # Compute next transition and state. - state = HMCState(0, t, h.metric, κ, adaptor) - - # Take actual first step. - return AbstractMCMC.step(rng, model, spl, state; kwargs...) -end - -function AbstractMCMC.step( - rng::AbstractRNG, - model::DifferentiableDensityModel, - spl::HMCSampler, - state::HMCState; - nadapts::Int = 0, - kwargs... -) - # Get step size - @debug "current ϵ" getstepsize(spl, state) - - # Compute transition. - i = state.i + 1 - t_old = state.transition - adaptor = state.adaptor - κ = state.κ - metric = state.metric - - # Reconstruct hamiltonian. - h = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) - - # Make new transition. - t = transition(rng, h, κ, t_old.z) - - # Adapt h and spl. - tstat = stat(t) - h, κ, isadapted = adapt!(h, κ, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) - tstat = merge(tstat, (is_adapt=isadapted,)) - - # Compute next transition and state. - newstate = HMCState(i, t, h.metric, κ, adaptor) - - # Return `Transition` with additional stats added. - return Transition(t.z, tstat), newstate -end - - -################ -### Callback ### -################ -struct HMCProgressCallback{P} - pm::P - progress::Bool - verbose::Bool -end - -function HMCProgressCallback(n_samples; progress=true, verbose=false) - pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing - HMCProgressCallback(pm, progress, verbose) -end - -function (cb::HMCProgressCallback)( - rng, model, spl, t, state, i; - nadapts = 0, - kwargs... -) - progress = cb.progress - verbose = cb.verbose - pm = cb.pm - - metric = state.metric - adaptor = state.adaptor - κ = state.κ - tstat = t.stat - isadapted = tstat.is_adapt - - # Update progress meter - if progress - # Do include current iteration and mass matrix - pm_next!( - pm, - (iterations=i, tstat..., mass_matrix=metric) - ) - # Report finish of adapation - elseif verbose && isadapted && i == nadapts - @info "Finished $nadapts adapation steps" adaptor κ.τ.integrator metric - end -end From dad71a85f9df121cc198498a22ccb1bf5799f27e Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 12 Apr 2021 15:32:44 +0200 Subject: [PATCH 10/36] version bump --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 20fd50bc..24822690 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedHMC" uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" -version = "0.2.28" +version = "0.3.0" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" From 07be7ac49344a6618f8f99344f56e3be9e89651c Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 12 Apr 2021 16:00:45 +0200 Subject: [PATCH 11/36] only bump patch version to see if tests succeed --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 24822690..e31c4f23 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedHMC" uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" -version = "0.3.0" +version = "0.2.29" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" From 486e9d19f65a57bf601585c1d86e6890f7c10128 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 01:15:17 +0100 Subject: [PATCH 12/36] move away from using extras in Project.toml --- .gitignore | 1 - Project.toml | 18 ------------------ test/Project.toml | 18 ++++++++++++++++++ 3 files changed, 18 insertions(+), 19 deletions(-) create mode 100644 test/Project.toml diff --git a/.gitignore b/.gitignore index cae5cb87..7f554567 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,3 @@ .history .DS_Store Manifest.toml -test/Project.toml diff --git a/Project.toml b/Project.toml index 039bba1c..89250b2c 100644 --- a/Project.toml +++ b/Project.toml @@ -27,21 +27,3 @@ StatsBase = "0.31, 0.32, 0.33" StatsFuns = "0.8, 0.9" UnPack = "1" julia = "1" - -[extras] -Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" -Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" -Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -MCMCDebugging = "6d524b87-5f90-4494-b601-374a5b87a94b" -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" -Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" - -[targets] -test = ["CUDA", "Distributed", "Distributions", "ComponentArrays", "ForwardDiff", "Plots", "MCMCDebugging", "Test", "Turing", "UnicodePlots", "Bijectors", "OrdinaryDiffEq", "Zygote"] diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 00000000..7404991f --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,18 @@ +[deps] +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" +Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MCMCDebugging = "6d524b87-5f90-4494-b601-374a5b87a94b" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" +UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" From 608442bb8843361c7a6df27f41f6dafc13a1613a Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:21:00 +0100 Subject: [PATCH 13/36] added integration tests for Turing.jl --- test/turing/Project.toml | 12 ++++++++++++ test/turing/runtests.jl | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+) create mode 100644 test/turing/Project.toml create mode 100644 test/turing/runtests.jl diff --git a/test/turing/Project.toml b/test/turing/Project.toml new file mode 100644 index 00000000..d10dd06e --- /dev/null +++ b/test/turing/Project.toml @@ -0,0 +1,12 @@ +[deps] +AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" + +[compat] +StatsFuns = "0.9.5" +Turing = "0.16" +julia = "1.3" \ No newline at end of file diff --git a/test/turing/runtests.jl b/test/turing/runtests.jl new file mode 100644 index 00000000..b8eacb33 --- /dev/null +++ b/test/turing/runtests.jl @@ -0,0 +1,19 @@ +using Random +using Test +using LinearAlgebra + +using Turing +using Turing.DynamicPPL + +using StatsFuns: logistic + +Turing.setprogress!(false) + +Random.seed!(100) + +# Load test utilities. +testdir(args...) = joinpath(pathof(Turing), "..", "..", "test", args...) +include(testdir("test_utils", "AllUtils.jl")) + +# Test HMC. +include(testdir("inference", "hmc.jl")) From cb77c9db70104030cbb19b6a89d393d759e59ea7 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:21:38 +0100 Subject: [PATCH 14/36] removed usage of Turing.jl and MCMCDebugging.jl in main testsuite --- test/Project.toml | 2 -- test/common.jl | 39 --------------------------------------- test/sampler.jl | 6 ------ 3 files changed, 47 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 7404991f..e6e0aa58 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,13 +6,11 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -MCMCDebugging = "6d524b87-5f90-4494-b601-374a5b87a94b" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/test/common.jl b/test/common.jl index 3af536a4..e6caa99f 100644 --- a/test/common.jl +++ b/test/common.jl @@ -79,45 +79,6 @@ function ℓπ_gdemo(θ) return logprior + loglikelihood end -using Distributions: MvNormal -import Turing - -Turing.@model function mvntest(θ=missing, x=missing) - θ ~ MvNormal(zeros(D), 2) - x ~ Normal(sum(θ), 1) - return θ, x -end - -function get_primitives(x, modelgen) - spl_prior = Turing.SampleFromPrior() - function ℓπ(θ) - vi = Turing.VarInfo(model) - vi[spl_prior] = θ - model(vi, spl_prior) - Turing.getlogp(vi) - end - adbackend = Turing.Core.ForwardDiffAD{40} - alg_ad = Turing.HMC{adbackend}(0.1, 1) - model = modelgen(missing, x) - vi = Turing.VarInfo(model) - spl = Turing.Sampler(alg_ad, model) - Turing.Core.link!(vi, spl) - ∂ℓπ∂θ = θ -> Turing.Core.gradient_logp(adbackend(), θ, vi, model, spl) - θ₀ = Turing.VarInfo(model)[Turing.SampleFromPrior()] - return ℓπ, ∂ℓπ∂θ, θ₀ -end - -function rand_θ_given(x, modelgen, metric, κ; n_samples=20) - ℓπ, ∂ℓπ∂θ, θ₀ = get_primitives(x, modelgen) - h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ) - samples, stats = sample(h, κ, θ₀, n_samples; verbose=false, progress=false) - s = samples[end] - return length(s) == 1 ? s[1] : s -end - -# Test function -geweke_g(θ, x) = cat(θ, x; dims=1) - test_show(x) = test_show(s -> length(s) > 0, x) function test_show(pred, x) io = IOBuffer(; append = true) diff --git a/test/sampler.jl b/test/sampler.jl index 937ea3f6..a15548d8 100644 --- a/test/sampler.jl +++ b/test/sampler.jl @@ -61,12 +61,6 @@ end Random.seed!(1) samples, stats = sample(h, HMCKernel(τ), θ_init, n_samples; verbose=false, progress=PROGRESS) @test mean(samples[n_adapts+1:end]) ≈ zeros(D) atol=RNDATOL - if "GEWEKE_TEST" in keys(ENV) && ENV["GEWEKE_TEST"] == "1" - res = perform(GewekeTest(5_000), mvntest, x -> rand_θ_given(x, mvntest, metric, HMCKernel(τ)); g=geweke_g, progress=false) - p = plot(res, mvntest()) - display(p) - println() - end end # Skip adaptation tests with tempering From 75e2404d1bd7ae42ac524a00b748d536dd60e00b Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:21:54 +0100 Subject: [PATCH 15/36] fixed bug in deprecated HMCDA constructor --- src/AdvancedHMC.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index a7cdf05a..9c90bff0 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -78,8 +78,8 @@ struct StaticTrajectory{TS} end struct HMCDA{TS} end @deprecate HMCDA{TS}(int::AbstractIntegrator, λ) where {TS} HMCKernel(Trajectory{TS}(int, FixedIntegrationTime(λ))) -@deprecate HMCDA(int::AbstractIntegrator, λ) HMCKernel(Trajectory{MetropolisTS}(int, FixedIntegrationTime(λ))) -@deprecate HMCDA(ϵ::AbstractScalarOrVec{<:Real}, λ) HMCKernel(Trajectory{MetropolisTS}(Leapfrog(ϵ), FixedIntegrationTime(λ))) +@deprecate HMCDA(int::AbstractIntegrator, λ) HMCKernel(Trajectory{EndPointTS}(int, FixedIntegrationTime(λ))) +@deprecate HMCDA(ϵ::AbstractScalarOrVec{<:Real}, λ) HMCKernel(Trajectory{EndPointTS}(Leapfrog(ϵ), FixedIntegrationTime(λ))) @deprecate find_good_eps find_good_stepsize From 68884d74d4b75480d677d8961074e90ddfc5a47b Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:42:26 +0100 Subject: [PATCH 16/36] allow specification of which testing suites to run --- test/runtests.jl | 71 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 51 insertions(+), 20 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index d0bc48f4..faf880f9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,31 +3,62 @@ using Distributed, Test, CUDA println("Envronment variables for testing") println(ENV) +const DIRECTORY_AdvancedHMC = dirname(dirname(pathof(AdvancedHMC))) +const DIRECTORY_Turing_tests = joinpath(DIRECTORY_AdvancedHMC, "test", "turing") +const GROUP = get(ENV, "GROUP", "All") + @testset "AdvancedHMC" begin - tests = [ - "metric", - "hamiltonian", - "integrator", - "trajectory", - "adaptation", - "sampler", - "sampler-vec", - "demo", - "models", - ] - - if CUDA.functional() - @eval module TestCUDA + if GROUP == "All" || GROUP == "AdvancedHMC" + tests = [ + "metric", + "hamiltonian", + "integrator", + "trajectory", + "adaptation", + "sampler", + "sampler-vec", + "demo", + "models", + ] + + if CUDA.functional() + @eval module TestCUDA include("cuda.jl") + end + else + @warn "Skipping GPU tests because no GPU available." end - else - @warn "Skipping GPU tests because no GPU available." - end - res = map(tests) do t - @eval module $(Symbol("Test_", t)) + res = map(tests) do t + @eval module $(Symbol("Test_", t)) include($t * ".jl") + end + return + end + end + + if GROUP == "All" || GROUP == "Downstream" + @testset "turing" begin + try + # activate separate test environment + Pkg.activate(DIRECTORY_Turing_tests) + Pkg.develop(PackageSpec(; path=DIRECTORY_AdvancedHMC)) + Pkg.instantiate() + + # make sure that the new environment is considered `using` and `import` statements + # (not added automatically on Julia 1.3, see e.g. PR #209) + if !(joinpath(DIRECTORY_Turing_tests, "Project.toml") in Base.load_path()) + pushfirst!(LOAD_PATH, DIRECTORY_Turing_tests) + end + + include(joinpath("turing", "runtests.jl")) + catch err + err isa Pkg.Resolve.ResolverError || rethrow() + # If we can't resolve that means this is incompatible by SemVer and this is fine + # It means we marked this as a breaking change, so we don't need to worry about + # Mistakenly introducing a breaking change, as we have intentionally made one + @info "Not compatible with this release. No problem." exception = err + end end - return end end From bfc8f4d84e1638a7b32108753ed0ae49e6acd2a2 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:42:46 +0100 Subject: [PATCH 17/36] added Turing.jl integration tests to CI --- .github/workflows/CI.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c26e8d04..c34bcd04 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -39,4 +39,9 @@ jobs: - name: Run tests uses: julia-actions/julia-runtest@latest env: - GEWEKE_TEST: 1 + GROUP: AdvancedHMC + - name: Run tests + uses: julia-actions/julia-runtest@latest + continue-on-error: true + env: + GROUP: Downstream From d002bfc99f136185defce211cf3118a93504025d Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:43:55 +0100 Subject: [PATCH 18/36] fixed name for integration tests --- .github/workflows/CI.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index c34bcd04..7a9584d7 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -40,7 +40,7 @@ jobs: uses: julia-actions/julia-runtest@latest env: GROUP: AdvancedHMC - - name: Run tests + - name: Run integration tests uses: julia-actions/julia-runtest@latest continue-on-error: true env: From 1ec826250851aee4bd0b863d8ae2b59a3957842d Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:50:06 +0100 Subject: [PATCH 19/36] added using AdvancedHMC in runtests.jl --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index faf880f9..f3fdc8b3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,4 @@ -using Distributed, Test, CUDA +using Distributed, Test, CUDA, AdvancedHMC println("Envronment variables for testing") println(ENV) From f7f91e2719cbecad63988f6e502f92400a78d22a Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 02:59:20 +0100 Subject: [PATCH 20/36] removed some now unnecessary usings --- test/integrator.jl | 11 ----------- test/sampler.jl | 2 +- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/test/integrator.jl b/test/integrator.jl index acffd8f0..3430f21b 100644 --- a/test/integrator.jl +++ b/test/integrator.jl @@ -26,17 +26,6 @@ n_steps = 10 @test z_step.r ≈ z_steps.r atol=DETATOL end -# using Turing: Inference -# @testset "step(::Leapfrog) against Turing.Inference._leapfrog()" begin -# z = AdvancedHMC.phasepoint(h, θ_init, r_init) -# t_Turing = @elapsed θ_Turing, r_Turing, _ = Inference._leapfrog(θ_init, r_init, n_steps, ϵ, x -> (nothing, ∂logπ∂θ(x))) -# t_AHMC = @elapsed z_AHMC = AdvancedHMC.step(lf, h, z, n_steps) -# @info "Performance of leapfrog of AdvancedHMC v.s. Turing" n_steps t_Turing t_AHMC t_Turing / t_AHMC -# -# @test θ_Turing ≈ z_AHMC.θ atol=DETATOL -# @test r_Turing ≈ z_AHMC.r atol=DETATOL -# end - @testset "jitter" begin @testset "Leapfrog" begin ϵ0 = 0.1 diff --git a/test/sampler.jl b/test/sampler.jl index a15548d8..40177c0c 100644 --- a/test/sampler.jl +++ b/test/sampler.jl @@ -1,7 +1,7 @@ # Allow pass --progress when running this script individually to turn on progress meter const PROGRESS = length(ARGS) > 0 && ARGS[1] == "--progress" ? true : false -using Test, AdvancedHMC, LinearAlgebra, Random, MCMCDebugging, Plots +using Test, AdvancedHMC, LinearAlgebra, Random, Plots using AdvancedHMC: StaticTerminationCriterion, DynamicTerminationCriterion using Setfield using Statistics: mean, var, cov From 639776a41885e139fa1aca33b21cb3d9b028050e Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 03:33:46 +0100 Subject: [PATCH 21/36] fixed a bug in the downstream testing --- test/Project.toml | 1 + test/runtests.jl | 8 +++++++- test/turing/Project.toml | 2 +- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index e6e0aa58..8927e4ba 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" diff --git a/test/runtests.jl b/test/runtests.jl index f3fdc8b3..c1f66b96 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,6 @@ -using Distributed, Test, CUDA, AdvancedHMC +using Distributed, Test, CUDA, Pkg + +using AdvancedHMC: AdvancedHMC println("Envronment variables for testing") println(ENV) @@ -51,7 +53,11 @@ const GROUP = get(ENV, "GROUP", "All") pushfirst!(LOAD_PATH, DIRECTORY_Turing_tests) end + # Avoids conflicting namespaces, e.g. `NUTS` used in Turing.jl's tests + # refers to `Turing.NUTS` not `AdvancedHMC.NUTS`. + @eval module TuringIntegrationTests include(joinpath("turing", "runtests.jl")) + end catch err err isa Pkg.Resolve.ResolverError || rethrow() # If we can't resolve that means this is incompatible by SemVer and this is fine diff --git a/test/turing/Project.toml b/test/turing/Project.toml index d10dd06e..7241c9c6 100644 --- a/test/turing/Project.toml +++ b/test/turing/Project.toml @@ -9,4 +9,4 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [compat] StatsFuns = "0.9.5" Turing = "0.16" -julia = "1.3" \ No newline at end of file +julia = "1.3" From e47a5a4008e75b46449059f372d6c27ec5cf68cf Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 03:39:18 +0100 Subject: [PATCH 22/36] give integration tests a separate CI --- .github/workflows/CI.yml | 5 ---- .github/workflows/IntegrationTests.yml | 41 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 5 deletions(-) create mode 100644 .github/workflows/IntegrationTests.yml diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 7a9584d7..747729a7 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -40,8 +40,3 @@ jobs: uses: julia-actions/julia-runtest@latest env: GROUP: AdvancedHMC - - name: Run integration tests - uses: julia-actions/julia-runtest@latest - continue-on-error: true - env: - GROUP: Downstream diff --git a/.github/workflows/IntegrationTests.yml b/.github/workflows/IntegrationTests.yml new file mode 100644 index 00000000..c65c7ca0 --- /dev/null +++ b/.github/workflows/IntegrationTests.yml @@ -0,0 +1,41 @@ +name: IntegrationTests + +on: + push: + branches: + - master + pull_request: + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + version: + - '1' + os: + - ubuntu-latest + - macOS-latest + - windows-latest + arch: + - x86 + - x64 + exclude: + - os: ubuntu-latest + arch: x86 + - os: macOS-latest + arch: x86 + - os: windows-latest + arch: x86 + steps: + - uses: actions/checkout@v2 + - uses: julia-actions/setup-julia@v1 + with: + version: ${{ matrix.version }} + arch: ${{ matrix.arch }} + - uses: julia-actions/julia-buildpkg@latest + - name: Run integration tests + uses: julia-actions/julia-runtest@latest + continue-on-error: true + env: + GROUP: Downstream From 678196c7fc042cf3805ce96f25f6a6ad7090dd74 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 03:48:05 +0100 Subject: [PATCH 23/36] forgot to remove the continue-on-error from CI --- .github/workflows/IntegrationTests.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/IntegrationTests.yml b/.github/workflows/IntegrationTests.yml index c65c7ca0..95888585 100644 --- a/.github/workflows/IntegrationTests.yml +++ b/.github/workflows/IntegrationTests.yml @@ -36,6 +36,5 @@ jobs: - uses: julia-actions/julia-buildpkg@latest - name: Run integration tests uses: julia-actions/julia-runtest@latest - continue-on-error: true env: GROUP: Downstream From 5aeefe61b19db957958d472b14b0736e5b6f917f Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 15:08:06 +0100 Subject: [PATCH 24/36] added convenient constructor for DifferentiableDensityModel using Hamiltonians defaults --- src/abstractmcmc.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 25a47203..652b65ae 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -10,6 +10,13 @@ struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.Abstra ∂ℓπ∂θ::T∂logπ∂θ end +struct DummyMetric <: AbstractMetric end + +function DifferentiableDensityModel(ℓπ, m::Module=ForwardDiff) + h = Hamiltonian(DummyMetric(), ℓπ, m) + return DifferentiableDensityModel(h.ℓπ, h.∂ℓπ∂θ) +end + struct HMCState{ TTrans<:Transition, TMetric<:AbstractMetric, From 2b599137df766e61f306779933771a9c4702e450 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 15:08:40 +0100 Subject: [PATCH 25/36] fixed tests for AbstractMCMC interface --- test/Project.toml | 1 + test/abstractmcmc.jl | 17 ++++++++++++----- test/runtests.jl | 1 + 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 8927e4ba..20d589ee 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index ac79ebed..eec65ec7 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -10,20 +10,27 @@ include("common.jl") θ_init = randn(rng, 2) - metric = DiagEuclideanMetric(2) - h = Hamiltonian(metric, ℓπ_gdemo, ForwardDiff) + model = AdvancedHMC.DifferentiableDensityModel(ℓπ_gdemo, ForwardDiff) init_eps = Leapfrog(0.1) κ = NUTS(init_eps) + metric = DiagEuclideanMetric(2) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) - # samples, _ = sample(rng, h, κ, θ_init, n_samples, adaptor, n_adapts) samples = AbstractMCMC.sample( model, κ, metric, adaptor, n_samples; nadapts = n_adapts, init_params = θ_init - ) + ); - m_est = mean(map(invlink_gdemo, samples[1000:end])) + # Transform back to original space. + # NOTE: We're not correctly for the `logabsdetjac` here since, but + # we're only interested in the mean it doesn't matter. + for t in samples + t.z.θ .= invlink_gdemo(t.z.θ) + end + m_est = mean(samples[1000:end]) do t + t.z.θ + end @test m_est ≈ [49 / 24, 7 / 6] atol=RNDATOL end diff --git a/test/runtests.jl b/test/runtests.jl index c1f66b96..6525d5e7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -21,6 +21,7 @@ const GROUP = get(ENV, "GROUP", "All") "sampler-vec", "demo", "models", + "abstractmcmc" ] if CUDA.functional() From d164e22add0bbf91a62f9c44411693b7a4bc37b7 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 15:41:49 +0100 Subject: [PATCH 26/36] added a bunch of docstrings --- src/AdvancedHMC.jl | 1 + src/abstractmcmc.jl | 76 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index c42384d5..b7814c66 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -132,6 +132,7 @@ include("sampler.jl") export sample include("abstractmcmc.jl") +export DifferentiableDensityModel include("contrib/ad.jl") diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 652b65ae..eb50cd8e 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -1,35 +1,95 @@ +""" + HMCSampler + +A `AbstractMCMC.AbstractSampler` for kernels in AdvancedHMC.jl. + +# Fields + +$(FIELDS) + +# Notes + +Note that all the fields have the prefix `initial_` to indicate +that these will not necessarily correspond to the `kernel`, `metric`, +and `adaptor` after sampling. + +To access the updated fields use the resulting [`HMCState`](@ref). +""" struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler + "Initial [`AbstractMCMCKernel`](@ref)." initial_kernel::K + "Initial [`AbstractMetric`](@ref)." initial_metric::M + "Initial [`AbstractAdaptor`](@ref)." initial_adaptor::A end HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation()) +""" + DifferentiableDensityModel(ℓπ, ∂ℓπ∂θ) + DifferentiableDensityModel(ℓπ, m::Module) + +A `AbstractMCMC.AbstractMCMCModel` representing a differentiable log-density. + +If a module `m` is given as the second argument, then `m` is assumed to be an +automatic-differentiation package and this will be used to compute the gradients. + +Note that the module `m` must be imported before usage, e.g. +```julia +using Zygote: Zygote +model = DifferentiableDensityModel(ℓπ, Zygote) +``` +results in a `model` which will use Zygote.jl as its AD-backend. + +# Fields +$(FIELDS) +""" struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel + "Log-density. Maps `AbstractArray` to value of the log-density." ℓπ::Tlogπ + "Gradient of log-density. Returns a tuple of `ℓπ` and the gradient evaluated at the given point." ∂ℓπ∂θ::T∂logπ∂θ end struct DummyMetric <: AbstractMetric end - -function DifferentiableDensityModel(ℓπ, m::Module=ForwardDiff) +function DifferentiableDensityModel(ℓπ, m::Module) h = Hamiltonian(DummyMetric(), ℓπ, m) return DifferentiableDensityModel(h.ℓπ, h.∂ℓπ∂θ) end +""" + HMCState + +Represents the state of a [`HMCSampler`](@ref). + +# Fields + +$(FIELDS) + +""" struct HMCState{ TTrans<:Transition, TMetric<:AbstractMetric, TKernel<:AbstractMCMCKernel, TAdapt<:Adaptation.AbstractAdaptor } + "Index of current iteration." i::Int + "Current [`Transition`](@ref)." transition::TTrans + "Current [`AbstractMetric`](@ref), possibly adapted." metric::TMetric + "Current [`AbstractMCMCKernel`](@ref)." κ::TKernel + "Current [`AbstractAdaptor`](@ref)." adaptor::TAdapt end +""" + $(TYPEDSIGNATURES) + +A convenient wrapper around `AbstractMCMC.sample` avoiding explicit construction of [`HMCSampler`](@ref). +""" function AbstractMCMC.sample( model::DifferentiableDensityModel, kernel::AbstractMCMCKernel, @@ -181,9 +241,21 @@ end ################ ### Callback ### ################ +""" + HMCProgressCallback + +A callback to be used with AbstractMCMC.jl's interface, replicating the +logging behavior of the non-AbstractMCMC [`sample`](@ref). + +# Fields +$(FIELDS) +""" struct HMCProgressCallback{P} + "`Progress` meter from ProgressMeters.jl." pm::P + "Specifies whether or not to use display a progress bar." progress::Bool + "If `progress` is not specified and this is `true` some information will be logged upon completion of adaptation." verbose::Bool end From de4e33a755095d11bebb009e3f432af5fed9a1e8 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Wed, 14 Jul 2021 15:42:02 +0100 Subject: [PATCH 27/36] bumped minor version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index b933a3b2..cbbf6ef2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedHMC" uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" -version = "0.2.29" +version = "0.3.0" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" From fa17e41a606717cdf9060d6f2603efa0145af28c Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 00:35:43 +0100 Subject: [PATCH 28/36] increased number of samples used in abstractmcmc tests --- test/abstractmcmc.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index eec65ec7..de6017ef 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -5,7 +5,7 @@ include("common.jl") @testset "`gdemo`" begin rng = MersenneTwister(0) - n_samples = 5_000 + n_samples = 10_000 n_adapts = 1_000 θ_init = randn(rng, 2) @@ -28,7 +28,7 @@ include("common.jl") for t in samples t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples[1000:end]) do t + m_est = mean(samples[1000:10:end]) do t t.z.θ end From 8f8e64cecf98a5960964c5f846f34d4e5e706733 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 01:21:13 +0100 Subject: [PATCH 29/36] remove thinning from tests --- test/abstractmcmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index de6017ef..55aedb26 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -28,7 +28,7 @@ include("common.jl") for t in samples t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples[1000:10:end]) do t + m_est = mean(samples[1000:end]) do t t.z.θ end From 8d13ff5a7e25e0e3f7d31ca95fb79046e46b2bce Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 01:25:05 +0100 Subject: [PATCH 30/36] make initial Leapfrog step size smaller --- test/abstractmcmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 55aedb26..7c6d144a 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -11,7 +11,7 @@ include("common.jl") θ_init = randn(rng, 2) model = AdvancedHMC.DifferentiableDensityModel(ℓπ_gdemo, ForwardDiff) - init_eps = Leapfrog(0.1) + init_eps = Leapfrog(1e-3) κ = NUTS(init_eps) metric = DiagEuclideanMetric(2) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) From ccdc832f9969aa8535598b34aeb24b647b380e00 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 01:50:29 +0100 Subject: [PATCH 31/36] mistakenly removed AbstractMCMC as a test dep in previous commit --- test/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/test/Project.toml b/test/Project.toml index 8927e4ba..20d589ee 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" From 5a417bcb81634c8933e2205051aaa5eb1581f620 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 02:17:41 +0100 Subject: [PATCH 32/36] increase adaptation to see if it helps --- test/abstractmcmc.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 7c6d144a..ab188304 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -5,8 +5,8 @@ include("common.jl") @testset "`gdemo`" begin rng = MersenneTwister(0) - n_samples = 10_000 - n_adapts = 1_000 + n_samples = 5_000 + n_adapts = 5_000 θ_init = randn(rng, 2) From 32f6ff84b4763932d4ed499013e074578a72c1c5 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 02:18:46 +0100 Subject: [PATCH 33/36] ensure we drop the adaptation samples in the test --- test/abstractmcmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index ab188304..7fa0f645 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -28,7 +28,7 @@ include("common.jl") for t in samples t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples[1000:end]) do t + m_est = mean(samples[n_adapts + 1:end]) do t t.z.θ end From 13a07e8fc638123ee563dfbc3424f519c9285c55 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 03:53:40 +0100 Subject: [PATCH 34/36] made a mistake apparently --- test/abstractmcmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 7fa0f645..170483dc 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -28,7 +28,7 @@ include("common.jl") for t in samples t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples[n_adapts + 1:end]) do t + m_est = mean(samples) do t t.z.θ end From b82a965da4c7ea6f718e28ce2ed6211321eb9167 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 04:58:12 +0100 Subject: [PATCH 35/36] think I finally fixed the tests --- test/abstractmcmc.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 170483dc..84153074 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -17,7 +17,7 @@ include("common.jl") adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) samples = AbstractMCMC.sample( - model, κ, metric, adaptor, n_samples; + model, κ, metric, adaptor, n_adapts + n_samples; nadapts = n_adapts, init_params = θ_init ); @@ -28,7 +28,7 @@ include("common.jl") for t in samples t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples) do t + m_est = mean(samples[n_adapts + 1:end]) do t t.z.θ end From 38408cf7ddeba46dae09bd478572854347e97a9b Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 15 Jul 2021 13:21:59 +0100 Subject: [PATCH 36/36] disable progress in test --- test/abstractmcmc.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 84153074..61ae4f0e 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -19,11 +19,13 @@ include("common.jl") samples = AbstractMCMC.sample( model, κ, metric, adaptor, n_adapts + n_samples; nadapts = n_adapts, - init_params = θ_init + init_params = θ_init, + progress=false, + verbose=false ); # Transform back to original space. - # NOTE: We're not correctly for the `logabsdetjac` here since, but + # NOTE: We're not correcting for the `logabsdetjac` here since, but # we're only interested in the mean it doesn't matter. for t in samples t.z.θ .= invlink_gdemo(t.z.θ)