TuringLang / Turing.jl

Bayesian inference with probabilistic programming.
https://turinglang.org
MIT License
2k stars 216 forks source link

Model using LKJCholesky fails with large covariance matrices #2081

Closed juehang closed 9 months ago

juehang commented 9 months ago

I am having some issues with larger covariance matrices sampled using LKJCholesky, in Turing 0.29.1. I find that I reliably get domain errors with covariance matrices larger than 10x10; with smaller covariance matrix sizes the errors are sporadic and sampling with NUTS works sometimes.

I think a minimal example to demonstrate this is:

using LinearAlgebra, PDMats, Random, Turing, ReverseDiff

Random.seed!(121)

N = 100
J = 10

# generate data
sigma = [i for i in 1:J]
Omega = PDMat(Cholesky(rand(LKJCholesky(J, 1.0)).L*Diagonal(sigma) + I(J)*eps()))

# Sigma = Diagonal(sigma) * Omega * Diagonal(sigma)

y = rand(MvNormal(Omega), N)

Turing.setadbackend(:forwarddiff)

@model function correlation_chol(J, N, y)
    sigma ~ filldist(Exponential(0.5), J)
    F ~ LKJCholesky(J, 1)
    Σ_L = Diagonal(sigma) * F.L
    Sigma = PDMat(Cholesky(Σ_L + I*1e-8))
    # print(sigma[1])
    for i in 1:N
        y[:, i] ~ MvNormal(Sigma)
    end
end

# sampler = NUTS(100, 0.8)
sampler = HMC(0.05, 10)
chain1 = sample(correlation_chol(J, N, float.(y)), sampler, 1000);

This works just fine using HMC where the parameters are manually entered, but not with the NUTS sampler, so I suspect something is happening during tuning. This code is based on a code sample in the comments of #1870 by @joshualeond.

Edit: Actually, that's just because my HMC was tuned so poorly in this thrown-together script that the acceptance rate was essentially zero; when I tune HMC better the same error occurs so it is not related to NUTS tuning at all.

The error message I get is:

DomainError with -2.220446049250313e-14:
sqrt will only return a complex result if called with a complex argument. Try sqrt(Complex(x)).
DomainError detected in the user `f` function. This occurs when the domain of a function is violated.
For example, `log(-1.0)` is undefined because `log` of a real number is defined to only output real
numbers, but `log` of a negative number is complex valued and therefore Julia throws a DomainError
by default. Cases to be aware of include:

* `log(x)`, `sqrt(x)`, `cbrt(x)`, etc. where `x<0`
* `x^y` for `x<0` floating point `y` (example: `(-1.0)^(1/2) == im`)

Within the context of SciML, this error can occur within the solver process even if the domain constraint
would not be violated in the solution due to adaptivity. For example, an ODE solver or optimization
routine may check a step at `new_u` which violates the domain constraint, and if violated reject the
step and use a smaller `dt`. However, the throwing of this error will have halted the solving process.

