<|start_filename|>src/hamiltonian.jl<|end_filename|>
#####
##### Building blocks for traversing a Hamiltonian deterministically, using the leapfrog
##### integrator.
#####

####
#### kinetic energy
####

"""
$(TYPEDEF)

Kinetic energy specifications. Implements the methods

- `Base.size`

- [`kinetic_energy`](@ref)

- [`calculate_p♯`](@ref)

- [`∇kinetic_energy`](@ref)

- [`rand_p`](@ref)

For all subtypes, it is implicitly assumed that kinetic energy is symmetric in
the momentum `p`,

```julia
kinetic_energy(κ, p, q) == kinetic_energy(κ, .-p, q)
```

When the above is violated, the consequences are undefined.
"""
abstract type KineticEnergy end

"""
$(TYPEDEF)

Euclidean kinetic energies (position independent).
"""
abstract type EuclideanKineticEnergy <: KineticEnergy end

"""
$(TYPEDEF)

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.

!!! note
    Making ``M⁻¹`` approximate the posterior variance is a reasonable starting point.
"""
struct GaussianKineticEnergy{T <: AbstractMatrix,
                             S <: AbstractMatrix} <: EuclideanKineticEnergy
    "M⁻¹"
    M⁻¹::T
    "W such that W*W'=M. Used for generating random draws."
    W::S
    function GaussianKineticEnergy(M⁻¹::T, W::S) where {T, S}
        @argcheck checksquare(M⁻¹) == checksquare(W)
        new{T,S}(M⁻¹, W)
    end
end

"""
$(SIGNATURES)

Gaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.
"""
GaussianKineticEnergy(M⁻¹::AbstractMatrix) = GaussianKineticEnergy(M⁻¹, cholesky(inv(M⁻¹)).L)

"""
$(SIGNATURES)

Gaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.
"""
GaussianKineticEnergy(M⁻¹::Diagonal) = GaussianKineticEnergy(M⁻¹, Diagonal(.√inv.(diag(M⁻¹))))

"""
$(SIGNATURES)

Gaussian kinetic energy with a diagonal inverse covariance matrix `M⁻¹=m⁻¹*I`.
"""
GaussianKineticEnergy(N::Integer, m⁻¹ = 1.0) = GaussianKineticEnergy(Diagonal(fill(m⁻¹, N)))

function Base.show(io::IO, κ::GaussianKineticEnergy{T}) where {T}
    print(io::IO, "Gaussian kinetic energy ($(nameof(T))), √diag(M⁻¹): $(.√(diag(κ.M⁻¹)))")
end

## NOTE about implementation: the 3 methods are callable without a third argument (`q`)
## because they are defined for Gaussian (Euclidean) kinetic energies.

Base.size(κ::GaussianKineticEnergy, args...) = size(κ.M⁻¹, args...)

"""
$(SIGNATURES)

Return kinetic energy `κ`, at momentum `p`.
"""
kinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = dot(p, κ.M⁻¹ * p) / 2

"""
$(SIGNATURES)

Return ``p♯ = M⁻¹⋅p``, used for turn diagnostics.
"""
calculate_p♯(κ::GaussianKineticEnergy, p, q = nothing) = κ.M⁻¹ * p

"""
$(SIGNATURES)

Calculate the gradient of the logarithm of kinetic energy in momentum `p`.
"""
∇kinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = calculate_p♯(κ, p)

"""
$(SIGNATURES)

Generate a random momentum from a kinetic energy at position `q`.
"""
rand_p(rng::AbstractRNG, κ::GaussianKineticEnergy, q = nothing) = κ.W * randn(rng, size(κ.W, 1))

####
#### Hamiltonian
####

struct Hamiltonian{K,L}
    "The kinetic energy specification."
    κ::K
    """
    The (log) density we are sampling from. Supports the `LogDensityProblem` API.
    Technically, it is the negative of the potential energy.
    """
    ℓ::L
    """
    $(SIGNATURES)

    Construct a Hamiltonian from the log density `ℓ`, and the kinetic energy specification
    `κ`. `ℓ` with a vector are expected to support the `LogDensityProblems` API, with
    gradients.
    """
    function Hamiltonian(κ::K, ℓ::L) where {K <: KineticEnergy,L}
        @argcheck capabilities(ℓ) ≥ LogDensityOrder(1)
        @argcheck dimension(ℓ) == size(κ, 1)
        new{K,L}(κ, ℓ)
    end
end

Base.show(io::IO, H::Hamiltonian) = print(io, "Hamiltonian with $(H.κ)")

"""
$(TYPEDEF)

A log density evaluated at position `q`. The log densities and gradient are saved, so that
they are not calculated twice for every leapfrog step (both as start- and endpoints).

Because of caching, a `EvaluatedLogDensity` should only be used with a specific Hamiltonian,
preferably constructed with the `evaluate_ℓ` constructor.

In composite types and arguments, `Q` is usually used for this type.
"""
struct EvaluatedLogDensity{T,S}
    "Position."
    q::T
    "ℓ(q). Saved for reuse in sampling."
    ℓq::S
    "∇ℓ(q). Cached for reuse in sampling."
    ∇ℓq::T
    function EvaluatedLogDensity(q::T, ℓq::S, ∇ℓq::T) where {T <: AbstractVector,S <: Real}
        @argcheck length(q) == length(∇ℓq)
        new{T,S}(q, ℓq, ∇ℓq)
    end
