Improve Hamiltonian constructors (#307)

* Improve Hamiltonian constructors

* Update Project.toml

* Simplify code

* Update Project.toml

* Update integrator.jl

* Update demo.jl

* Update common.jl

* Update README.md

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-version = "0.4"
+version = "0.4.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

README.md

@@ -60,6 +60,7 @@ struct LogTargetDensity
 end
 LogDensityProblems.logdensity(p::LogTargetDensity, θ) = -sum(abs2, θ) / 2  # standard multivariate normal
 LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
+LogDensityProblems.capabilities(::Type{LogTargetDensity}) = LogDensityProblems.LogDensityOrder{0}()
 
 # Choose parameter dimensionality and initial parameter value
 D = 10; initial_θ = rand(D)

src/AdvancedHMC.jl

@@ -144,34 +144,43 @@ export sample
 
 include("abstractmcmc.jl")
 
-function Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel)
-    ℓπ = ℓ.logdensity
-
-    # Check we're capable of computing gradients.
-    cap = LogDensityProblems.capabilities(ℓπ)
+## Without explicit AD backend
+function Hamiltonian(metric::AbstractMetric, ℓ::LogDensityModel; kwargs...)
+    return Hamiltonian(metric, ℓ.logdensity; kwargs...)
+end
+function Hamiltonian(metric::AbstractMetric, ℓ; kwargs...)
+    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"))
+    # Check if we're capable of computing gradients.
+    ℓπ = if cap === LogDensityProblems.LogDensityOrder{0}()
+        # In this case ℓ does not support evaluation of the gradient of the log density function
+        # We use ForwardDiff to compute the gradient
+        LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; kwargs...)
+    else
+        # In this case ℓ already supports evaluation of the gradient of the log density function
+        ℓ
     end
-
     return Hamiltonian(
         metric,
-        Base.Fix1(LogDensityProblems.logdensity, ℓ.logdensity),
-        Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ.logdensity),
+        Base.Fix1(LogDensityProblems.logdensity, ℓπ),
+        Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓπ),
     )
 end
-function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val}; kwargs...)
-    ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ.logdensity; kwargs...))
-    return Hamiltonian(metric, ℓ)
+
+## With explicit AD specification
+function Hamiltonian(metric::AbstractMetric, ℓπ::LogDensityModel, kind::Union{Symbol,Val,Module}; kwargs...)
+    return Hamiltonian(metric, ℓπ.logdensity, kind; kwargs...)
 end
-function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val} = Val{:ForwardDiff}(); kwargs...)
-    ℓ = LogDensityModel(LogDensityProblemsAD.ADgradient(kind, ℓπ; kwargs...))
+Hamiltonian(metric::AbstractMetric, ℓπ, m::Module; kwargs...) = Hamiltonian(metric, ℓπ, Val(Symbol(m)); kwargs...)
+function Hamiltonian(metric::AbstractMetric, ℓπ, kind::Union{Symbol,Val}; kwargs...)
+    if LogDensityProblems.capabilities(ℓπ) === nothing
+        throw(ArgumentError("The log density function does not support the LogDensityProblems.jl interface"))
+    end
+    ℓ = LogDensityProblemsAD.ADgradient(kind, ℓπ; kwargs...)
     return Hamiltonian(metric, ℓ)
 end
-Hamiltonian(metric::AbstractMetric, ℓπ, m::Module; kwargs...) = Hamiltonian(metric, ℓπ, Val(Symbol(m)); kwargs...)
 
 ### Init
 

test/common.jl

@@ -18,13 +18,15 @@ using Bijectors: Bijectors
 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
 end
+LogDensityProblems.capabilities(::Type{<:LogDensityDistribution}) = LogDensityProblems.LogDensityOrder{0}()
 
 # Hand-coded multivariate Gaussian
 
@@ -41,6 +43,7 @@ end
 
 LogDensityProblems.dimension(g::Gaussian) = dim(g.m)
 LogDensityProblems.logdensity(g::Gaussian, x) = ℓπ_gaussian(g.m. g.s, x)
+LogDensityProblems.capabilities(::Type{<:Gaussian}) = LogDensityProblems.LogDensityOrder{0}()
 
 function ∇ℓπ_gaussianl(m, s, x)
     g = m .- x
@@ -97,6 +100,7 @@ end
 # Make compat with `LogDensityProblems`.
 LogDensityProblems.dimension(::typeof(ℓπ_gdemo)) = 2
 LogDensityProblems.logdensity(::typeof(ℓπ_gdemo), θ) = ℓπ_gdemo(θ)
+LogDensityProblems.capabilities(::Type{typeof(ℓπ_gdemo)}) = LogDensityProblems.LogDensityOrder{0}()
 
 test_show(x) = test_show(s -> length(s) > 0, x)
 function test_show(pred, x)

test/demo.jl

@@ -7,8 +7,10 @@ using LinearAlgebra
     struct DemoProblem
         dim::Int
     end
+
     LogDensityProblems.logdensity(p::DemoProblem, θ) = logpdf(MvNormal(zeros(p.dim), I), θ)
     LogDensityProblems.dimension(p::DemoProblem) = p.dim
+    LogDensityProblems.capabilities(::Type{DemoProblem}) = LogDensityProblems.LogDensityOrder{0}()
 
     # Choose parameter dimensionality and initial parameter value
     D = 10
@@ -50,10 +52,13 @@ end
     # target distribution parametrized by ComponentsArray
     p1 = ComponentVector(μ=2.0, σ=1)
     struct DemoProblemComponentArrays end
+
     function LogDensityProblems.logdensity(::DemoProblemComponentArrays, p::ComponentArray)
         return -((1 - p.μ) / p.σ)^2
     end
     LogDensityProblems.dimension(::DemoProblemComponentArrays) = 2
+    LogDensityProblems.capabilities(::Type{DemoProblemComponentArrays}) = LogDensityProblems.LogDensityOrder{0}()
+
     ℓπ = DemoProblemComponentArrays()
 
     # Define a Hamiltonian system

test/integrator.jl

@@ -110,8 +110,11 @@ using Statistics: mean
         struct NegU
             dim::Int
         end
-        LogDensityProblems.logdensity(d::NegU, x) = -dot(x, x) / 2
+
+        LogDensityProblems.logdensity(::NegU, x) = -dot(x, x) / 2
         LogDensityProblems.dimension(d::NegU) = d.dim
+        LogDensityProblems.capabilities(::Type{NegU}) = LogDensityProblems.LogDensityOrder{0}()
+
         negU = NegU(1)
 
         ϵ = 0.01