Thus the recommended fix is to replace this function with the equivalent ones from NaNMath.jl
(https://github.com/JuliaMath/NaNMath.jl) which returns a NaN instead of an error. The solver will then
effectively use the NaN within the error control routines to reject the out of bounds step. Additionally,
one could perform a domain transformation on the variables so that such an issue does not occur in the
definition of `f`.

For more information, check out the following FAQ page:
https://docs.sciml.ai/Optimization/stable/API/FAQ/#The-Solver-Seems-to-Violate-Constraints-During-the-Optimization,-Causing-DomainErrors,-What-Can-I-Do-About-That?

Stacktrace:
  [1] throw_complex_domainerror(f::Symbol, x::Float64)
    @ Base.Math ./math.jl:33
  [2] sqrt
    @ ./math.jl:677 [inlined]
  [3] sqrt
    @ ~/.julia/packages/ForwardDiff/PcZ48/src/dual.jl:240 [inlined]
  [4] _link_chol_lkj_from_upper(W::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}})
    @ Bijectors ~/.julia/packages/Bijectors/QhObI/src/bijectors/corr.jl:319
  [5] _link_chol_lkj_from_lower
    @ ~/.julia/packages/Bijectors/QhObI/src/bijectors/corr.jl:328 [inlined]
  [6] transform(b::Bijectors.VecCholeskyBijector, X::Cholesky{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}})
    @ Bijectors ~/.julia/packages/Bijectors/QhObI/src/bijectors/corr.jl:220
  [7] Transform
    @ ~/.julia/packages/Bijectors/QhObI/src/interface.jl:81 [inlined]
  [8] logabsdetjac
    @ ~/.julia/packages/Bijectors/QhObI/src/bijectors/corr.jl:225 [inlined]
  [9] with_logabsdet_jacobian(ib::Inverse{Bijectors.VecCholeskyBijector}, y::Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}})
    @ Bijectors ~/.julia/packages/Bijectors/QhObI/src/interface.jl:214
 [10] with_logabsdet_jacobian_and_reconstruct
    @ ~/.julia/packages/DynamicPPL/YThRW/src/abstract_varinfo.jl:604 [inlined]
 [11] invlink_with_logpdf(vi::DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, vn::AbstractPPL.VarName{:F, Setfield.IdentityLens}, dist::LKJCholesky{Float64}, y::SubArray{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}, 1, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Tuple{UnitRange{Int64}}, true})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/YThRW/src/abstract_varinfo.jl:684
 [12] invlink_with_logpdf
    @ ~/.julia/packages/DynamicPPL/YThRW/src/abstract_varinfo.jl:679 [inlined]
 [13] assume
    @ ~/.julia/packages/DynamicPPL/YThRW/src/context_implementations.jl:197 [inlined]
 [14] assume
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/hmc.jl:461 [inlined]
 [15] tilde_assume
    @ ~/.julia/packages/DynamicPPL/YThRW/src/context_implementations.jl:49 [inlined]
 [16] tilde_assume
    @ ~/.julia/packages/DynamicPPL/YThRW/src/context_implementations.jl:46 [inlined]
 [17] tilde_assume
    @ ~/.julia/packages/DynamicPPL/YThRW/src/context_implementations.jl:31 [inlined]
 [18] tilde_assume!!(context::DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}, right::LKJCholesky{Float64}, vn::AbstractPPL.VarName{:F, Setfield.IdentityLens}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/YThRW/src/context_implementations.jl:117
 [19] correlation_chol(__model__::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, __varinfo__::DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, __context__::DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}, J::Int64, N::Int64, y::Matrix{Float64})
    @ Main ./In[147]:3
 [20] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/YThRW/src/model.jl:963 [inlined]
 [21] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/YThRW/src/model.jl:936 [inlined]
 [22] evaluate!!(model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, varinfo::DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, Vector{Set{DynamicPPL.Selector}}}}}, ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}, context::DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/YThRW/src/model.jl:889
 [23] logdensity
    @ ~/.julia/packages/DynamicPPL/YThRW/src/logdensityfunction.jl:94 [inlined]
 [24] Fix1
    @ ./operators.jl:1108 [inlined]
 [25] chunk_mode_gradient!(result::DiffResults.MutableDiffResult{1, Float64, Tuple{Vector{Float64}}}, f::Base.Fix1{typeof(LogDensityProblems.logdensity), LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}}, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:123
 [26] gradient!
    @ ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:39 [inlined]
 [27] gradient!
    @ ~/.julia/packages/ForwardDiff/PcZ48/src/gradient.jl:35 [inlined]
 [28] logdensity_and_gradient
    @ ~/.julia/packages/LogDensityProblemsAD/JoNjv/ext/LogDensityProblemsADForwardDiffExt.jl:113 [inlined]
 [29] ∂logπ∂θ
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/hmc.jl:160 [inlined]
 [30] ∂H∂θ
    @ ~/.julia/packages/AdvancedHMC/dgxuI/src/hamiltonian.jl:38 [inlined]
 [31] step(lf::AdvancedHMC.Leapfrog{Float64}, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, AdvancedHMC.GaussianKinetic, Base.Fix1{typeof(LogDensityProblems.logdensity), LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}, Turing.Inference.var"#∂logπ∂θ#36"{LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}}, z::AdvancedHMC.PhasePoint{Vector{Float64}, AdvancedHMC.DualValue{Float64, Vector{Float64}}}, n_steps::Int64; fwd::Bool, full_trajectory::Val{false})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/dgxuI/src/integrator.jl:229
 [32] step (repeats 2 times)
    @ ~/.julia/packages/AdvancedHMC/dgxuI/src/integrator.jl:199 [inlined]
 [33] A
    @ ~/.julia/packages/AdvancedHMC/dgxuI/src/trajectory.jl:743 [inlined]
 [34] find_good_stepsize(rng::TaskLocalRNG, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, AdvancedHMC.GaussianKinetic, Base.Fix1{typeof(LogDensityProblems.logdensity), LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}, Turing.Inference.var"#∂logπ∂θ#36"{LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}}, θ::Vector{Float64}; max_n_iters::Int64)
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/dgxuI/src/trajectory.jl:776
 [35] find_good_stepsize(rng::TaskLocalRNG, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, AdvancedHMC.GaussianKinetic, Base.Fix1{typeof(LogDensityProblems.logdensity), LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}, Turing.Inference.var"#∂logπ∂θ#36"{LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{LogDensityFunction{DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.SamplingContext{DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.DefaultContext, TaskLocalRNG}}, ForwardDiff.Chunk{11}, ForwardDiff.Tag{Turing.TuringTag, Float64}, ForwardDiff.GradientConfig{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11, Vector{ForwardDiff.Dual{ForwardDiff.Tag{Turing.TuringTag, Float64}, Float64, 11}}}}}}, θ::Vector{Float64})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/dgxuI/src/trajectory.jl:749
 [36] initialstep(rng::TaskLocalRNG, model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, spl::DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:sigma, :F), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}, Int64}, Vector{Product{Continuous, Exponential{Float64}, FillArrays.Fill{Exponential{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{AbstractPPL.VarName{:sigma, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:F, Setfield.IdentityLens}, Int64}, Vector{LKJCholesky{Float64}}, Vector{AbstractPPL.VarName{:F, Setfield.IdentityLens}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/CP3Ic/src/mcmc/hmc.jl:191
 [37] step(rng::TaskLocalRNG, model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, spl::DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, init_params::Nothing, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/YThRW/src/sampler.jl:111
 [38] step
    @ ~/.julia/packages/DynamicPPL/YThRW/src/sampler.jl:84 [inlined]
 [39] macro expansion
    @ ~/.julia/packages/AbstractMCMC/fWWW0/src/sample.jl:125 [inlined]
 [40] macro expansion
    @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
 [41] (::AbstractMCMC.var"#21#22"{Bool, String, Nothing, Int64, Int64, Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}}, TaskLocalRNG, DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, Int64, Int64})()
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/fWWW0/src/logging.jl:12
 [42] with_logstate(f::Function, logstate::Any)
    @ Base.CoreLogging ./logging.jl:514
 [43] with_logger
    @ ./logging.jl:626 [inlined]
 [44] with_progresslogger(f::Function, _module::Module, logger::Logging.ConsoleLogger)
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/fWWW0/src/logging.jl:36
 [45] macro expansion
    @ ~/.julia/packages/AbstractMCMC/fWWW0/src/logging.jl:11 [inlined]
 [46] mcmcsample(rng::TaskLocalRNG, model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/fWWW0/src/sample.jl:116
 [47] sample(rng::TaskLocalRNG, model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; chain_type::Type, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/CP3Ic/src/mcmc/hmc.jl:121
 [48] sample
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/hmc.jl:91 [inlined]
 [49] #sample#2
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/Inference.jl:194 [inlined]
 [50] sample
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/Inference.jl:187 [inlined]
 [51] #sample#1
    @ ~/.julia/packages/Turing/CP3Ic/src/mcmc/Inference.jl:184 [inlined]
 [52] sample(model::DynamicPPL.Model{typeof(correlation_chol), (:J, :N, :y), (), (), Tuple{Int64, Int64, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, alg::NUTS{Turing.Essential.ForwardDiffAD{0}, (), AdvancedHMC.DiagEuclideanMetric}, N::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/CP3Ic/src/mcmc/Inference.jl:178
juehang commented 9 months ago

Figured out what the issue is. For others encountering this, it has something to do with sigma being too small, resulting in floating point errors. Adding a small offset to sigma fixes this.