end

# general constructors below are necessary to sanitize input from eg Diagnostics, or an
# initial position given as integers, etc

function EvaluatedLogDensity(q::AbstractVector, ℓq::Real, ∇ℓq::AbstractVector)
    q, ∇ℓq = promote(q, ∇ℓq)
    EvaluatedLogDensity(q, ℓq, ∇ℓq)
end

EvaluatedLogDensity(q, ℓq::Real, ∇ℓq) = EvaluatedLogDensity(collect(q), ℓq, collect(∇ℓq))

"""
$(SIGNATURES)

Evaluate log density and gradient and save with the position. Preferred interface for
creating `EvaluatedLogDensity` instances.
"""
function evaluate_ℓ(ℓ, q)
    ℓq, ∇ℓq = logdensity_and_gradient(ℓ, q)
    if isfinite(ℓq)
        EvaluatedLogDensity(q, ℓq, ∇ℓq)
    else
        EvaluatedLogDensity(q, oftype(ℓq, -Inf), q) # second q used just as a placeholder
    end
end

"""
$(TYPEDEF)

A point in phase space, consists of a position (in the form of an evaluated log density `ℓ`
at `q`) and a momentum.
"""
struct PhasePoint{T <: EvaluatedLogDensity,S}
    "Evaluated log density."
    Q::T
    "Momentum."
    p::S
    function PhasePoint(Q::T, p::S) where {T,S}
        @argcheck length(p) == length(Q.q)
        new{T,S}(Q, p)
    end
end

"""
$(SIGNATURES)

Log density for Hamiltonian `H` at point `z`.

If `ℓ(q) == -Inf` (rejected), skips the kinetic energy calculation.

Non-finite values (incl `NaN`, `Inf`) are automatically converted to `-Inf`. This can happen
if

1. the log density is not a finite value,

2. the kinetic energy is not a finite value (which usually happens when `NaN` or `Inf` got
mixed in the leapfrog step, leading to an invalid position).
"""
function logdensity(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)
    @unpack ℓq = z.Q
    isfinite(ℓq) || return oftype(ℓq, -Inf)
    K = kinetic_energy(H.κ, z.p)
    ℓq - (isfinite(K) ? K : oftype(K, Inf))
end

function calculate_p♯(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)
    calculate_p♯(H.κ, z.p)
end

"""
    leapfrog(H, z, ϵ)

Take a leapfrog step of length `ϵ` from `z` along the Hamiltonian `H`.

Return the new phase point.

The leapfrog algorithm uses the gradient of the next position to evolve the momentum. If
this is not finite, the momentum won't be either, `logdensity` above will catch this and
return an `-Inf`, making the point divergent.
"""
function leapfrog(H::Hamiltonian{<: EuclideanKineticEnergy}, z::PhasePoint, ϵ)
    @unpack ℓ, κ = H
    @unpack p, Q = z
    @argcheck isfinite(Q.ℓq) "Internal error: leapfrog called from non-finite log density"
    pₘ = p + ϵ/2 * Q.∇ℓq
    q′ = Q.q + ϵ * ∇kinetic_energy(κ, pₘ)
    Q′ = evaluate_ℓ(H.ℓ, q′)
    p′ = pₘ + ϵ/2 * Q′.∇ℓq
    PhasePoint(Q′, p′)
end


<|start_filename|>src/trees.jl<|end_filename|>
#####
##### Abstract tree/trajectory interface
#####

####
#### Directions
####

"Maximum number of iterations [`next_direction`](@ref) supports."
const MAX_DIRECTIONS_DEPTH = 32

"""
Internal type implementing random directions. Draw a new value with `rand`, see
[`next_direction`](@ref).

Serves two purposes: a fixed value of `Directions` is useful for unit testing, and drawing a
single bit flag collection economizes on the RNG cost.
"""
struct Directions
    flags::UInt32
end

Base.rand(rng::AbstractRNG, ::Type{Directions}) = Directions(rand(rng, UInt32))

"""
$(SIGNATURES)

Return the next direction flag and the new state of directions. Results are undefined for
more than [`MAX_DIRECTIONS_DEPTH`](@ref) updates.
"""
function next_direction(directions::Directions)
    @unpack flags = directions
    Bool(flags & 0x01), Directions(flags >>> 1)
end

####
#### Trajectory interface
####

"""
    $(FUNCTIONNAME)(trajectory, z, is_forward)

Move along the trajectory in the specified direction. Return the new position.
"""
function move end

"""
    $(FUNCTIONNAME)(trajectory, τ)

Test if the turn statistics `τ` indicate that the corresponding tree is turning.

Will only be called on nontrivial trees (at least two nodes).
"""
function is_turning end

