Update no-U-turn criteria and added some unit tests

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

README.md

@@ -14,8 +14,8 @@ using Distributions: logpdf, MvNormal
 using DiffResults: GradientResult, value, gradient
 using ForwardDiff: gradient!
 
-const D = 10
-const target = MvNormal(zeros(D), ones(D))
+D = 10
+target = MvNormal(zeros(D), ones(D))
 ℓπ(θ) = logpdf(target, θ)
 
 function ∂ℓπ∂θ(θ)
@@ -28,17 +28,17 @@ end
 using AdvancedHMC
 
 # Sampling parameter settings
-n_samples = 100_000
-n_adapts = 2_000
+n_samples, n_adapts = 10_000, 2_000
 
 # Draw a random starting points
 θ_init = randn(D)
 
 # Define metric space, Hamiltonian, sampling method and adaptor
 metric = DiagEuclideanMetric(D)
 h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
-prop = NUTS(Leapfrog(find_good_eps(h, θ_init)))
-adaptor = StanHMCAdaptor(n_adapts, Preconditioner(metric), NesterovDualAveraging(0.8, prop.integrator.ϵ))
+int = Leapfrog(find_good_eps(h, θ_init))
+prop = NUTS{MultinomialTrajectorySampler,GeneralisedNoUTurn}(int)
+adaptor = StanHMCAdaptor(n_adapts, Preconditioner(metric), NesterovDualAveraging(0.8, int.ϵ))
 
 # Draw samples via simulating Hamiltonian dynamics
 # - `samples` will store the samples

src/AdvancedHMC.jl

@@ -28,7 +28,7 @@ export Hamiltonian
 include("integrator.jl")
 export Leapfrog
 include("trajectory.jl")
-export StaticTrajectory, find_good_eps, HMCDA, NUTS
+export StaticTrajectory, find_good_eps, HMCDA, NUTS, SliceTrajectorySampler, MultinomialTrajectorySampler, NoUTurn, GeneralisedNoUTurn
 
 include("diagnosis.jl")
 include("sampler.jl")

src/hamiltonian.jl

@@ -98,6 +98,8 @@ function neg_energy(
     return -dot(r, h.metric._temp) / 2
 end
 
+energy(args...) = -neg_energy(args...)
+
 ####
 #### Momentum refreshment
 ####