use enums instead of strings

Author: vandenman

src/bgmCompare_parallel.cpp

@@ -9,6 +9,7 @@
 #include "progress_manager.h"
 #include "sampler_output.h"
 #include "mcmc_adaptation.h"
+#include "common_helpers.h"
 
 using namespace Rcpp;
 using namespace RcppParallel;
@@ -131,7 +132,7 @@ struct GibbsCompareChainRunner : public Worker {
   const arma::mat& inclusion_probability_master;
   // RNG seeds
   const std::vector<SafeRNG>& chain_rngs;
-  const std::string& update_method;
+  const UpdateMethod update_method;
   const int hmc_num_leapfrogs;
   ProgressManager& pm;
   // output
@@ -169,7 +170,7 @@ struct GibbsCompareChainRunner : public Worker {
     const arma::imat& interaction_index_matrix,
     const arma::mat& inclusion_probability_master,
     const std::vector<SafeRNG>& chain_rngs,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const int hmc_num_leapfrogs,
     ProgressManager& pm,
     std::vector<ChainResultCompare>& results
@@ -394,8 +395,10 @@ Rcpp::List run_bgmCompare_parallel(
     chain_rngs[c] = SafeRNG(seed + c);
   }
 
+  UpdateMethod update_method_enum = update_method_from_string(update_method);
+
   // only used to determine the total no. warmup iterations, a bit hacky
-  WarmupSchedule warmup_schedule_temp(warmup, difference_selection, (update_method != "adaptive-metropolis"));
+  WarmupSchedule warmup_schedule_temp(warmup, difference_selection, (update_method_enum != adaptive_metropolis));
   int total_warmup = warmup_schedule_temp.total_warmup;
   ProgressManager pm(num_chains, iter, total_warmup, 50, progress_type);
 
@@ -408,7 +411,7 @@ Rcpp::List run_bgmCompare_parallel(
       baseline_category, difference_selection, main_effect_indices,
       pairwise_effect_indices, target_accept, nuts_max_depth, learn_mass_matrix,
       projection, group_membership, group_indices, interaction_index_matrix,
-      inclusion_probability, chain_rngs, update_method, hmc_num_leapfrogs,
+      inclusion_probability, chain_rngs, update_method_enum, hmc_num_leapfrogs,
       pm, results
   );
 

src/bgmCompare_sampler.cpp

@@ -1389,7 +1389,7 @@ void gibbs_update_step_bgmcompare (
     SafeRNG& rng,
     arma::mat& inclusion_probability,
     int hmc_nuts_leapfrogs,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     arma::mat& proposal_sd_main,
     arma::mat& proposal_sd_pair,
     const arma::imat& index
@@ -1416,7 +1416,7 @@ void gibbs_update_step_bgmcompare (
   }
 
   // Step 2: Update parameters
-  if(update_method == "adaptive-metropolis") {
+  if(update_method == adaptive_metropolis) {
     update_main_effects_metropolis_bgmcompare (
         main_effects, pairwise_effects, main_effect_indices,
         pairwise_effect_indices, inclusion_indicator, projection,
@@ -1434,7 +1434,7 @@ void gibbs_update_step_bgmcompare (
         pairwise_scale, difference_scale, iteration, rwm_adapt_pair, rng,
         proposal_sd_pair
     );
-  } else if (update_method == "hamiltonian-mc") {
+  } else if (update_method == hamiltonian_mc) {
     update_hmc_bgmcompare(
       main_effects, pairwise_effects, main_effect_indices,
       pairwise_effect_indices, inclusion_indicator, projection, num_categories,
@@ -1444,7 +1444,7 @@ void gibbs_update_step_bgmcompare (
       main_beta, hmc_nuts_leapfrogs, iteration, hmc_adapt, learn_mass_matrix,
       schedule.selection_enabled(iteration), rng
     );
-  } else if (update_method == "nuts") {
+  } else if (update_method == nuts) {
     SamplerResult result = update_nuts_bgmcompare(
       main_effects, pairwise_effects, main_effect_indices,
       pairwise_effect_indices, inclusion_indicator, projection, num_categories,
@@ -1577,7 +1577,7 @@ SamplerOutput run_gibbs_sampler_bgmCompare(
     const arma::imat& interaction_index_matrix,
     arma::mat inclusion_probability,
     SafeRNG& rng,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const int hmc_num_leapfrogs,
     ProgressManager& pm
 ) {
@@ -1618,7 +1618,7 @@ SamplerOutput run_gibbs_sampler_bgmCompare(
 
   // --- Optional HMC/NUTS warmup stage
   double initial_step_size = 1.0;
-  if (update_method == "hamiltonian-mc" || update_method == "nuts") {
+  if (update_method == hamiltonian_mc || update_method == nuts) {
     initial_step_size = find_initial_stepsize_bgmcompare(
       main_effects, pairwise_effects, main_effect_indices,
       pairwise_effect_indices, inclusion_indicator, projection, num_categories,

src/bgmCompare_sampler.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <RcppArmadillo.h>
+#include "common_helpers.h"
 #include <string>
 
 struct SamplerOutput;
@@ -40,7 +41,7 @@ SamplerOutput run_gibbs_sampler_bgmCompare(
     const arma::imat& interaction_index_matrix,
     arma::mat inclusion_probability,
     SafeRNG& rng,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const int hmc_num_leapfrogs,
     ProgressManager& pm
 );

src/bgm_parallel.cpp

@@ -8,6 +8,7 @@
 #include <string>
 #include "progress_manager.h"
 #include "mcmc_adaptation.h"
+#include "common_helpers.h"
 
 using namespace Rcpp;
 using namespace RcppParallel;
@@ -61,7 +62,7 @@ struct GibbsChainRunner : public Worker {
   const arma::imat& observations;
   const arma::ivec& num_categories;
   double  pairwise_scale;
-  const std::string& edge_prior;
+  const EdgePrior edge_prior;
   const arma::mat& inclusion_probability;
   double beta_bernoulli_alpha;
   double beta_bernoulli_beta;
@@ -79,7 +80,7 @@ struct GibbsChainRunner : public Worker {
   const arma::uvec& is_ordinal_variable;
   const arma::ivec& baseline_category;
   bool edge_selection;
-  const std::string& update_method;
+  const UpdateMethod update_method;
   const arma::imat& pairwise_effect_indices;
   double target_accept;
   const arma::imat& pairwise_stats;
@@ -98,7 +99,7 @@ struct GibbsChainRunner : public Worker {
     const arma::imat& observations,
     const arma::ivec& num_categories,
     double  pairwise_scale,
-    const std::string& edge_prior,
+    const EdgePrior edge_prior,
     const arma::mat& inclusion_probability,
     double beta_bernoulli_alpha,
     double beta_bernoulli_beta,
@@ -116,7 +117,7 @@ struct GibbsChainRunner : public Worker {
     const arma::uvec& is_ordinal_variable,
     const arma::ivec& baseline_category,
     bool edge_selection,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const arma::imat& pairwise_effect_indices,
     double target_accept,
     const arma::imat& pairwise_stats,
@@ -308,18 +309,20 @@ Rcpp::List run_bgm_parallel(
     chain_rngs[c] = SafeRNG(seed + c);
   }
 
+  UpdateMethod update_method_enum = update_method_from_string(update_method);
+  EdgePrior edge_prior_enum = edge_prior_from_string(edge_prior);
   // only used to determine the total no. warmup iterations, a bit hacky
-  WarmupSchedule warmup_schedule_temp(warmup, edge_selection, (update_method != "adaptive-metropolis"));
+  WarmupSchedule warmup_schedule_temp(warmup, edge_selection, (update_method_enum != adaptive_metropolis));
   int total_warmup = warmup_schedule_temp.total_warmup;
   ProgressManager pm(num_chains, iter, total_warmup, 50, progress_type);
 
   GibbsChainRunner worker(
-      observations, num_categories,  pairwise_scale, edge_prior,
+      observations, num_categories,  pairwise_scale, edge_prior_enum,
       inclusion_probability, beta_bernoulli_alpha, beta_bernoulli_beta,
       dirichlet_alpha, lambda, interaction_index_matrix, iter, warmup,
       counts_per_category, blume_capel_stats, main_alpha, main_beta,
       na_impute, missing_index, is_ordinal_variable, baseline_category,
-      edge_selection, update_method, pairwise_effect_indices, target_accept,
+      edge_selection, update_method_enum, pairwise_effect_indices, target_accept,
       pairwise_stats, hmc_num_leapfrogs, nuts_max_depth, learn_mass_matrix,
       chain_rngs, pm, results
   );

src/bgm_sampler.cpp

@@ -1015,7 +1015,7 @@ void gibbs_update_step_bgm (
     const arma::uvec& is_ordinal_variable,
     const arma::ivec& baseline_category,
     const int iteration,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const arma::imat& pairwise_effect_indices,
     arma::imat& pairwise_stats,
     const int hmc_num_leapfrogs,
@@ -1051,7 +1051,7 @@ void gibbs_update_step_bgm (
   }
 
   // Step 2a: Update interaction weights for active edges
-  if (update_method == "adaptive-metropolis") {
+  if (update_method == adaptive_metropolis) {
     update_pairwise_effects_metropolis_bgm (
         pairwise_effects, main_effects, inclusion_indicator, observations,
         num_categories, proposal_sd_pairwise, adapt_pairwise, pairwise_scale,
@@ -1061,7 +1061,7 @@ void gibbs_update_step_bgm (
   }
 
   // Step 2b: Update main effect (main_effect) parameters
-  if (update_method == "adaptive-metropolis") {
+  if (update_method == adaptive_metropolis) {
     update_main_effects_metropolis_bgm (
         main_effects, observations, num_categories, counts_per_category,
         blume_capel_stats, baseline_category, is_ordinal_variable,
@@ -1072,7 +1072,7 @@ void gibbs_update_step_bgm (
   }
 
   // Step 2: Update joint parameters if applicable
-  if (update_method == "hamiltonian-mc") {
+  if (update_method == hamiltonian_mc) {
     update_hmc_bgm(
       main_effects, pairwise_effects, inclusion_indicator, observations,
       num_categories, counts_per_category, blume_capel_stats,
@@ -1081,7 +1081,7 @@ void gibbs_update_step_bgm (
       iteration, adapt, learn_mass_matrix, schedule.selection_enabled(iteration),
       rng
     );
-  } else if (update_method == "nuts") {
+  } else if (update_method == nuts) {
     SamplerResult result = update_nuts_bgm(
       main_effects, pairwise_effects, inclusion_indicator,
       observations, num_categories, counts_per_category, blume_capel_stats,
@@ -1171,7 +1171,7 @@ Rcpp::List run_gibbs_sampler_bgm(
     arma::imat observations,
     const arma::ivec& num_categories,
     const double pairwise_scale,
-    const std::string& edge_prior,
+    const EdgePrior edge_prior,
     arma::mat inclusion_probability,
     const double beta_bernoulli_alpha,
     const double beta_bernoulli_beta,
@@ -1189,7 +1189,7 @@ Rcpp::List run_gibbs_sampler_bgm(
     const arma::uvec& is_ordinal_variable,
     const arma::ivec& baseline_category,
     bool edge_selection,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const arma::imat pairwise_effect_indices,
     const double target_accept,
     arma::imat pairwise_stats,
@@ -1223,7 +1223,7 @@ Rcpp::List run_gibbs_sampler_bgm(
   if (edge_selection) {
     indicator_samples.set_size(iter, num_pairwise);
   }
-  if (edge_selection && edge_prior == "Stochastic-Block") {
+  if (edge_selection && edge_prior == Stochastic_Block) {
     allocation_samples.set_size(iter, num_variables);
   }
 
@@ -1245,7 +1245,7 @@ Rcpp::List run_gibbs_sampler_bgm(
   arma::vec log_Vn(1);
 
   // --- Initialize SBM prior if applicable
-  if (edge_prior == "Stochastic-Block") {
+  if (edge_prior == Stochastic_Block) {
     cluster_allocations[0] = 0;
     cluster_allocations[1] = 1;
     for (int i = 2; i < num_variables; i++) {
@@ -1273,7 +1273,7 @@ Rcpp::List run_gibbs_sampler_bgm(
 
   // --- Optional HMC/NUTS warmup stage
   double initial_step_size_joint = 1.0;
-  if (update_method == "hamiltonian-mc" || update_method == "nuts") {
+  if (update_method == hamiltonian_mc || update_method == nuts) {
     initial_step_size_joint = find_initial_stepsize_bgm(
       main_effects, pairwise_effects, inclusion_indicator, observations,
       num_categories, counts_per_category, blume_capel_stats,
@@ -1283,7 +1283,7 @@ Rcpp::List run_gibbs_sampler_bgm(
   }
 
   // --- Warmup scheduling + adaptation controller
-  WarmupSchedule warmup_schedule(warmup, edge_selection, (update_method != "adaptive-metropolis"));
+  WarmupSchedule warmup_schedule(warmup, edge_selection, (update_method != adaptive_metropolis));
   HMCAdaptationController adapt_joint(
       num_main + num_pairwise, initial_step_size_joint, target_accept,
       warmup_schedule, learn_mass_matrix
@@ -1339,7 +1339,7 @@ Rcpp::List run_gibbs_sampler_bgm(
 
     // --- Update edge probabilities under the prior (if edge selection is active)
     if (warmup_schedule.selection_enabled(iteration)) {
-      if (edge_prior == "Beta-Bernoulli") {
+      if (edge_prior == Beta_Bernoulli) {
         int num_edges_included = 0;
         for (int i = 0; i < num_variables - 1; i++)
           for (int j = i + 1; j < num_variables; j++)
@@ -1354,7 +1354,7 @@ Rcpp::List run_gibbs_sampler_bgm(
           for (int j = i + 1; j < num_variables; j++)
             inclusion_probability(i, j) = inclusion_probability(j, i) = prob;
 
-      } else if (edge_prior == "Stochastic-Block") {
+      } else if (edge_prior == Stochastic_Block) {
         cluster_allocations = block_allocations_mfm_sbm(
           cluster_allocations, num_variables, log_Vn, cluster_prob,
           arma::conv_to<arma::umat>::from(inclusion_indicator), dirichlet_alpha,
@@ -1396,7 +1396,7 @@ Rcpp::List run_gibbs_sampler_bgm(
         }
       }
 
-      if (edge_selection && edge_prior == "Stochastic-Block") {
+      if (edge_selection && edge_prior == Stochastic_Block) {
         for (int v = 0; v < num_variables; v++) {
           allocation_samples(sample_index, v) = cluster_allocations[v] + 1;
         }
@@ -1408,7 +1408,7 @@ Rcpp::List run_gibbs_sampler_bgm(
   out["main_samples"] = main_effect_samples;
   out["pairwise_samples"] = pairwise_effect_samples;
 
-  if (update_method == "nuts") {
+  if (update_method == nuts) {
     out["treedepth__"] = treedepth_samples;
     out["divergent__"] = divergent_samples;
     out["energy__"] = energy_samples;
@@ -1418,7 +1418,7 @@ Rcpp::List run_gibbs_sampler_bgm(
     out["indicator_samples"] = indicator_samples;
   }
 
-  if (edge_selection && edge_prior == "Stochastic-Block") {
+  if (edge_selection && edge_prior == Stochastic_Block) {
     out["allocations"] = allocation_samples;
   }
 

src/bgm_sampler.h

@@ -1,5 +1,6 @@
 #pragma once
 #include <RcppArmadillo.h>
+#include "common_helpers.h"
 // forward declaration
 struct SafeRNG;
 class ProgressManager;
@@ -9,7 +10,7 @@ Rcpp::List run_gibbs_sampler_bgm(
     arma::imat observations,
     const arma::ivec& num_categories,
     const double pairwise_scale,
-    const std::string& edge_prior,
+    const EdgePrior edge_prior,
     arma::mat inclusion_probability,
     const double beta_bernoulli_alpha,
     const double beta_bernoulli_beta,
@@ -27,7 +28,7 @@ Rcpp::List run_gibbs_sampler_bgm(
     const arma::uvec& is_ordinal_variable,
     const arma::ivec& baseline_category,
     bool edge_selection,
-    const std::string& update_method,
+    const UpdateMethod update_method,
     const arma::imat pairwise_effect_indices,
     const double target_accept,
     arma::imat pairwise_stats,

src/common_helpers.h

@@ -25,3 +25,30 @@ inline int count_num_main_effects(const arma::ivec& num_categories,
   }
   return n_params;
 }
+
+enum UpdateMethod { adaptive_metropolis, hamiltonian_mc, nuts };
+
+inline UpdateMethod update_method_from_string(const std::string& update_method) {
+  if (update_method == "adaptive-metropolis")
+    return adaptive_metropolis;
+
+  if (update_method == "hamiltonian-mc")
+    return hamiltonian_mc;
+
+  if (update_method == "nuts")
+    return nuts;
+
+  throw std::invalid_argument("Invalid update_method: " + update_method);
+}
+
+enum EdgePrior { Stochastic_Block, Beta_Bernoulli };
+
+inline EdgePrior edge_prior_from_string(const std::string& edge_prior) {
+  if (edge_prior == "stochastic-block")
+    return Stochastic_Block;
+
+  if (edge_prior == "Beta-Bernoulli")
+    return Beta_Bernoulli;
+
+  throw std::invalid_argument("Invalid edge_prior: " + edge_prior);
+}