diff --git a/Project.toml b/Project.toml index ca096669..48205f22 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedHMC" uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" -version = "0.4.6" +version = "0.5.0" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" diff --git a/README.md b/README.md index a9f8cf46..dc3d31b5 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,10 @@ If you are interested in using AdvancedHMC.jl through a probabilistic programmin - We presented a poster for AdvancedHMC.jl at [StanCon 2019](https://mc-stan.org/events/stancon2019Cambridge/) in Cambridge, UK. ([pdf](https://github.com/TuringLang/AdvancedHMC.jl/files/3730367/StanCon-AHMC.pdf)) **API CHANGES** +- [v0.5.0] **Breaking!** Convinience constructors for common samplers changed to: + - `HMC(init_ϵ::Float64=init_ϵ, n_leapfrog::Int=n_leapfrog)` + - `NUTS(n_adapts::Int=n_adapts, δ::Float64=δ)` + - `HMCDA(n_adapts::Int=n_adapts, δ::Float64=δ, λ::Float64=λ)` - [v0.2.22] Three functions are renamed. - `Preconditioner(metric::AbstractMetric)` -> `MassMatrixAdaptor(metric)` and - `NesterovDualAveraging(δ, integrator::AbstractIntegrator)` -> `StepSizeAdaptor(δ, integrator)` diff --git a/docs/src/api.md b/docs/src/api.md index 7b9b9a84..8754eb2f 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -8,15 +8,14 @@ Documentation for AdvancedHMC.jl ## Structs ```@docs ClassicNoUTurn +HMCSampler +HMC +NUTS +HMCDA ``` ## Functions ```@docs sample -``` - -## Index - -```@index ``` \ No newline at end of file diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 28bc440f..40b409b9 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -66,28 +66,6 @@ export Trajectory, # Useful defaults -struct NUTS{TS,TC} end - -""" -$(SIGNATURES) - -Convenient constructor for the no-U-turn sampler (NUTS). -This falls back to `HMCKernel(Trajectory{TS}(int, TC(args...; kwargs...)))` where - -- `TS<:Union{MultinomialTS, SliceTS}` is the type for trajectory sampler -- `TC<:Union{ClassicNoUTurn, GeneralisedNoUTurn, StrictGeneralisedNoUTurn}` is the type for termination criterion. - -See [`ClassicNoUTurn`](@ref), [`GeneralisedNoUTurn`](@ref) and [`StrictGeneralisedNoUTurn`](@ref) for details in parameters. -""" -NUTS{TS,TC}(int::AbstractIntegrator, args...; kwargs...) where {TS,TC} = - HMCKernel(Trajectory{TS}(int, TC(args...; kwargs...))) -NUTS(int::AbstractIntegrator, args...; kwargs...) = - HMCKernel(Trajectory{MultinomialTS}(int, GeneralisedNoUTurn(args...; kwargs...))) -NUTS(ϵ::AbstractScalarOrVec{<:Real}) = - HMCKernel(Trajectory{MultinomialTS}(Leapfrog(ϵ), GeneralisedNoUTurn())) - -export NUTS - # Deprecations for trajectory.jl abstract type AbstractTrajectory end @@ -103,25 +81,14 @@ struct StaticTrajectory{TS} end Trajectory{EndPointTS}(Leapfrog(ϵ), FixedNSteps(L)), ) -struct HMCDA{TS} end -@deprecate HMCDA{TS}(int::AbstractIntegrator, λ) where {TS} HMCKernel( - Trajectory{TS}(int, 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 -export StaticTrajectory, HMCDA, find_good_eps +export StaticTrajectory, find_good_eps include("adaptation/Adaptation.jl") using .Adaptation import .Adaptation: - StepSizeAdaptor, MassMatrixAdaptor, StanHMCAdaptor, NesterovDualAveraging + StepSizeAdaptor, MassMatrixAdaptor, StanHMCAdaptor, NesterovDualAveraging, NoAdaptation # Helpers for initializing adaptors via AHMC structs @@ -161,13 +128,17 @@ export StepSizeAdaptor, WelfordVar, WelfordCov, NaiveHMCAdaptor, - StanHMCAdaptor + StanHMCAdaptor, + NoAdaptation include("diagnosis.jl") include("sampler.jl") export sample +include("constructors.jl") +export HMCSampler, HMC, NUTS, HMCDA + include("abstractmcmc.jl") ## Without explicit AD backend diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index e491b53b..31bdf999 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -1,30 +1,3 @@ -""" - 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()) - """ HMCState @@ -58,38 +31,17 @@ end A convenient wrapper around `AbstractMCMC.sample` avoiding explicit construction of [`HMCSampler`](@ref). """ -function AbstractMCMC.sample( - model::LogDensityModel, - 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::LogDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, + sampler::AbstractHMCSampler, 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 @@ -107,34 +59,10 @@ function AbstractMCMC.sample( ) end -function AbstractMCMC.sample( - model::LogDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, - parallel::AbstractMCMC.AbstractMCMCEnsemble, - 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::LogDensityModel, - kernel::AbstractMCMCKernel, - metric::AbstractMetric, - adaptor::AbstractAdaptor, + sampler::AbstractHMCSampler, parallel::AbstractMCMC.AbstractMCMCEnsemble, N::Integer, nchains::Integer; @@ -143,7 +71,7 @@ function AbstractMCMC.sample( 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 @@ -166,27 +94,36 @@ end function AbstractMCMC.step( rng::AbstractRNG, model::LogDensityModel, - spl::HMCSampler; + spl::AbstractHMCSampler; init_params = nothing, kwargs..., ) - metric = spl.initial_metric - κ = spl.initial_kernel - adaptor = spl.initial_adaptor + # Unpack model + logdensity = model.logdensity - if init_params === nothing - init_params = randn(rng, size(metric, 1)) - end + # Define metric + metric = make_metric(spl, logdensity) # Construct the hamiltonian using the initial metric hamiltonian = Hamiltonian(metric, model) + # Define integration algorithm + # Find good eps if not provided one + init_params = make_init_params(spl, logdensity, init_params) + ϵ = make_step_size(rng, spl, hamiltonian, init_params) + integrator = make_integrator(spl, ϵ) + + # Make kernel + κ = make_kernel(spl, integrator) + + # Make adaptor + adaptor = make_adaptor(spl, metric, integrator) + # 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) - + state = HMCState(0, t, metric, κ, adaptor) # Take actual first step. return AbstractMCMC.step(rng, model, spl, state; kwargs...) end @@ -194,9 +131,8 @@ end function AbstractMCMC.step( rng::AbstractRNG, model::LogDensityModel, - spl::HMCSampler, + spl::AbstractHMCSampler, state::HMCState; - nadapts::Int = 0, kwargs..., ) # Compute transition. @@ -214,7 +150,8 @@ function AbstractMCMC.step( # Adapt h and spl. tstat = stat(t) - h, κ, isadapted = adapt!(h, κ, adaptor, i, nadapts, t.z.θ, tstat.acceptance_rate) + n_adapts = get_nadapts(spl) + h, κ, isadapted = adapt!(h, κ, adaptor, i, n_adapts, t.z.θ, tstat.acceptance_rate) tstat = merge(tstat, (is_adapt = isadapted,)) # Compute next transition and state. @@ -224,7 +161,6 @@ function AbstractMCMC.step( return Transition(t.z, tstat), newstate end - ################ ### Callback ### ################ @@ -303,3 +239,120 @@ function (cb::HMCProgressCallback)(rng, model, spl, t, state, i; nadapts = 0, kw @info "Finished $nadapts adapation steps" adaptor κ.τ.integrator metric end end + +############# +### Utils ### +############# + +function get_type_of_spl(::AbstractHMCSampler{T}) where {T<:Real} + return T +end + +######### + +function make_init_params(spl::AbstractHMCSampler, logdensity, init_params) + T = get_type_of_spl(spl) + if init_params == nothing + d = LogDensityProblems.dimension(logdensity) + init_params = randn(rng, d) + end + return T.(init_params) +end + +######### + +function make_step_size( + rng::Random.AbstractRNG, + spl::AbstractHMCSampler, + hamiltonian::Hamiltonian, + init_params, +) + ϵ = spl.init_ϵ + if iszero(ϵ) + ϵ = find_good_stepsize(rng, hamiltonian, init_params) + T = get_type_of_spl(spl) + ϵ = T(ϵ) + @info string("Found initial step size ", ϵ) + end + return ϵ +end + +function make_step_size( + rng::Random.AbstractRNG, + spl::HMCSampler, + hamiltonian::Hamiltonian, + init_params, +) + return spl.κ.τ.integrator.ϵ +end + +make_integrator(spl::HMCSampler, ϵ::Real) = spl.κ.τ.integrator +make_integrator(spl::AbstractHMCSampler, ϵ::Real) = make_integrator(spl.integrator, ϵ) +make_integrator(i::AbstractIntegrator, ϵ::Real) = i +make_integrator(i::Type{<:AbstractIntegrator}, ϵ::Real) = i +make_integrator(i::Symbol, ϵ::Real) = make_integrator(Val(i), ϵ) +make_integrator(i...) = error("Integrator $(typeof(i)) not supported.") +make_integrator(i::Val{:leapfrog}, ϵ::Real) = Leapfrog(ϵ) +make_integrator(i::Val{:jitteredleapfrog}, ϵ::Real) = JitteredLeapfrog(ϵ) +make_integrator(i::Val{:temperedleapfrog}, ϵ::Real) = TemperedLeapfrog(ϵ) + +######### + +make_metric(i...) = error("Metric $(typeof(i)) not supported.") +make_metric(i::Symbol, T::Type, d::Int) = make_metric(Val(i), T, d) +make_metric(i::AbstractMetric, T::Type, d::Int) = i +make_metric(i::Type{AbstractMetric}, T::Type, d::Int) = i +make_metric(i::Val{:diagonal}, T::Type, d::Int) = DiagEuclideanMetric(T, d) +make_metric(i::Val{:unit}, T::Type, d::Int) = UnitEuclideanMetric(T, d) +make_metric(i::Val{:dense}, T::Type, d::Int) = DenseEuclideanMetric(T, d) + +function make_metric(spl::AbstractHMCSampler, logdensity) + d = LogDensityProblems.dimension(logdensity) + T = get_type_of_spl(spl) + return make_metric(spl.metric, T, d) +end + +######### + +function make_adaptor( + spl::Union{NUTS,HMCDA}, + metric::AbstractMetric, + integrator::AbstractIntegrator, +) + return StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(spl.δ, integrator)) +end + +function make_adaptor(spl::HMC, metric::AbstractMetric, integrator::AbstractIntegrator) + return NoAdaptation() +end + +function make_adaptor( + spl::HMCSampler, + metric::AbstractMetric, + integrator::AbstractIntegrator, +) + return spl.adaptor +end + +######### + +get_nadapts(spl::Union{HMCSampler,NUTS,HMCDA}) = spl.n_adapts +get_nadapts(spl::HMC) = 0 + +######### + +function make_kernel(spl::NUTS, integrator::AbstractIntegrator) + return HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) +end + +function make_kernel(spl::HMC, integrator::AbstractIntegrator) + return HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(spl.n_leapfrog))) +end + +function make_kernel(spl::HMCDA, integrator::AbstractIntegrator) + return HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(spl.λ))) +end + +function make_kernel(spl::HMCSampler, integrator::AbstractIntegrator) + return spl.κ +end diff --git a/src/constructors.jl b/src/constructors.jl new file mode 100644 index 00000000..6ff05b60 --- /dev/null +++ b/src/constructors.jl @@ -0,0 +1,166 @@ +abstract type AbstractHMCSampler{T<:Real} <: AbstractMCMC.AbstractSampler end + +############## +### Custom ### +############## + +""" + 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{T<:Real} <: AbstractHMCSampler{T} + "[`AbstractMCMCKernel`](@ref)." + κ::AbstractMCMCKernel + "[`AbstractMetric`](@ref)." + metric::AbstractMetric + "[`AbstractAdaptor`](@ref)." + adaptor::AbstractAdaptor + "Adaptation steps if any" + n_adapts::Int +end + +function HMCSampler(κ, metric, adaptor; n_adapts = 0) + T = collect(typeof(metric).parameters)[1] + return HMCSampler{T}(κ, metric, adaptor, n_adapts) +end + +############ +### NUTS ### +############ +""" + NUTS(n_adapts::Int, δ::Real; max_depth::Int=10, Δ_max::Real=1000, init_ϵ::Real=0) + +No-U-Turn Sampler (NUTS) sampler. + +# Fields + +$(FIELDS) + +# Usage: + +```julia +NUTS(n_adapts=1000, δ=0.65) # Use 1000 adaption steps, and target accept ratio 0.65. +``` +""" +struct NUTS{T<:Real} <: AbstractHMCSampler{T} + "Number of adaptation steps." + n_adapts::Int + "Target acceptance rate for dual averaging." + δ::T + "Maximum doubling tree depth." + max_depth::Int + "Maximum divergence during doubling tree." + Δ_max::T + "Initial step size; 0 means it is automatically chosen." + init_ϵ::T + "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" + integrator::Union{Symbol,AbstractIntegrator} + "Choice of metric, specified either using a `Symbol` or `AbstractMetric`" + metric::Union{Symbol,AbstractMetric} +end + +function NUTS( + n_adapts, + δ; + max_depth = 10, + Δ_max = 1000.0, + init_ϵ = 0.0, + integrator = :leapfrog, + metric = :diagonal, +) + T = typeof(δ) + return NUTS(n_adapts, δ, max_depth, T(Δ_max), T(init_ϵ), integrator, metric) +end + +########### +### HMC ### +########### +""" + HMC(ϵ::Real, n_leapfrog::Int) + +Hamiltonian Monte Carlo sampler with static trajectory. + +# Fields + +$(FIELDS) + +# Usage: + +```julia +HMC(init_ϵ=0.05, n_leapfrog=10) +``` +""" +struct HMC{T<:Real} <: AbstractHMCSampler{T} + "Initial step size; 0 means automatically searching using a heuristic procedure." + init_ϵ::T + "Number of leapfrog steps." + n_leapfrog::Int + "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" + integrator::Union{Symbol,AbstractIntegrator} + "Choice of metric, specified either using a `Symbol` or `AbstractMetric`" + metric::Union{Symbol,AbstractMetric} +end + +function HMC(init_ϵ, n_leapfrog; integrator = :leapfrog, metric = :diagonal) + return HMC(init_ϵ, n_leapfrog, integrator, metric) +end + +############# +### HMCDA ### +############# +""" + HMCDA(n_adapts::Int, δ::Real, λ::Real; ϵ::Real=0) + +Hamiltonian Monte Carlo sampler with Dual Averaging algorithm. + +# Fields + +$(FIELDS) + +# Usage: + +```julia +HMCDA(n_adapts=200, δ=0.65, λ=0.3) +``` + +For more information, please view the following paper ([arXiv link](https://arxiv.org/abs/1111.4246)): + +- Hoffman, Matthew D., and Andrew Gelman. "The No-U-turn sampler: adaptively + setting path lengths in Hamiltonian Monte Carlo." Journal of Machine Learning + Research 15, no. 1 (2014): 1593-1623. +""" +struct HMCDA{T<:Real} <: AbstractHMCSampler{T} + "`Number of adaptation steps." + n_adapts::Int + "Target acceptance rate for dual averaging." + δ::T + "Target leapfrog length." + λ::T + "Initial step size; 0 means automatically searching using a heuristic procedure." + init_ϵ::T + "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" + integrator::Union{Symbol,AbstractIntegrator} + "Choice of metric, specified either using a `Symbol` or `AbstractMetric`" + metric::Union{Symbol,AbstractMetric} +end + +function HMCDA(n_adapts, δ, λ; init_ϵ = 0.0, integrator = :leapfrog, metric = :diagonal) + if typeof(δ) != typeof(λ) + @warn "typeof(δ) != typeof(λ) --> using typeof(δ)" + end + T = typeof(δ) + return HMCDA(n_adapts, δ, T(λ), T(init_ϵ), integrator, metric) +end diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index f14dbf2f..c0ea04e0 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -4,27 +4,74 @@ include("common.jl") @testset "AbstractMCMC w/ gdemo" begin rng = MersenneTwister(0) - n_samples = 5_000 n_adapts = 5_000 - θ_init = randn(rng, 2) + nuts = NUTS(n_adapts, 0.8) + hmc = HMC(0.05, 100) + hmcda = HMCDA(n_adapts, 0.8, 0.1) + + integrator = Leapfrog(1e-3) + κ = AdvancedHMC.make_kernel(nuts, integrator) + metric = DiagEuclideanMetric(2) + adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) + custom = HMCSampler(κ, metric, adaptor) + model = AdvancedHMC.LogDensityModel( LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ_gdemo), ) - init_eps = Leapfrog(1e-3) - κ = NUTS(init_eps) - metric = DiagEuclideanMetric(2) - adaptor = - StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) - samples = AbstractMCMC.sample( + samples_nuts = AbstractMCMC.sample( + rng, + model, + nuts, + n_adapts + n_samples; + nadapts = n_adapts, + init_params = θ_init, + progress = false, + verbose = false, + ) + + # Transform back to original space. + # 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_nuts + t.z.θ .= invlink_gdemo(t.z.θ) + end + m_est_nuts = mean(samples_nuts[n_adapts+1:end]) do t + t.z.θ + end + + @test m_est_nuts ≈ [49 / 24, 7 / 6] atol = RNDATOL + + samples_hmc = AbstractMCMC.sample( + rng, + model, + hmc, + n_adapts + n_samples; + nadapts = n_adapts, + init_params = θ_init, + progress = false, + verbose = false, + ) + + # Transform back to original space. + # 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_hmc + t.z.θ .= invlink_gdemo(t.z.θ) + end + m_est_hmc = mean(samples_hmc) do t + t.z.θ + end + + @test m_est_hmc ≈ [49 / 24, 7 / 6] atol = RNDATOL + + samples_custom = AbstractMCMC.sample( rng, model, - κ, - metric, - adaptor, + custom, n_adapts + n_samples; nadapts = n_adapts, init_params = θ_init, @@ -35,14 +82,14 @@ include("common.jl") # Transform back to original space. # 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 + for t in samples_custom t.z.θ .= invlink_gdemo(t.z.θ) end - m_est = mean(samples[n_adapts+1:end]) do t + m_est_custom = mean(samples_custom[n_adapts+1:end]) do t t.z.θ end - @test m_est ≈ [49 / 24, 7 / 6] atol = RNDATOL + @test m_est_custom ≈ [49 / 24, 7 / 6] atol = RNDATOL # Test that using the same AbstractRNG results in the same chain rng1 = MersenneTwister(42) @@ -50,22 +97,20 @@ include("common.jl") samples1 = AbstractMCMC.sample( rng1, model, - κ, - metric, - adaptor, + custom, 10; nadapts = 0, + init_params = θ_init, progress = false, verbose = false, ) samples2 = AbstractMCMC.sample( rng2, model, - κ, - metric, - adaptor, + custom, 10; nadapts = 0, + init_params = θ_init, progress = false, verbose = false, ) diff --git a/test/adaptation.jl b/test/adaptation.jl index 3e5422ad..0c968873 100644 --- a/test/adaptation.jl +++ b/test/adaptation.jl @@ -5,13 +5,15 @@ using AdvancedHMC.Adaptation: function runnuts(ℓπ, metric; n_samples = 3_000) D = size(metric, 1) n_adapts = 1_500 - θ_init = rand(D) + rng = MersenneTwister(0) + nuts = NUTS(n_adapts, 0.8) h = Hamiltonian(metric, ℓπ, ForwardDiff) - κ = NUTS(Leapfrog(find_good_stepsize(h, θ_init))) - adaptor = - StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) + step_size = AdvancedHMC.make_step_size(rng, nuts, h, θ_init) + integrator = AdvancedHMC.make_integrator(nuts, step_size) + κ = AdvancedHMC.make_kernel(nuts, integrator) + adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) samples, stats = sample(h, κ, θ_init, n_samples, adaptor, n_adapts; verbose = false) return (samples = samples, stats = stats, adaptor = adaptor) end diff --git a/test/constructors.jl b/test/constructors.jl new file mode 100644 index 00000000..88e69750 --- /dev/null +++ b/test/constructors.jl @@ -0,0 +1,92 @@ +using AdvancedHMC, AbstractMCMC, Random +include("common.jl") + +# Initalize samplers +nuts = NUTS(1000, 0.8) +nuts_32 = NUTS(1000, 0.8f0) +hmc = HMC(0.1, 25) +hmcda = HMCDA(1000, 0.8, 1.0) +hmcda_32 = HMCDA(1000, 0.8f0, 1.0) + +integrator = Leapfrog(1e-3) +kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) +metric = DiagEuclideanMetric(2) +adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) +custom = HMCSampler(kernel, metric, adaptor) + +# Check that everything is initalized correctly +@testset "Constructors" begin + # Types + @test typeof(nuts) == NUTS{Float64} + @test typeof(nuts_32) == NUTS{Float32} + @test typeof(hmc) == HMC{Float64} + @test typeof(hmcda) == HMCDA{Float64} + @test typeof(nuts) <: AdvancedHMC.AbstractHMCSampler + @test typeof(nuts) <: AbstractMCMC.AbstractSampler + + # NUTS + @test nuts.n_adapts == 1000 + @test nuts.δ == 0.8 + @test nuts.max_depth == 10 + @test nuts.Δ_max == 1000.0 + @test nuts.init_ϵ == 0.0 + @test nuts.integrator == :leapfrog + @test nuts.metric == :diagonal + + # NUTS Float32 + @test nuts_32.n_adapts == 1000 + @test nuts_32.δ == 0.8f0 + @test nuts_32.max_depth == 10 + @test nuts_32.Δ_max == 1000.0f0 + @test nuts_32.init_ϵ == 0.0f0 + + # HMC + @test hmc.n_leapfrog == 25 + @test hmc.init_ϵ == 0.1 + @test hmc.integrator == :leapfrog + @test hmc.metric == :diagonal + + # HMCDA + @test hmcda.n_adapts == 1000 + @test hmcda.δ == 0.8 + @test hmcda.λ == 1.0 + @test hmcda.init_ϵ == 0.0 + @test hmcda.integrator == :leapfrog + @test hmcda.metric == :diagonal + + # HMCDA Float32 + @test hmcda_32.n_adapts == 1000 + @test hmcda_32.δ == 0.8f0 + @test hmcda_32.λ == 1.0f0 + @test hmcda_32.init_ϵ == 0.0f0 +end + +@testset "First step" begin + rng = MersenneTwister(0) + θ_init = randn(rng, 2) + logdensitymodel = AbstractMCMC.LogDensityModel(ℓπ_gdemo) + _, nuts_state = AbstractMCMC.step(rng, logdensitymodel, nuts; init_params = θ_init) + _, hmc_state = AbstractMCMC.step(rng, logdensitymodel, hmc; init_params = θ_init) + _, nuts_32_state = + AbstractMCMC.step(rng, logdensitymodel, nuts_32; init_params = θ_init) + _, custom_state = AbstractMCMC.step(rng, logdensitymodel, custom; init_params = θ_init) + + # Metric + @test typeof(nuts_state.metric) == DiagEuclideanMetric{Float64,Vector{Float64}} + @test typeof(nuts_32_state.metric) == DiagEuclideanMetric{Float32,Vector{Float32}} + @test custom_state.metric == metric + + # Integrator + @test typeof(nuts_state.κ.τ.integrator) == Leapfrog{Float64} + @test typeof(nuts_32_state.κ.τ.integrator) == Leapfrog{Float32} + @test custom_state.κ.τ.integrator == integrator + + # Kernel + @test nuts_state.κ == AdvancedHMC.make_kernel(nuts, nuts_state.κ.τ.integrator) + @test custom_state.κ == kernel + + # Adaptor + @test typeof(nuts_state.adaptor) <: StanHMCAdaptor + @test hmc_state.adaptor == NoAdaptation() + @test custom_state.adaptor == adaptor +end diff --git a/test/cuda.jl b/test/cuda.jl index 5610e0f9..fb9655a9 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -24,6 +24,8 @@ using CUDA samples, stats = sample(hamiltonian, proposal, θ₀, n_samples) end +#= +Broken! See https://github.com/JuliaTesting/ReTest.jl/issues/50 @testset "PhasePoint GPU" begin for T in [Float32, Float64] init_z1() = PhasePoint( @@ -55,3 +57,4 @@ end @test z1.ℓκ.value == z2.ℓκ.value end end +=# diff --git a/test/demo.jl b/test/demo.jl index 0139336e..068b82dc 100644 --- a/test/demo.jl +++ b/test/demo.jl @@ -33,7 +33,7 @@ using LinearAlgebra # - multinomial sampling scheme, # - generalised No-U-Turn criteria, and # - windowed adaption for step-size and diagonal mass matrix - proposal = NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator) + proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # Run the sampler to draw samples from the specified Gaussian, where @@ -84,7 +84,7 @@ end integrator = Leapfrog(initial_ϵ) # Define an HMC sampler, with the following components - proposal = NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator) + proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # -- run sampler diff --git a/test/mcmcchains.jl b/test/mcmcchains.jl index 3547f2d0..1c896884 100644 --- a/test/mcmcchains.jl +++ b/test/mcmcchains.jl @@ -13,18 +13,16 @@ include("common.jl") model = AdvancedHMC.LogDensityModel( LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ_gdemo), ) - init_eps = Leapfrog(1e-3) - κ = NUTS(init_eps) + integrator = Leapfrog(1e-3) + kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) metric = DiagEuclideanMetric(2) - adaptor = - StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) + adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) + sampler = HMCSampler(kernel, metric, adaptor) samples = AbstractMCMC.sample( rng, model, - κ, - metric, - adaptor, + sampler, n_adapts + n_samples; nadapts = n_adapts, init_params = θ_init, diff --git a/test/models.jl b/test/models.jl index de24e4ce..08be9197 100644 --- a/test/models.jl +++ b/test/models.jl @@ -14,10 +14,10 @@ include("common.jl") metric = DiagEuclideanMetric(2) h = Hamiltonian(metric, ℓπ_gdemo, ForwardDiff) - init_eps = Leapfrog(0.1) - κ = NUTS(init_eps) + integrator = Leapfrog(0.1) + κ = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) adaptor = - StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) + StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) samples, _ = sample( rng, diff --git a/test/runtests.jl b/test/runtests.jl index 0b243b01..0a58c56a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -25,6 +25,7 @@ if GROUP == "All" || GROUP == "AdvancedHMC" include("models.jl") include("abstractmcmc.jl") include("mcmcchains.jl") + include("constructors.jl") if CUDA.functional() include("cuda.jl") diff --git a/test/sampler.jl b/test/sampler.jl index 522d598e..177cafde 100644 --- a/test/sampler.jl +++ b/test/sampler.jl @@ -159,11 +159,12 @@ end end end @testset "drop_warmup" begin + nuts = NUTS(n_adapts, 0.8) metric = DiagEuclideanMetric(D) h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ) - κ = NUTS(Leapfrog(ϵ)) - adaptor = - StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, κ.τ.integrator)) + integrator = Leapfrog(ϵ) + κ = AdvancedHMC.make_kernel(nuts, integrator) + adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) samples, stats = sample( h, κ,