"""
    $(FUNCTIONNAME)(trajectory, τ₁, τ₂)

Combine turn statistics on trajectory. Implementation can assume that the trees that
correspond to the turn statistics have the same ordering.

When
```julia
τ = combine_turn_statistics(trajectory, τ₁, τ₂)
is_turning(trajectory, τ)
```
the combined turn statistic `τ` is guaranteed not to escape the caller, so it can eg change
type.
"""
function combine_turn_statistics end

"""
    $(FUNCTIONNAME)(trajectory, v₁, v₂)

Combine visited node statistics for adjacent trees trajectory. Implementation should be
invariant to the ordering of `v₁` and `v₂` (ie the operation is commutative).
"""
function combine_visited_statistics end

"""
    $(FUNCTIONNAME)(trajectory, is_doubling::Bool, ω₁, ω₂, ω)

Calculate the log probability if selecting the subtree corresponding to `ω₂`. Being the log
of a probability, it is always `≤ 0`, but implementations are allowed to return and accept
values `> 0` and treat them as `0`.

When `is_doubling`, the tree corresponding to `ω₂` was obtained from a doubling step (this
can be relevant eg for biased progressive sampling).

The value `ω = logaddexp(ω₁, ω₂)` is provided for avoiding redundant calculations.

See [`biased_progressive_logprob2`](@ref) for an implementation.
"""
function calculate_logprob2 end

"""
    $(FUNCTIONNAME)(rng, trajectory, ζ₁, ζ₂, logprob2::Real, is_forward::Bool)

Combine two proposals `ζ₁, ζ₂` on `trajectory`, with log probability `logprob2` for
selecting `ζ₂`.

 `ζ₁` is before `ζ₂` iff `is_forward`.
"""
function combine_proposals end

"""
    ζωτ_or_nothing, v = $(FUNCTIONNAME)(trajectory, z, is_initial)

Information for a tree made of a single node. When `is_initial == true`, this is the first
node.

The first value is either

1. `nothing` for a divergent node,

2. a tuple containing the proposal `ζ`, the log weight (probability) of the node `ω`, the
turn statistics `τ` (never tested with `is_turning` for leaf nodes).

The second value is the visited node information.
"""
function leaf end

####
#### utilities
####

"""
$(SIGNATURES)

Combine turn statistics with the given direction. When `is_forward`, `τ₁` is before `τ₂`,
otherwise after.

Internal helper function.
"""
@inline function combine_turn_statistics_in_direction(trajectory, τ₁, τ₂, is_forward::Bool)
    if is_forward
        combine_turn_statistics(trajectory, τ₁, τ₂)
    else
        combine_turn_statistics(trajectory, τ₂, τ₁)
    end
end

function combine_proposals_and_logweights(rng, trajectory, ζ₁, ζ₂, ω₁::Real, ω₂::Real,
                                          is_forward::Bool, is_doubling::Bool)
    ω = logaddexp(ω₁, ω₂)
    logprob2 = calculate_logprob2(trajectory, is_doubling, ω₁, ω₂, ω)
    ζ = combine_proposals(rng, trajectory, ζ₁, ζ₂, logprob2, is_forward)
    ζ, ω
end

"""
$(SIGNATURES)

Given (relative) log probabilities `ω₁` and `ω₂`, return the log probabiliy of
drawing a sample from the second (`logprob2`).

When `bias`, biases towards the second argument, introducing anti-correlations.
"""
function biased_progressive_logprob2(bias::Bool, ω₁::Real, ω₂::Real, ω = logaddexp(ω₁, ω₂))
    ω₂ - (bias ? ω₁ : ω)
end

####
#### abstract trajectory interface
####

"""
$(SIGNATURES)

Information about an invalid (sub)tree, using positions relative to the starting node.

1. When `left < right`, this tree was *turning*.

2. When `left == right`, this is a *divergent* node.

3. `left == 1 && right == 0` is used as a sentinel value for reaching maximum depth without
encountering any invalid trees (see [`REACHED_MAX_DEPTH`](@ref). All other `left > right`
values are disallowed.
"""
struct InvalidTree
    left::Int
    right::Int
end

InvalidTree(i::Integer) = InvalidTree(i, i)

is_divergent(invalid_tree::InvalidTree) = invalid_tree.left == invalid_tree.right

function Base.show(io::IO, invalid_tree::InvalidTree)
    msg = if is_divergent(invalid_tree)
        "divergence at position $(invalid_tree.left)"
    elseif invalid_tree == REACHED_MAX_DEPTH
        "reached maximum depth without divergence or turning"
    else
        @unpack left, right = invalid_tree
        "turning at positions $(left):$(right)"
    end
    print(io, msg)
end

"Sentinel value for reaching maximum depth."
const REACHED_MAX_DEPTH = InvalidTree(1, 0)

