Module Norn.Diag

Diagnostics of draws.

A chain diagnostic of draws of 'u is a value of 'u holding one number per element, at the element's dtype: (rhat schools post).tau is the R-hat of tau. The importance sampling diagnostics take one tensor, of log ratios or log likelihoods, and return numbers at its dtype. Diagnostics read their draws on the host, compute at float64, and take float tensors.

The chain diagnostics follow Vehtari, Gelman, Simpson, Carpenter and Bürkner (2021). Each splits every chain into its two halves, dropping the middle draw of an odd count, and ranks the pooled draws, ties at their average rank r; the normal scores ndtri ((r - 3/8) / (S + 1/4)), S the number of draws, make them insensitive to the draws' scale and tails. An element whose draws are all equal, or not all finite, or a chain of fewer than four draws, has no diagnostic: it is NaN.

Raises Invalid_argument naming the path for a tensor that is not of a float dtype.

Convergence

#val rhat : ('u, 'd) Nx.Ptree.t -> 'u Draws.t -> 'u

rhat u d is the rank-normalised split R-hat of each element: the larger of the R-hat of the split chains' normal scores (bulk) and of the normal scores of their distance to the median (tails). Near 1 the chains agree; above 1.01 they have not mixed.

#val nested_rhat : ('u, 'd) Nx.Ptree.t -> superchains:int -> 'u Draws.t -> 'u

nested_rhat u ~superchains d is the rank-normalised nested R-hat (Margossian, Hoffman, Sountsov, Riou-Durand, Vehtari and Gelman 2024) of each element, for many short chains: the chains form superchains consecutive groups of equal size, each started from one point, and the statistic compares the variance between groups with the variance within them, chains split and scored as for rhat.

Raises Invalid_argument if superchains < 2 or if it does not divide the number of chains.

#val ess_bulk : ('u, 'd) Nx.Ptree.t -> 'u Draws.t -> 'u

ess_bulk u d is the effective sample size of each element's normal scores, by Geyer's initial monotone sequence over the autocorrelations of the split chains. It measures how well the draws estimate the centre of the distribution.

#val ess_tail : ('u, 'd) Nx.Ptree.t -> 'u Draws.t -> 'u

ess_tail u d is the smaller of the effective sample sizes of the indicators of each element being at most its 5% and its 95% quantile. It measures how well the draws estimate the tails.

#val mcse_mean : ('u, 'd) Nx.Ptree.t -> 'u Draws.t -> 'u

mcse_mean u d is the Monte Carlo standard error of each element's mean: its standard deviation over the square root of the effective sample size of the split chains' draws.

Hamiltonian transitions

#val ebfmi : ('f, 'd) Stats.t Draws.t -> (float, 'f, 'd) Nx.t

ebfmi s is each chain's energy Bayesian fraction of missing information, shape [chain]: the mean squared change of the energy between draws over the energy's variance. Below 0.3 the momentum resampling explores the energy poorly.

Raises Invalid_argument if the chains have fewer than two draws.

val divergent_shift : 
  ('u, 'd) Nx.Ptree.t ->
  ('f, 'd) Stats.t Draws.t ->
  'u Draws.t ->
  'u

divergent_shift u s d is, for each element, the mean of the draws whose transition diverged minus the mean of all draws, in standard deviations of all draws. Where divergences gather, it is far from 0; with no divergence it is NaN.

Raises Invalid_argument naming both shapes if s and d differ in their chains or draws.

Calibration

#val rank : ('u, 'd) Nx.Ptree.t -> truth:'u -> 'u Draws.t -> 'u

rank u ~truth d is, for each element, the number of draws of d below the element of truth: from 0 to the number of draws. With truth drawn from the prior, data simulated from it and d exact posterior draws, each rank is uniform (Talts, Betancourt, Simpson, Vehtari and Gelman 2018).

Raises Invalid_argument naming the path if there are more draws than the element's dtype holds exact integers.

#val rank_uniformity : ('u, 'd) Nx.Ptree.t -> draws:int -> 'u list -> 'u

rank_uniformity u ~draws ranks is, for each element, the simultaneous p-value of its ranks against the uniform distribution on 0, ..., draws (Säilynoja, Bürkner and Vehtari 2022), computed exactly, with no simulation: the probability, under uniformity, that the count of ranks below some point i / (draws + 1) strays from its expectation at least as far, in pointwise p-value, as the ranks' farthest count. Below alpha rejects uniformity at level alpha.

