Things to discuss

• [x] AFAIK, the way to customize the logging in AbstractMCMC.jl is to pass progress=false to the underlying AbstractMCMC.mcmcsample and then use the callback keyword argument to log the progress. So the question is: should we do this so as to preserve the current logging functionality?
• [x] To replicate the current summarization functionality (e.g. inform the user of average acceptance rates and EBFMI) as a post-sample step, we can overload StatsBase.sample and then perform this step after the call to AbstractMCMC.mcmcsample. Should we do this?

Turing.jl has somewhat "recently" started using LogDensityProblems.jl under the hood to simply all the handling of AD. This PR does the same for AdvancedHMC.jl, which in turn means that we can just plug a Turing.LogDensityFunction into AHMC and sample, in theory making the HMC implementation, etc. in Turing.jl redundant + we get additional AD-backends for free, e.g. Enzyme.

IMO we might even want to consider making a bridge-package between AbstractMCMC.jl and LogDensityProblems.jl simply containing this LogDensityModel.jl, given how this implementation would be very useful in other packages implementing samplers, e.g. AdvancedMH.jl. Thoughts on this @devmotion @yebai @xukai92 ?

A couple of (fixable) caveats:

• We're loosing the ForwardDiff implementation for AbstractMatrix.
• The DiffResults.jl helpers in LogDensityProblems.jl seems to be me a bit overly restrictive, e.g. https://github.com/tpapp/LogDensityProblems.jl/blob/a6a570751d0ee79345e92efd88062a0e6d59ef1b/src/DiffResults_helpers.jl#L14-L18 I believe will convert a ComponentVector into Vector, thus dropping the named dimensions. (@tpapp is there a reason why we can't just use similar(x)?)

EDIT: This now depends on https://github.com/TuringLang/AbstractMCMC.jl/pull/110.

Current design of type system for HMC:

• A type Metric to indicate which metric space we are working on
• A type Hamiltonian which contains the metric space, the log density function of the target distribution and its gradient function
• A type Leapfrog (a sub-type of AbstractIntegrator) to indicate we are using leapfrog integrator
  • We may want other integrators later
• A type Trajectory to indicate the way we build trajectory
  • StaticTrajectory is just HMC and NoUTurnTrajectory will be NUTS
• A type Proposal to wrap Trajectory with what kind of trajectory sampler is used
  • ~~Question: do we want this type, or may be an alternative type called Proposal, which build abstraction on proposal distribution directly.~~
• HMC samplers are just different combination of these. This package is not aiming to provide integrated samplers but those building blocks. However, some simple examples like normal HMC and NUTS will be there.
  • [x] HMC
  • [x] NUTS

Side note:

• All functions are supposed to be immutable
• During adaptation, we are meant to keep constructing new instances for Metric and Leapfrog types as their parameters are changing. However, after adaptation, all instances should be fixed.
• I hope this implementation is type stable on its own.

Time used for collecting 2 million samples

• Turing.NUTS: 3318.293540 seconds (33.35 G allocations: 1.299 TiB, 29.35% gc time)
• Turing.DynamicNUTS: 848.741643 seconds (7.86 G allocations: 251.076 GiB, 32.13% gc time)

using the following

using DynamicHMC, Turing, Test

@model gdemo(x, y) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  x ~ Normal(m, sqrt(s))
  y ~ Normal(m, sqrt(s))
  return s, m
end

mf = gdemo(1.5, 2.0)

@time chn1 = sample(mf, DynamicNUTS(2000000));

@time chn2 = sample(mf, Turing.NUTS(2000000, 0.65));

