Dont hard dep on MCMCLogDensityProblems

Author: ErikQQY

Project.toml

@@ -6,7 +6,6 @@ version = "0.7.1"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -16,11 +15,6 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-VecTargets = "8a639fad-7908-4fe4-8003-906e9297f002"
-
-[sources]
-VecTargets = {url = "https://github.com/chalk-lab/VecTargets.jl", rev = "main"}
-
 
 [weakdeps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -43,7 +37,6 @@ ArgCheck = "1, 2"
 ComponentArrays = "0.15"
 CUDA = "3, 4, 5"
 DocStringExtensions = "0.8, 0.9"
-ForwardDiff = "0.10.38"
 LinearAlgebra = "<0.1, 1"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1"

src/AdvancedHMC.jl

@@ -30,12 +30,8 @@ using LogDensityProblemsAD: LogDensityProblemsAD
 
 using AbstractMCMC: AbstractMCMC, LogDensityModel
 
-using VecTargets: VecTargets
-
 import StatsBase: sample
 
-using ForwardDiff: ForwardDiff
-
 const DEFAULT_FLOAT_TYPE = typeof(float(0))
 
 include("utilities.jl")

src/riemannian/metric.jl

@@ -43,57 +43,6 @@ function DenseRiemannianMetric(size, G, ∂G∂θ, map=IdentityMap())
     return DenseRiemannianMetric(size, G, ∂G∂θ, map, _temp)
 end
 
-# Convenient constructor
-function DenseRiemannianMetric(size, ℓπ, initial_θ, λ, map=IdentityMap())
-    _Hfunc = VecTargets.gen_hess(x -> -ℓπ(x), initial_θ) # x -> (value, gradient, hessian)
-    Hfunc = x -> copy.(_Hfunc(x)) # _Hfunc do in-place computation, copy to avoid bug
-
-    fstabilize = H -> H + λ * I
-    Gfunc = x -> begin
-        H = fstabilize(Hfunc(x)[3])
-        all(isfinite, H) ? H : diagm(ones(length(x)))
-    end
-    _∂G∂θfunc = gen_∂G∂θ_fwd(x -> -ℓπ(x), initial_θ; f=fstabilize)
-    ∂G∂θfunc = x -> reshape_∂G∂θ(_∂G∂θfunc(x))
-
-    _temp = Vector{Float64}(undef, first(size))
-
-    return DenseRiemannianMetric(size, Gfunc, ∂G∂θfunc, map, _temp)
-end
-
-function gen_hess_fwd(func, x::AbstractVector)
-    function hess(x::AbstractVector)
-        return nothing, nothing, ForwardDiff.hessian(func, x)
-    end
-    return hess
-end
-
-#= possible integrate DI for AD-independent fisher information metric
-function gen_∂G∂θ_rev(Vfunc, x; f=identity)
-    _Hfunc = VecTargets.gen_hess(Vfunc, ReverseDiff.track.(x))
-    Hfunc = x -> _Hfunc(x)[3]
-    # QUES What's the best output format of this function?
-    return x -> ReverseDiff.jacobian(x -> f(Hfunc(x)), x) # default output shape [∂H∂x₁; ∂H∂x₂; ...]
-end
-=#
-
-# Fisher information metric
-function gen_∂G∂θ_fwd(Vfunc, x; f=identity)
-    _Hfunc = gen_hess_fwd(Vfunc, x)
-    Hfunc = x -> _Hfunc(x)[3]
-    # QUES What's the best output format of this function?
-    cfg = ForwardDiff.JacobianConfig(Hfunc, x)
-    d = length(x)
-    out = zeros(eltype(x), d^2, d)
-    return x -> ForwardDiff.jacobian!(out, Hfunc, x, cfg)
-    return out # default output shape [∂H∂x₁; ∂H∂x₂; ...]
-end
-
-function reshape_∂G∂θ(H)
-    d = size(H, 2)
-    return cat((H[((i - 1) * d + 1):(i * d), :] for i in 1:d)...; dims=3)
-end
-
 Base.size(e::DenseRiemannianMetric) = e.size
 Base.size(e::DenseRiemannianMetric, dim::Int) = e.size[dim]
 function Base.show(io::IO, drm::DenseRiemannianMetric)

test/riemannian.jl

@@ -2,16 +2,60 @@ using ReTest, Random
 using AdvancedHMC, ForwardDiff, AbstractMCMC
 using LinearAlgebra
 
+using Pkg
+Pkg.develop(; url="https://github.com/chalk-lab/MCMCLogDensityProblems.jl")
+using MCMCLogDensityProblems
+
+# Fisher information metric
+function gen_∂G∂θ_fwd(Vfunc, x; f=identity)
+    _Hfunc = gen_hess_fwd(Vfunc, x)
+    Hfunc = x -> _Hfunc(x)[3]
+    # QUES What's the best output format of this function?
+    cfg = ForwardDiff.JacobianConfig(Hfunc, x)
+    d = length(x)
+    out = zeros(eltype(x), d^2, d)
+    return x -> ForwardDiff.jacobian!(out, Hfunc, x, cfg)
+    return out # default output shape [∂H∂x₁; ∂H∂x₂; ...]
+end
+
+function gen_hess_fwd(func, x::AbstractVector)
+    function hess(x::AbstractVector)
+        return nothing, nothing, ForwardDiff.hessian(func, x)
+    end
+    return hess
+end
+
+function reshape_∂G∂θ(H)
+    d = size(H, 2)
+    return cat((H[((i - 1) * d + 1):(i * d), :] for i in 1:d)...; dims=3)
+end
+
+function prepare_sample(ℓπ, initial_θ, λ)
+    _Hfunc = MCMCLogDensityProblems.gen_hess(x -> -ℓπ(x), initial_θ) # x -> (value, gradient, hessian)
+    Hfunc = x -> copy.(_Hfunc(x)) # _Hfunc do in-place computation, copy to avoid bug
+
+    fstabilize = H -> H + λ * I
+    Gfunc = x -> begin
+        H = fstabilize(Hfunc(x)[3])
+        all(isfinite, H) ? H : diagm(ones(length(x)))
+    end
+    _∂G∂θfunc = gen_∂G∂θ_fwd(x -> -ℓπ(x), initial_θ; f=fstabilize)
+    ∂G∂θfunc = x -> reshape_∂G∂θ(_∂G∂θfunc(x))
+
+    return Gfunc, ∂G∂θfunc
+end
+
 @testset "Multi variate Normal with Riemannian HMC" begin
     # Set the number of samples to draw and warmup iterations
     n_samples = 2_000
     rng = MersenneTwister(1110)
     initial_θ = rand(rng, D)
     λ = 1e-2
+    G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
     # Define a Hamiltonian system
-    metric = DenseRiemannianMetric((D,), ℓπ, initial_θ, λ)
+    metric = DenseRiemannianMetric((D,), G, ∂G∂θ)
     kinetic = GaussianKinetic()
-    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∇ℓπ)
+    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
 
     # Define a leapfrog solver, with the initial step size chosen heuristically
     initial_ϵ = 0.01
@@ -35,11 +79,13 @@ end
     n_samples = 2_000
     rng = MersenneTwister(1110)
     initial_θ = rand(rng, D)
+    λ = 1e-2
+    G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
 
     # Define a Hamiltonian system
-    metric = DenseRiemannianMetric((D,), ℓπ, initial_θ, λSoftAbsMap(20.0))
+    metric = DenseRiemannianMetric((D,), G, ∂G∂θ, λSoftAbsMap(20.0))
     kinetic = GaussianKinetic()
-    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∇ℓπ)
+    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
 
     # Define a leapfrog solver, with the initial step size chosen heuristically
     initial_ϵ = 0.01