Fix na_impute bug bgmCompare.

Author: MaartenMarsman

src/bgmCompare_parallel.cpp

@@ -58,7 +58,7 @@ struct ChainResultCompare {
  *  - Store results into the shared `results` vector at the chain index.
  *
  * Inputs (stored as const references or values):
- *  - observations: Input observation matrix (persons × variables).
+ *  - observations_master: Input observation matrix (persons × variables).
  *  - num_groups: Number of groups.
  *  - counts_per_category_master: Group-level category counts.
  *  - blume_capel_stats_master: Group-level Blume–Capel sufficient statistics.
@@ -99,7 +99,7 @@ struct ChainResultCompare {
  *  - Errors are caught locally so one failing chain does not crash the entire run.
  */
 struct GibbsCompareChainRunner : public Worker {
-  const arma::imat& observations;
+  const arma::imat& observations_master;
   const int num_groups;
   const std::vector<arma::imat>& counts_per_category_master;
   const std::vector<arma::imat>& blume_capel_stats_master;
@@ -138,7 +138,7 @@ struct GibbsCompareChainRunner : public Worker {
   std::vector<ChainResultCompare>& results;
 
   GibbsCompareChainRunner(
-    const arma::imat& observations,
+    const arma::imat& observations_master,
     int num_groups,
     const std::vector<arma::imat>& counts_per_category_master,
     const std::vector<arma::imat>& blume_capel_stats_master,
@@ -174,7 +174,7 @@ struct GibbsCompareChainRunner : public Worker {
     ProgressManager& pm,
     std::vector<ChainResultCompare>& results
   ) :
-    observations(observations),
+    observations_master(observations_master),
     num_groups(num_groups),
     counts_per_category_master(counts_per_category_master),
     blume_capel_stats_master(blume_capel_stats_master),
@@ -226,12 +226,12 @@ struct GibbsCompareChainRunner : public Worker {
         std::vector<arma::imat> blume_capel_stats = blume_capel_stats_master;
         std::vector<arma::mat>  pairwise_stats = pairwise_stats_master;
         arma::mat inclusion_probability = inclusion_probability_master;
-        arma::imat observations_copy = observations;
+        arma::imat observations = observations_master;
 
         // run sampler (pure C++)
         SamplerOutput result = run_gibbs_sampler_bgmCompare(
           out.chain_id,
-          observations_copy,
+          observations,
           num_groups,
           counts_per_category,
           blume_capel_stats,

src/bgmCompare_sampler.cpp

@@ -79,15 +79,15 @@ void impute_missing_bgmcompare(
     std::vector<arma::imat>& counts_per_category,
     std::vector<arma::imat>& blume_capel_stats,
     std::vector<arma::mat>& pairwise_stats,
-    const arma::imat& num_categories,
+    const arma::ivec& num_categories,
     const arma::imat& missing_data_indices,
     const arma::uvec& is_ordinal_variable,
     const arma::ivec& baseline_category,
     SafeRNG& rng
 ) {
   const int num_variables = observations.n_cols;
   const int num_missings = missing_data_indices.n_rows;
-  const int max_num_categories = arma::max(arma::vectorise(num_categories));
+  const int max_num_categories = arma::max(num_categories);
 
   arma::vec category_response_probabilities(max_num_categories + 1);
   double exponent, cumsum, u;
@@ -120,11 +120,12 @@ void impute_missing_bgmcompare(
 
     double rest_score =
       arma::as_scalar(observations.row(person) * group_pairwise_effects.col(variable));
+
     if(is_ordinal_variable[variable] == true) {
       // For regular binary or ordinal variables
       cumsum = 1.0;
       category_response_probabilities[0] = 1.0;
-      for(int category = 1; category <= num_categories(variable, group); category++) {
+      for(int category = 1; category <= num_categories(variable); category++) {
         exponent = group_main_effects(category - 1);
         exponent += category * rest_score;
         cumsum += MY_EXP(exponent);
@@ -133,7 +134,7 @@ void impute_missing_bgmcompare(
     } else {
       // For Blume-Capel variables
       cumsum = 0.0;
-      for(int category = 0; category <= num_categories(variable, group); category++) {
+      for(int category = 0; category <= num_categories(variable); category++) {
         exponent = group_main_effects[0] * category;
         exponent += group_main_effects[1] *
           (category - baseline_category[variable]) *
@@ -161,9 +162,9 @@ void impute_missing_bgmcompare(
       if(is_ordinal_variable[variable] == true) {
         arma::imat counts_per_category_group = counts_per_category[group];
         if(old_observation > 0)
-          counts_per_category_group(old_observation, variable)--;
+          counts_per_category_group(old_observation-1, variable)--;
         if(new_observation > 0)
-          counts_per_category_group(new_observation, variable)++;
+          counts_per_category_group(new_observation-1, variable)++;
         counts_per_category[group] = counts_per_category_group;
       } else {
         arma::imat blume_capel_stats_group = blume_capel_stats[group];