From 2cb27ec2956870dcee6e43f25a4fa40ea3d52030 Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 4 Jun 2026 22:00:02 +0200 Subject: [PATCH] refactor(mcmc): extract store_nuts_diagnostics_if_present helper The NUTS per-iteration diagnostics store in run_mcmc_chain's sampling-phase block nested a dynamic_cast, a null check, and a six-argument store call inline. Move it to a file-local helper so the main loop reads as one call. Per the design-review critique this stays a chain_runner-local helper: no DiagnosticsBase::store_into (that would invert the algorithm/storage layering bgmCompare relies on), and the reserve/convert ladders stay explicit. Behavior-preserving: same guard, same dynamic_cast, same stored fields. Verified bitwise: NUTS GGM raw draws AND the stored NUTS diagnostics (tree depth, divergent, energy, accept prob, ...) are digest-identical pre/post (edge-sel + no-edge); adaptive-Metropolis draws unaffected. identical(main, branch) == TRUE. --- src/mcmc/execution/chain_runner.cpp | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/src/mcmc/execution/chain_runner.cpp b/src/mcmc/execution/chain_runner.cpp index fdd186a8..4ea9c008 100644 --- a/src/mcmc/execution/chain_runner.cpp +++ b/src/mcmc/execution/chain_runner.cpp @@ -6,6 +6,24 @@ #include "mcmc/samplers/metropolis_sampler.h" +namespace { + +// Store this sample's NUTS per-iteration diagnostics, if the chain is collecting +// them and the step produced a NUTSDiagnostics object. Keeps the dynamic_cast +// and the wide store_nuts_diagnostics call out of the main sampling loop. +void store_nuts_diagnostics_if_present(ChainResult& chain_result, int sample_index, + const SamplerBase& sampler, const StepResult& result) { + if (!chain_result.has_nuts_diagnostics || !sampler.has_nuts_diagnostics()) return; + auto* diag = dynamic_cast(result.diagnostics.get()); + if (diag) { + chain_result.store_nuts_diagnostics(sample_index, diag->tree_depth, diag->divergent, + diag->non_reversible, diag->energy, diag->accept_prob); + } +} + +} // namespace + + std::unique_ptr create_sampler(const SamplerConfig& config, WarmupSchedule& schedule) { if (config.sampler_type == "nuts") { return std::make_unique(config, schedule); @@ -78,12 +96,7 @@ void run_mcmc_chain( if (schedule.sampling(iter)) { int sample_index = iter - config.no_warmup; - if (chain_result.has_nuts_diagnostics && sampler->has_nuts_diagnostics()) { - auto* diag = dynamic_cast(result.diagnostics.get()); - if (diag) { - chain_result.store_nuts_diagnostics(sample_index, diag->tree_depth, diag->divergent, diag->non_reversible, diag->energy, diag->accept_prob); - } - } + store_nuts_diagnostics_if_present(chain_result, sample_index, *sampler, result); if (chain_result.has_am_diagnostics) { chain_result.store_am_diagnostics(sample_index, result.accept_prob);