"""
    result, v = adjacent_tree(rng, trajectory, z, i, depth, is_forward)

Traverse the tree of given `depth` adjacent to point `z` in `trajectory`.

`is_forward` specifies the direction, `rng` is used for random numbers in
[`combine_proposals`](@ref). `i` is an integer position relative to the initial node (`0`).

The *first value* is either

1. an `InvalidTree`, indicating the first divergent node or turning subtree that was
encounteted and invalidated this tree.

2. a tuple of `(ζ, ω, τ, z′, i′), with

    - `ζ`: the proposal from the tree.

    - `ω`: the log weight of the subtree that corresponds to the proposal

    - `τ`: turn statistics

    - `z′`: the last node of the tree

    - `i′`: the position of the last node relative to the initial node.

The *second value* is always the visited node statistic.
"""
function adjacent_tree(rng, trajectory, z, i, depth, is_forward)
    i′ = i + (is_forward ? 1 : -1)
    if depth == 0
        z′ = move(trajectory, z, is_forward)
        ζωτ, v = leaf(trajectory, z′, false)
        if ζωτ ≡ nothing
            InvalidTree(i′), v
        else
            (ζωτ..., z′, i′), v
        end
    else
        # “left” tree
        t₋, v₋ = adjacent_tree(rng, trajectory, z, i, depth - 1, is_forward)
        t₋ isa InvalidTree && return t₋, v₋
        ζ₋, ω₋, τ₋, z₋, i₋ = t₋

        # “right” tree — visited information from left is kept even if invalid
        t₊, v₊ = adjacent_tree(rng, trajectory, z₋, i₋, depth - 1, is_forward)
        v = combine_visited_statistics(trajectory, v₋, v₊)
        t₊ isa InvalidTree && return t₊, v
        ζ₊, ω₊, τ₊, z₊, i₊ = t₊

        # turning invalidates
        τ = combine_turn_statistics_in_direction(trajectory, τ₋, τ₊, is_forward)
        is_turning(trajectory, τ) && return InvalidTree(i′, i₊), v

        # valid subtree, combine proposals
        ζ, ω = combine_proposals_and_logweights(rng, trajectory, ζ₋, ζ₊, ω₋, ω₊,
                                                is_forward, false)
        (ζ, ω, τ, z₊, i₊), v
    end
end

"""
$(SIGNATURES)

Sample a `trajectory` starting at `z`, up to `max_depth`. `directions` determines the tree
expansion directions.

Return the following values

- `ζ`: proposal from the tree

- `v`: visited node statistics

- `termination`: an `InvalidTree` (this includes the last doubling step turning, which is
  technically a valid tree) or `REACHED_MAX_DEPTH` when all subtrees were valid and no
  turning happens.

- `depth`: the depth of the tree that was sampled from. Doubling steps that lead to an
  invalid adjacent tree do not contribute to `depth`.
"""
function sample_trajectory(rng, trajectory, z, max_depth::Integer, directions::Directions)
    @argcheck max_depth ≤ MAX_DIRECTIONS_DEPTH
    (ζ, ω, τ), v = leaf(trajectory, z, true)
    z₋ = z₊ = z
    depth = 0
    termination = REACHED_MAX_DEPTH
    i₋ = i₊ = 0
    while depth < max_depth
        is_forward, directions = next_direction(directions)
        t′, v′ = adjacent_tree(rng, trajectory, is_forward ? z₊ : z₋, is_forward ? i₊ : i₋,
                               depth, is_forward)
        v = combine_visited_statistics(trajectory, v, v′)

        # invalid adjacent tree: stop
        t′ isa InvalidTree && (termination = t′; break)

        # extract information from adjacent tree
        ζ′, ω′, τ′, z′, i′ = t′

        # update edges and combine proposals
        if is_forward
            z₊, i₊ = z′, i′
        else
            z₋, i₋ = z′, i′
        end

        # tree has doubled successfully
        ζ, ω = combine_proposals_and_logweights(rng, trajectory, ζ, ζ′, ω, ω′,
                                                is_forward, true)
        depth += 1

        # when the combined tree is turning, stop
        τ = combine_turn_statistics_in_direction(trajectory, τ, τ′, is_forward)
        is_turning(trajectory, τ) && (termination = InvalidTree(i₋, i₊); break)
    end
    ζ, v, termination, depth
end


<|start_filename|>test/test-stepsize.jl<|end_filename|>
#####
##### Stepsize and adaptation tests
#####

isinteractive() && include("common.jl")

