diff --git a/src/mcmc/execution/chain_runner.cpp b/src/mcmc/execution/chain_runner.cpp index fdd186a8..fef5ade8 100644 --- a/src/mcmc/execution/chain_runner.cpp +++ b/src/mcmc/execution/chain_runner.cpp @@ -6,16 +6,26 @@ #include "mcmc/samplers/metropolis_sampler.h" -std::unique_ptr create_sampler(const SamplerConfig& config, WarmupSchedule& schedule) { - if (config.sampler_type == "nuts") { - return std::make_unique(config, schedule); - } else if (config.sampler_type == "adaptive-metropolis") { - return std::make_unique(config, schedule); +SamplerSpec resolve_sampler_spec(const std::string& sampler_type) { + if (sampler_type == "nuts") { + return SamplerSpec{SamplerKind::NUTS, /*learn_sd=*/true, /*nuts_diag=*/true, /*am_diag=*/false}; + } else if (sampler_type == "adaptive-metropolis") { + return SamplerSpec{SamplerKind::AdaptiveMetropolis, /*learn_sd=*/false, /*nuts_diag=*/false, /*am_diag=*/true}; } else { - Rcpp::stop("Unknown sampler_type: '%s'", config.sampler_type.c_str()); + Rcpp::stop("Unknown sampler_type: '%s'", sampler_type.c_str()); } } +std::unique_ptr create_sampler(SamplerKind kind, const SamplerConfig& config, WarmupSchedule& schedule) { + switch (kind) { + case SamplerKind::NUTS: + return std::make_unique(config, schedule); + case SamplerKind::AdaptiveMetropolis: + return std::make_unique(config, schedule); + } + Rcpp::stop("Unhandled SamplerKind"); // unreachable: kind comes from resolve_sampler_spec +} + void run_mcmc_chain( ChainResult& chain_result, @@ -28,10 +38,10 @@ void run_mcmc_chain( chain_result.chain_id = chain_id + 1; // Construct warmup schedule (shared by runner and sampler) - const bool learn_sd = (config.sampler_type == "nuts"); - WarmupSchedule schedule(config.no_warmup, config.edge_selection, learn_sd); + const SamplerSpec spec = resolve_sampler_spec(config.sampler_type); + WarmupSchedule schedule(config.no_warmup, config.edge_selection, spec.learn_sd); - auto sampler = create_sampler(config, schedule); + auto sampler = create_sampler(spec.kind, config, schedule); // Initialize sampler (step-size heuristic) before the main loop sampler->initialize(model); @@ -138,8 +148,9 @@ std::vector run_mcmc_sampler( const int no_threads, ProgressManager& pm ) { - const bool has_nuts_diag = (config.sampler_type == "nuts"); - const bool has_am_diag = (config.sampler_type == "adaptive-metropolis"); + const SamplerSpec spec = resolve_sampler_spec(config.sampler_type); + const bool has_nuts_diag = spec.nuts_diag; + const bool has_am_diag = spec.am_diag; const bool has_sbm_alloc = edge_prior.has_allocations() || (config.edge_selection && dynamic_cast(&edge_prior) != nullptr); diff --git a/src/mcmc/execution/chain_runner.h b/src/mcmc/execution/chain_runner.h index fca3cc27..84ae6cd6 100644 --- a/src/mcmc/execution/chain_runner.h +++ b/src/mcmc/execution/chain_runner.h @@ -14,14 +14,38 @@ #include "mcmc/execution/warmup_schedule.h" +/** Which concrete sampler a run uses. */ +enum class SamplerKind { NUTS, AdaptiveMetropolis }; + +/** + * Behavioral descriptor for a sampler type. resolve_sampler_spec is the single + * place that decodes the config.sampler_type string; all downstream code reads + * these fields instead of re-comparing the string. + */ +struct SamplerSpec { + SamplerKind kind; + bool learn_sd; ///< NUTS dual-averaging step-size adaptation during warmup + bool nuts_diag; ///< reserve/store NUTS per-iteration diagnostics + bool am_diag; ///< reserve/store adaptive-Metropolis diagnostics +}; + +/** + * Decode a sampler-type string into a SamplerSpec. + * + * @param sampler_type "nuts" or "adaptive-metropolis". + * @return Descriptor with the concrete kind and its derived behavior flags. + */ +SamplerSpec resolve_sampler_spec(const std::string& sampler_type); + /** - * Create a sampler matching config.sampler_type + * Create a sampler of the given kind. * - * @param config Sampler configuration (type, step size, tree depth, etc.) + * @param kind Concrete sampler kind (from resolve_sampler_spec). + * @param config Sampler configuration (step size, tree depth, etc.) * @param schedule Shared warmup schedule for adaptation staging * @return Owning pointer to a concrete SamplerBase subclass */ -std::unique_ptr create_sampler(const SamplerConfig& config, WarmupSchedule& schedule); +std::unique_ptr create_sampler(SamplerKind kind, const SamplerConfig& config, WarmupSchedule& schedule); /** * Run a single MCMC chain (warmup + sampling)