Using DynamicHMC to calculate paths of Quantum Simple Harmonic Oscillator

Specific Domains Modelling & Simulations
1

Hello, I need some help getting sampling to work using DynamicHMC. The model thus far is:

using LogDensityProblems
using Statistics
using Distributions
using UnPack
using TransformVariables
using DynamicHMC
using MCMCDiagnostics
import ForwardDiff
using Random

struct SHO_Model
    N::Int64
    ε::Float64
    ω::Float64
    m::Float64
    ħ::Float64
end;

function SHO_Model(x::AbstractVector)
    SHO_Model(x[1], x[2], x[3], x[4], x[5])
end;

function Energy(φ, N, ε, ω, m, ħ)::Vector{Float64}
    H = Vector{Float64}(undef, N+1)

    @inbounds for i in 1:N
        H[i] = (φ[i+1] - φ[i])^2 + 0.25*ω^2*(φ[i] + φ[i+1])^2
    end
    
    H[N+1] = H[1]
    
    return 0.5*H*m   
end;

function Action(H, ε)
    return sum(H)
end;

function (Problem::SHO_Model)(φ)
    @unpack N, ε, ω, m, ħ, = Problem
    
    E = Energy(φ, N, ε, ω, m, ħ)
    
    -Action(E, ε)
end

Initializing the Problem:

Problem = SHO_Model([500, 0.01, 1., 1., 1.])

and calculating that the log-probability works correctly:

Fields = rand(Uniform(-1., 1.), 500);
p = Problem((Fields))

And then I do, what I believe to be, the usual:

T = as(Array, Problem.N);
P = TransformedLogDensity(T, Problem);  
∇P = ADgradient(:ForwardDiff, P);
mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 10)

The error message which results is thus:

MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12})
Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  (::Type{T})(::T) where T<:Number at boot.jl:760
  (::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number} at char.jl:50
  ...

Stacktrace:
  [1] convert(#unused#::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12})
    @ Base .\number.jl:7
  [2] setindex!(A::Vector{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}, i1::Int64)
    @ Base .\array.jl:839
  [3] Energy(φ::Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}, N::Int64, ε::Float64, ω::Float64, m::Float64, ħ::Float64)
    @ Main .\In[217]:5
  [4] (::SHO_Model)(φ::Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}})
    @ Main .\In[218]:6
  [5] transform_logdensity(t::TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, f::SHO_Model, x::Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}})
    @ TransformVariables ~\.julia\packages\TransformVariables\N5VYB\src\generic.jl:179
  [6] logdensity
    @ ~\.julia\packages\LogDensityProblems\R15tV\src\LogDensityProblems.jl:168 [inlined]
  [7] #28
    @ ~\.julia\packages\LogDensityProblems\R15tV\src\AD_ForwardDiff.jl:21 [inlined]
  [8] chunk_mode_gradient!(result::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, f::LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}})
    @ ForwardDiff ~\.julia\packages\ForwardDiff\QOqCN\src\gradient.jl:150
  [9] gradient!
    @ ~\.julia\packages\ForwardDiff\QOqCN\src\gradient.jl:39 [inlined]
 [10] gradient!
    @ ~\.julia\packages\ForwardDiff\QOqCN\src\gradient.jl:35 [inlined]
 [11] logdensity_and_gradient(fℓ::LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}, ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}}}, x::Vector{Float64})
    @ LogDensityProblems ~\.julia\packages\LogDensityProblems\R15tV\src\AD_ForwardDiff.jl:51
 [12] evaluate_ℓ(ℓ::LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}, ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}}}, q::Vector{Float64})
    @ DynamicHMC ~\.julia\packages\DynamicHMC\lYo0t\src\hamiltonian.jl:193
 [13] #initialize_warmup_state#25
    @ ~\.julia\packages\DynamicHMC\lYo0t\src\mcmc.jl:109 [inlined]
 [14] initialize_warmup_state
    @ ~\.julia\packages\DynamicHMC\lYo0t\src\mcmc.jl:109 [inlined]
 [15] mcmc_keep_warmup(rng::Random._GLOBAL_RNG, ℓ::LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}, ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}}}, N::Int64; initialization::Tuple{}, warmup_stages::Tuple{InitialStepsizeSearch, TuningNUTS{Nothing, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{Nothing, DualAveraging{Float64}}}, algorithm::DynamicHMC.NUTS{Val{:generalized}}, reporter::LogProgressReport{Nothing})
    @ DynamicHMC ~\.julia\packages\DynamicHMC\lYo0t\src\mcmc.jl:464
 [16] macro expansion
    @ ~\.julia\packages\UnPack\EkESO\src\UnPack.jl:100 [inlined]
 [17] mcmc_with_warmup(rng::Random._GLOBAL_RNG, ℓ::LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}, ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}}}, N::Int64; initialization::Tuple{}, warmup_stages::Tuple{InitialStepsizeSearch, TuningNUTS{Nothing, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{LinearAlgebra.Diagonal, DualAveraging{Float64}}, TuningNUTS{Nothing, DualAveraging{Float64}}}, algorithm::DynamicHMC.NUTS{Val{:generalized}}, reporter::LogProgressReport{Nothing})
    @ DynamicHMC ~\.julia\packages\DynamicHMC\lYo0t\src\mcmc.jl:516
 [18] mcmc_with_warmup(rng::Random._GLOBAL_RNG, ℓ::LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}, ForwardDiff.GradientConfig{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12, Vector{ForwardDiff.Dual{ForwardDiff.Tag{LogDensityProblems.var"#28#29"{TransformedLogDensity{TransformVariables.ArrayTransform{TransformVariables.Identity, 1}, SHO_Model}}, Float64}, Float64, 12}}}}, N::Int64)
    @ DynamicHMC ~\.julia\packages\DynamicHMC\lYo0t\src\mcmc.jl:516
 [19] top-level scope
    @ In[240]:1
 [20] eval
    @ .\boot.jl:360 [inlined]
 [21] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base .\loading.jl:1094

I’ve been looking through the examples and source code but I am unable to find the issue. Any help would be greatly appreciated and any additional performance/code issues tips would be appreciated too. Thank you very much!

1 Like
2

I am no expert with DynamicHMC, but usually you need to change the strict type annotation you have at the moment to allow for Dual numbers, i.e., can you try:

struct SHO_Model{I<:Integer,F<:Real}
    N::I
    ε::F
    ω::F
    m::F
    ħ::F
end;

instead of your current version? This one as well I guess:

function Energy(φ, N, ε, ω, m, ħ)
    H = Vector{typeof(φ)}(undef, N+1)
3

As @mrVeng suggested, you should not impose those type annotations as they will prevent ForwardDiff from working. I would recommend using something like

function Energy(φ, N, ε, ω, m, ħ)
    T = typeof(promote(φ,ω))
    H = Vector{T}(undef, N+1)

    @inbounds for i in 1:N
        H[i] = (φ[i+1] - φ[i])^2 + 0.25*ω^2*(φ[i] + φ[i+1])^2
    end
    
    H[N+1] = H[1]
    
    return 0.5*H*m   
end

Note that for testing this, you can just call ∇P directly with a vector, as described in the manual.