@testset "stepsize general rootfinding" begin
    Δ = 3.0                   # shift exponential so that ϵ = 1 is not in interval
    A = ϵ -> exp(-ϵ*Δ)
    invA = A -> -log(A) / Δ
    @test 0.1 ≈ invA(A(0.1)) atol = 1e-10 # inverse is OK
    params = InitialStepsizeSearch()
    @test A(params.ϵ₀) ≤ params.a_min || A(params.ϵ₀) ≥ params.a_max # outside interval
    constantA(ϵ) = params.a_max + 0.1
    # parameters and defaults
    @test params.a_min == 0.25
    @test params.a_max ≤ 1.0
    @test params.ϵ₀ == 1.0
    @test params.C == 2.0
    @test params.maxiter_crossing == 400
    @test params.maxiter_bisect == 400
    @test_throws ArgumentError InitialStepsizeSearch(; a_min = 0.9, a_max = 0.1) # not <
    @test_throws ArgumentError InitialStepsizeSearch(; C = 0.5) # not > 1
    @test_throws ArgumentError InitialStepsizeSearch(; maxiter_crossing = 2)
    @test_throws ArgumentError InitialStepsizeSearch(; maxiter_bisect = 2)
    # crossing from below
    ϵ₀, Aϵ₀, ϵ₁, Aϵ₁ = find_crossing_stepsize(params, A, invA(params.a_min-0.1))
    @test ϵ₀ ≥ invA(params.a_min) ≥ ϵ₁
    @test Aϵ₀ ≤ params.a_min ≤ Aϵ₁
    @test_throws ErrorException find_crossing_stepsize(params, constantA, 100.0)
    # crossing from above
    ϵ₀, Aϵ₀, ϵ₁, Aϵ₁ = find_crossing_stepsize(params, A, invA(params.a_max+0.1))
    @test ϵ₀ ≤ invA(params.a_max) ≤ ϵ₁
    @test Aϵ₀ ≥ params.a_max ≥ Aϵ₁
    @test_throws ErrorException find_crossing_stepsize(params, constantA, 0.1)
    # bisection
    ϵ = bisect_stepsize(params, A, invA(params.a_max + 0.1), invA(params.a_min - 0.1))
    @test params.a_min ≤ A(ϵ) ≤ params.a_max
    @test_throws ArgumentError bisect_stepsize(params, A, 0.4, 0.3) # order
    @test_throws ArgumentError bisect_stepsize(params, A, # already in interval
                                               invA(params.a_max - 0.1),
                                               invA(params.a_min + 0.1))
    # combined algorithm - single test
    ϵ = find_initial_stepsize(params, A)
    @test params.a_min ≤ A(ϵ) ≤ params.a_max
    # combined algorithm - random tests
    for _ in 1:1000
        α = abs(randn())
        β = -abs(randn())
        A = ϵ -> exp(-(α*ϵ+β)*ϵ)  # A(0)=1, has a hump, goes to A(∞)→0
        Aϵ₀ = A(params.ϵ₀)
        # uncomment line belows for coverage debugging or inverse
        # println("α = $α, β = $β")
        # println(Aϵ₀ ≤ params.a_min ? "below" : (Aϵ₀ ≥ params.a_max ? "above" : "in"))
        # invA(A) = -(log(A)+β)/A
        ϵ = find_initial_stepsize(params, A)
        @test params.a_min ≤ A(ϵ) ≤ params.a_max
    end
end

"""
$(SIGNATURES)

A parametric random acceptance rate that depends on the stepsize. For unit
testing acceptance rate tuning.
"""
dummy_acceptance_rate(ϵ, σ = 0.05) = min(1/ϵ * exp(randn()*σ - σ^2/2), 1)

mean_dummy_acceptance_rate(ϵ, σ = 0.05) = mean(dummy_acceptance_rate(ϵ, σ) for _ in 1:10000)

@testset "dual averaging far" begin
    ϵ₀ = 100.0                # way off
    δ = 0.65
    dual_averaging = DualAveraging(; δ = δ)
    A = initial_adaptation_state(dual_averaging, ϵ₀)
    @test A.logϵ̄ == 0           # ϵ̄₀ = 1 in Gelman and Hoffman (2014)
    @test A.m == 0
    @test A.H̄ == 0
    for _ in 1:500
        A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A)))
    end
    @test mean_dummy_acceptance_rate(final_ϵ(A)) ≈ δ atol = 0.02
end

@testset "dual averaging close" begin
    ϵ₀ = 2.0
    δ = 0.65
    dual_averaging = DualAveraging(; δ = δ)
    A = initial_adaptation_state(dual_averaging, ϵ₀)
    for _ in 1:2000
        A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A)))
    end
    @test mean_dummy_acceptance_rate(final_ϵ(A)) ≈ δ atol = 0.01
end

@testset "dual averaging far and noisy" begin
    ϵ₀ = 20.0
    δ = 0.65
    dual_averaging = DualAveraging(; δ = δ)
    A = initial_adaptation_state(dual_averaging, ϵ₀)
    for _ in 1:2000
        A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A), 2.0))
    end
    @test mean_dummy_acceptance_rate(final_ϵ(A), 2.0) ≈ δ atol = 0.03
end

@testset "fixed stepsize sanity checks" begin
    fs = FixedStepsize()
    ϵ = 1.0
    A = initial_adaptation_state(fs, ϵ)
    @test A == adapt_stepsize(fs, A, ϵ)
    @test current_ϵ(A) == ϵ
    @test final_ϵ(A) == ϵ
end

@testset "find reasonable stepsize - random H, z" begin
    p = InitialStepsizeSearch()
    for _ in 1:100
        @unpack H, z = rand_Hz(rand(3:5))
        ϵ = find_initial_stepsize(p, H, z)
        logA = logdensity(H, leapfrog(H, z, ϵ)) - logdensity(H, z)
        @test p.a_min ≤ exp(logA) ≤ p.a_max
    end
end

