Skip to content

fix(diagnostics,summary): NA-safe criteria, AM acceptance, summary contracts #213

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MaartenMarsman merged 1 commit into from
Oct 21, 2026
Merged

fix(diagnostics,summary): NA-safe criteria, AM acceptance, summary contracts #213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 68 additions & 23 deletions R/diagnostics_am.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@
# necessarily lives in its own file because the underlying sampler
# organisation is different: NUTS does one accept/reject per iteration
# over the whole parameter vector, whereas adaptive-Metropolis does
# one accept/reject per parameter per iteration. The natural per-parameter
# acceptance summary therefore differs in shape (parameter x chain)
# from the per-trajectory NUTS summary (chain x iteration).
# one accept/reject per parameter per iteration. The per-parameter move
# summary therefore has a shape (parameter x chain) with no NUTS analogue,
# alongside the per-sweep acceptance trace (chain x iteration) that does
# mirror the NUTS one.
# ==============================================================================


Expand All @@ -18,29 +19,48 @@
# ------------------------------------------------------------------------------
# Combine and summarize adaptive-Metropolis diagnostics across chains.
#
# The acceptance probability for a componentwise Metropolis update of a
# single parameter is estimated as the empirical *move rate*: the
# fraction of post-warmup iterations on which that parameter's value
# changed from the previous iteration. Since the sampler proposes from
# a continuous Normal proposal, "no change" is essentially equivalent
# to "rejected" (a probability-zero exact-match accept is ignored).
# Two quantities are reported:
#
# accept_prob -- the Metropolis acceptance probability the sampler itself
# records. The C++ sampler averages the acceptance probability over all
# components updated in a sweep and stores that per iteration as the
# `am_accept_prob__` trace (metropolis_sampler.h -> ChainResult ->
# build_output.R). This is the statistic Robbins-Monro adaptation drives
# towards `target_accept`, so it is the one to compare against it.
#
# move_rate -- the empirical *move rate*: the fraction of post-warmup
# iterations on which a parameter's stored value differed from the
# previous iteration. It is not an acceptance rate whenever a parameter
# can go un-proposed: under edge selection an excluded pairwise effect is
# held at exactly 0 and never moves, so its move rate is approximately
# P(included) * P(accept) and sits well below `target_accept` even when
# the sampler is tuned correctly. Reported per parameter, which the
# acceptance trace is not, so it still shows which parameters are moving.
#
# @param out List of chain outputs. Each element is a named list
# that must contain "main_samples" and
# "pairwise_samples" matrices (iterations x params).
# "pairwise_samples" matrices (iterations x params),
# and optionally the "am_accept_prob__" trace.
# @param names_main Character vector of main-effect parameter names.
# @param names_pairwise Character vector of pairwise interaction names.
# @param target_accept Target acceptance rate the sampler was tuned to
# (default: 0.44 for componentwise random-walk MH).
#
# Returns: An invisible named list with:
# - accept_prob: Numeric matrix (parameter x chain) of per-parameter
# per-chain empirical Metropolis acceptance rates. Rows are
# labelled by parameter (main effects first, then pairwise);
# columns are labelled "chain 1", "chain 2", ...
# - target_accept: Numeric scalar; the target acceptance rate.
# - summary: List with mean_accept_prob (mean across all
# parameters and chains).
# - accept_prob: Numeric matrix (chains x iterations) of the sampler's
# own per-sweep mean Metropolis acceptance probability, or NULL when
# the fit carries no `am_accept_prob__` trace (pre-0.2.0.0 fits, and
# models that take no Metropolis step). Rows are labelled
# "chain 1", "chain 2", ...
# - move_rate: Numeric matrix (parameter x chain) of per-parameter
# per-chain empirical move rates. Rows are labelled by parameter
# (main effects first, then pairwise); columns are labelled
# "chain 1", "chain 2", ...
# - target_accept: Numeric scalar; the target acceptance rate. Comparable
# with accept_prob, not with move_rate.
# - summary: List with mean_accept_prob (mean of accept_prob over
# all iterations and chains, NA when the trace is absent) and
# mean_move_rate (mean of move_rate over all parameters and chains).
# ------------------------------------------------------------------------------
summarize_am_diagnostics = function(out, names_main, names_pairwise,
target_accept = 0.44) {
Expand All @@ -52,7 +72,7 @@ summarize_am_diagnostics = function(out, names_main, names_pairwise,
return(NULL)
}

per_chain_accept = function(chain) {
per_chain_move_rate = function(chain) {
main_acc = apply(
chain$main_samples, 2,
function(col) mean(diff(col) != 0)
Expand All @@ -64,19 +84,44 @@ summarize_am_diagnostics = function(out, names_main, names_pairwise,
c(main_acc, pair_acc)
}

accept_prob_mat = vapply(
am_chains, per_chain_accept,
move_rate_mat = vapply(
am_chains, per_chain_move_rate,
numeric(length(names_main) + length(names_pairwise))
)

rownames(accept_prob_mat) = c(names_main, names_pairwise)
colnames(accept_prob_mat) = paste0("chain ", seq_along(am_chains))
rownames(move_rate_mat) = c(names_main, names_pairwise)
colnames(move_rate_mat) = paste0("chain ", seq_along(am_chains))

# accept_prob is built only when every chain has an am_accept_prob__ trace
# with at least one finite value. ChainResult pre-fills the trace with NaN,
# so an all-NaN trace means the sampler never wrote one.
has_accept_trace = all(vapply(
am_chains,
function(chain) {
trace = chain[["am_accept_prob__"]]
!is.null(trace) && any(is.finite(trace))
},
logical(1)
))

if(has_accept_trace) {
accept_prob_mat = do.call(rbind, lapply(
am_chains, function(chain) as.numeric(chain[["am_accept_prob__"]])
))
rownames(accept_prob_mat) = paste0("chain ", seq_along(am_chains))
mean_accept_prob = mean(accept_prob_mat, na.rm = TRUE)
} else {
accept_prob_mat = NULL
mean_accept_prob = NA_real_
}

invisible(list(
accept_prob = accept_prob_mat,
move_rate = move_rate_mat,
target_accept = target_accept,
summary = list(
mean_accept_prob = mean(accept_prob_mat, na.rm = TRUE)
mean_accept_prob = mean_accept_prob,
mean_move_rate = mean(move_rate_mat, na.rm = TRUE)
)
))
}
51 changes: 42 additions & 9 deletions R/diagnostics_nuts.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,19 @@ check_warmup_complete = function(energy_mat) {
time_idx = seq_len(n_chain)
trend_lm = stats::lm(energy ~ time_idx)
slope = stats::coef(trend_lm)[2]
slope_se = summary(trend_lm)$coefficients[2, 2]
slope_significant = abs(slope / slope_se) > 2.58
# Constant trace: slope and its SE are both ~0, so the t-ratio is
# undefined. Set slope = 0, SE = 0, not significant, and skip summary.lm(),
# which warns about the exact fit.
if(stats::var(energy) == 0) {
slope[] = 0
slope_se = 0
slope_significant = FALSE
} else {
slope_se = summary(trend_lm)$coefficients[2, 2]
# isTRUE(): an NA comparison counts as not triggered, so the flag stays
# logical.
slope_significant = isTRUE(abs(slope / slope_se) > 2.58)
}

# Reported alongside the flag, not part of it: the same statistic with its
# standard error corrected for autocorrelation. tau comes from the second
Expand All @@ -142,8 +153,11 @@ check_warmup_complete = function(energy_mat) {
# Variance ratio
var_ratio = stats::var(first_half) / stats::var(second_half)

# Flag if any criterion triggered
warmup_incomplete = slope_significant || ebfmi_first < 0.3 || var_ratio > 2.0
# Flag if any criterion triggered. A constant half makes ebfmi_first and
# var_ratio NaN; isTRUE() keeps the flag logical instead of NA.
warmup_incomplete = slope_significant ||
isTRUE(ebfmi_first < 0.3) ||
isTRUE(var_ratio > 2.0)

list(
warmup_incomplete = warmup_incomplete,
Expand Down Expand Up @@ -193,10 +207,12 @@ check_warmup_complete = function(energy_mat) {
# - energy: Numeric matrix (chains x iterations).
# - accept_prob: Numeric matrix (chains x iterations) of mean
# per-trajectory Metropolis acceptance (Stan's accept_stat__).
# - ebfmi: Numeric vector of per-chain E-BFMI values.
# - ebfmi: Numeric vector of per-chain E-BFMI values, computed from the
# finite energy draws only and NA for a chain with fewer than two.
# - warmup_check: Output of check_warmup_complete().
# - summary: List with total_divergences, max_tree_depth_hits,
# min_ebfmi, mean_accept_prob, and warmup_incomplete (logical).
# min_ebfmi (the minimum over the chains with a computable E-BFMI, NA if
# no chain has one), mean_accept_prob, and warmup_incomplete (logical).
# ------------------------------------------------------------------------------
summarize_nuts_diagnostics = function(out, nuts_max_depth = 10, verbose = TRUE) {
nuts_chains = Filter(function(chain) {
Expand Down Expand Up @@ -225,8 +241,13 @@ summarize_nuts_diagnostics = function(out, nuts_max_depth = 10, verbose = TRUE)
matrix(NA_real_, nrow = nrow(divergent_mat), ncol = ncol(divergent_mat))
}

# E-BFMI per chain
# E-BFMI per chain, from the finite energy draws only (as in
# check_warmup_complete()). NA with fewer than two finite draws.
compute_ebfmi = function(energy) {
energy = energy[is.finite(energy)]
if(length(energy) < 2L) {
return(NA_real_)
}
mean(diff(energy)^2) / stats::var(energy)
}
ebfmi_per_chain = apply(energy_mat, 1, compute_ebfmi)
Expand All @@ -237,8 +258,12 @@ summarize_nuts_diagnostics = function(out, nuts_max_depth = 10, verbose = TRUE)
n_total = nrow(divergent_mat) * ncol(divergent_mat)
total_divergences = sum(divergent_mat)
max_tree_depth_hits = sum(treedepth_mat == nuts_max_depth)
min_ebfmi = min(ebfmi_per_chain)
low_ebfmi_chains = which(ebfmi_per_chain < 0.2)
# min and threshold over the chains with a computable E-BFMI; the others get
# their own issue line below.
ebfmi_computed = is.finite(ebfmi_per_chain)
min_ebfmi = if(any(ebfmi_computed)) min(ebfmi_per_chain[ebfmi_computed]) else NA_real_
low_ebfmi_chains = which(ebfmi_computed & ebfmi_per_chain < 0.2)
unusable_ebfmi_chains = which(!ebfmi_computed)

divergence_rate = total_divergences / n_total
depth_hit_rate = max_tree_depth_hits / n_total
Expand Down Expand Up @@ -285,6 +310,14 @@ summarize_nuts_diagnostics = function(out, nuts_max_depth = 10, verbose = TRUE)
))
}

if(length(unusable_ebfmi_chains) > 0) {
issues = c(issues, sprintf(
"E-BFMI: not computable in chain%s %s - fewer than two finite energy draws",
if(length(unusable_ebfmi_chains) > 1) "s" else "",
paste(unusable_ebfmi_chains, collapse = ", ")
))
}

incomplete_chains = which(warmup_check$warmup_incomplete)
if(length(incomplete_chains) > 0) {
# Name the criterion that actually fired. A low first-half E-BFMI is a
Expand Down
Loading
Loading