Implementation of the No U-Turn Sampler for MCMC.

Please read the documentation. For the impatient: you probably want to

1. define a log density problem (eg for Bayesian inference) using the LogDensityProblems

package, then

1. use it with mcmc_with_warmup.

mcmc_with_warmup(
    rng,
    ℓ,
    N;
    initialization,
    warmup_stages,
    algorithm,
    reporter
)

Perform MCMC with NUTS, including warmup which is not returned. Return a NamedTuple of

• posterior_matrix, a matrix of position vectors, indexes by [parameter_index, draw_index]

• tree_statistics, a vector of tree statistics for each sample

• κ and ϵ, the adapted metric and stepsize.

Arguments

• rng: the random number generator, eg Random.default_rng().

• ℓ: the log density, supporting the API of the LogDensityProblems package

• N: the number of samples for inference, after the warmup.

Keyword arguments

• initialization: see below.

• warmup_stages: a sequence of warmup stages. See default_warmup_stages and fixed_stepsize_warmup_stages; the latter requires an ϵ in initialization.

• algorithm: see NUTS. It is very unlikely you need to modify this, except perhaps for the maximum depth.

• reporter: how progress is reported. By default, verbosely for interactive sessions using the log message mechanism (see LogProgressReport, and no reporting for non-interactive sessions (see NoProgressReport).

Initialization

The initialization keyword argument should be a NamedTuple which can contain the following fields (all of them optional and provided with reasonable defaults):

• q: initial position. Default: random (uniform [-2,2] for each coordinate).

• κ: kinetic energy specification. Default: Gaussian with identity matrix.

• ϵ: a scalar for initial stepsize, or nothing for heuristic finders.

Usage examples

Using a fixed stepsize:

mcmc_with_warmup(rng, ℓ, N;
                 initialization = (ϵ = 0.1, ),
                 warmup_stages = fixed_stepsize_warmup_stages())

Starting from a given position q₀ and kinetic energy scaled down (will still be adapted):

mcmc_with_warmup(rng, ℓ, N;
                 initialization = (q = q₀, κ = GaussianKineticEnergy(5, 0.1)))

Using a dense metric:

mcmc_with_warmup(rng, ℓ, N;
                 warmup_stages = default_warmup_stages(; M = Symmetric))

Disabling the initial stepsize search (provided explicitly, still adapted):

mcmc_with_warmup(rng, ℓ, N;
                 initialization = (ϵ = 1.0, ),
                 warmup_stages = default_warmup_stages(; stepsize_search = nothing))

stack_posterior_matrices(results)

Given a vector of results, each containing a property posterior_matrix (eg obtained from mcmc_with_warmup with the same sample length), return a lazy view as an array indexed by [draw_index, chain_index, parameter_index].

This is useful as an input for eg MCMCDiagnosticTools.ess_rhat.

pool_posterior_matrices(results)

Given a vector of results, each containing a property posterior_matrix (eg obtained from mcmc_with_warmup with the same sample length), return a lazy view as an array indexed by [parameter_index, pooled_draw_index].

This is useful for posterior analysis after diagnostics (see eg Base.eachcol).

default_warmup_stages(
;
    stepsize_search,
    M,
    stepsize_adaptation,
    init_steps,
    middle_steps,
    doubling_stages,
    terminating_steps
)

A sequence of warmup stages:

1. select an initial stepsize using stepsize_search (default: based on a heuristic),

2. tuning stepsize with init_steps steps

3. tuning stepsize and covariance: first with middle_steps steps, then repeat with twice the steps doubling_stages times

4. tuning stepsize with terminating_steps steps.

M (Diagonal, the default or Symmetric) determines the type of the metric adapted from the sample.

This is the suggested tuner of most applications.

Use nothing for stepsize_adaptation to skip the corresponding step.

fixed_stepsize_warmup_stages(
;
    M,
    middle_steps,
    doubling_stages
)

A sequence of warmup stages for fixed stepsize, only tuning covariance: first with middle_steps steps, then repeat with twice the steps doubling_stages times

Very similar to default_warmup_stages, but omits the warmup stages with just stepsize tuning.

struct InitialStepsizeSearch

Parameters for the search algorithm for the initial stepsize.

The algorithm finds an initial stepsize $ϵ$ so that the local log acceptance ratio $A(ϵ)$ is near params.log_threshold.

• initial_ϵ: The stepsize where the search is started.

• log_threshold: Log of the threshold that needs to be crossed.

• maxiter_crossing: Maximum number of iterations for crossing the threshold.

!!! NOTE

The algorithm is from Hoffman and Gelman (2014), default threshold modified to `0.8` following later practice in Stan.

struct DualAveraging{T}

Parameters for the dual averaging algorithm of Gelman and Hoffman (2014, Algorithm 6).

To get reasonable defaults, initialize with DualAveraging().

Fields

• δ: target acceptance rate

• γ: regularization scale

• κ: relaxation exponent

• t₀: offset

struct TuningNUTS{M, D}

Tune the step size ϵ during sampling, and the metric of the kinetic energy at the end of the block. The method for the latter is determined by the type parameter M, which can be

1. Diagonal for diagonal metric (the default),

2. Symmetric for a dense metric,

3. Nothing for an unchanged metric.

Results

A NamedTuple with the following fields:

• posterior_matrix, a matrix of position vectors, indexes by [parameter_index, draw_index]

• tree_statistics, a vector of tree statistics for each sample

• ϵs, a vector of step sizes for each sample

Fields

• N: Number of samples.

• stepsize_adaptation: Dual averaging parameters.

• λ: Regularization factor for normalizing variance. An estimated covariance matrix Σ is rescaled by λ towards $σ²I$, where $σ²$ is the median of the diagonal. The constructor has a reasonable default.

struct GaussianKineticEnergy{T<:(AbstractMatrix), S<:(AbstractMatrix)} <: DynamicHMC.EuclideanKineticEnergy

Gaussian kinetic energy, with $K(q,p) = p ∣ q ∼ 1/2 pᵀ⋅M⁻¹⋅p + log|M|$ (without constant), which is independent of $q$.

The inverse covariance $M⁻¹$ is stored.

struct NoProgressReport

A placeholder type for not reporting any information.

struct LogProgressReport{T}

Report progress into the Logging framework, using @info.

For the information reported, a step is a NUTS transition, not a leapfrog step.

Fields

• chain_id: ID of chain. Can be an arbitrary object, eg nothing.

• step_interval: Always report progress past step_interval of the last report.

• time_interval_s: Always report progress past this much time (in seconds) after the last report.

struct ProgressMeterReport

Report progress via a progress bar, using ProgressMeter.jl.

Example usage:

julia> ProgressMeterReport()

struct NUTS{S}

Implementation of the “No-U-Turn Sampler” of Hoffman and Gelman (2014), including subsequent developments, as described in Betancourt (2017).

Fields

• max_depth: Maximum tree depth.

• min_Δ: Threshold for negative energy relative to starting point that indicates divergence.

• turn_statistic_configuration: Turn statistic configuration. Currently only Val(:generalized) (the default) is supported.

mcmc_keep_warmup(
    rng,
    ℓ,
    N;
    initialization,
    warmup_stages,
    algorithm,
    reporter
)

Perform MCMC with NUTS, keeping the warmup results. Returns a NamedTuple of

• initial_warmup_state, which contains the initial warmup state

• warmup, an iterable of NamedTuples each containing fields

  • stage: the relevant warmup stage

  • results: results returned by that warmup stage (may be nothing if not applicable, or a chain, with tree statistics, etc; see the documentation of stages)

  • warmup_state: the warmup state after the corresponding stage.

• final_warmup_state, which contains the final adaptation after all the warmup

• inference, which has posterior_matrix and tree_statistics, see mcmc_with_warmup.

• sampling_logdensity, which contains information that is invariant to warmup

Arguments

• rng: the random number generator, eg Random.default_rng().

• ℓ: the log density, supporting the API of the LogDensityProblems package

• N: the number of samples for inference, after the warmup.

Keyword arguments

• initialization: see below.

• warmup_stages: a sequence of warmup stages. See default_warmup_stages and fixed_stepsize_warmup_stages; the latter requires an ϵ in initialization.

• algorithm: see NUTS. It is very unlikely you need to modify this, except perhaps for the maximum depth.

• reporter: how progress is reported. By default, verbosely for interactive sessions using the log message mechanism (see LogProgressReport, and no reporting for non-interactive sessions (see NoProgressReport).

Initialization

The initialization keyword argument should be a NamedTuple which can contain the following fields (all of them optional and provided with reasonable defaults):

• q: initial position. Default: random (uniform [-2,2] for each coordinate).

• κ: kinetic energy specification. Default: Gaussian with identity matrix.

• ϵ: a scalar for initial stepsize, or nothing for heuristic finders.

mcmc_steps(rng, algorithm, κ, ℓ, ϵ)

Return a value which can be used to perform MCMC stepwise, eg until some criterion is satisfied about the sample. See mcmc_next_step.

Two constructors are available:

1. Explicitly providing

   • rng, the random number generator,
   • algorithm, see mcmc_with_warmup,
   • κ, the (adapted) metric,
   • ℓ, the log density callable (see mcmc_with_warmup,
   • ϵ, the stepsize.

2. Using the fields sampling_logdensity and warmup_state, eg from mcmc_keep_warmup (make sure you use eg final_warmup_state).

Example

# initialization
results = DynamicHMC.mcmc_keep_warmup(RNG, ℓ, 0; reporter = NoProgressReport())
steps = mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
Q = results.final_warmup_state.Q

# a single update step
Q, tree_stats = mcmc_next_step(steps, Q)

# extract the position
Q.q

mcmc_next_step(mcmc_steps, Q)

Given Q (an evaluated log density at a position), return the next Q and tree statistics.

explore_log_acceptance_ratios(ℓ, q, log2ϵs; rng, κ, N, ps)

From an initial position, calculate the uncapped log acceptance ratio for the given log2 stepsizes and momentums ps, N of which are generated randomly by default.

summarize_tree_statistics(tree_statistics)

Summarize tree statistics. Mostly useful for NUTS diagnostics.

leapfrog_trajectory(ℓ, q, ϵ, positions; rng, κ, p)

Calculate a leapfrog trajectory visiting positions (specified as a UnitRange, eg -5:5) relative to the starting point q, with stepsize ϵ. positions has to contain 0, and the trajectories are only tracked up to the first non-finite log density in each direction.

Returns a vector of NamedTuples, each containin

• z, a PhasePoint object,

• position, the corresponding position,

• Δ, the log density + the kinetic energy relative to position 0.

EBFMI(tree_statistics)

Energy Bayesian fraction of missing information. Useful for diagnosing poorly chosen kinetic energies.

Low values (≤ 0.3) are considered problematic. See Betancourt (2016).

struct PhasePoint{T<:DynamicHMC.EvaluatedLogDensity, S}

A point in phase space, consists of a position (in the form of an evaluated log density ℓ at q) and a momentum.