@testset "error for non-finite initial density" begin
    p = InitialStepsizeSearch()
    @unpack H, z = rand_Hz(2)
    z = DynamicHMC.PhasePoint(z.Q, [NaN, NaN])
    @test_throws DomainError find_initial_stepsize(p, H, z)
end


<|start_filename|>docs/make.jl<|end_filename|>
using Documenter, DynamicHMC

makedocs(modules = [DynamicHMC],
         format = Documenter.HTML(; prettyurls = get(ENV, "CI", nothing) == "true"),
         clean = true,
         sitename = "DynamicHMC.jl",
         authors = "<NAME>",
         checkdocs = :exports,
         # strict = true,
         pages = ["Introduction" => "index.md",
                  "A worked example" => "worked_example.md",
                  "Documentation" => "interface.md"])

deploydocs(repo = "github.com/tpapp/DynamicHMC.jl.git",
           push_preview = true)


<|start_filename|>test/runtests.jl<|end_filename|>
include("common.jl")

macro include_testset(filename)
    @assert filename isa AbstractString
    quote
        @testset $(filename) begin
            include($(filename))
        end
    end
end

####
#### unit tests
####

@include_testset("test-trees.jl")
@include_testset("test-hamiltonian.jl")
@include_testset("test-NUTS.jl")
@include_testset("test-stepsize.jl")
@include_testset("test-mcmc.jl")
@include_testset("test-diagnostics.jl")
@include_testset("test-logging.jl")

####
#### sample correctness tests
####

@include_testset("sample-correctness-tests.jl")


<|start_filename|>test/test-NUTS.jl<|end_filename|>
isinteractive() && include("common.jl")

####
#### utilities
####

"""
    $SIGNATURES

Recursive comparison by fields; types and values should match by `==`.
"""
function ≂(x::T, y::T) where T
    fn = fieldnames(T)
    isempty(fn) ? x == y : all(getfield(x, f) ≂ getfield(y, f) for f in fn)
end

≂(x, y) = x == y

"""
Type for testing `≂`.
"""
struct Foo{T}
    a::T
    b::T
end

@testset "≂ comparisons" begin
    @test 1 ≂ 1
    @test 1 ≂ 1.0
    @test [1,2] ≂ [1,2]
    @test [1.0,2.0] ≂ [1,2]
    @test !(1 ≂ 2)
    @test Foo(1,2) ≂ Foo(1,2)
    @test !(Foo(1,2) ≂ Foo(1,3))
    @test !(Foo{Any}(1,2) ≂ Foo(1,2))
end


###
### random booleans
###

@testset "random booleans" begin
    for prob in (1:9) ./ 10
        logprob = log(prob)
        @test abs(mean(rand_bool_logprob(RNG, logprob) for _ in 1:10000) - prob) ≤ 0.01
    end

    # these operations don't call the RNG, this is checked
    RNG′ = copy(RNG)
    @test all(rand_bool_logprob(RNG, 0) for _ in 1:10000)
    @test all(rand_bool_logprob(RNG, 10) for _ in 1:10000)
    @test rand(RNG′) == rand(RNG)
end

###
### test turn statistics
###

@testset "low-level turn statistics" begin
    trajectory = TrajectoryNUTS(nothing, 0, 1, -1000, Val{:generalized})
    p = ones(3)                 # unit vector
    c = 0.1                     # a constant, just for consistency checking of combination
    # turn statistics constructed so that τ₁ + τ₂ won't be turning, τ₁ + τ₃ will be
    τ₁ = GeneralizedTurnStatistic(p, p .- c, p, p .- c, p)
    τ₂ = GeneralizedTurnStatistic(3 .* p, 3 .* p .+ c, 3 .* p, 3 .* p .+ c, 3 .* p)
    τ₃ = GeneralizedTurnStatistic(2 .* p, 2 .* p .+ c, 2 .* p, 2 .* p .+ c, -2 .* p)
    τ = combine_turn_statistics(trajectory, τ₁, τ₂)
    # test mechanics of combination
    @test τ ≂ GeneralizedTurnStatistic(τ₁.p₋, τ₁.p♯₋, τ₂.p₊, τ₂.p♯₊, τ₁.ρ .+ τ₂.ρ)
    # test non-turning
    @test !is_turning(trajectory, τ)
    # test turning
    @test is_turning(trajectory, combine_turn_statistics(trajectory, τ₁, τ₃))
end

@testset "low-level visited statistics" begin
    trajectory = TrajectoryNUTS(nothing, 0, 1, -1000, Val{:generalized})
    vs(p, is_initial = false) = leaf_acceptance_statistic(log(p), is_initial)
    x = vs(0.3)
    @test acceptance_rate(x) ≈ 0.3
    y = vs(0.6)
    @test acceptance_rate(y) ≈ 0.6
    x0 = vs(10, true)    # initial node, does not count
    z = reduce((x, y) -> combine_visited_statistics(trajectory, x, y),
               [x, x, y, x0])
    @test acceptance_rate(z) ≈ 0.4
end

# define a distribution which is divergent everywhere except at 0
struct AlwaysDivergentTest
    K::Int
end

function LogDensityProblems.capabilities(::Type{AlwaysDivergentTest})
    LogDensityProblems.LogDensityOrder{1}()