Raises Invalid_argument if ranks is empty, if draws < 1, or if a rank is not an integer in 0, ..., draws.

Importance sampling

Draws θ_s of a proposal g estimate expectations under a target p once each is weighted by its importance ratio r_s = p(θ_s) / g(θ_s), known up to a constant factor. When the ratios have a heavy right tail, a few draws dominate the estimate. Pareto smoothed importance sampling (Vehtari, Simpson, Gelman, Yao and Gabry 2024) replaces the M largest ratios by the values a generalized Pareto tail fitted to them predicts, and the tail's shape k̂ says whether the estimate can be trusted: the ratios have about 1 / k̂ finite moments, and from S draws the estimate is reliable when k̂ <= min (1 - 1 / log10 S, 0.7), 0.7 from about 2000 draws.

k̂ is -inf when the M largest ratios are equal, a bounded tail, as when every ratio is equal; they are left as they are. It is NaN, and the ratios are left as they are, when no tail can be fitted: fewer than 25 draws, a ratio NaN or +inf, every ratio -inf, or the tail's first quartile at the largest of the other ratios. A ratio of -inf is a weight of zero, which stays zero unless it is among the M largest.

#type ('f, 'd) psis = {

#log_weights : (float, 'f, 'd) Nx.t;
(*
The smoothed log weights, of the ratios' shape, not normalised.
*)
#khat : (float, 'f, 'd) Nx.t;
(*
Each set's k̂.
*)

}

The type for smoothed importance weights.

#val psis : (float, 'f, 'd) Nx.t -> ('f, 'd) psis

psis lr smooths the log importance ratios lr. Axis 0 of lr holds S independent draws, and each index of the other axes is a set of S ratios smoothed on its own: khat has the other axes' shape. { w with log_weights = (psis w.log_weights).log_weights } smooths Weighted.t draws.

Of each set, the M = floor (min (S / 5, 3 sqrt S)) largest ratios are replaced. A generalized Pareto distribution is fitted to their excesses over u, the largest of the others, by Zhang and Stephens' (2009) posterior mean, its shape k then shrunk toward 1/2 as k̂ = (M k + 5) / (M + 10). The z-th smallest of the M becomes u plus the fitted quantile of (z - 1/2) / M, capped at the largest ratio. The others are unchanged, the order of the ratios is kept, and adding a constant to lr adds it to the weights.

Raises Invalid_argument if lr is a scalar.

#type ('f, 'd) loo = {

#elpd : (float, 'f, 'd) Nx.t;
(*
elpd_loo, the sum of pointwise, a scalar.
*)
#se : (float, 'f, 'd) Nx.t;
(*
The standard error of elpd from the spread of the n points' terms, sqrt (n var elpd_i) with the variance of n - 1 degrees of freedom: the points taken as a sample of the data. It omits the Monte Carlo error. NaN for fewer than two points.
*)
#p : (float, 'f, 'd) Nx.t;
(*
p_loo, the effective number of parameters, lpd - elpd: lpd = Σ_i log ((1/S) Σ_s p(y_i | θ_s)) is the log density of each point given all of the data, summed.
*)
#pointwise : (float, 'f, 'd) Nx.t;
(*
elpd_i, the estimate of log p(y_i | y_-i) of each point given the others, on the points' axes.
*)
#khat : (float, 'f, 'd) Nx.t;
(*
The k̂ of each point's ratios (psis), on the points' axes.
*)

}

The type for leave-one-out cross-validations.

#val loo : (float, 'f, 'd) Nx.t Draws.t -> ('f, 'd) loo

loo ll is the Pareto-smoothed leave-one-out cross-validation (Vehtari, Gelman and Gabry 2017) of the log likelihoods ll, of shape [chain; draw; ...]: log p(y_i | θ_s) of each point y_i at each posterior draw θ_s, as Draws.map of Norn_model.pointwise gives them.

The draws are importance draws of each posterior that leaves one point out, the points independent given θ: point i's ratios 1 / p(y_i | θ_s) are smoothed as psis does into w_s, and elpd_i = log (Σ_s w_s p(y_i | θ_s) / Σ_s w_s). Autocorrelated chains widen the tail to M = floor (min (S / 5, 3 sqrt (S / r))), r the ratios' bulk effective sample size (ess_bulk) over S, or 1 where they have none; the threshold on k̂ stays that of S draws.

A point's elpd_i is unreliable when its k̂ is above the threshold. A point whose log likelihood is NaN or infinite at some draw has NaN elpd_i, k̂ and term of p.