From e7489a60ad81047ff3d46b7aa4cb1f43b44c5b6c Mon Sep 17 00:00:00 2001 From: jaimerz Date: Fri, 21 Jul 2023 14:00:43 +0100 Subject: [PATCH 01/29] bug + docs --- src/abstractmcmc.jl | 4 ++-- test/constructors.jl | 23 +++++++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 31bdf999..55f3ea26 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -293,8 +293,8 @@ 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_integrator(i::Val{:jitteredleapfrog}, ϵ::Real) = JitteredLeapfrog(ϵ, 0.1ϵ) +make_integrator(i::Val{:temperedleapfrog}, ϵ::Real) = TemperedLeapfrog(ϵ, 1.0) ######### diff --git a/test/constructors.jl b/test/constructors.jl index 88e69750..e827362a 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -14,6 +14,14 @@ metric = DiagEuclideanMetric(2) adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) custom = HMCSampler(kernel, metric, adaptor) +nuts_metric1 = NUTS(1000, 0.8; metric=:unit) +nuts_metric2 = NUTS(1000, 0.8; metric=:dense) +hmc_metric1 = HMC(0.1, 25; metric=metric) + +nuts_integrator1 = NUTS(1000, 0.8, integrator=:jitteredleapfrog) +nuts_integrator2 = NUTS(1000, 0.8, integrator=:temperedleapfrog) +hmc_integrator1 = HMC(0.1, 25, integrator=integrator) + # Check that everything is initalized correctly @testset "Constructors" begin # Types @@ -59,6 +67,7 @@ custom = HMCSampler(kernel, metric, adaptor) @test hmcda_32.δ == 0.8f0 @test hmcda_32.λ == 1.0f0 @test hmcda_32.init_ϵ == 0.0f0 + end @testset "First step" begin @@ -71,14 +80,28 @@ end AbstractMCMC.step(rng, logdensitymodel, nuts_32; init_params = θ_init) _, custom_state = AbstractMCMC.step(rng, logdensitymodel, custom; init_params = θ_init) + _, nuts_metric1_state = AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; init_params = θ_init) + _, nuts_metric2_state = AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; init_params = θ_init) + _, hmc_metric1_state = AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; init_params = θ_init) + + _, nuts_integrator1_state = AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; init_params = θ_init) + _, nuts_integrator2_state = AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; init_params = θ_init) + _, hmc_integrator1_state = AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; init_params = θ_init) + # Metric @test typeof(nuts_state.metric) == DiagEuclideanMetric{Float64,Vector{Float64}} @test typeof(nuts_32_state.metric) == DiagEuclideanMetric{Float32,Vector{Float32}} + @test typeof(nuts_metric1_state.metric) <: UnitEuclideanMetric + @test typeof(nuts_metric2_state.metric) <: DenseEuclideanMetric + @test hmc_metric1_state.metric == metric @test custom_state.metric == metric # Integrator @test typeof(nuts_state.κ.τ.integrator) == Leapfrog{Float64} @test typeof(nuts_32_state.κ.τ.integrator) == Leapfrog{Float32} + @test typeof(nuts_integrator1_state.κ.τ.integrator) <: JitteredLeapfrog + @test typeof(nuts_integrator2_state.κ.τ.integrator) <: TemperedLeapfrog + @test hmc_integrator1_state.κ.τ.integrator == integrator @test custom_state.κ.τ.integrator == integrator # Kernel From bca300c5c39686d3706ce694f5701e87a1d961ff Mon Sep 17 00:00:00 2001 From: jaimerz Date: Fri, 21 Jul 2023 14:02:52 +0100 Subject: [PATCH 02/29] read me --- README.md | 128 ++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 120 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index dc3d31b5..53a15ec8 100644 --- a/README.md +++ b/README.md @@ -47,10 +47,10 @@ In this section we demonstrate a minimal example of sampling from a multivariate Suppose ${\bf x}$ and ${\bf v}$ are the position and velocity of an individual particle respectively; $i$ and $i+1$ are the indices for time values $t_i$ and $t_{i+1}$ respectively; $dt = t_{i+1} - t_i$ is the time step size (constant and regularly spaced intervals); and ${\bf a}$ is the acceleration induced on a particle by the forces of all other particles. Furthermore, suppose positions are defined at times $t_i, t_{i+1}, t_{i+2}, \dots $, spaced at constant intervals $dt$, the velocities are defined at halfway times in between, denoted by $t_{i-1/2}, t_{i+1/2}, t_{i+3/2}, \dots $, where $t_{i+1} - t_{i + 1/2} = t_{i + 1/2} - t_i = dt / 2$, and the accelerations ${\bf a}$ are defined only on integer times, just like the positions. Then the leapfrog integration scheme is given as: $x_{i} = x_{i-1} + v_{i-1/2} dt; \quad v_{i+1/2} = v_{i-1/2} + a_i dt$. For available integrators refer Integrator. -- **Proposal for trajectories (static or dynamic)**: Different types of proposals can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle. +- **kernel for trajectories (static or dynamic)**: Different types of kernels can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle.
- More about the proposals - In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. Whereas, during the second step, a Metropolis update is performed, using Hamiltonian dynamics to provide a new state. For available proposals refer Proposal. + More about the kernels + In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. Whereas, during the second step, a Metropolis update is performed, using Hamiltonian dynamics to provide a new state. For available kernels refer kernel.
```julia @@ -85,13 +85,13 @@ integrator = Leapfrog(initial_ϵ) # - multinomial sampling scheme, # - generalised No-U-Turn criteria, and # - windowed adaption for step-size and diagonal mass matrix -proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) +kernel = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # Run the sampler to draw samples from the specified Gaussian, where # - `samples` will store the samples # - `stats` will store diagnostic statistics for each sample -samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true) +samples, stats = sample(hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; progress=true) ``` ### Parallel sampling @@ -114,11 +114,123 @@ chains = Vector{Any}(undef, nchains) # The `samples` from each parallel chain is stored in the `chains` vector # Adjust the `verbose` flag as per need Threads.@threads for i in 1:nchains - samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; verbose=false) + samples, stats = sample(hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; verbose=false) chains[i] = samples end ``` +### Using the `AbstractMCMC` interface + +Users can also make use of the `AbstractMCMC` interface to sample which employs the same API as other popular Bayesian inference libraries in Julia such as `Turing`. +In order to show how this is done let us start from our previous example where we defined a `LogTargetDensity`, ℓπ. + +```julia +# Wrap the previous LogTargetDensity as LogDensityModel +# where ℓπ::LogTargetDensity +model = AdvancedHMC.LogDensityModel(LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ)) + +# Wrap the previous sampler as a HMCSampler <: AbstractMCMC.AbstractSampler +D = 10; initial_θ = rand(D) +n_samples, n_adapts, δ = 1_000, 2_000, 0.8 +sampler = HMCSampler(kernel, metric, adaptor) + +# Now just sample +samples = AbstractMCMC.sample( + model, + sampler, + n_adapts + n_samples; + nadapts = n_adapts, + init_params = initial_θ, + ) +``` + +### Covenience Constructors + +In the previous examples we built the sampler by manually specifying the integrator, metric, kernel and adaptor to build our own sampler. However, in many cases users might want to simply initialize a standard NUTS sampler. In such cases having to manually define each of these aspects is tedious and error prone. For these reasons `AdvancedHMC` also provides users with a series of convenience constructors for standard samplers. We will now show how to use them. + +- HMC: + ```julia + # HMC Sampler + # step size, number of leapfrog steps + ϵ, n_leapfrogs = 0.1, 0.25 + hmc = HMC(ϵ, n_leapfrogs) + ``` + + Equivalent to: + + ```julia + metric = DiagEuclideanMetric(D) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = Leapfrog(initial_ϵ) + kernel = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) + adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(δ, integrator)) + nuts = HMCSampler(kernel, metric, adaptor) + ``` + +- NUTS: + ```julia + # NUTS Sampler + # adaptation steps, target acceptance probability, + n_adapt, δ = 1000, 0.8 + nuts = NUTS(n_adapt, δ) + ``` + + Equivalent to: + + ```julia + metric = DiagEuclideanMetric(D) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = Leapfrog(initial_ϵ) + kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(n_leapfrog))) + adaptor = NoAdaptation() + hmc = HMCSampler(kernel, metric, adaptor) + ``` + + +- HMCDA: + ```julia + #HMCDA (dual averaging) + # adaptation steps, target acceptance probability, target trajectory length + n_adapt, δ, λ = 1000, 0.8 + hmcda = HMCDA(1000, 0.8, 1.0) + ``` + + Equivalent to: + + ```julia + metric = DiagEuclideanMetric(D) + hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) + integrator = Leapfrog(initial_ϵ) + kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(λ))) + adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(δ, integrator)) + hmcda = HMCSampler(kernel, metric, adaptor) + ``` + +Moreover, there's some flexibility in how these samplers can be initialized. +For example, a user can initialize a NUTS (as well as HMC and HMCDA) sampler with their own metric and integrator. +This can be done as follows: + ```julia + nuts = NUTS(n_adapt, δ, metric = :diagonal) #metric = DiagEuclideanMetric(D) (Default!) + nuts = NUTS(n_adapt, δ, metric = :unit) #metric = UnitEuclideanMetric(D) + nuts = NUTS(n_adapt, δ, metric = :dense) #metric = DenseEuclideanMetric(D) + # Provide your own :AbstractMetric + metric = DiagEuclideanMetric(10) + nuts = NUTS(n_adapt, δ, metric = metric) + + nuts = NUTS(n_adapt, δ, integrator = :leapfrog) #integrator = Leapfrog(ϵ) (Default!) + nuts = NUTS(n_adapt, δ, integrator = :jitteredleapfrog) #integrator = JitteredLeapfrog(ϵ, 0.1ϵ) + nuts = NUTS(n_adapt, δ, integrator = :temperedleapfrog) #integrator = TemperedLeapfrog(ϵ, 1.0) + + # Provide your own :AbstractIntegrator + integrator = JitteredLeapfrog(ϵ, 0.2ϵ) + nuts = NUTS(n_adapt, δ, integrator = integrator) + ``` + +Finally, bare in mind that the convinience constructors return `AbstractMCMC.AbstractSampler` and therefore they must be used using the `AbstractMCMC` interface as in the case of `HMCSampler`. + ### GPU Sampling with CUDA There is experimental support for running static HMC on the GPU using CUDA. @@ -147,7 +259,7 @@ where `dim` is the dimensionality of the sampling space. where `ϵ` is the step size of leapfrog integration. -### Proposal (`proposal`) +### Kernel (`kernel`) - Static HMC with a fixed number of steps (`n_steps`) (Neal, R. M. (2011)): `StaticTrajectory(integrator, n_steps)` - HMC with a fixed total trajectory length (`trajectory_length`) (Neal, R. M. (2011)): `HMCDA(integrator, trajectory_length)` @@ -187,7 +299,7 @@ function sample( ) ``` -Draw `n_samples` samples using the proposal `κ` under the Hamiltonian system `h` +Draw `n_samples` samples using the kernel `κ` under the Hamiltonian system `h` - The randomness is controlled by `rng`. - If `rng` is not provided, `GLOBAL_RNG` will be used. From 1c698a3213ec1bd851d524df3489a18f5e1312ae Mon Sep 17 00:00:00 2001 From: jaimerz Date: Fri, 21 Jul 2023 14:04:22 +0100 Subject: [PATCH 03/29] format --- test/constructors.jl | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index e827362a..fe6853f8 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -14,13 +14,13 @@ metric = DiagEuclideanMetric(2) adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) custom = HMCSampler(kernel, metric, adaptor) -nuts_metric1 = NUTS(1000, 0.8; metric=:unit) -nuts_metric2 = NUTS(1000, 0.8; metric=:dense) -hmc_metric1 = HMC(0.1, 25; metric=metric) +nuts_metric1 = NUTS(1000, 0.8; metric = :unit) +nuts_metric2 = NUTS(1000, 0.8; metric = :dense) +hmc_metric1 = HMC(0.1, 25; metric = metric) -nuts_integrator1 = NUTS(1000, 0.8, integrator=:jitteredleapfrog) -nuts_integrator2 = NUTS(1000, 0.8, integrator=:temperedleapfrog) -hmc_integrator1 = HMC(0.1, 25, integrator=integrator) +nuts_integrator1 = NUTS(1000, 0.8, integrator = :jitteredleapfrog) +nuts_integrator2 = NUTS(1000, 0.8, integrator = :temperedleapfrog) +hmc_integrator1 = HMC(0.1, 25, integrator = integrator) # Check that everything is initalized correctly @testset "Constructors" begin @@ -80,13 +80,19 @@ end AbstractMCMC.step(rng, logdensitymodel, nuts_32; init_params = θ_init) _, custom_state = AbstractMCMC.step(rng, logdensitymodel, custom; init_params = θ_init) - _, nuts_metric1_state = AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; init_params = θ_init) - _, nuts_metric2_state = AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; init_params = θ_init) - _, hmc_metric1_state = AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; init_params = θ_init) - - _, nuts_integrator1_state = AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; init_params = θ_init) - _, nuts_integrator2_state = AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; init_params = θ_init) - _, hmc_integrator1_state = AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; init_params = θ_init) + _, nuts_metric1_state = + AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; init_params = θ_init) + _, nuts_metric2_state = + AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; init_params = θ_init) + _, hmc_metric1_state = + AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; init_params = θ_init) + + _, nuts_integrator1_state = + AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; init_params = θ_init) + _, nuts_integrator2_state = + AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; init_params = θ_init) + _, hmc_integrator1_state = + AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; init_params = θ_init) # Metric @test typeof(nuts_state.metric) == DiagEuclideanMetric{Float64,Vector{Float64}} From 484461986c6ae051e4570cbd318d8d215c5cb6ff Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Fri, 21 Jul 2023 14:13:04 +0100 Subject: [PATCH 04/29] Update abstractmcmc.jl --- src/abstractmcmc.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 55f3ea26..bb092fbf 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -291,14 +291,14 @@ make_integrator(spl::AbstractHMCSampler, ϵ::Real) = make_integrator(spl.integra 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(@nospecialize(i), ::Real) = error("Integrator $i not supported.") make_integrator(i::Val{:leapfrog}, ϵ::Real) = Leapfrog(ϵ) make_integrator(i::Val{:jitteredleapfrog}, ϵ::Real) = JitteredLeapfrog(ϵ, 0.1ϵ) make_integrator(i::Val{:temperedleapfrog}, ϵ::Real) = TemperedLeapfrog(ϵ, 1.0) ######### -make_metric(i...) = error("Metric $(typeof(i)) not supported.") +make_metric(@nospecialize(i), T::Type, d::Int) = 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 From 108f98c1ac6b13626c0ab769fc57e95b10d6b7fd Mon Sep 17 00:00:00 2001 From: Jaime RZ Date: Fri, 21 Jul 2023 14:17:04 +0100 Subject: [PATCH 05/29] Apply suggestions from code review Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 53a15ec8..a959ae41 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,7 @@ integrator = Leapfrog(initial_ϵ) # - multinomial sampling scheme, # - generalised No-U-Turn criteria, and # - windowed adaption for step-size and diagonal mass matrix -kernel = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) +kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # Run the sampler to draw samples from the specified Gaussian, where From 75333d5b606b4fe7fae1b3dbef51b1622bf6e8d5 Mon Sep 17 00:00:00 2001 From: jaimerz Date: Fri, 21 Jul 2023 14:34:14 +0100 Subject: [PATCH 06/29] API --- README.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 53a15ec8..0b63a879 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,7 @@ integrator = Leapfrog(initial_ϵ) # - multinomial sampling scheme, # - generalised No-U-Turn criteria, and # - windowed adaption for step-size and diagonal mass matrix -kernel = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) +kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # Run the sampler to draw samples from the specified Gaussian, where @@ -261,12 +261,12 @@ where `ϵ` is the step size of leapfrog integration. ### Kernel (`kernel`) -- Static HMC with a fixed number of steps (`n_steps`) (Neal, R. M. (2011)): `StaticTrajectory(integrator, n_steps)` -- HMC with a fixed total trajectory length (`trajectory_length`) (Neal, R. M. (2011)): `HMCDA(integrator, trajectory_length)` -- Original NUTS with slice sampling (Hoffman, M. D., & Gelman, A. (2014)): `NUTS{SliceTS,ClassicNoUTurn}(integrator)` -- Generalised NUTS with slice sampling (Betancourt, M. (2017)): `NUTS{SliceTS,GeneralisedNoUTurn}(integrator)` -- Original NUTS with multinomial sampling (Betancourt, M. (2017)): `NUTS{MultinomialTS,ClassicNoUTurn}(integrator)` -- Generalised NUTS with multinomial sampling (Betancourt, M. (2017)): `NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator)` +- Static HMC with a fixed number of steps (`n_steps`) (Neal, R. M. (2011)): `HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(integrator)))` +- HMC with a fixed total trajectory length (`trajectory_length`) (Neal, R. M. (2011)): `HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(trajectory_length)))` +- Original NUTS with slice sampling (Hoffman, M. D., & Gelman, A. (2014)): `HMCKernel(Trajectory{SliceTS}(integrator, ClassicNoUTurn()))` +- Generalised NUTS with slice sampling (Betancourt, M. (2017)): `HMCKernel(Trajectory{SliceTS}(integrator, GeneralisedNoUTurn()))` +- Original NUTS with multinomial sampling (Betancourt, M. (2017)): `HMCKernel(Trajectory{MultinomialTS}(integrator, ClassicNoUTurn()))` +- Generalised NUTS with multinomial sampling (Betancourt, M. (2017)): `HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))` ### Adaptor (`adaptor`) From a7055cd74ce0159f7161ed3f68429959bb95f41c Mon Sep 17 00:00:00 2001 From: Jaime RZ Date: Tue, 25 Jul 2023 09:09:48 +0100 Subject: [PATCH 07/29] Apply suggestions from code review Co-authored-by: Tor Erlend Fjelde --- README.md | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 0b63a879..115db176 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ In this section we demonstrate a minimal example of sampling from a multivariate Suppose ${\bf x}$ and ${\bf v}$ are the position and velocity of an individual particle respectively; $i$ and $i+1$ are the indices for time values $t_i$ and $t_{i+1}$ respectively; $dt = t_{i+1} - t_i$ is the time step size (constant and regularly spaced intervals); and ${\bf a}$ is the acceleration induced on a particle by the forces of all other particles. Furthermore, suppose positions are defined at times $t_i, t_{i+1}, t_{i+2}, \dots $, spaced at constant intervals $dt$, the velocities are defined at halfway times in between, denoted by $t_{i-1/2}, t_{i+1/2}, t_{i+3/2}, \dots $, where $t_{i+1} - t_{i + 1/2} = t_{i + 1/2} - t_i = dt / 2$, and the accelerations ${\bf a}$ are defined only on integer times, just like the positions. Then the leapfrog integration scheme is given as: $x_{i} = x_{i-1} + v_{i-1/2} dt; \quad v_{i+1/2} = v_{i-1/2} + a_i dt$. For available integrators refer Integrator. -- **kernel for trajectories (static or dynamic)**: Different types of kernels can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle. +- **Kernel for trajectories (static or dynamic)**: Different types of kernels can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle.
More about the kernels In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. Whereas, during the second step, a Metropolis update is performed, using Hamiltonian dynamics to provide a new state. For available kernels refer kernel. @@ -121,8 +121,8 @@ end ### Using the `AbstractMCMC` interface -Users can also make use of the `AbstractMCMC` interface to sample which employs the same API as other popular Bayesian inference libraries in Julia such as `Turing`. -In order to show how this is done let us start from our previous example where we defined a `LogTargetDensity`, ℓπ. +Users can also make use of the `AbstractMCMC` interface to sample, which is also what is used in Turing.jl. +In order to show how this is done let us start from our previous example where we defined a `LogTargetDensity`, `ℓπ`. ```julia # Wrap the previous LogTargetDensity as LogDensityModel @@ -216,7 +216,7 @@ This can be done as follows: nuts = NUTS(n_adapt, δ, metric = :diagonal) #metric = DiagEuclideanMetric(D) (Default!) nuts = NUTS(n_adapt, δ, metric = :unit) #metric = UnitEuclideanMetric(D) nuts = NUTS(n_adapt, δ, metric = :dense) #metric = DenseEuclideanMetric(D) - # Provide your own :AbstractMetric + # Provide your own AbstractMetric metric = DiagEuclideanMetric(10) nuts = NUTS(n_adapt, δ, metric = metric) @@ -224,13 +224,11 @@ This can be done as follows: nuts = NUTS(n_adapt, δ, integrator = :jitteredleapfrog) #integrator = JitteredLeapfrog(ϵ, 0.1ϵ) nuts = NUTS(n_adapt, δ, integrator = :temperedleapfrog) #integrator = TemperedLeapfrog(ϵ, 1.0) - # Provide your own :AbstractIntegrator + # Provide your own AbstractIntegrator integrator = JitteredLeapfrog(ϵ, 0.2ϵ) nuts = NUTS(n_adapt, δ, integrator = integrator) ``` -Finally, bare in mind that the convinience constructors return `AbstractMCMC.AbstractSampler` and therefore they must be used using the `AbstractMCMC` interface as in the case of `HMCSampler`. - ### GPU Sampling with CUDA There is experimental support for running static HMC on the GPU using CUDA. From 7085642d5d69cad470ffc4548a0641deb4eb503a Mon Sep 17 00:00:00 2001 From: Jaime RZ Date: Tue, 25 Jul 2023 12:26:47 +0100 Subject: [PATCH 08/29] Apply suggestions from code review Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --- README.md | 2 +- src/abstractmcmc.jl | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index 115db176..bc11f226 100644 --- a/README.md +++ b/README.md @@ -194,7 +194,7 @@ In the previous examples we built the sampler by manually specifying the integra #HMCDA (dual averaging) # adaptation steps, target acceptance probability, target trajectory length n_adapt, δ, λ = 1000, 0.8 - hmcda = HMCDA(1000, 0.8, 1.0) + hmcda = HMCDA(n_adapt, δ, λ) ``` Equivalent to: diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index bb092fbf..e68112cf 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -289,7 +289,6 @@ 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(@nospecialize(i), ::Real) = error("Integrator $i not supported.") make_integrator(i::Val{:leapfrog}, ϵ::Real) = Leapfrog(ϵ) From 8e62dc933e7577f9945dbafbd97bbafe828ff60a Mon Sep 17 00:00:00 2001 From: jaimerz Date: Tue, 25 Jul 2023 12:29:36 +0100 Subject: [PATCH 09/29] Hong s comments --- README.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index bc11f226..fa31c0bd 100644 --- a/README.md +++ b/README.md @@ -163,9 +163,9 @@ In the previous examples we built the sampler by manually specifying the integra hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) integrator = Leapfrog(initial_ϵ) - kernel = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) - adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(δ, integrator)) - nuts = HMCSampler(kernel, metric, adaptor) + kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(n_leapfrog))) + adaptor = NoAdaptation() + hmc = HMCSampler(kernel, metric, adaptor) ``` - NUTS: @@ -183,9 +183,9 @@ In the previous examples we built the sampler by manually specifying the integra hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) integrator = Leapfrog(initial_ϵ) - kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(n_leapfrog))) - adaptor = NoAdaptation() - hmc = HMCSampler(kernel, metric, adaptor) + kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn())) + adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(δ, integrator)) + nuts = HMCSampler(kernel, metric, adaptor) ``` @@ -193,7 +193,7 @@ In the previous examples we built the sampler by manually specifying the integra ```julia #HMCDA (dual averaging) # adaptation steps, target acceptance probability, target trajectory length - n_adapt, δ, λ = 1000, 0.8 + n_adapt, δ, λ = 1000, 0.8, 1.0 hmcda = HMCDA(n_adapt, δ, λ) ``` From 306ad4d07c9a98a95678e0524b9b60a7c6d48859 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Tue, 25 Jul 2023 13:46:01 +0100 Subject: [PATCH 10/29] Fix typo and simplify arguments (#331) --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index fa31c0bd..cfb6f638 100644 --- a/README.md +++ b/README.md @@ -17,10 +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.5.0] **Breaking!** Convenience constructors for common samplers changed to: + - `HMC(init_ϵ, n_leapfrog)` + - `NUTS(n_adapts, target_acceptance)` + - `HMCDA(n_adapts, target_acceptance, integration_time)` - [v0.2.22] Three functions are renamed. - `Preconditioner(metric::AbstractMetric)` -> `MassMatrixAdaptor(metric)` and - `NesterovDualAveraging(δ, integrator::AbstractIntegrator)` -> `StepSizeAdaptor(δ, integrator)` From 021de4f22c2e24d4e9922122f0111a49bb045335 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 10:59:26 +0100 Subject: [PATCH 11/29] Removed `n_adapts` from sampler constructors and some fixes. (#333) * Update README.md * Update README.md Co-authored-by: Tor Erlend Fjelde --------- Co-authored-by: Jaime RZ Co-authored-by: Tor Erlend Fjelde --- README.md | 64 +++++++++++++++++++++++++++---------------------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index cfb6f638..d4dfcc77 100644 --- a/README.md +++ b/README.md @@ -6,21 +6,21 @@ [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/stable/) [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/dev/) -AdvancedHMC.jl provides a robust, modular and efficient implementation of advanced HMC algorithms. An illustrative example for AdvancedHMC's usage is given below. AdvancedHMC.jl is part of [Turing.jl](https://github.com/TuringLang/Turing.jl), a probabilistic programming library in Julia. +AdvancedHMC.jl provides a robust, modular, and efficient implementation of advanced HMC algorithms. An illustrative example of AdvancedHMC's usage is given below. AdvancedHMC.jl is part of [Turing.jl](https://github.com/TuringLang/Turing.jl), a probabilistic programming library in Julia. If you are interested in using AdvancedHMC.jl through a probabilistic programming language, please check it out! **Interfaces** - [`IMP.hmc`](https://github.com/salilab/hmc): an experimental Python module for the Integrative Modeling Platform, which uses AdvancedHMC in its backend to sample protein structures. **NEWS** -- We presented a paper for AdvancedHMC.jl at [AABI](http://approximateinference.org/) 2019 in Vancouver, Canada. ([abs](http://proceedings.mlr.press/v118/xu20a.html), [pdf](http://proceedings.mlr.press/v118/xu20a/xu20a.pdf), [OpenReview](https://openreview.net/forum?id=rJgzckn4tH)) +- We presented a paper for AdvancedHMC.jl at [AABI](http://approximateinference.org/) in 2019 in Vancouver, Canada. ([abs](http://proceedings.mlr.press/v118/xu20a.html), [pdf](http://proceedings.mlr.press/v118/xu20a/xu20a.pdf), [OpenReview](https://openreview.net/forum?id=rJgzckn4tH)) - 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!** Convenience constructors for common samplers changed to: - `HMC(init_ϵ, n_leapfrog)` - - `NUTS(n_adapts, target_acceptance)` - - `HMCDA(n_adapts, target_acceptance, integration_time)` + - `NUTS(target_acceptance)` + - `HMCDA(target_acceptance, integration_time)` - [v0.2.22] Three functions are renamed. - `Preconditioner(metric::AbstractMetric)` -> `MassMatrixAdaptor(metric)` and - `NesterovDualAveraging(δ, integrator::AbstractIntegrator)` -> `StepSizeAdaptor(δ, integrator)` @@ -33,24 +33,24 @@ If you are interested in using AdvancedHMC.jl through a probabilistic programmin ## A minimal example - sampling from a multivariate Gaussian using NUTS -In this section we demonstrate a minimal example of sampling from a multivariate Gaussian (10 dimensional) using the no U-turn sampler (NUTS). Below we describe the major components of the Hamiltonian system which are essential to sample using this approach: +This section demonstrates a minimal example of sampling from a multivariate Gaussian (10-dimensional) using the no U-turn sampler (NUTS). Below we describe the major components of the Hamiltonian system which are essential to sample using this approach: -- **Metric**: In many sampling problems the sample space is usually associated with a metric, that allows us to measure the distance between any two points, and other similar quantities. In the example in this section, we use a special metric called the **Euclidean Metric**, represented with a `D × D` matrix from which we can compute distances. +- **Metric**: In many sampling problems the sample space is usually associated with a metric that allows us to measure the distance between any two points, and other similar quantities. In the example in this section, we use a special metric called the **Euclidean Metric**, represented with a `D × D` matrix from which we can compute distances.
Further details about the Metric component The Euclidean metric is also known as the mass matrix in the physical perspective. For available metrics refer Hamiltonian mass matrix.
-- **Leapfrog integration**: Leapfrog integration is a second-order numerical method for integrating differential equations (In this case they are, equations of motion for the relative position of one particle with respect to the other). The order of this integration signifies its rate of convergence. Any alogrithm with a finite time step size will have numerical errors and the order is related to this error. For a second-order algorithm, this error scales as the second power of the time step, hence, the name second-order. High-order intergrators are usually complex to code and have a limited region of convergence, hence they do not allow arbitrarily large time steps. A second-order integrator is suitable for our purpose, hence we opt for the leapfrog integrator. It is called `leapfrog` due to the ways this algorithm is written, where the positions and velocities of particles `leap over` each other. +- **Leapfrog integration**: Leapfrog integration is a second-order numerical method for integrating differential equations (In this case they are equations of motion for the relative position of one particle with respect to the other). The order of this integration signifies its rate of convergence. Any algorithm with a finite time step size will have numerical errors, and the order is related to this error. For a second-order algorithm, this error scales as the second power of the time step, hence, the name second-order. High-order integrators are usually complex to code and have a limited region of convergence; hence they do not allow arbitrarily large time steps. A second-order integrator is suitable for our purpose. Hence we opt for the leapfrog integrator. It is called `leapfrog` due to the ways this algorithm is written, where the positions and velocities of particles `leap over` each other.
About the leapfrog integration scheme - Suppose ${\bf x}$ and ${\bf v}$ are the position and velocity of an individual particle respectively; $i$ and $i+1$ are the indices for time values $t_i$ and $t_{i+1}$ respectively; $dt = t_{i+1} - t_i$ is the time step size (constant and regularly spaced intervals); and ${\bf a}$ is the acceleration induced on a particle by the forces of all other particles. Furthermore, suppose positions are defined at times $t_i, t_{i+1}, t_{i+2}, \dots $, spaced at constant intervals $dt$, the velocities are defined at halfway times in between, denoted by $t_{i-1/2}, t_{i+1/2}, t_{i+3/2}, \dots $, where $t_{i+1} - t_{i + 1/2} = t_{i + 1/2} - t_i = dt / 2$, and the accelerations ${\bf a}$ are defined only on integer times, just like the positions. Then the leapfrog integration scheme is given as: $x_{i} = x_{i-1} + v_{i-1/2} dt; \quad v_{i+1/2} = v_{i-1/2} + a_i dt$. For available integrators refer Integrator. + Suppose ${\bf x}$ and ${\bf v}$ are the position and velocity of an individual particle respectively; $i$ and $i+1$ are the indices for time values $t_i$ and $t_{i+1}$ respectively; $dt = t_{i+1} - t_i$ is the time step size (constant and regularly spaced intervals), and ${\bf a}$ is the acceleration induced on a particle by the forces of all other particles. Furthermore, suppose positions are defined at times $t_i, t_{i+1}, t_{i+2}, \dots $, spaced at constant intervals $dt$, the velocities are defined at halfway times in between, denoted by $t_{i-1/2}, t_{i+1/2}, t_{i+3/2}, \dots $, where $t_{i+1} - t_{i + 1/2} = t_{i + 1/2} - t_i = dt / 2$, and the accelerations ${\bf a}$ are defined only on integer times, just like the positions. Then the leapfrog integration scheme is given as: $x_{i} = x_{i-1} + v_{i-1/2} dt; \quad v_{i+1/2} = v_{i-1/2} + a_i dt$. For available integrators refer Integrator.
-- **Kernel for trajectories (static or dynamic)**: Different types of kernels can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle. +- **Kernel for trajectories (static or dynamic)**: Different kernels, which may be static or dynamic, can be used. At each iteration of any variant of the HMC algorithm, there are two main steps - the first step changes the momentum and the second step may change both the position and the momentum of a particle.
More about the kernels - In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. Whereas, during the second step, a Metropolis update is performed, using Hamiltonian dynamics to provide a new state. For available kernels refer kernel. + In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. A Metropolis update is performed during the second step, using Hamiltonian dynamics to provide a new state. For available kernels refer kernel.
```julia @@ -77,11 +77,11 @@ n_samples, n_adapts = 2_000, 1_000 metric = DiagEuclideanMetric(D) hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) -# Define a leapfrog solver, with initial step size chosen heuristically +# Define a leapfrog solver, with the initial step size chosen heuristically initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) integrator = Leapfrog(initial_ϵ) -# Define an HMC sampler, with the following components +# Define an HMC sampler with the following components # - multinomial sampling scheme, # - generalised No-U-Turn criteria, and # - windowed adaption for step-size and diagonal mass matrix @@ -102,7 +102,7 @@ It also supports vectorized sampling for static HMC and has been discussed in mo The below example utilizes the `@threads` macro to sample 4 chains across 4 threads. ```julia -# Ensure that julia was launched with appropriate number of threads +# Ensure that Julia was launched with an appropriate number of threads println(Threads.nthreads()) # Number of chains to sample @@ -121,7 +121,7 @@ end ### Using the `AbstractMCMC` interface -Users can also make use of the `AbstractMCMC` interface to sample, which is also what is used in Turing.jl. +Users can also use the `AbstractMCMC` interface to sample, which is also used in Turing.jl. In order to show how this is done let us start from our previous example where we defined a `LogTargetDensity`, `ℓπ`. ```julia @@ -144,9 +144,9 @@ samples = AbstractMCMC.sample( ) ``` -### Covenience Constructors +### Convenience Constructors -In the previous examples we built the sampler by manually specifying the integrator, metric, kernel and adaptor to build our own sampler. However, in many cases users might want to simply initialize a standard NUTS sampler. In such cases having to manually define each of these aspects is tedious and error prone. For these reasons `AdvancedHMC` also provides users with a series of convenience constructors for standard samplers. We will now show how to use them. +In the previous examples, we built the sampler by manually specifying the integrator, metric, kernel, and adaptor to build our own sampler. However, in many cases, users might want to initialize a standard NUTS sampler. In such cases having to define each of these aspects manually is tedious and error-prone. For these reasons `AdvancedHMC` also provides users with a series of convenience constructors for standard samplers. We will now show how to use them. - HMC: ```julia @@ -172,8 +172,8 @@ In the previous examples we built the sampler by manually specifying the integra ```julia # NUTS Sampler # adaptation steps, target acceptance probability, - n_adapt, δ = 1000, 0.8 - nuts = NUTS(n_adapt, δ) + δ = 0.8 + nuts = NUTS(δ) ``` Equivalent to: @@ -193,8 +193,8 @@ In the previous examples we built the sampler by manually specifying the integra ```julia #HMCDA (dual averaging) # adaptation steps, target acceptance probability, target trajectory length - n_adapt, δ, λ = 1000, 0.8, 1.0 - hmcda = HMCDA(n_adapt, δ, λ) + δ, λ = 0.8, 1.0 + hmcda = HMCDA(δ, λ) ``` Equivalent to: @@ -210,35 +210,35 @@ In the previous examples we built the sampler by manually specifying the integra ``` Moreover, there's some flexibility in how these samplers can be initialized. -For example, a user can initialize a NUTS (as well as HMC and HMCDA) sampler with their own metric and integrator. +For example, a user can initialize a NUTS (HMC and HMCDA) sampler with their own metrics and integrators. This can be done as follows: ```julia - nuts = NUTS(n_adapt, δ, metric = :diagonal) #metric = DiagEuclideanMetric(D) (Default!) - nuts = NUTS(n_adapt, δ, metric = :unit) #metric = UnitEuclideanMetric(D) - nuts = NUTS(n_adapt, δ, metric = :dense) #metric = DenseEuclideanMetric(D) + nuts = NUTS(δ, metric = :diagonal) #metric = DiagEuclideanMetric(D) (Default!) + nuts = NUTS(δ, metric = :unit) #metric = UnitEuclideanMetric(D) + nuts = NUTS(δ, metric = :dense) #metric = DenseEuclideanMetric(D) # Provide your own AbstractMetric metric = DiagEuclideanMetric(10) nuts = NUTS(n_adapt, δ, metric = metric) - nuts = NUTS(n_adapt, δ, integrator = :leapfrog) #integrator = Leapfrog(ϵ) (Default!) - nuts = NUTS(n_adapt, δ, integrator = :jitteredleapfrog) #integrator = JitteredLeapfrog(ϵ, 0.1ϵ) - nuts = NUTS(n_adapt, δ, integrator = :temperedleapfrog) #integrator = TemperedLeapfrog(ϵ, 1.0) + nuts = NUTS(δ, integrator = :leapfrog) #integrator = Leapfrog(ϵ) (Default!) + nuts = NUTS(δ, integrator = :jitteredleapfrog) #integrator = JitteredLeapfrog(ϵ, 0.1ϵ) + nuts = NUTS(δ, integrator = :temperedleapfrog) #integrator = TemperedLeapfrog(ϵ, 1.0) # Provide your own AbstractIntegrator integrator = JitteredLeapfrog(ϵ, 0.2ϵ) - nuts = NUTS(n_adapt, δ, integrator = integrator) + nuts = NUTS(δ, integrator = integrator) ``` ### GPU Sampling with CUDA There is experimental support for running static HMC on the GPU using CUDA. -To do so the user needs to have [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) installed, ensure the logdensity of the `Hamiltonian` can be executed on the GPU and that the initial points are a `CuArray`. +To do so, the user needs to have [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) installed, ensure the logdensity of the `Hamiltonian` can be executed on the GPU and that the initial points are a `CuArray`. A small working example can be found at `test/cuda.jl`. ## API and supported HMC algorithms An important design goal of AdvancedHMC.jl is modularity; we would like to support algorithmic research on HMC. -This modularity means that different HMC variants can be easily constructed by composing various components, such as preconditioning metric (i.e. mass matrix), leapfrog integrators, trajectories (static or dynamic), and adaption schemes etc. +This modularity means that different HMC variants can be easily constructed by composing various components, such as preconditioning metric (i.e., mass matrix), leapfrog integrators, trajectories (static or dynamic), and adaption schemes, etc. The minimal example above can be modified to suit particular inference problems by picking components from the list below. ### Hamiltonian mass matrix (`metric`) @@ -276,9 +276,9 @@ where `ϵ` is the step size of leapfrog integration. - Combine the first two using Stan's windowed adaptation: `StanHMCAdaptor(mma, ssa)` ### Gradients -`AdvancedHMC` supports both AD-based (`Zygote`, `Tracker` and `ForwardDiff`) and user-specified gradients. In order to use user-specified gradients, please replace `ForwardDiff` with `ℓπ_grad` in the `Hamiltonian` constructor, where the gradient function `ℓπ_grad` should return a tuple containing both the log-posterior and its gradient. +`AdvancedHMC` supports AD-based using [`LogDensityProblemsAD`](https://github.com/tpapp/LogDensityProblemsAD.jl) and user-specified gradients. In order to use user-specified gradients, please replace `ForwardDiff` with `ℓπ_grad` in the `Hamiltonian` constructor, where the gradient function `ℓπ_grad` should return a tuple containing both the log-posterior and its gradient. -All the combinations are tested in [this file](https://github.com/TuringLang/AdvancedHMC.jl/blob/master/test/sampler.jl) except from using tempered leapfrog integrator together with adaptation, which we found unstable empirically. +All the combinations are tested in [this file](https://github.com/TuringLang/AdvancedHMC.jl/blob/master/test/sampler.jl) except for using tempered leapfrog integrator together with adaptation, which we found unstable empirically. ## The `sample` function signature in detail From c207e818b196d395c08cb44c5dcc2211e3e6e2b9 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 11:13:05 +0100 Subject: [PATCH 12/29] Minor tweaks to the metric field comments. --- src/constructors.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/constructors.jl b/src/constructors.jl index 93933570..ccde73f9 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -66,7 +66,7 @@ struct NUTS{T<:Real} <: AbstractHMCSampler{T} init_ϵ::T "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" integrator::Union{Symbol,AbstractIntegrator} - "Choice of initial metric, specified using a `Symbol` or `AbstractMetric`. The metric type will be preserved during adaption." + "Choice of initial metric; `Symbol` means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption." metric::Union{Symbol,AbstractMetric} end @@ -107,7 +107,7 @@ struct HMC{T<:Real} <: AbstractHMCSampler{T} n_leapfrog::Int "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" integrator::Union{Symbol,AbstractIntegrator} - "Choice of initial metric, specified using a `Symbol` or `AbstractMetric`. The metric type will be preserved during adaption." + "Choice of initial metric; `Symbol` means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption." metric::Union{Symbol,AbstractMetric} end @@ -148,7 +148,7 @@ struct HMCDA{T<:Real} <: AbstractHMCSampler{T} init_ϵ::T "Choice of integrator, specified either using a `Symbol` or [`AbstractIntegrator`](@ref)" integrator::Union{Symbol,AbstractIntegrator} - "Choice of initial metric, specified using a `Symbol` or `AbstractMetric`. The metric type will be preserved during adaption." + "Choice of initial metric; `Symbol` means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption." metric::Union{Symbol,AbstractMetric} end From 4e9ed2682875648e0bb218ef4a3d45debe8ad95d Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 11:15:46 +0100 Subject: [PATCH 13/29] Removed redundant make_metric function --- src/abstractmcmc.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index bd856d9d..8adb459c 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -304,7 +304,6 @@ make_integrator(i::Val{:temperedleapfrog}, ϵ::Real) = TemperedLeapfrog(ϵ, 1.0) make_metric(@nospecialize(i), T::Type, d::Int) = 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) From 0ef0ccfb8b287bfc095586ac71f48c3ee12d10bb Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 11:32:57 +0100 Subject: [PATCH 14/29] Fix typos in constructor tests --- test/constructors.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index 7684e223..4f40dc35 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -14,12 +14,12 @@ metric = DiagEuclideanMetric(2) adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator) custom = HMCSampler(kernel, metric, adaptor) -nuts_metric1 = NUTS(1000, 0.8; metric = :unit) -nuts_metric2 = NUTS(1000, 0.8; metric = :dense) +nuts_metric1 = NUTS(0.8; metric = :unit) +nuts_metric2 = NUTS(0.8; metric = :dense) hmc_metric1 = HMC(0.1, 25; metric = metric) -nuts_integrator1 = NUTS(1000, 0.8, integrator = :jitteredleapfrog) -nuts_integrator2 = NUTS(1000, 0.8, integrator = :temperedleapfrog) +nuts_integrator1 = NUTS(0.8, integrator = :jitteredleapfrog) +nuts_integrator2 = NUTS(0.8, integrator = :temperedleapfrog) hmc_integrator1 = HMC(0.1, 25, integrator = integrator) # Check that everything is initalized correctly From b002f8567ba910b392031ba578b6dd3359cb8fc6 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 11:42:33 +0100 Subject: [PATCH 15/29] More fixes. --- src/constructors.jl | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/constructors.jl b/src/constructors.jl index ccde73f9..35367963 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -28,20 +28,18 @@ struct HMCSampler{T<:Real} <: AbstractHMCSampler{T} metric::AbstractMetric "[`AbstractAdaptor`](@ref)." adaptor::AbstractAdaptor - "Adaptation steps if any" - n_adapts::Int end -function HMCSampler(κ, metric, adaptor; n_adapts = 0) +function HMCSampler(κ, metric, adaptor) T = collect(typeof(metric).parameters)[1] - return HMCSampler{T}(κ, metric, adaptor, n_adapts) + return HMCSampler{T}(κ, metric, adaptor) end ############ ### NUTS ### ############ """ - NUTS(n_adapts::Int, δ::Real; max_depth::Int=10, Δ_max::Real=1000, init_ϵ::Real=0) + NUTS(δ::Real; max_depth::Int=10, Δ_max::Real=1000, init_ϵ::Real=0, init_ϵ = 0.0, integrator = :leapfrog, metric = :diagonal) No-U-Turn Sampler (NUTS) sampler. @@ -52,7 +50,7 @@ $(FIELDS) # Usage: ```julia -NUTS(n_adapts=1000, δ=0.65) # Use 1000 adaption steps, and target accept ratio 0.65. +NUTS(δ=0.65) # Use target accept ratio 0.65. ``` """ struct NUTS{T<:Real} <: AbstractHMCSampler{T} @@ -97,7 +95,7 @@ $(FIELDS) # Usage: ```julia -HMC(init_ϵ=0.05, n_leapfrog=10) +HMC(init_ϵ=0.05, n_leapfrog=10, integrator = :leapfrog, metric = :diagonal) ``` """ struct HMC{T<:Real} <: AbstractHMCSampler{T} @@ -119,7 +117,7 @@ end ### HMCDA ### ############# """ - HMCDA(n_adapts::Int, δ::Real, λ::Real; ϵ::Real=0) + HMCDA(δ::Real, λ::Real; ϵ::Real=0, integrator = :leapfrog, metric = :diagonal) Hamiltonian Monte Carlo sampler with Dual Averaging algorithm. @@ -130,7 +128,7 @@ $(FIELDS) # Usage: ```julia -HMCDA(n_adapts=200, δ=0.65, λ=0.3) +HMCDA(δ=0.65, λ=0.3) ``` For more information, please view the following paper ([arXiv link](https://arxiv.org/abs/1111.4246)): From cc0ae8649069d50ca4fc11936cbcbf4208d34dc0 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 11:57:47 +0100 Subject: [PATCH 16/29] Typofix. --- src/constructors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/constructors.jl b/src/constructors.jl index 35367963..ac4c45f4 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -39,7 +39,7 @@ end ### NUTS ### ############ """ - NUTS(δ::Real; max_depth::Int=10, Δ_max::Real=1000, init_ϵ::Real=0, init_ϵ = 0.0, integrator = :leapfrog, metric = :diagonal) + NUTS(δ::Real; max_depth::Int=10, Δ_max::Real=1000, init_ϵ::Real=0, integrator = :leapfrog, metric = :diagonal) No-U-Turn Sampler (NUTS) sampler. From a5c4983b023f5e0815381a23c309aed9819a2c52 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 12:25:18 +0100 Subject: [PATCH 17/29] More test fixes. --- test/constructors.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index 4f40dc35..9de58030 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -80,18 +80,18 @@ end AbstractMCMC.step(rng, logdensitymodel, custom; n_adapts = 0, init_params = θ_init) _, nuts_metric1_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; n_adapts = 0, init_params = θ_init) _, nuts_metric2_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; n_adapts = 0, init_params = θ_init) _, hmc_metric1_state = - AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; n_adapts = 0, init_params = θ_init) _, nuts_integrator1_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; n_adapts = 0, init_params = θ_init) _, nuts_integrator2_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; n_adapts = 0, init_params = θ_init) _, hmc_integrator1_state = - AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; init_params = θ_init) + AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; n_adapts = 0, init_params = θ_init) # Metric @test typeof(nuts_state.metric) == DiagEuclideanMetric{Float64,Vector{Float64}} From 032cf35281c893bf6da12cf24aad43c4bbdeedb2 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 12:28:03 +0100 Subject: [PATCH 18/29] Update test/constructors.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/constructors.jl | 56 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 43 insertions(+), 13 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index 9de58030..e6099feb 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -79,19 +79,49 @@ end _, custom_state = AbstractMCMC.step(rng, logdensitymodel, custom; n_adapts = 0, init_params = θ_init) - _, nuts_metric1_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_metric1; n_adapts = 0, init_params = θ_init) - _, nuts_metric2_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_metric2; n_adapts = 0, init_params = θ_init) - _, hmc_metric1_state = - AbstractMCMC.step(rng, logdensitymodel, hmc_metric1; n_adapts = 0, init_params = θ_init) - - _, nuts_integrator1_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_integrator1; n_adapts = 0, init_params = θ_init) - _, nuts_integrator2_state = - AbstractMCMC.step(rng, logdensitymodel, nuts_integrator2; n_adapts = 0, init_params = θ_init) - _, hmc_integrator1_state = - AbstractMCMC.step(rng, logdensitymodel, hmc_integrator1; n_adapts = 0, init_params = θ_init) + _, nuts_metric1_state = AbstractMCMC.step( + rng, + logdensitymodel, + nuts_metric1; + n_adapts = 0, + init_params = θ_init, + ) + _, nuts_metric2_state = AbstractMCMC.step( + rng, + logdensitymodel, + nuts_metric2; + n_adapts = 0, + init_params = θ_init, + ) + _, hmc_metric1_state = AbstractMCMC.step( + rng, + logdensitymodel, + hmc_metric1; + n_adapts = 0, + init_params = θ_init, + ) + + _, nuts_integrator1_state = AbstractMCMC.step( + rng, + logdensitymodel, + nuts_integrator1; + n_adapts = 0, + init_params = θ_init, + ) + _, nuts_integrator2_state = AbstractMCMC.step( + rng, + logdensitymodel, + nuts_integrator2; + n_adapts = 0, + init_params = θ_init, + ) + _, hmc_integrator1_state = AbstractMCMC.step( + rng, + logdensitymodel, + hmc_integrator1; + n_adapts = 0, + init_params = θ_init, + ) # Metric @test typeof(nuts_state.metric) == DiagEuclideanMetric{Float64,Vector{Float64}} From 3f7fbcb7f94ba62619a2110a0fb3f3d065fba072 Mon Sep 17 00:00:00 2001 From: Jaime RZ Date: Wed, 26 Jul 2023 15:48:21 +0100 Subject: [PATCH 19/29] no init_e (#335) * no init_e * format * bug * Bugfix. (#337) * Bugfix. * Update src/constructors.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Update test/abstractmcmc.jl * Update src/abstractmcmc.jl Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --------- Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/abstractmcmc.jl | 36 ++++++++++++++++++++++++++---------- src/constructors.jl | 30 +++++++++++------------------- src/integrator.jl | 1 + test/abstractmcmc.jl | 2 +- test/constructors.jl | 12 +++--------- 5 files changed, 42 insertions(+), 39 deletions(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 8adb459c..75fbe483 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -265,29 +265,45 @@ end ######### +function make_step_size( + rng::Random.AbstractRNG, + spl::HMCSampler, + hamiltonian::Hamiltonian, + init_params, +) + return spl.κ.τ.integrator.ϵ +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 ϵ + T = get_type_of_spl(spl) + return make_step_size(rng, spl.integrator, T, hamiltonian, init_params) end function make_step_size( rng::Random.AbstractRNG, - spl::HMCSampler, + integrator::AbstractIntegrator, + T::Type, hamiltonian::Hamiltonian, init_params, ) - return spl.κ.τ.integrator.ϵ + return integrator.ϵ +end + +function make_step_size( + rng::Random.AbstractRNG, + integrator::Symbol, + T::Type, + hamiltonian::Hamiltonian, + init_params, +) + ϵ = find_good_stepsize(rng, hamiltonian, init_params) + @info string("Found initial step size ", ϵ) + return T(ϵ) end make_integrator(spl::HMCSampler, ϵ::Real) = spl.κ.τ.integrator diff --git a/src/constructors.jl b/src/constructors.jl index ac4c45f4..a37227d6 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -60,24 +60,15 @@ struct NUTS{T<:Real} <: AbstractHMCSampler{T} 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 initial metric; `Symbol` means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption." metric::Union{Symbol,AbstractMetric} end -function NUTS( - δ; - max_depth = 10, - Δ_max = 1000.0, - init_ϵ = 0.0, - integrator = :leapfrog, - metric = :diagonal, -) +function NUTS(δ; max_depth = 10, Δ_max = 1000.0, integrator = :leapfrog, metric = :diagonal) T = typeof(δ) - return NUTS(δ, max_depth, T(Δ_max), T(init_ϵ), integrator, metric) + return NUTS(δ, max_depth, T(Δ_max), integrator, metric) end ########### @@ -99,8 +90,6 @@ HMC(init_ϵ=0.05, n_leapfrog=10, integrator = :leapfrog, metric = :diagonal) ``` """ 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)" @@ -109,8 +98,13 @@ struct HMC{T<:Real} <: AbstractHMCSampler{T} metric::Union{Symbol,AbstractMetric} end -function HMC(init_ϵ, n_leapfrog; integrator = :leapfrog, metric = :diagonal) - return HMC(init_ϵ, n_leapfrog, integrator, metric) +function HMC(n_leapfrog; integrator = :leapfrog, metric = :diagonal) + if integrator isa Symbol + T = typeof(0.0) # current default float type + else + T = integrator_eltype(integrator) + end + return HMC{T}(n_leapfrog, integrator, metric) end ############# @@ -142,18 +136,16 @@ struct HMCDA{T<:Real} <: AbstractHMCSampler{T} δ::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 initial metric; `Symbol` means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption." metric::Union{Symbol,AbstractMetric} end -function HMCDA(δ, λ; init_ϵ = 0.0, integrator = :leapfrog, metric = :diagonal) +function HMCDA(δ, λ; integrator = :leapfrog, metric = :diagonal) if typeof(δ) != typeof(λ) @warn "typeof(δ) != typeof(λ) --> using typeof(δ)" end T = typeof(δ) - return HMCDA(δ, T(λ), T(init_ϵ), integrator, metric) + return HMCDA(δ, T(λ), integrator, metric) end diff --git a/src/integrator.jl b/src/integrator.jl index 8391881f..2a2a26c2 100644 --- a/src/integrator.jl +++ b/src/integrator.jl @@ -70,6 +70,7 @@ struct Leapfrog{T<:AbstractScalarOrVec{<:AbstractFloat}} <: AbstractLeapfrog{T} ϵ::T end Base.show(io::IO, l::Leapfrog) = print(io, "Leapfrog(ϵ=$(round.(l.ϵ; sigdigits=3)))") +integrator_eltype(i::AbstractLeapfrog{T}) where {T<:AbstractFloat} = T ### Jittering diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 207eb21f..52d6c35c 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -9,7 +9,7 @@ include("common.jl") θ_init = randn(rng, 2) nuts = NUTS(0.8) - hmc = HMC(0.05, 100) + hmc = HMC(100; integrator = Leapfrog(0.05)) hmcda = HMCDA(0.8, 0.1) integrator = Leapfrog(1e-3) diff --git a/test/constructors.jl b/test/constructors.jl index e6099feb..147d400d 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -4,7 +4,7 @@ include("common.jl") # Initalize samplers nuts = NUTS(0.8) nuts_32 = NUTS(0.8f0) -hmc = HMC(0.1, 25) +hmc = HMC(25) hmcda = HMCDA(0.8, 1.0) hmcda_32 = HMCDA(0.8f0, 1.0) @@ -16,11 +16,11 @@ custom = HMCSampler(kernel, metric, adaptor) nuts_metric1 = NUTS(0.8; metric = :unit) nuts_metric2 = NUTS(0.8; metric = :dense) -hmc_metric1 = HMC(0.1, 25; metric = metric) +hmc_metric1 = HMC(25; integrator = Leapfrog(0.1), metric = metric) nuts_integrator1 = NUTS(0.8, integrator = :jitteredleapfrog) nuts_integrator2 = NUTS(0.8, integrator = :temperedleapfrog) -hmc_integrator1 = HMC(0.1, 25, integrator = integrator) +hmc_integrator1 = HMC(25, integrator = integrator) # Check that everything is initalized correctly @testset "Constructors" begin @@ -36,7 +36,6 @@ hmc_integrator1 = HMC(0.1, 25, integrator = integrator) @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 @@ -44,26 +43,21 @@ hmc_integrator1 = HMC(0.1, 25, integrator = integrator) @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.δ == 0.8 @test hmcda.λ == 1.0 - @test hmcda.init_ϵ == 0.0 @test hmcda.integrator == :leapfrog @test hmcda.metric == :diagonal # HMCDA Float32 @test hmcda_32.δ == 0.8f0 @test hmcda_32.λ == 1.0f0 - @test hmcda_32.init_ϵ == 0.0f0 - end @testset "First step" begin From 5e4659167bf8b43d58e5aaa6f23fbe6501a082c5 Mon Sep 17 00:00:00 2001 From: jaimerz Date: Wed, 26 Jul 2023 16:41:57 +0100 Subject: [PATCH 20/29] format --- test/constructors.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index db62fb39..c411e197 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -57,14 +57,13 @@ include("common.jl") ), ), ( - NUTS(T(0.8); integrator = :temperedleapfrog), + NUTS(T(0.8); integrator = :temperedleapfrog), ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, integrator_type = TemperedLeapfrog{T}, ), ), - ] # Make sure the sampler element type is preserved. @test AdvancedHMC.sampler_eltype(sampler) == T @@ -73,7 +72,7 @@ include("common.jl") rng = Random.default_rng() transition, state = AbstractMCMC.step(rng, model, sampler; n_adapts = 0, init_params = θ_init) - + # Verify that the types are preserved in the transition. @test eltype(transition.z.θ) == T @test eltype(transition.z.r) == T From 795fefbe9451f4e8fdac23ca166e063286d34540 Mon Sep 17 00:00:00 2001 From: jaimerz Date: Wed, 26 Jul 2023 16:50:20 +0100 Subject: [PATCH 21/29] docs update for init_e --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index d4dfcc77..19d31b99 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ If you are interested in using AdvancedHMC.jl through a probabilistic programmin **API CHANGES** - [v0.5.0] **Breaking!** Convenience constructors for common samplers changed to: - - `HMC(init_ϵ, n_leapfrog)` + - `HMC(n_leapfrog)` - `NUTS(target_acceptance)` - `HMCDA(target_acceptance, integration_time)` - [v0.2.22] Three functions are renamed. @@ -152,8 +152,8 @@ In the previous examples, we built the sampler by manually specifying the integr ```julia # HMC Sampler # step size, number of leapfrog steps - ϵ, n_leapfrogs = 0.1, 0.25 - hmc = HMC(ϵ, n_leapfrogs) + n_leapfrogs = 0.25 + hmc = HMC(n_leapfrogs) ``` Equivalent to: From f5679ec71e07cc968b387efa3b88dcf316eb939e Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 17:15:27 +0100 Subject: [PATCH 22/29] Update constructors.jl --- test/constructors.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index c411e197..8a288c39 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -8,16 +8,16 @@ include("common.jl") @testset "$T" for T in [Float32, Float64] @testset "$(nameof(typeof(sampler)))" for (sampler, expected) in [ ( - HMC(25), + HMC(25, integrator = Leapfrog(0.1)), ( adaptor_type = NoAdaptation, metric_type = DiagEuclideanMetric{T}, integrator_type = Leapfrog{T}, ), ), - # This should peform the correct promotion for the 2nd argument. + # This should perform the correct promotion for the 2nd argument. ( - HMCDA(T(0.1), 1), + HMCDA(T(0.8), 1, integrator = Leapfrog(0.1)), ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, From d0673a20e353aeadeae80c213afe83440ede64c2 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 17:15:44 +0100 Subject: [PATCH 23/29] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 19d31b99..dfdc9e31 100644 --- a/README.md +++ b/README.md @@ -152,8 +152,8 @@ In the previous examples, we built the sampler by manually specifying the integr ```julia # HMC Sampler # step size, number of leapfrog steps - n_leapfrogs = 0.25 - hmc = HMC(n_leapfrogs) + n_leapfrogs, lf_integrator = 0.25, Leapfrog(0.1) + hmc = HMC(n_leapfrogs, integrator = lf_integrator) ``` Equivalent to: From 4e90876086172c0e6ac4bd2bcf14d9a0bd22b4af Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 17:18:25 +0100 Subject: [PATCH 24/29] Update constructors.jl --- test/constructors.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/constructors.jl b/test/constructors.jl index 8a288c39..00ee5083 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -8,7 +8,7 @@ include("common.jl") @testset "$T" for T in [Float32, Float64] @testset "$(nameof(typeof(sampler)))" for (sampler, expected) in [ ( - HMC(25, integrator = Leapfrog(0.1)), + HMC(25, integrator = Leapfrog(T(0.1))), ( adaptor_type = NoAdaptation, metric_type = DiagEuclideanMetric{T}, @@ -17,7 +17,7 @@ include("common.jl") ), # This should perform the correct promotion for the 2nd argument. ( - HMCDA(T(0.8), 1, integrator = Leapfrog(0.1)), + HMCDA(T(0.8), 1, integrator = Leapfrog(T(0.1))), ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, From 0857774b858598a80a205dd4daed4d18e49c7dc6 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 26 Jul 2023 19:15:17 +0100 Subject: [PATCH 25/29] More bugfixes. --- src/abstractmcmc.jl | 6 ++---- src/integrator.jl | 2 +- test/constructors.jl | 4 ++-- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index f5998591..4f7b6582 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -28,8 +28,6 @@ end getadaptor(state::HMCState) = state.adaptor getmetric(state::HMCState) = state.metric - -getintegrator(state::HMCState) = state.κ.τ.integrator getintegrator(state::HMCState) = state.κ.τ.integrator """ @@ -319,8 +317,8 @@ make_integrator(i::AbstractIntegrator, ϵ::Real) = i make_integrator(i::Symbol, ϵ::Real) = make_integrator(Val(i), ϵ) make_integrator(@nospecialize(i), ::Real) = error("Integrator $i not supported.") make_integrator(i::Val{:leapfrog}, ϵ::Real) = Leapfrog(ϵ) -make_integrator(i::Val{:jitteredleapfrog}, ϵ::Real) = JitteredLeapfrog(ϵ, 0.1ϵ) -make_integrator(i::Val{:temperedleapfrog}, ϵ::Real) = TemperedLeapfrog(ϵ, 1.0) +make_integrator(i::Val{:jitteredleapfrog}, ϵ::T) where T<:Real = JitteredLeapfrog(ϵ, T(0.1ϵ)) +make_integrator(i::Val{:temperedleapfrog}, ϵ::T) where T<:Real = TemperedLeapfrog(ϵ, T(1)) ######### diff --git a/src/integrator.jl b/src/integrator.jl index 2a2a26c2..5dd69ac4 100644 --- a/src/integrator.jl +++ b/src/integrator.jl @@ -132,7 +132,7 @@ function _jitter( lf::JitteredLeapfrog{FT,T}, ) where {FT<:AbstractFloat,T<:AbstractScalarOrVec{FT}} ϵ = lf.ϵ0 .* (1 .+ lf.jitter .* (2 .* rand(rng) .- 1)) - return @set lf.ϵ = ϵ + return @set lf.ϵ = FT.(ϵ) end jitter(rng::AbstractRNG, lf::JitteredLeapfrog) = _jitter(rng, lf) diff --git a/test/constructors.jl b/test/constructors.jl index 00ee5083..673838de 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -53,7 +53,7 @@ include("common.jl") ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, - integrator_type = JitterdLeapfrog{T}, + integrator_type = JitteredLeapfrog{T, T}, ), ), ( @@ -61,7 +61,7 @@ include("common.jl") ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, - integrator_type = TemperedLeapfrog{T}, + integrator_type = TemperedLeapfrog{T, T}, ), ), ] From f27f5a7147020d7a1b8514dbe2c5c192b8ae968c Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 19:25:52 +0100 Subject: [PATCH 26/29] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/abstractmcmc.jl | 5 +++-- test/constructors.jl | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 4f7b6582..bdb27b5b 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -317,8 +317,9 @@ make_integrator(i::AbstractIntegrator, ϵ::Real) = i make_integrator(i::Symbol, ϵ::Real) = make_integrator(Val(i), ϵ) make_integrator(@nospecialize(i), ::Real) = error("Integrator $i not supported.") make_integrator(i::Val{:leapfrog}, ϵ::Real) = Leapfrog(ϵ) -make_integrator(i::Val{:jitteredleapfrog}, ϵ::T) where T<:Real = JitteredLeapfrog(ϵ, T(0.1ϵ)) -make_integrator(i::Val{:temperedleapfrog}, ϵ::T) where T<:Real = TemperedLeapfrog(ϵ, T(1)) +make_integrator(i::Val{:jitteredleapfrog}, ϵ::T) where {T<:Real} = + JitteredLeapfrog(ϵ, T(0.1ϵ)) +make_integrator(i::Val{:temperedleapfrog}, ϵ::T) where {T<:Real} = TemperedLeapfrog(ϵ, T(1)) ######### diff --git a/test/constructors.jl b/test/constructors.jl index 673838de..5deb2df3 100644 --- a/test/constructors.jl +++ b/test/constructors.jl @@ -53,7 +53,7 @@ include("common.jl") ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, - integrator_type = JitteredLeapfrog{T, T}, + integrator_type = JitteredLeapfrog{T,T}, ), ), ( @@ -61,7 +61,7 @@ include("common.jl") ( adaptor_type = StanHMCAdaptor, metric_type = DiagEuclideanMetric{T}, - integrator_type = TemperedLeapfrog{T, T}, + integrator_type = TemperedLeapfrog{T,T}, ), ), ] From 6cd1d408a5f8719a126e8bb2ec518a172fb4b643 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 19:44:51 +0100 Subject: [PATCH 27/29] Update src/constructors.jl --- src/constructors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/constructors.jl b/src/constructors.jl index 3173618b..e54580ef 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -86,7 +86,7 @@ $(FIELDS) # Usage: ```julia -HMC(init_ϵ=0.05, n_leapfrog=10, integrator = :leapfrog, metric = :diagonal) +HMC(n_leapfrog=10, integrator = Leapfrog(0.05), metric = :diagonal) ``` """ struct HMC{T<:Real} <: AbstractHMCSampler{T} From 70cd531777417c5c554e9876cce0a36cf0bbb75f Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 19:46:11 +0100 Subject: [PATCH 28/29] Update src/constructors.jl --- src/constructors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/constructors.jl b/src/constructors.jl index e54580ef..f2238224 100644 --- a/src/constructors.jl +++ b/src/constructors.jl @@ -86,7 +86,7 @@ $(FIELDS) # Usage: ```julia -HMC(n_leapfrog=10, integrator = Leapfrog(0.05), metric = :diagonal) +HMC(10, integrator = Leapfrog(0.05), metric = :diagonal) ``` """ struct HMC{T<:Real} <: AbstractHMCSampler{T} From b05759dbc35d93deb9232a591cccc6c90b69671d Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Wed, 26 Jul 2023 20:00:33 +0100 Subject: [PATCH 29/29] Update README.md --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index dfdc9e31..f96e6e70 100644 --- a/README.md +++ b/README.md @@ -161,8 +161,7 @@ In the previous examples, we built the sampler by manually specifying the integr ```julia metric = DiagEuclideanMetric(D) hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) - initial_ϵ = find_good_stepsize(hamiltonian, initial_θ) - integrator = Leapfrog(initial_ϵ) + integrator = Leapfrog(0.1) kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(n_leapfrog))) adaptor = NoAdaptation() hmc = HMCSampler(kernel, metric, adaptor)