end
LogDensityProblems.dimension(d::AlwaysDivergentTest) = d.K
function LogDensityProblems.logdensity_and_gradient(d::AlwaysDivergentTest, x)
    if all(iszero.(x))
        0.0, zeros(length(x))
    else
        -Inf, nothing
    end
end

@testset "unconditional divergence" begin

    # test NUTS sampler where all movements are divergent
    K = 3
    ℓ = AlwaysDivergentTest(K)
    Q, tree_statistics = sample_tree(RNG, NUTS(), Hamiltonian(GaussianKineticEnergy(K), ℓ),
                                     evaluate_ℓ(ℓ, zeros(K)), 1.0)
    @test is_divergent(tree_statistics.termination)
    @test iszero(tree_statistics.acceptance_rate)
    @test iszero(tree_statistics.depth)
    @test tree_statistics.steps == 1
end

@testset "normal NUTS HMC transition mean and cov" begin
    # A test for sample_tree with a fixed ϵ and κ, which is perfectly adapted and should
    # provide excellent mixing
    for _ in 1:10
        K = rand(2:8)
        N = 10000
        μ = randn(K)
        Σ = rand_Σ(K)
        L = cholesky(Σ).L
        ℓ = multivariate_normal(μ, L)
        Q = evaluate_ℓ(ℓ, randn(K))
        H = Hamiltonian(GaussianKineticEnergy(Σ), ℓ)
        qs = Array{Float64}(undef, N, K)
        ϵ = 0.5
        algorithm = NUTS()
        for i in 1:N
            Q = first(sample_tree(RNG, algorithm, H, Q, ϵ))
            qs[i, :] = Q.q
        end
        m, C = mean_and_cov(qs, 1)
        @test vec(m) ≈ μ atol = 0.1 rtol = maximum(diag(C))*0.02 norm = x -> norm(x,1)
        @test cov(qs, dims = 1) ≈ L*L' atol = 0.1 rtol = 0.1
    end
end

###
### test proposals
###

# NOTE superseded by separating ω and ζ
# @testset "proposal" begin
#     trajectory = TrajectoryNUTS(nothing, 0, 1, -1000)
#     function test_sample(rng, prop1, prop2, bias, prob_prob2; atol = 0.02, N = 10000)
#         count = 0
#         for _ in 1:N
#             prop = combine_proposals(rng, trajectory, prop1, prop2,
#                                      # direction irrelevant for thie method
#                                      rand(Bool),
#                                      # test with given bias
#                                      bias)
#             if prop.z ≂ prop2.z
#                 count += 1
#             else
#                 @test prop.z ≂ prop1.z
#             end
#             @test prop.ω ≈ logaddexp(prop1.ω, prop2.ω)
#         end
#         @test count / N ≈ prob_prob2 atol = atol
#     end
#     prop1 = Proposal(1, log(1.0))
#     prop2 = Proposal(2, log(3.0))
#     prop3 = Proposal(3, log(1/3))
#     test_sample(RNG, prop1, prop2, true, 1; atol = 0, N = 100)
#     test_sample(RNG, prop1, prop2, false, 0.75)
#     test_sample(RNG, prop1, prop3, true, 1/3)
#     test_sample(RNG, prop1, prop3, false, 0.25)
# end


<|start_filename|>src/reporting.jl<|end_filename|>
import ProgressMeter

#####
##### Reporting progress.
#####

"""
$(TYPEDEF)

A placeholder type for not reporting any information.
"""
struct NoProgressReport end

"""
$(SIGNATURES)

Report to the given `reporter`.

The second argument can be

1. a string, which is displayed as is (this is supported by all reporters).

2. or a step in an MCMC chain with a known number of steps for progress reporters (see
[`make_mcmc_reporter`](@ref)).

`meta` arguments are key-value pairs.

In this context, a *step* is a NUTS transition, not a leapfrog step.
"""
report(reporter::NoProgressReport, step::Union{AbstractString,Integer}; meta...) = nothing

"""
$(SIGNATURES)

Return a reporter which can be used for progress reports with a known number of
`total_steps`. May return the same reporter, or a related object. Will display `meta` as
key-value pairs.

## Arguments:
- `reporter::NoProgressReport`: the original reporter
- `total_steps`: total number of steps

## Keyword arguments:
- `currently_warmup::Bool`: `true` if we are currently doing warmup; `false` if we are currently doing MCMC
- `meta`: key-value pairs that will be displayed by the reporter
"""
make_mcmc_reporter(reporter::NoProgressReport, total_steps; currently_warmup::Bool = false, meta...) = reporter

"""
$(TYPEDEF)

Report progress into the `Logging` framework, using `@info`.

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

# Fields

$(FIELDS)
"""
@with_kw struct LogProgressReport{T}
    "ID of chain. Can be an arbitrary object, eg `nothing`."
    chain_id::T = nothing
    "Always report progress past `step_interval` of the last report."
    step_interval::Int = 100
    "Always report progress past this much time (in seconds) after the last report."
    time_interval_s::Float64 = 1000.0
end

"""
$(SIGNATURES)

Assemble log message metadata.

Currently, it adds `chain_id` *iff* it is not `nothing`.
"""
_log_meta(chain_id::Nothing, meta) = meta

_log_meta(chain_id, meta) = (chain_id = chain_id, meta...)

function report(reporter::LogProgressReport, message::AbstractString; meta...)
    @info message _log_meta(reporter.chain_id, meta)...
    nothing
end

"""
$(TYPEDEF)

A composite type for tracking the state for which the last log message was emitted, for MCMC
reporting with a given total number of steps (see [`make_mcmc_reporter`](@ref).

# Fields

$(FIELDS)
"""
mutable struct LogMCMCReport{T}
    "The progress report sink."
    log_progress_report::T
    "Total steps for this stage."
    total_steps::Int
    "Index of the last reported step."
    last_reported_step::Int
    "The last time a report was logged (determined using `time_ns`)."
    last_reported_time_ns::UInt64
end

function report(reporter::LogMCMCReport, message::AbstractString; meta...)
    @info message _log_meta(reporter.log_progress_report.chain_id, meta)...
    nothing
end

function make_mcmc_reporter(reporter::LogProgressReport, total_steps::Integer; currently_warmup::Bool = false, meta...)
    @info "Starting MCMC" total_steps = total_steps meta...
    LogMCMCReport(reporter, total_steps, -1, time_ns())
end

function report(reporter::LogMCMCReport, step::Integer; meta...)
    @unpack (log_progress_report, total_steps, last_reported_step,
             last_reported_time_ns) = reporter
    @unpack chain_id, step_interval, time_interval_s = log_progress_report
    @argcheck 1 ≤ step ≤ total_steps
    Δ_steps = step - last_reported_step
    t_ns = time_ns()
    Δ_time_s = (t_ns - last_reported_time_ns) / 1_000_000_000
    if last_reported_step < 0 || Δ_steps ≥ step_interval || Δ_time_s ≥ time_interval_s
        seconds_per_step = Δ_time_s / Δ_steps
        meta_progress = (step = step,
                         seconds_per_step = round(seconds_per_step; sigdigits = 2),
                         estimated_seconds_left = round((total_steps - step) *
                                                        seconds_per_step; sigdigits = 2))
        @info "MCMC progress" merge(_log_meta(chain_id, meta_progress), meta)...
        reporter.last_reported_step = step
        reporter.last_reported_time_ns = t_ns
    end
    nothing
end

"""
$(TYPEDEF)

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

Example usage:
```julia
julia> ProgressMeterReport()
```
"""
struct ProgressMeterReport
end

struct ProgressMeterReportMCMC{T}
    currently_warmup::Bool
    progress_meter::T
end

function make_mcmc_reporter(reporter::ProgressMeterReport, total_steps::Integer; currently_warmup::Bool=false, meta...)
    description = currently_warmup ? "Warmup: " : "MCMC: "
    return ProgressMeterReportMCMC(currently_warmup, ProgressMeter.Progress(total_steps, 1, description))
end

function report(reporter::ProgressMeterReport, message::AbstractString; meta...)
    return nothing
end

function report(reporter::ProgressMeterReportMCMC, message::AbstractString; meta...)
    return nothing
end

function report(reporter::ProgressMeterReport, step::Integer; meta...)
    return nothing
end

function report(reporter::ProgressMeterReportMCMC, step::Integer; meta...)
    ProgressMeter.next!(reporter.progress_meter)
    return nothing
end

"""
$(SIGNATURES)

Return a default reporter, taking the environment into account. Keyword arguments are passed
to constructors when applicable.
"""
function default_reporter(; kwargs...)
    if isinteractive()
        LogProgressReport(; kwargs...)
    else
        NoProgressReport()
    end
end


<|start_filename|>test/test-logging.jl<|end_filename|>
ℓ = multivariate_normal(ones(1))
κ = GaussianKineticEnergy(1)
Q = evaluate_ℓ(ℓ, [1.0])

reporters_1 = [
    NoProgressReport(),
    ProgressMeterReport(),
]

reporters_2 = [
    LogProgressReport(),
]

for reporter in deepcopy(vcat(reporters_1, reporters_2))
    results = mcmc_with_warmup(RNG, ℓ, 10_000; reporter = reporter)
end

for reporter in deepcopy(vcat(reporters_1, reporters_2))
    DynamicHMC.report(reporter, "")

    mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)
    DynamicHMC.report(mcmc_reporter_1, "")
    DynamicHMC.report(mcmc_reporter_1, 1)

    mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)
    DynamicHMC.report(mcmc_reporter_2, "")
    DynamicHMC.report(mcmc_reporter_2, 1)
end

for reporter in deepcopy(reporters_1)
    DynamicHMC.report(reporter, "")
    DynamicHMC.report(reporter, 1)

    mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)
    DynamicHMC.report(mcmc_reporter_1, "")
    DynamicHMC.report(mcmc_reporter_1, 1)

    mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)
    DynamicHMC.report(mcmc_reporter_2, "")
    DynamicHMC.report(mcmc_reporter_2, 1)
end
