Objects of type LogTargetDensity are not callable

[arnold-pdev]

Hi, I’m trying to execute the following example for AdvancedHMC:

using AdvancedHMC, ForwardDiff
using LogDensityProblems
using LinearAlgebra

# Define the target distribution using the `LogDensityProblem` interface
struct LogTargetDensity
    dim::Int
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)
ℓπ = LogTargetDensity(D)

# Set the number of samples to draw and warmup iterations
n_samples, n_adapts = 2_000, 1_000

# Define a Hamiltonian system
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)

# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)

# Define an HMC sampler, with the following components
#   - multinomial sampling scheme,
#   - generalised No-U-Turn criteria, and
#   - windowed adaption for step-size and diagonal mass matrix
proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

# Run the sampler to draw samples from the specified Gaussian, where
#   - `samples` will store the samples
#   - `stats` will store diagnostic statistics for each sample
samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)

However, ForwardDiff does not seem equipped to handle the LogTargetDensity type. Here’s the stacktrace:

ERROR: MethodError: objects of type LogTargetDensity are not callable
Stacktrace:
  [1] vector_mode_dual_eval!(f::LogTargetDensity, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10}}}, x::Vector{Float64})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/apiutils.jl:24
  [2] vector_mode_gradient!(result::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, f::LogTargetDensity, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/gradient.jl:96
  [3] gradient!(result::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, f::LogTargetDensity, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10}}}, ::Val{true})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/gradient.jl:37
  [4] gradient!(result::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, f::LogTargetDensity, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogTargetDensity, Float64}, Float64, 10}}}) (repeats 2 times)
    @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/gradient.jl:35
  [5] ∂ℓπ∂θ_forwarddiff(ℓπ::LogTargetDensity, θ::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/contrib/forwarddiff.jl:5
  [6] (::AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity})(θ::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/contrib/forwarddiff.jl:46
  [7] ∂H∂θ(h::Hamiltonian{DiagEuclideanMetric{Float64, Vector{Float64}}, LogTargetDensity, AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity}}, θ::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/hamiltonian.jl:31
  [8] phasepoint(h::Hamiltonian{DiagEuclideanMetric{Float64, Vector{Float64}}, LogTargetDensity, AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity}}, θ::Vector{Float64}, r::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/hamiltonian.jl:76
  [9] find_good_stepsize(rng::Random._GLOBAL_RNG, h::Hamiltonian{DiagEuclideanMetric{Float64, Vector{Float64}}, LogTargetDensity, AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity}}, θ::Vector{Float64}; max_n_iters::Int64)
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/trajectory.jl:710
 [10] find_good_stepsize(h::Hamiltonian{DiagEuclideanMetric{Float64, Vector{Float64}}, LogTargetDensity, AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity}}, θ::Vector{Float64}; max_n_iters::Int64)
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/trajectory.jl:770
 [11] find_good_stepsize(h::Hamiltonian{DiagEuclideanMetric{Float64, Vector{Float64}}, LogTargetDensity, AdvancedHMC.var"#∂ℓπ∂θ#66"{LogTargetDensity}}, θ::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/51xgc/src/trajectory.jl:765
 [12] top-level scope
    @ Untitled-1:25

Any suggestions? Thanks.

[yebai]

@JaimeRZP can you see what's the issue here please?

[torfjelde]

What version are you using @arnold-pdev ?

[ErikQQY]

With the latest version of AdvancedHMC.jl, this issue is now fixed. Below is a working example based on OP's example:

using AdvancedHMC, ForwardDiff
using LogDensityProblems
using LinearAlgebra

# Define the target distribution using the `LogDensityProblem` interface
struct LogTargetDensity
    dim::Int
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)
ℓπ = LogTargetDensity(D)

# Set the number of samples to draw and warmup iterations
n_samples, n_adapts = 2_000, 1_000

# Define a Hamiltonian system
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)

# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)

# Define an HMC sampler, with the following components
#   - multinomial sampling scheme,
#   - generalised No-U-Turn criteria, and
#   - windowed adaption for step-size and diagonal mass matrix
proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))

adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

# Run the sampler to draw samples from the specified Gaussian, where
#   - `samples` will store the samples
#   - `stats` will store diagnostic statistics for each sample
samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)