From 632669193eb37616ee887648e53ab0d9a834b1a9 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 17 Nov 2022 13:48:41 +0000 Subject: [PATCH 01/22] removed AD compats and added implementation of LogDensityProblems --- Project.toml | 1 + README.md | 5 +++- src/AdvancedHMC.jl | 57 ++++++++++++++++++++++++++++++++------ src/abstractmcmc.jl | 48 ++++++-------------------------- src/contrib/ad.jl | 18 ------------ src/contrib/forwarddiff.jl | 50 --------------------------------- src/contrib/zygote.jl | 18 ------------ 7 files changed, 61 insertions(+), 136 deletions(-) delete mode 100644 src/contrib/ad.jl delete mode 100644 src/contrib/forwarddiff.jl delete mode 100644 src/contrib/zygote.jl diff --git a/Project.toml b/Project.toml index 63895f4d..6f8710a7 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" InplaceOps = "505f98c9-085e-5b2c-8e89-488be7bf1f34" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Requires = "ae029012-a4dd-5104-9daa-d747884805df" diff --git a/README.md b/README.md index da701fa5..bdb3df6c 100644 --- a/README.md +++ b/README.md @@ -51,13 +51,16 @@ In this section we demonstrate a minimal example of sampling from a multivariate ```julia using AdvancedHMC, Distributions, ForwardDiff +using LogDensityProblems using LinearAlgebra # Choose parameter dimensionality and initial parameter value D = 10; initial_θ = rand(D) -# Define the target distribution +# Define the target distribution using the `LogDensityProblem` interface ℓπ(θ) = logpdf(MvNormal(zeros(D), I), θ) +LogDensityProblems.logdensity(::typeof(ℓπ), x) = ℓπ(x) +LogDensityProblems.dimension(::typeof(ℓπ)) = D # Set the number of samples to draw and warmup iterations n_samples, n_adapts = 2_000, 1_000 diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 6ca83048..96f343e1 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -17,6 +17,8 @@ using ArgCheck: @argcheck using DocStringExtensions +using LogDensityProblems + import AbstractMCMC import StatsBase: sample @@ -49,6 +51,40 @@ export Trajectory, HMCKernel, # Useful defaults +""" + $(TYPEDEF) + +Wrapper around something that implements the `LogDensityProblem` interface. + +This itself then implements the `LogDensityProblem` interface by simply deferring to the wrapped object. + +Since this is a sub-type of `AbstractMCMC.AbstractModel`, it is also compatible with the `AbstractMCMC` interface. +""" +struct LogDensityModel{L} <: AbstractMCMC.AbstractModel + logdensity::L +end + +LogDensityModel(logdensity, ad) = LogDensityModel(ADgradient(ad, logdensity)) + +function LogDensityProblems.ADgradient(kind::Symbol, ℓ::LogDensityModel) + return LogDensityModel(LogDensityProblems.ADgradient(kind, ℓ.logdensity)) +end + +for kind in [:ForwardDiff, :ReverseDiff, :Zygote, :Tracker, :Enzyme] + @eval function LogDensityProblems.ADgradient( + ::Val{$(QuoteNode(kind))}, ℓ::LogDensityModel + ) + return LogDensityModel(LogDensityProblems.ADgradient(Val($(QuoteNode(kind))), ℓ.logdensity)) + end +end + +LogDensityProblems.dimension(model::LogDensityModel) = LogDensityProblems.dimension(model.logdensity) +LogDensityProblems.capabilities(model::LogDensityModel) = LogDensityProblems.capabilities(model.logdensity) +LogDensityProblems.logdensity(model::LogDensityModel, x) = LogDensityProblems.logdensity(model.logdensity, x) +function LogDensityProblems.logdensity_and_gradient(model::LogDensityModel, x) + return LogDensityProblems.logdensity_and_gradient(model.logdensity, x) +end + struct NUTS{TS, TC} end """ @@ -135,7 +171,18 @@ export sample include("abstractmcmc.jl") export DifferentiableDensityModel -include("contrib/ad.jl") +# include("contrib/ad.jl") +Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( + metric, + Base.Fix1(LogDensityProblems.logdensity, ℓ), + Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ) +) +function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val}) + ℓ = LogDensityModel(LogDensityProblems.ADgradient(kind, ℓπ)) + return Hamiltonian(metric, ℓ) +end +Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = Hamiltonian(metric, ℓπ, Val(Symbol(m))) +Hamiltonian(metric::AbstractMetric, ℓπ) = Hamiltonian(metric, ℓπ, Val{:ForwardDiff}()) ### Init @@ -147,14 +194,6 @@ function __init__() include("contrib/diffeq.jl") end - @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin - include("contrib/forwarddiff.jl") - end - - @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin - include("contrib/zygote.jl") - end - @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin include("contrib/cuda.jl") end diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 037f0ab0..d0adad76 100644 --- a/src/abstractmcmc.jl +++ b/src/abstractmcmc.jl @@ -25,38 +25,6 @@ struct HMCSampler{K, M, A} <: AbstractMCMC.AbstractSampler end HMCSampler(kernel, metric) = HMCSampler(kernel, metric, Adaptation.NoAdaptation()) -""" - DifferentiableDensityModel(ℓπ, ∂ℓπ∂θ) - DifferentiableDensityModel(ℓπ, m::Module) - -A `AbstractMCMC.AbstractMCMCModel` representing a differentiable log-density. - -If a module `m` is given as the second argument, then `m` is assumed to be an -automatic-differentiation package and this will be used to compute the gradients. - -Note that the module `m` must be imported before usage, e.g. -```julia -using Zygote: Zygote -model = DifferentiableDensityModel(ℓπ, Zygote) -``` -results in a `model` which will use Zygote.jl as its AD-backend. - -# Fields -$(FIELDS) -""" -struct DifferentiableDensityModel{Tlogπ, T∂logπ∂θ} <: AbstractMCMC.AbstractModel - "Log-density. Maps `AbstractArray` to value of the log-density." - ℓπ::Tlogπ - "Gradient of log-density. Returns a tuple of `ℓπ` and the gradient evaluated at the given point." - ∂ℓπ∂θ::T∂logπ∂θ -end - -struct DummyMetric <: AbstractMetric end -function DifferentiableDensityModel(ℓπ, m::Module) - h = Hamiltonian(DummyMetric(), ℓπ, m) - return DifferentiableDensityModel(h.ℓπ, h.∂ℓπ∂θ) -end - """ HMCState @@ -91,7 +59,7 @@ end A convenient wrapper around `AbstractMCMC.sample` avoiding explicit construction of [`HMCSampler`](@ref). """ function AbstractMCMC.sample( - model::DifferentiableDensityModel, + model::LogDensityModel, kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, @@ -103,7 +71,7 @@ end function AbstractMCMC.sample( rng::Random.AbstractRNG, - model::DifferentiableDensityModel, + model::LogDensityModel, kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, @@ -129,7 +97,7 @@ function AbstractMCMC.sample( end function AbstractMCMC.sample( - model::DifferentiableDensityModel, + model::LogDensityModel, kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, @@ -146,7 +114,7 @@ end function AbstractMCMC.sample( rng::Random.AbstractRNG, - model::DifferentiableDensityModel, + model::LogDensityModel, kernel::AbstractMCMCKernel, metric::AbstractMetric, adaptor::AbstractAdaptor, @@ -175,7 +143,7 @@ end function AbstractMCMC.step( rng::AbstractRNG, - model::DifferentiableDensityModel, + model::LogDensityModel, spl::HMCSampler; init_params = nothing, kwargs... @@ -189,7 +157,7 @@ function AbstractMCMC.step( end # Construct the hamiltonian using the initial metric - hamiltonian = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) + hamiltonian = Hamiltonian(metric, model) # Get an initial sample. h, t = AdvancedHMC.sample_init(rng, hamiltonian, init_params) @@ -203,7 +171,7 @@ end function AbstractMCMC.step( rng::AbstractRNG, - model::DifferentiableDensityModel, + model::LogDensityModel, spl::HMCSampler, state::HMCState; nadapts::Int = 0, @@ -220,7 +188,7 @@ function AbstractMCMC.step( metric = state.metric # Reconstruct hamiltonian. - h = Hamiltonian(metric, model.ℓπ, model.∂ℓπ∂θ) + h = Hamiltonian(metric, model) # Make new transition. t = transition(rng, h, κ, t_old.z) diff --git a/src/contrib/ad.jl b/src/contrib/ad.jl deleted file mode 100644 index 33d6bcc1..00000000 --- a/src/contrib/ad.jl +++ /dev/null @@ -1,18 +0,0 @@ -const ADSUPPORT = (:ForwardDiff, :Zygote) -const ADAVAILABLE = Dict{Module, Function}() - -Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = ADAVAILABLE[m](metric, ℓπ) - -function Hamiltonian(metric::AbstractMetric, ℓπ) - available = collect(keys(ADAVAILABLE)) - if length(available) == 0 - support_list_str = join(ADSUPPORT, " or ") - error("MethodError: no method matching Hamiltonian(metric::AbstractMetric, ℓπ) because no backend is loaded. Please load an AD package ($support_list_str) first.") - elseif length(available) > 1 - available_list_str = join(keys(ADAVAILABLE), " and ") - constructor_list_str = join(map(m -> "Hamiltonian(metric, ℓπ, $m)", available), "\n ") - error("MethodError: Hamiltonian(metric::AbstractMetric, ℓπ) is ambiguous because multiple AD pakcages are available ($available_list_str). Please use AD explictly. Candidates:\n $constructor_list_str") - else - return Hamiltonian(metric, ℓπ, first(available)) - end -end diff --git a/src/contrib/forwarddiff.jl b/src/contrib/forwarddiff.jl deleted file mode 100644 index 122dbd7a..00000000 --- a/src/contrib/forwarddiff.jl +++ /dev/null @@ -1,50 +0,0 @@ -import .ForwardDiff, .ForwardDiff.DiffResults - -function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractVector) - res = DiffResults.GradientResult(θ) - ForwardDiff.gradient!(res, ℓπ, θ) - return DiffResults.value(res), DiffResults.gradient(res) -end - -# Implementation 1 -function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix) - jacob = similar(θ) - res = DiffResults.JacobianResult(similar(θ, size(θ, 2)), jacob) - ForwardDiff.jacobian!(res, ℓπ, θ) - jacob_full = DiffResults.jacobian(res) - - d, n = size(jacob) - for i in 1:n - jacob[:,i] = jacob_full[i,1+(i-1)*d:i*d] - end - return DiffResults.value(res), jacob -end - -# Implementation 2 -# function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix) -# local densities -# f(x) = (densities = ℓπ(x); sum(densities)) -# res = DiffResults.GradientResult(θ) -# ForwardDiff.gradient!(res, f, θ) -# return ForwardDiff.value.(densities), DiffResults.gradient(res) -# end - -# Implementation 3 -# function ∂ℓπ∂θ_forwarddiff(ℓπ, θ::AbstractMatrix) -# v = similar(θ, size(θ, 2)) -# g = similar(θ) -# for i in 1:size(θ, 2) -# res = GradientResult(θ[:,i]) -# gradient!(res, ℓπ, θ[:,i]) -# v[i] = value(res) -# g[:,i] = gradient(res) -# end -# return v, g -# end - -function ForwardDiffHamiltonian(metric::AbstractMetric, ℓπ) - ∂ℓπ∂θ(θ::AbstractVecOrMat) = ∂ℓπ∂θ_forwarddiff(ℓπ, θ) - return Hamiltonian(metric, ℓπ, ∂ℓπ∂θ) -end - -ADAVAILABLE[ForwardDiff] = ForwardDiffHamiltonian diff --git a/src/contrib/zygote.jl b/src/contrib/zygote.jl deleted file mode 100644 index d89511d3..00000000 --- a/src/contrib/zygote.jl +++ /dev/null @@ -1,18 +0,0 @@ -import .Zygote - -function ∂ℓπ∂θ_zygote(ℓπ, θ::AbstractVector) - res, back = Zygote.pullback(ℓπ, θ) - return res, first(back(Zygote.sensitivity(res))) -end - -function ∂ℓπ∂θ_zygote(ℓπ, θ::AbstractMatrix) - res, back = Zygote.pullback(ℓπ, θ) - return res, first(back(ones(eltype(res), size(res)))) -end - -function ZygoteADHamiltonian(metric::AbstractMetric, ℓπ) - ∂ℓπ∂θ(θ::AbstractVecOrMat) = ∂ℓπ∂θ_zygote(ℓπ, θ) - return Hamiltonian(metric, ℓπ, ∂ℓπ∂θ) -end - -ADAVAILABLE[Zygote] = ZygoteADHamiltonian From 364dfb4178bd47a8096bd18a1c7fb918b96bd559 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 17 Nov 2022 13:57:35 +0000 Subject: [PATCH 02/22] improved README a bit --- README.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index bdb3df6c..396590d7 100644 --- a/README.md +++ b/README.md @@ -54,13 +54,16 @@ using AdvancedHMC, Distributions, ForwardDiff using LogDensityProblems using LinearAlgebra +# Define the target distribution using the `LogDensityProblem` interface +struct Problem + dim::Int +end +LogDensityProblems.logdensity(p::Problem, θ) = logpdf(MvNormal(zeros(p.dim), I), θ) +LogDensityProblems.dimension(p::Problem) = p.dim + # Choose parameter dimensionality and initial parameter value D = 10; initial_θ = rand(D) - -# Define the target distribution using the `LogDensityProblem` interface -ℓπ(θ) = logpdf(MvNormal(zeros(D), I), θ) -LogDensityProblems.logdensity(::typeof(ℓπ), x) = ℓπ(x) -LogDensityProblems.dimension(::typeof(ℓπ)) = D +ℓπ = Problem(D) # Set the number of samples to draw and warmup iterations n_samples, n_adapts = 2_000, 1_000 From 4e6301bc1e49beafedb98838f0ec0144c2a6601b Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 17 Nov 2022 13:57:56 +0000 Subject: [PATCH 03/22] added LogDensityProblems to test deps --- test/Project.toml | 1 + test/integrator.jl | 10 ++++++++-- test/runtests.jl | 1 + 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index fb0fe35c..6c311902 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" diff --git a/test/integrator.jl b/test/integrator.jl index fe1279e5..c228f3c1 100644 --- a/test/integrator.jl +++ b/test/integrator.jl @@ -106,8 +106,14 @@ using Statistics: mean end - @testset "Analyitcal solution to Eq (2.11) of Neal (2011)" begin - negU = q -> -dot(q, q) / 2 + @testset "Analytical solution to Eq (2.11) of Neal (2011)" begin + struct NegU + dim::Int + end + (::NegU)(q) = -dot(q, q) / 2 + LogDensityProblems.logdensity(::NegU, x) = (x) + LogDensityProblems.dimension(d::NegU) = d.dim + negU = NegU(1) ϵ = 0.01 for lf in [ diff --git a/test/runtests.jl b/test/runtests.jl index 42bbeb8b..a2bc5d7c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,6 @@ using Comonicon using AdvancedHMC: AdvancedHMC +using LogDensityProblems: LogDensityProblems println("Environment variables for testing") println(ENV) From 7dea4dbbe0b0cdc64fd61bad73a3112343da12f8 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 17 Nov 2022 14:55:34 +0000 Subject: [PATCH 04/22] updated tests --- test/abstractmcmc.jl | 2 +- test/adaptation.jl | 9 +++------ test/common.jl | 17 +++++++++++++++++ test/integrator.jl | 2 +- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 3ffb2084..696f15e6 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -10,7 +10,7 @@ include("common.jl") θ_init = randn(rng, 2) - model = AdvancedHMC.DifferentiableDensityModel(ℓπ_gdemo, ForwardDiff) + model = AdvancedHMC.LogDensityModel(ℓπ_gdemo, ForwardDiff) init_eps = Leapfrog(1e-3) κ = NUTS(init_eps) metric = DiagEuclideanMetric(2) diff --git a/test/adaptation.jl b/test/adaptation.jl index 91dc941e..ae991c30 100644 --- a/test/adaptation.jl +++ b/test/adaptation.jl @@ -124,8 +124,7 @@ end σ² = 1 .+ abs.(randn(D)) # Diagonal Gaussian - target = MvNormal(zeros(D), Diagonal(σ²)) - ℓπ = θ -> logpdf(target, θ) + ℓπ = MvNormal(zeros(D), Diagonal(σ²)) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ σ² rtol=0.2 @@ -142,8 +141,7 @@ end Σ = m' * m # Correlated Gaussian - target = MvNormal(zeros(D), Σ) - ℓπ = θ -> logpdf(target, θ) + ℓπ = MvNormal(zeros(D), Σ) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ diag(Σ) rtol=0.2 @@ -156,8 +154,7 @@ end end @testset "Initialisation adaptor by metric" begin - target = MvNormal(zeros(D), I) - ℓπ = θ -> logpdf(target, θ) + ℓπ = MvNormal(zeros(D), I) mass_init = fill(0.5, D) res = runnuts(ℓπ, DiagEuclideanMetric(mass_init); n_samples=1) diff --git a/test/common.jl b/test/common.jl index e5382732..dd2da58a 100644 --- a/test/common.jl +++ b/test/common.jl @@ -11,6 +11,16 @@ const DETATOL = 1e-3 * D * TRATIO # Random tolerance const RNDATOL = 5e-2 * D * TRATIO * 2 +# Convenience +using Distributions: Distributions +using Bijectors: Bijectors +LogDensityProblems.dimension(d::Distributions.Distribution) = length(d) +function LogDensityProblems.logdensity(d::Distributions.Distribution, y) + b = Bijectors.inverse(Bijectors.bijector(d)) + x, logjac = Bijectors.with_logabsdet_jacobian(b, y) + return logpdf(d, x) + logjac +end + # Hand-coded multivariate Gaussian struct Gaussian{Tm, Ts} @@ -24,6 +34,9 @@ end ℓπ_gaussian(m, s, x) = ℓπ_gaussian(m .- x, s) +LogDensityProblems.dimension(g::Gaussian) = dim(g.m) +LogDensityProblems.logdensity(g::Gaussian, x) = ℓπ_gaussian(g.m. g.s, x) + function ∇ℓπ_gaussianl(m, s, x) g = m .- x v = ℓπ_gaussian(g, s) @@ -76,6 +89,10 @@ function ℓπ_gdemo(θ) return logprior + loglikelihood end +# Make compat with `LogDensityProblems`. +LogDensityProblems.dimension(::typeof(ℓπ_gdemo)) = 2 +LogDensityProblems.logdensity(::typeof(ℓπ_gdemo), θ) = ℓπ_gdemo(θ) + test_show(x) = test_show(s -> length(s) > 0, x) function test_show(pred, x) io = IOBuffer(; append = true) diff --git a/test/integrator.jl b/test/integrator.jl index c228f3c1..8553493f 100644 --- a/test/integrator.jl +++ b/test/integrator.jl @@ -111,7 +111,7 @@ using Statistics: mean dim::Int end (::NegU)(q) = -dot(q, q) / 2 - LogDensityProblems.logdensity(::NegU, x) = (x) + LogDensityProblems.logdensity(d::NegU, x) = d(x) LogDensityProblems.dimension(d::NegU) = d.dim negU = NegU(1) From 9300af936fd9355bee347f5068ce7adc946251a7 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Thu, 17 Nov 2022 16:59:55 +0000 Subject: [PATCH 05/22] Apply suggestions from code review Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 396590d7..c089d008 100644 --- a/README.md +++ b/README.md @@ -55,15 +55,15 @@ using LogDensityProblems using LinearAlgebra # Define the target distribution using the `LogDensityProblem` interface -struct Problem +struct LogTargetDensity dim::Int end -LogDensityProblems.logdensity(p::Problem, θ) = logpdf(MvNormal(zeros(p.dim), I), θ) -LogDensityProblems.dimension(p::Problem) = p.dim +LogDensityProblems.logdensity(p::LogTargetDensity, θ) = logpdf(MvNormal(zeros(p.dim), I), θ) +LogDensityProblems.dimension(p::LogTargetDensity) = p.dim # Choose parameter dimensionality and initial parameter value D = 10; initial_θ = rand(D) -ℓπ = Problem(D) +ℓπ = LogTargetDensity(D) # Set the number of samples to draw and warmup iterations n_samples, n_adapts = 2_000, 1_000 From c9a67921eaf628d8ede766a3d6d2c80fc652f451 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sat, 19 Nov 2022 04:21:56 +0000 Subject: [PATCH 06/22] fixed tests --- src/AdvancedHMC.jl | 3 +++ test/demo.jl | 25 ++++++++++++++++++------- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 96f343e1..42621071 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -66,10 +66,13 @@ end LogDensityModel(logdensity, ad) = LogDensityModel(ADgradient(ad, logdensity)) +# Apply the `ADgradient` wrapper to the wrapped `logdensity` function since we +# need to ensure that the outermost "wrapper" is a sub-type of `AbstractModel`. function LogDensityProblems.ADgradient(kind::Symbol, ℓ::LogDensityModel) return LogDensityModel(LogDensityProblems.ADgradient(kind, ℓ.logdensity)) end +# Otherwise we'll run into method ambiguities. for kind in [:ForwardDiff, :ReverseDiff, :Zygote, :Tracker, :Enzyme] @eval function LogDensityProblems.ADgradient( ::Val{$(QuoteNode(kind))}, ℓ::LogDensityModel diff --git a/test/demo.jl b/test/demo.jl index a38eaf25..825c4381 100644 --- a/test/demo.jl +++ b/test/demo.jl @@ -3,11 +3,17 @@ using AdvancedHMC, Distributions, ForwardDiff, ComponentArrays using LinearAlgebra @testset "Demo" begin - # Choose parameter dimensionality and initial parameter value - D = 10; initial_θ = rand(D) + # Define the target distribution using the `LogDensityProblem` interface + struct DemoProblem + dim::Int + end + LogDensityProblems.logdensity(p::DemoProblem, θ) = logpdf(MvNormal(zeros(p.dim), I), θ) + LogDensityProblems.dimension(p::DemoProblem) = p.dim - # Define the target distribution - ℓπ(θ) = logpdf(MvNormal(zeros(D), I), θ) + # Choose parameter dimensionality and initial parameter value + D = 10 + initial_θ = rand(D) + ℓπ = DemoProblem(D) # Set the number of samples to draw and warmup iterations n_samples, n_adapts = 2_000, 1_000 @@ -24,7 +30,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 = NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)) # Run the sampler to draw samples from the specified Gaussian, where @@ -42,8 +48,13 @@ end @testset "Demo ComponentsArrays" begin # target distribution parametrized by ComponentsArray - p1 = ComponentVector(μ = 2.0, σ=1) - ℓπ(p::ComponentArray) = -(1 .- p.μ)^2/(p.σ)^2 + p1 = ComponentVector(μ=2.0, σ=1) + struct DemoProblemComponentArrays end + function LogDensityProblems.logdensity(::DemoProblemComponentArrays, p::ComponentArray) + return -(1 .- p.μ)^2 / (p.σ)^2 + end + LogDensityProblems.dimension(::DemoProblemComponentArrays) = 2 + ℓπ = DemoProblemComponentArrays() # Define a Hamiltonian system D = length(p1) # number of parameters From ee10ad308559cf541e983778ca0df845d082115e Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Fri, 9 Dec 2022 08:35:28 +0000 Subject: [PATCH 07/22] depend on LogDensityProblemsAD and new version of AbstractMCMC --- Project.toml | 1 + src/AdvancedHMC.jl | 45 +++++++-------------------------------------- 2 files changed, 8 insertions(+), 38 deletions(-) diff --git a/Project.toml b/Project.toml index 6f8710a7..cca6e68e 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" InplaceOps = "505f98c9-085e-5b2c-8e89-488be7bf1f34" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Requires = "ae029012-a4dd-5104-9daa-d747884805df" diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 42621071..256a9ed9 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -18,8 +18,10 @@ using ArgCheck: @argcheck using DocStringExtensions using LogDensityProblems +using LogDensityProblemsAD: LogDensityProblemsAD import AbstractMCMC +using AbstractMCMC: LogDensityModel import StatsBase: sample @@ -51,43 +53,6 @@ export Trajectory, HMCKernel, # Useful defaults -""" - $(TYPEDEF) - -Wrapper around something that implements the `LogDensityProblem` interface. - -This itself then implements the `LogDensityProblem` interface by simply deferring to the wrapped object. - -Since this is a sub-type of `AbstractMCMC.AbstractModel`, it is also compatible with the `AbstractMCMC` interface. -""" -struct LogDensityModel{L} <: AbstractMCMC.AbstractModel - logdensity::L -end - -LogDensityModel(logdensity, ad) = LogDensityModel(ADgradient(ad, logdensity)) - -# Apply the `ADgradient` wrapper to the wrapped `logdensity` function since we -# need to ensure that the outermost "wrapper" is a sub-type of `AbstractModel`. -function LogDensityProblems.ADgradient(kind::Symbol, ℓ::LogDensityModel) - return LogDensityModel(LogDensityProblems.ADgradient(kind, ℓ.logdensity)) -end - -# Otherwise we'll run into method ambiguities. -for kind in [:ForwardDiff, :ReverseDiff, :Zygote, :Tracker, :Enzyme] - @eval function LogDensityProblems.ADgradient( - ::Val{$(QuoteNode(kind))}, ℓ::LogDensityModel - ) - return LogDensityModel(LogDensityProblems.ADgradient(Val($(QuoteNode(kind))), ℓ.logdensity)) - end -end - -LogDensityProblems.dimension(model::LogDensityModel) = LogDensityProblems.dimension(model.logdensity) -LogDensityProblems.capabilities(model::LogDensityModel) = LogDensityProblems.capabilities(model.logdensity) -LogDensityProblems.logdensity(model::LogDensityModel, x) = LogDensityProblems.logdensity(model.logdensity, x) -function LogDensityProblems.logdensity_and_gradient(model::LogDensityModel, x) - return LogDensityProblems.logdensity_and_gradient(model.logdensity, x) -end - struct NUTS{TS, TC} end """ @@ -180,8 +145,12 @@ Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( Base.Fix1(LogDensityProblems.logdensity, ℓ), Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ) ) +function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}) + ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity)) + return Hamiltonian(metric, ℓ) +end function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val}) - ℓ = LogDensityModel(LogDensityProblems.ADgradient(kind, ℓπ)) + ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ)) return Hamiltonian(metric, ℓ) end Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = Hamiltonian(metric, ℓπ, Val(Symbol(m))) From 1f918e8be49cd643e1e83f958071a70cd9f0f71b Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sat, 17 Dec 2022 15:22:22 +0000 Subject: [PATCH 08/22] removed unnecessary comment and export of no-longer-existing model --- src/AdvancedHMC.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 256a9ed9..a364fe2d 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -137,9 +137,7 @@ include("sampler.jl") export sample include("abstractmcmc.jl") -export DifferentiableDensityModel -# include("contrib/ad.jl") Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( metric, Base.Fix1(LogDensityProblems.logdensity, ℓ), From 15c41df73ab9f160a02bfa97104c979a12d5625f Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sat, 17 Dec 2022 15:25:36 +0000 Subject: [PATCH 09/22] updated to work with new AbstractMCMC --- src/AdvancedHMC.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index a364fe2d..68bc8e5c 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -140,8 +140,8 @@ include("abstractmcmc.jl") Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( metric, - Base.Fix1(LogDensityProblems.logdensity, ℓ), - Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ) + Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity), + Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ.logdensity), ) function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}) ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity)) From 574427fd9f5b57512e0a533bfcee7d57ad9cb028 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Mon, 19 Dec 2022 16:04:49 +0000 Subject: [PATCH 10/22] bump compat entry for AbstractMCMC --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 94c109e4..8e5f891d 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" [compat] -AbstractMCMC = "3.2, 4" +AbstractMCMC = "4.2" ArgCheck = "1, 2" DocStringExtensions = "0.8, 0.9" InplaceOps = "0.3" From afaea4f74e94871c2a9a2302e678e5f37ea98bdb Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 20 Dec 2022 19:35:38 +0000 Subject: [PATCH 11/22] fixed bug in tests --- test/Project.toml | 1 + test/abstractmcmc.jl | 2 +- test/runtests.jl | 1 + 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index 6c311902..6ab625b3 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,6 +8,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" diff --git a/test/abstractmcmc.jl b/test/abstractmcmc.jl index 696f15e6..76e7c362 100644 --- a/test/abstractmcmc.jl +++ b/test/abstractmcmc.jl @@ -10,7 +10,7 @@ include("common.jl") θ_init = randn(rng, 2) - model = AdvancedHMC.LogDensityModel(ℓπ_gdemo, ForwardDiff) + model = AdvancedHMC.LogDensityModel(LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ_gdemo)) init_eps = Leapfrog(1e-3) κ = NUTS(init_eps) metric = DiagEuclideanMetric(2) diff --git a/test/runtests.jl b/test/runtests.jl index f54e9ba2..e780ae67 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,7 @@ using Comonicon using AdvancedHMC: AdvancedHMC using LogDensityProblems: LogDensityProblems +using LogDensityProblemsAD: LogDensityProblemsAD println("Environment variables for testing") println(ENV) From 45327256f66dc2f4355275cfad623d792afd54cf Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 20 Dec 2022 19:35:44 +0000 Subject: [PATCH 12/22] added compat entries for LogDensityProblems and AD --- Project.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Project.toml b/Project.toml index 8e5f891d..1ae0eb8a 100644 --- a/Project.toml +++ b/Project.toml @@ -24,6 +24,8 @@ AbstractMCMC = "4.2" ArgCheck = "1, 2" DocStringExtensions = "0.8, 0.9" InplaceOps = "0.3" +LogDensityProblems = "2" +LogDensityProblemsAD = "1" ProgressMeter = "1" Requires = "0.5, 1" Setfield = "0.7, 0.8, 1" From c7afe423c33ea21336e3791bd461f0ef813dc053 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Tue, 20 Dec 2022 19:41:07 +0000 Subject: [PATCH 13/22] removed type-piracy in tests --- test/adaptation.jl | 6 +++--- test/common.jl | 8 ++++++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/test/adaptation.jl b/test/adaptation.jl index ae991c30..90404839 100644 --- a/test/adaptation.jl +++ b/test/adaptation.jl @@ -124,7 +124,7 @@ end σ² = 1 .+ abs.(randn(D)) # Diagonal Gaussian - ℓπ = MvNormal(zeros(D), Diagonal(σ²)) + ℓπ = LogDensityDistribution(MvNormal(zeros(D), Diagonal(σ²))) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ σ² rtol=0.2 @@ -141,7 +141,7 @@ end Σ = m' * m # Correlated Gaussian - ℓπ = MvNormal(zeros(D), Σ) + ℓπ = LogDensityDistribution(MvNormal(zeros(D), Σ)) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ diag(Σ) rtol=0.2 @@ -154,7 +154,7 @@ end end @testset "Initialisation adaptor by metric" begin - ℓπ = MvNormal(zeros(D), I) + ℓπ = LogDensityDistribution(MvNormal(zeros(D), I)) mass_init = fill(0.5, D) res = runnuts(ℓπ, DiagEuclideanMetric(mass_init); n_samples=1) diff --git a/test/common.jl b/test/common.jl index dd2da58a..b13fc56c 100644 --- a/test/common.jl +++ b/test/common.jl @@ -14,8 +14,12 @@ const RNDATOL = 5e-2 * D * TRATIO * 2 # Convenience using Distributions: Distributions using Bijectors: Bijectors -LogDensityProblems.dimension(d::Distributions.Distribution) = length(d) -function LogDensityProblems.logdensity(d::Distributions.Distribution, y) +struct LogDensityDistribution{D<:Distributions.Distribution} + dist::D +end +LogDensityProblems.dimension(d::LogDensityDistribution) = length(d.dist) +function LogDensityProblems.logdensity(ld::LogDensityDistribution, y) + d = ld.dist b = Bijectors.inverse(Bijectors.bijector(d)) x, logjac = Bijectors.with_logabsdet_jacobian(b, y) return logpdf(d, x) + logjac From 7d88cee9b1714714ee523350fdbb68fe4b159e94 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Fri, 23 Dec 2022 09:21:00 +0000 Subject: [PATCH 14/22] Apply suggestions from code review Co-authored-by: David Widmann --- src/AdvancedHMC.jl | 3 +-- test/adaptation.jl | 6 +++--- test/demo.jl | 2 +- test/integrator.jl | 3 +-- 4 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index a681dab2..a16ec822 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -153,12 +153,11 @@ function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity)) return Hamiltonian(metric, ℓ) end -function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val}) +function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val} = Val{:ForwardDiff}()) ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ)) return Hamiltonian(metric, ℓ) end Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = Hamiltonian(metric, ℓπ, Val(Symbol(m))) -Hamiltonian(metric::AbstractMetric, ℓπ) = Hamiltonian(metric, ℓπ, Val{:ForwardDiff}()) ### Init diff --git a/test/adaptation.jl b/test/adaptation.jl index 90404839..3aad5982 100644 --- a/test/adaptation.jl +++ b/test/adaptation.jl @@ -124,7 +124,7 @@ end σ² = 1 .+ abs.(randn(D)) # Diagonal Gaussian - ℓπ = LogDensityDistribution(MvNormal(zeros(D), Diagonal(σ²))) + ℓπ = LogDensityDistribution(MvNormal(Diagonal(σ²))) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ σ² rtol=0.2 @@ -141,7 +141,7 @@ end Σ = m' * m # Correlated Gaussian - ℓπ = LogDensityDistribution(MvNormal(zeros(D), Σ)) + ℓπ = LogDensityDistribution(MvNormal(Σ)) res = runnuts(ℓπ, DiagEuclideanMetric(D)) @test res.adaptor.pc.var ≈ diag(Σ) rtol=0.2 @@ -154,7 +154,7 @@ end end @testset "Initialisation adaptor by metric" begin - ℓπ = LogDensityDistribution(MvNormal(zeros(D), I)) + ℓπ = LogDensityDistribution(MvNormal(Eye(D))) mass_init = fill(0.5, D) res = runnuts(ℓπ, DiagEuclideanMetric(mass_init); n_samples=1) diff --git a/test/demo.jl b/test/demo.jl index 825c4381..c5e93009 100644 --- a/test/demo.jl +++ b/test/demo.jl @@ -51,7 +51,7 @@ end p1 = ComponentVector(μ=2.0, σ=1) struct DemoProblemComponentArrays end function LogDensityProblems.logdensity(::DemoProblemComponentArrays, p::ComponentArray) - return -(1 .- p.μ)^2 / (p.σ)^2 + return -((1 - p.μ) / p.σ)^2 end LogDensityProblems.dimension(::DemoProblemComponentArrays) = 2 ℓπ = DemoProblemComponentArrays() diff --git a/test/integrator.jl b/test/integrator.jl index 8553493f..ee350aa6 100644 --- a/test/integrator.jl +++ b/test/integrator.jl @@ -110,8 +110,7 @@ using Statistics: mean struct NegU dim::Int end - (::NegU)(q) = -dot(q, q) / 2 - LogDensityProblems.logdensity(d::NegU, x) = d(x) + LogDensityProblems.logdensity(d::NegU, x) = -dot(x, x) / 2 LogDensityProblems.dimension(d::NegU) = d.dim negU = NegU(1) From 3b009cf4415c4fef46e501c71ed59df66d29d211 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Fri, 23 Dec 2022 11:43:45 +0000 Subject: [PATCH 15/22] added FillArrays to tests --- test/Project.toml | 1 + test/runtests.jl | 1 + 2 files changed, 2 insertions(+) diff --git a/test/Project.toml b/test/Project.toml index 6ab625b3..d8efcd8c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -5,6 +5,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" diff --git a/test/runtests.jl b/test/runtests.jl index e780ae67..df8ff95a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,5 @@ using Comonicon +using FillArrays using AdvancedHMC: AdvancedHMC using LogDensityProblems: LogDensityProblems using LogDensityProblemsAD: LogDensityProblemsAD From b655d504bf81dde8b2bb4275c6bbc6410a31e691 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Fri, 23 Dec 2022 23:59:42 +0000 Subject: [PATCH 16/22] added manual specification of the gradientconfig when working with ComponentArrays --- test/demo.jl | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/test/demo.jl b/test/demo.jl index c5e93009..73de3fa4 100644 --- a/test/demo.jl +++ b/test/demo.jl @@ -1,5 +1,5 @@ using ReTest -using AdvancedHMC, Distributions, ForwardDiff, ComponentArrays +using AdvancedHMC, Distributions, ForwardDiff, ComponentArrays, AbstractMCMC using LinearAlgebra @testset "Demo" begin @@ -61,7 +61,18 @@ end metric = DiagEuclideanMetric(D) # choose AD framework or provide a function manually - hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff) + ℓπ_with_grad = LogDensityProblemsAD.ADgradient( + Val(:ForwardDiff), + ℓπ; + # NOTE: Have to manually construct the `gradientconfig` to ensure we're not using the default + # buffer in the config (which won't be of type `ComponentArray`). + gradientconfig = ForwardDiff.GradientConfig( + Base.Fix1(LogDensityProblems.logdensity, ℓπ), + similar(p1), + LogDensityProblemsAD._default_chunk(ℓπ), + ) + ) + hamiltonian = Hamiltonian(metric, AbstractMCMC.LogDensityModel(ℓπ_with_grad)) # Define a leapfrog solver, with initial step size chosen heuristically initial_ϵ = find_good_stepsize(hamiltonian, p1) From aa0c11af1c1f144d86b9f373100eb25f802c3c89 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 10:59:05 +0000 Subject: [PATCH 17/22] Apply suggestions from code review Co-authored-by: David Widmann --- src/AdvancedHMC.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index a16ec822..09eb9af0 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -149,15 +149,15 @@ Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity), Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ.logdensity), ) -function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}) - ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity)) +function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}; kwargs...) + ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity; kwargs...)) return Hamiltonian(metric, ℓ) end -function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val} = Val{:ForwardDiff}()) - ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ)) +function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val} = Val{:ForwardDiff}(); kwargs...) + ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ; kwargs...)) return Hamiltonian(metric, ℓ) end -Hamiltonian(metric::AbstractMetric, ℓπ, m::Module) = Hamiltonian(metric, ℓπ, Val(Symbol(m))) +Hamiltonian(metric::AbstractMetric, ℓπ, m::Module; kwargs...) = Hamiltonian(metric, ℓπ, Val(Symbol(m)); kwargs...) ### Init From 0f7bbec2e5412e758421f7e5fb1d6708edf63d8d Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 11:46:56 +0000 Subject: [PATCH 18/22] updated tests to work with LogDensityProblemsAD --- test/Project.toml | 3 +++ test/demo.jl | 14 +------------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index d8efcd8c..e112465f 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -19,3 +19,6 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +LogDensityProblemsAD = "1.1.1" \ No newline at end of file diff --git a/test/demo.jl b/test/demo.jl index 73de3fa4..e1554189 100644 --- a/test/demo.jl +++ b/test/demo.jl @@ -61,18 +61,7 @@ end metric = DiagEuclideanMetric(D) # choose AD framework or provide a function manually - ℓπ_with_grad = LogDensityProblemsAD.ADgradient( - Val(:ForwardDiff), - ℓπ; - # NOTE: Have to manually construct the `gradientconfig` to ensure we're not using the default - # buffer in the config (which won't be of type `ComponentArray`). - gradientconfig = ForwardDiff.GradientConfig( - Base.Fix1(LogDensityProblems.logdensity, ℓπ), - similar(p1), - LogDensityProblemsAD._default_chunk(ℓπ), - ) - ) - hamiltonian = Hamiltonian(metric, AbstractMCMC.LogDensityModel(ℓπ_with_grad)) + hamiltonian = Hamiltonian(metric, ℓπ, Val(:ForwardDiff); x=p1) # Define a leapfrog solver, with initial step size chosen heuristically initial_ϵ = find_good_stepsize(hamiltonian, p1) @@ -92,5 +81,4 @@ end labels = ComponentArrays.labels(samples[1]) @test "μ" ∈ labels @test "σ" ∈ labels - end From 5cd6f4aa3cc76809f03c57f48f5517383ddad323 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 11:49:59 +0000 Subject: [PATCH 19/22] remove usage of Distributions in README --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index c089d008..94f43672 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ In this section we demonstrate a minimal example of sampling from a multivariate ```julia -using AdvancedHMC, Distributions, ForwardDiff +using AdvancedHMC, ForwardDiff using LogDensityProblems using LinearAlgebra @@ -58,7 +58,7 @@ using LinearAlgebra struct LogTargetDensity dim::Int end -LogDensityProblems.logdensity(p::LogTargetDensity, θ) = logpdf(MvNormal(zeros(p.dim), I), θ) +LogDensityProblems.logdensity(p::LogTargetDensity, θ) = -sum(abs2, θ) / 2 # standard multivariate normal LogDensityProblems.dimension(p::LogTargetDensity) = p.dim # Choose parameter dimensionality and initial parameter value From 0fff007fc1a8e572020d491d96f6c2545bc73cda Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 11:55:03 +0000 Subject: [PATCH 20/22] added a TODO --- test/common.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/common.jl b/test/common.jl index b13fc56c..f516eecf 100644 --- a/test/common.jl +++ b/test/common.jl @@ -12,6 +12,7 @@ const DETATOL = 1e-3 * D * TRATIO const RNDATOL = 5e-2 * D * TRATIO * 2 # Convenience +# TODO: Remove this if made available in some other package. using Distributions: Distributions using Bijectors: Bijectors struct LogDensityDistribution{D<:Distributions.Distribution} From ff14099776e28eb7318134b7355579bc66dc4488 Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 16:40:28 +0000 Subject: [PATCH 21/22] Update src/AdvancedHMC.jl Co-authored-by: David Widmann --- src/AdvancedHMC.jl | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index 09eb9af0..fa18f6ec 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -144,11 +144,20 @@ export sample include("abstractmcmc.jl") -Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) = Hamiltonian( - metric, - Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity), - Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ.logdensity), -) +function Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) + ℓπ = ℓ.logdensity + cap = LogDensityProblems.capabilities(ℓπ) + if cap === nothing + throw(ArgumentError("The log density function does not support the LogDensityProblems.jl interface")) + end + if cap === LogDensityProblems.LogDensityOrder{0}() + throw(ArgumentError("The gradient of the log density function is not defined: Implement `LogDensityProblems.logdensity_and_gradient` or use automatic differentiation by calling `Hamiltionian(metric, model, AD; kwargs...)` where AD is one of the backends supported by LogDensityProblemsAD.jl")) + return Hamiltonian( + metric, + Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity), + Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ.logdensity), + ) +end function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}; kwargs...) ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity; kwargs...)) return Hamiltonian(metric, ℓ) From b683501de82198d81d809c0e0def252dafd54d7d Mon Sep 17 00:00:00 2001 From: Tor Erlend Fjelde Date: Sun, 25 Dec 2022 16:43:31 +0000 Subject: [PATCH 22/22] fixed missing end --- src/AdvancedHMC.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index fa18f6ec..cf327ed2 100644 --- a/src/AdvancedHMC.jl +++ b/src/AdvancedHMC.jl @@ -146,12 +146,17 @@ include("abstractmcmc.jl") function Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel) ℓπ = ℓ.logdensity + + # Check we're capable of computing gradients. cap = LogDensityProblems.capabilities(ℓπ) if cap === nothing throw(ArgumentError("The log density function does not support the LogDensityProblems.jl interface")) end + if cap === LogDensityProblems.LogDensityOrder{0}() throw(ArgumentError("The gradient of the log density function is not defined: Implement `LogDensityProblems.logdensity_and_gradient` or use automatic differentiation by calling `Hamiltionian(metric, model, AD; kwargs...)` where AD is one of the backends supported by LogDensityProblemsAD.jl")) + end + return Hamiltonian( metric, Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity),