Documentation for internals in Chains object

[treigerm]

When sampling from a model with NUTS then the Chains object that is returned contains a lot of internal fields which seem to be undocumented. Are those fields documented somewhere?

MWE:

Model and sampling:

using Turing

y = Float64[28, 8, -3, 7, -1, 1, 18, 12]
sigma = Float64[15, 10, 16, 11, 9, 11, 10, 18]

@model function eight_schools(y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5), 0, Inf)
    theta ~ filldist(Normal(mu, tau), length(y))
    y ~ MvNormal(theta, sigma)
end

chain = sample(eight_schools(y, sigma), NUTS(500, 0.65), 1200)

Then chain contains the following fields and I am interested in the internals:

Iterations        = 1:700
Thinning interval = 1
Chains            = 1
Samples per chain = 700
internals         = acceptance_rate, hamiltonian_energy, hamiltonian_energy_error, is_accept, log_density, lp, max_hamiltonian_energy_error, n_steps, nom_step_size, numerical_error, step_size, tree_depth
parameters        = mu, tau, theta[1], theta[2], theta[3], theta[4], theta[5], theta[6], theta[7], theta[8]

[treigerm]

Reading through the source of AdvancedHMC.jl I assume the following for each field:

acceptance_rate = acceptance rate for all the leapfrog steps
hamiltonian_energy = value of the hamiltonian for the current sample
hamiltonian_energy_error = integration error in the hamiltonian between the last sample and the current sample
is_accept = whether the sample is accepted (is always true?)
log_density = log density of the current sample
lp = same as log_density
max_hamiltonian_energy_error = maximum error in hamiltonian encountered during integration
n_steps = number of leapfrog steps
nom_step_size = nominal integration step size
numerical_error = indicates whether max_hamiltonian_energy_error is above a certain threshold (default is 1000)
step_size = actual step_size of integrator (might differ from nom_step_size if jittered step sizes are used)
tree_depth = depth of the tree built during integration

Is that roughly correct?

[devmotion]

@xukai92 is that correct?

BTW a bit confusing that there's both log_density and lp - are they guaranteed to be equal for every sample?

[xukai92]

Is that roughly correct?

Yes. And also yes for NUTS regarding "is_accept = whether the sample is accepted (is always true?)"?

BTW a bit confusing that there's both log_density and lp - are they guaranteed to be equal for every sample?

I agree. I believe long time ago Turing simply extracted log_density as lp which avoids this issue. What do you think the best way to make it consistent between Turing and other sampler packages?

@Vaibhavdixit02 maybe include this in the doc plan?

[devmotion]

I agree. I believe long time ago Turing simply extracted log_density as lp which avoids this issue. What do you think the best way to make it consistent between Turing and other sampler packages?

Turing uses getlogp to extract the log density from the sampled transitions (https://github.com/TuringLang/Turing.jl/blob/6db59629f1f189f63350aef9ce4fe6c0bebdaba1/src/inference/Inference.jl#L118), so in principle you could use custom transition types (such as the one for HMC) and overload the function. I guess the problem here is that we save the log density as part of the stats as well, so it should be possible to fix it by overloading getlogp and not saving it as a separate lp field in the HMC transitions.

[torfjelde]

I just want to add that I've several times tried to access log_density and it's always been empty... Sooo is it even used for anything?

[xukai92]

AHMC for sure outputs log_density (here for static and here for NUTS), and Turing should have warpped them into chains as they are declared here. So it looks like a bug somewhere in the pipeline.