@ChrisRackauckas pointed out a paper (https://arxiv.org/abs/1912.03253) that uses the Calvo and Sanz-Sara methods in HMC. These integrators are already implemented in DiffEq (https://docs.juliadiffeq.org/latest/solvers/dynamical_solve/#Symplectic-Integrators-1). It would be nice to interface over them in AHMC.

See #123 for discussion

Adds the following statistics:

• [x] hamiltonian_energy_error
• [x] max_hamiltonian_energy_error
• [x] is_adapt
• [x] nom_step_size

To-do:

• [x] test that stats are returned, in particular that is_adapt starts as true and switches to false after adaptation is over
• [x] refactor code around jitter so that the step size used is accessible for stats. I'd appreciate input on this.

I'll submit a PR to Turing so that the new stats fields show up in internals instead of parameters. Looks like Turing doesn't use AdvancedHMC's sample function, so I'll need to add an is_adapt over there.

This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers. Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue. I'll open a PR within a few hours, please be patient!

To address

• Support of theta_init of Matrix type / GPU-level multiple chain support #92

Progress and goal

• [x] StaticTrajectory
• [x] HMCDA
• [x] UnitEuclideanMetric
• [x] DiagEuclideanMetric
• [x] UnitPreconditioner
• [x] DiagPreconditioner
• [x] NesterovDualAveraging
  • Does not work with HMCDA because it could yeild different number of steps for each chain.
• [x] Support a vector of random number generators
• [x] Try to merge transition

This PR does NOT aim to support

• NUTS: different chain could have different depth - not sure how to implement
• DenseEuclideanMetric: require more effort to deal with 3-dimensional tensors and not frequently used - low priority
• DensePreconditioner: same as above

Approximation error of leapfrog integration (i.e. accumulated Hamiltonian energy error) can sometimes explode, for example when the curvature of the current region is very high. This type of approximation error is sometimes called divergence [1] since it shifts a leapfrog simulation away from the correct solution.

In Turing, this type of errors is currently caught a relatively ad-hoc function called is_valid,

https://github.com/TuringLang/AdvancedHMC.jl/blob/734c0fa6d802a852d6b0bff7fc9f70a049a3f367/src/integrator.jl#L7

is_valid can catch cases where one or more elements of the parameter vector is either nan or inf. This has several drawbacks

• it's not general enough for catching all leapfrog approximation errors, e.g. when the parameter vector is valid but Hamiltonian energy is invalid.
• it may be delayed because numerical errors can appear in Hamiltonian energy earlier than in parameter vectors
• it's hard to propagate the exception around, i.e. at the moment we use a heuristic to find the previous valid point before approximation/numerical error happens
  • https://github.com/TuringLang/AdvancedHMC.jl/blob/734c0fa6d802a852d6b0bff7fc9f70a049a3f367/src/integrator.jl#L26
  • https://github.com/TuringLang/AdvancedHMC.jl/blob/734c0fa6d802a852d6b0bff7fc9f70a049a3f367/src/integrator.jl#L45

Therefore, we might want to refactor the current code a bit for a more robust mechanism for handling leapfrog approximation errors. Perhaps we can learn from the DynamicHMC implementation:

https://github.com/tpapp/DynamicHMC.jl/blob/master/src/buildingblocks.jl#L168

[1]: Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.

Hi @xukai92,

As discussed on Slack, adding support for complex numbers would allow us to use many models from physics. The code below is a quantum model of human judgment based on Wang et al. (2014). The model has a single parameter theta, which rotates the basis vectors. Thank you for looking into this. This feature would be very useful for me and others as well.

function quantum_model(S, θ)
    N = length(S)
    H = zeros(N, N) # Hamiltonian
    idx = CartesianIndex.(2:N,1:(N-1))
    dind = diagind(H)
    H[idx] .= 1.0
    H .+= H'
    H[dind] .= 1:N
    U = exp(-1im * θ * H) 
    na = 1
    PA = zeros(N, N)
    PA[dind] .= [ones(na); zeros(N - na)]
    PnA = I - PA
    Sa = PA * S
    Sa ./= sqrt(Sa' * Sa) #normalized projection A
    Sna = PnA * S; 
    Sna ./= sqrt(Sna' * Sna) # normalized projection B

    nb = 1 # dim of B subspace
    PB = zeros(ComplexF64, N, N)
    PB[dind] = [zeros(nb, 1); ones(N - nb, 1)] # projector for B in B coord
    PB .= U * PB * U' # projector for B in A coord
    PnB = I - PB

    Sb = PB * S
    Sb ./= sqrt(Sb' * Sb)
    Snb = PnB * S
    Snb ./= sqrt(Snb' * Snb)
    
    pAtB = (S' * PA * S) * (Sa' * PB * Sa) # prob A then B
    pAtnB = (S' * PA * S) * (Sa' * PnB * Sa)
    pnAtB = (S' * PnA * S)*(Sna' * PB * Sna)
    pnAtnB = (S' * PnA * S) * (Sna' * PnB * Sna)
    # pAtB + pAtnB + pnAtB + pnAtnB == 1
    pBtA = (S' * PB * S)*(Sb' * PA * Sb) # prob B then A
    pBtnA = (S' * PB * S) * (Sb' * PnA * Sb)
    pnBtA = (S' * PnB*S) * (Snb' * PA * Snb)
    pnBtnA = (S' * PnB * S) * (Snb' * PnA * Snb)
    # pBtA + pBtnA + pnBtA + pnBtnA == 1
    # order 1
    c_probs1 = [pAtB, pAtnB , pnAtB, pnAtnB]
    # order 2
    c_probs2 = [pBtA, pnBtA, pBtnA, pnBtnA]
    return map(real, c_probs1), map(real, c_probs2)
end

function simulate(S, θ, n_sim)
    p1,p2 = quantum_model(S, θ)
    y1 = rand(Multinomial(n_sim, p1))
    y2 = rand(Multinomial(n_sim, p2))
    return y1, y2
end

using Distributions, Turing, LinearAlgebra
import Distributions: logpdf, loglikelihood


"""
Simplified model based on 
    Wang, Z., Solloway, T., Shiffrin, R. M., & Busemeyer, J. R. (2014). 
    Context effects produced by question orders reveal quantum nature of human 
    judgments. Proceedings of the National Academy of Sciences, 111(26), 9431-9436.

"""
struct Quantum{T1,T2} <: ContinuousUnivariateDistribution
    θ::T1
    S::T2
    n::Int64
end

function logpdf(d::Quantum, data)
    p = quantum_model(d.S, d.θ)
    LL = @. logpdf(Multinomial(d.n, p), data)
    return sum(LL)
end

loglikelihood(d::Quantum, data::Tuple{Vector{Int64}, Vector{Int64}}) = logpdf(d, data)


# number of observations per condition
n_sim = 100
# dimensionality of Hilbert Space
N = 4
# state vector
S = fill(sqrt(.25), N)
# rotation
θ = 2.0
data = simulate(S, θ, n_sim)

@model model(data, S, n_sim) = begin
    θ ~ Truncated(Normal(2, 2), 0.0, Inf)
    data ~ Quantum(θ, S, n_sim)
end

# Settings of the NUTS sampler.
n_samples = 1000
delta = 0.85
n_adapt = 1000
n_chains = 4
specs = NUTS(n_adapt, delta)
# Start sampling.
chain = sample(model(data, S, n_sim), specs, MCMCThreads(), n_samples, n_chains, progress=true)

This PR aims at introducing some additional types that are discussed in https://github.com/TuringLang/AdvancedHMC.jl/issues/16. More specifically, the following types are introduced:

• PhasePoint: stores θ, r and cached Potential energy (and its gradient). https://github.com/TuringLang/AdvancedHMC.jl/issues/17
• DualValue: stores log density logπ(θ), and cache its gradient.
• ~~LogDensityFunction: stores logπ and its gradient function ∂logπ∂θ.~~

TODOs

• [x] Refactor numerical error handling code leapfrog using PhasePoint
• [x] Remove DualFunction?
• [x] Refactor build_tree using PhasePoint
• ~~Return more information for each step #59~~
• ~~Add a Termination type.~~

This is a very minor pull request for conveniently bundling samples from AdvancedHMC.jl into an Chains object from MCMCChains.jl. Very similar to the interface in AdvancedMH.jl, it should make it easier to switch from one to the other.

https://github.com/TuringLang/AdvancedHMC.jl/blob/6a55a3f3f341c90a70491267635acb51dd989463/src/trajectory.jl#L770

This caused an actual reproducibility problem, see Discourse thread: https://discourse.julialang.org/t/stablerng-in-turing-is-not-producing-reproducible-output/92032

To get RNG from caller, it might need to be perculated down from callers, with a few other changes needed.

This is a draft to implement Riemannian HMC. There are many things to discuss. I put the high-level points here while leaving more specific ones in the code.

• Do we want to unify all function signatures for Hamiltonian to a non-separable one (i.e. position dependent)?
• Do we need to update our metric abstraction? Looks like we should decouple unit/diagonal/dense vs Euclidean/Riemannian.
• Where should SoftAbs best locate? I current implemented it as an external function to provide G to a generic Riemannian HMC implementation.
• The efficiency is pretty bad and there is fewer optimization opportunities due to our decoupled abstraction (e.g. compared to what derived in [1]).

To-dos

• The current SoftAbs implementation has numerical issues when running on Neal's funnel. I need to double check the if the manual Jacobian implementation is correct, or we could register softabs with ReverseDiff to see if this solves the issue. I could use someone else help on this.

How to play with this PR

I provided a notebook (which contains the same content as the test/experimental/riemannian_hmc.jl file) to play with the code. The notebook has some simple validation on the implementation and also shows the current numerical issue of SoftAbs. I highly recommend you to try this.

[1] Betancourt, M., 2013, August. A general metric for Riemannian manifold Hamiltonian Monte Carlo. In International Conference on Geometric Science of Information (pp. 327-334). Springer, Berlin, Heidelberg.

~~This PR limits warning messages to 10. Is that a good default?~~

~~Solution 1: we can probably disable these numerical messages, and instead print a summary message about total divergent transitions (see e.g. this comment) when the sample function is done.~~

A message about the percentage of divergent transitions is displayed in the progress bar. In addition, a warning message is shown if the percentage of divergent transitions is above a certain (30% by default) threshold.

It would be nice if the adapted mass matrix could be retrieved. @sethaxen pointed out that giving save_state=true to sample is possible, but that it currently only gives the metric before warm-up.