Allow specifying AD backend in config

Project.toml

@@ -1,7 +1,7 @@
 name = "SampleChainsDynamicHMC"
 uuid = "6d9fd711-e8b2-4778-9c70-c1dfb499d4c4"
 authors = ["Chad Scherrer <[email protected]> and contributors"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
@@ -32,7 +32,9 @@ TupleVectors = "0.1"
 julia = "1.5"
 
 [extras]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["LinearAlgebra", "ReverseDiff", "Test"]

docs/Project.toml

@@ -1,3 +1,4 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SampleChainsDynamicHMC = "6d9fd711-e8b2-4778-9c70-c1dfb499d4c4"

src/SampleChainsDynamicHMC.jl

@@ -59,6 +59,7 @@ end
     warmup_stages
     algorithm
     reporter
+    ad_backend
 end
 
 # Docs adapted from https://tamaspapp.eu/DynamicHMC.jl/stable/interface/
@@ -68,6 +69,7 @@ end
       , warmup_stages = DynamicHMC.default_warmup_stages()
       , algorithm     = DynamicHMC.NUTS()
       , reporter      = DynamicHMC.NoProgressReport()
+      , ad_backend    = Val(:ForwardDiff)
     )
 
 `init`: a `NamedTuple` that can contain the following fields (all of them
@@ -88,11 +90,16 @@ except perhaps for the maximum depth.
 `DynamicHMC.NoProgressReport`), but this default will likely change in future
 releases. 
 
+`ad_backend`: The automatic differentiation backend to use for gradient
+computation, specified as either a symbol or a `Val` type with a symbol that
+refers to an AD package. See [LogDensityProblems.jl](https://tamaspapp.eu/LogDensityProblems.jl/stable/#Automatic-differentiation)
+for supported packages, including `ForwardDiff`, `ReverseDiff`, `Zygote`, and `Tracker`.
+
 For more details see https://tamaspapp.eu/DynamicHMC.jl/stable/interface/
 # Example
 
 ```jldoctest
-julia> using LinearAlgebra
+julia> using LinearAlgebra, ReverseDiff
 
 julia> config = dynamichmc(
            warmup_stages=default_warmup_stages(
@@ -101,6 +108,7 @@ julia> config = dynamichmc(
                doubling_stages=7, # 7-stage metric adaptation
            ),
            reporter=LogProgressReport(), # log progress using Logging
+           ad_backend=Val(:ReverseDiff), # use ReverseDiff AD package
        );
 ```
 """
@@ -109,13 +117,14 @@ function dynamichmc(;
   , warmup_stages = DynamicHMC.default_warmup_stages()
   , algorithm     = DynamicHMC.NUTS()
   , reporter      = DynamicHMC.NoProgressReport()
+  , ad_backend    = Val(:ForwardDiff)
 )
-    DynamicHMCConfig(init, warmup_stages, algorithm, reporter)
+    DynamicHMCConfig(init, warmup_stages, algorithm, reporter, ad_backend)
 end
 
-function SampleChains.newchain(rng::Random.AbstractRNG, config::DynamicHMCConfig, ℓ, tr, ad_backend=Val(:ForwardDiff))
+function SampleChains.newchain(rng::Random.AbstractRNG, config::DynamicHMCConfig, ℓ, tr)
     P = LogDensityProblems.TransformedLogDensity(tr, ℓ)
-    ∇P = LogDensityProblems.ADgradient(ad_backend, P)
+    ∇P = LogDensityProblems.ADgradient(config.ad_backend, P)
     reporter = DynamicHMC.NoProgressReport()
 
     results = DynamicHMC.mcmc_keep_warmup(rng, ∇P, 0; 
@@ -135,9 +144,8 @@ function SampleChains.newchain(rng::Random.AbstractRNG, config::DynamicHMCConfig
     chain = DynamicHMCChain(tr, Q, tree_stats, steps)
 end
 
-function SampleChains.newchain(config::DynamicHMCConfig, ℓ, tr, ad_backend=Val(:ForwardDiff))
-    rng = Random.GLOBAL_RNG
-    return newchain(rng, config, ℓ, tr, ad_backend)
+function SampleChains.newchain(config::DynamicHMCConfig, ℓ, tr)
+    return newchain(Random.GLOBAL_RNG, config, ℓ, tr)
 end
 
 function SampleChains.sample!(chain::DynamicHMCChain, n::Int=1000)

test/notebook.jl

@@ -2,6 +2,7 @@
 
 using SampleChainsDynamicHMC
 using TransformVariables
+using Test
 
 function ℓ(nt)
 	z = nt.x/nt.σ
@@ -10,14 +11,57 @@ end
 
 t = as((x=asℝ, σ=asℝ₊))
 
-chain = newchain(dynamichmc(), ℓ, t)
+@testset "single chain" begin
+	chain = newchain(dynamichmc(), ℓ, t)
+	@test chain isa SampleChainsDynamicHMC.DynamicHMCChain
+	@test length(chain) == 1
 
-sample!(chain, 9)
-sample!(chain, 90)
+	sample!(chain, 9)
+	@test length(chain) == 10
 
-chains = newchain(4, dynamichmc(), ℓ, t)
+	sample!(chain, 90)
+	@test length(chain) == 100
+end
+
+@testset "multichain" begin
+	chains = newchain(4, dynamichmc(), ℓ, t)
+	@test chains isa MultiChain
+	chains_chains = getfield(chains, :chains)
+	@test length(chains_chains) == 4
+	@test all(x -> length(x) == 1, chains_chains)
+
+	sample!(chains, 9)
+	@test all(x -> length(x) == 10, chains_chains)
+	sample!(chains, 90)
+	@test all(x -> length(x) == 100, chains_chains)
+
+	samples(chains)
+end
 
-sample!(chains, 9)
-sample!(chains, 90)
+using LinearAlgebra
+using ReverseDiff
+using SampleChainsDynamicHMC.LogDensityProblems
 
-samples(chains)
+@testset "config options" begin
+	config = dynamichmc(
+		warmup_stages=default_warmup_stages(
+			M=Symmetric, # adapt dense positive definite metric
+			stepsize_adaptation=DualAveraging(δ=0.9), # target acceptance rate 0.9
+			doubling_stages=7, # 7-stage metric adaptation
+		),
+		reporter=LogProgressReport(), # log progress using Logging
+		ad_backend=Val(:ReverseDiff), # use ReverseDiff AD package
+	)
+
+	chain = newchain(config, ℓ, t)
+	@test length(chain) == 1
+
+	meta = getfield(chain, :meta)
+	@test meta.H.κ.M⁻¹ isa Symmetric
+	@test meta.H.ℓ isa LogDensityProblems.ReverseDiffLogDensity
+
+	sample!(chain, 9)
+	@test length(chain) == 10
+	sample!(chain, 90)
+	@test length(chain) == 100
+end