Merge branch 'main' into sparse_leaf_matrix

Author: andrewherren

debug/api_debug.cpp

@@ -169,10 +169,34 @@ void RunAPI() {
   double outcome_scale;
   OutcomeOffsetScale(residual, outcome_offset, outcome_scale);
 
-  // // Construct a random effects dataset
-  // RandomEffectsDataset rfx_dataset = RandomEffectsDataset();
-  // rfx_dataset.AddBasis(rfx_basis_raw.data(), n, rfx_basis_cols, row_major);
-  // rfx_dataset.AddGroupLabels(rfx_groups);
+  // Construct a random effects dataset
+  RandomEffectsDataset rfx_dataset = RandomEffectsDataset();
+  rfx_dataset.AddBasis(rfx_basis_raw.data(), n, rfx_basis_cols, true);
+  rfx_dataset.AddGroupLabels(rfx_groups);
+
+  // Construct random effects tracker / model / container
+  RandomEffectsTracker rfx_tracker = RandomEffectsTracker(rfx_groups);
+  MultivariateRegressionRandomEffectsModel rfx_model = MultivariateRegressionRandomEffectsModel(rfx_basis_cols, num_rfx_groups);
+  RandomEffectsContainer rfx_container = RandomEffectsContainer(rfx_basis_cols, num_rfx_groups);
+  LabelMapper label_mapper = LabelMapper(rfx_tracker.GetLabelMap());
+
+  // Set random effects model parameters
+  Eigen::VectorXd working_param_init(rfx_basis_cols);
+  Eigen::MatrixXd group_param_init(rfx_basis_cols, num_rfx_groups);
+  Eigen::MatrixXd working_param_cov_init(rfx_basis_cols, rfx_basis_cols);
+  Eigen::MatrixXd group_param_cov_init(rfx_basis_cols, rfx_basis_cols);
+  double variance_prior_shape = 1.;
+  double variance_prior_scale = 1.;
+  working_param_init << 1.;
+  group_param_init << 1., 1.;
+  working_param_cov_init << 1;
+  group_param_cov_init << 1;
+  rfx_model.SetWorkingParameter(working_param_init);
+  rfx_model.SetGroupParameters(group_param_init);
+  rfx_model.SetWorkingParameterCovariance(working_param_cov_init);
+  rfx_model.SetGroupParameterCovariance(group_param_cov_init);
+  rfx_model.SetVariancePriorShape(variance_prior_shape);
+  rfx_model.SetVariancePriorScale(variance_prior_scale);
 
   // Initialize an ensemble
   int num_trees = 100;
@@ -244,6 +268,10 @@ void RunAPI() {
     sampleGFR(tracker, tree_prior, forest_samples, dataset, residual, rng, feature_types, variable_weights, 
               leaf_model_type, leaf_scale_matrix, global_variance, leaf_scale, cutpoint_grid_size);
 
+    // Sample random effects
+    rfx_model.SampleRandomEffects(rfx_dataset, residual, rfx_tracker, global_variance, rng);
+    rfx_container.AddSample(rfx_model);
+
     // Sample leaf node variance
     leaf_variance_samples.push_back(leaf_var_model.SampleVarianceParameter(forest_samples.GetEnsemble(i), a_leaf, b_leaf, rng));
 
@@ -266,24 +294,36 @@ void RunAPI() {
     sampleMCMC(tracker, tree_prior, forest_samples, dataset, residual, rng, feature_types, variable_weights, 
                leaf_model_type, leaf_scale_matrix, global_variance, leaf_scale, cutpoint_grid_size);
 
+    // Sample random effects
+    rfx_model.SampleRandomEffects(rfx_dataset, residual, rfx_tracker, global_variance, rng);
+    rfx_container.AddSample(rfx_model);
+
     // Sample leaf node variance
     leaf_variance_samples.push_back(leaf_var_model.SampleVarianceParameter(forest_samples.GetEnsemble(i), a_leaf, b_leaf, rng));
 
     // Sample global variance
     global_variance_samples.push_back(global_var_model.SampleVarianceParameter(residual.GetData(), nu, nu*lamb, rng));
   }
 
-  // Write model to a file
-  std::string filename = "model.json";
-  forest_samples.SaveToJsonFile(filename);
+  // Predict from the tree ensemble
+  std::vector<double> pred_orig = forest_samples.Predict(dataset);
+
+  // Predict from the random effects dataset
+  int num_samples = num_gfr_samples + num_mcmc_samples;
+  std::vector<double> rfx_predictions(n*num_samples);
+  rfx_container.Predict(rfx_dataset, label_mapper, rfx_predictions);
 
-  // Read and parse json from file
-  ForestContainer forest_samples_parsed = ForestContainer(num_trees, output_dimension, is_leaf_constant);
-  forest_samples_parsed.LoadFromJsonFile(filename);
+  // // Write model to a file
+  // std::string filename = "model.json";
+  // forest_samples.SaveToJsonFile(filename);
+
+  // // Read and parse json from file
+  // ForestContainer forest_samples_parsed = ForestContainer(num_trees, output_dimension, is_leaf_constant);
+  // forest_samples_parsed.LoadFromJsonFile(filename);
 
-  // Make sure we can predict from both the original and parsed forest containers
-  std::vector<double> pred_orig = forest_samples.Predict(dataset);
-  std::vector<double> pred_parsed = forest_samples_parsed.Predict(dataset);
+  // // Make sure we can predict from both the original and parsed forest containers
+  // std::vector<double> pred_orig = forest_samples.Predict(dataset);
+  // std::vector<double> pred_parsed = forest_samples_parsed.Predict(dataset);
 }
 
 } // namespace StochTree

include/stochtree/category_tracker.h

@@ -96,6 +96,7 @@ class CategorySampleTracker {
       if (i > 0) new_group_cond = group_indices[indices_[i]] != group_indices[indices_[i-1]];
       if (start_cond || new_group_cond) {
         category_id_map_.insert({group_indices[indices_[i]], category_count_});
+        unique_category_ids_.push_back(group_indices[indices_[i]]);
         node_index_vector_.emplace_back();
         if (i == 0) {
           category_begin_.push_back(i);
@@ -115,6 +116,11 @@ class CategorySampleTracker {
     }
   }
 
+  /*! \brief Zero-indexed numeric index that category_id is remapped to internally */
+  inline int32_t CategoryNumber(int category_id) {
+    return category_id_map_[category_id];
+  }
+
   /*! \brief First index of data points contained in node_id */
   inline data_size_t CategoryBegin(int category_id) {return category_begin_[category_id_map_[category_id]];}
 
@@ -144,15 +150,28 @@ class CategorySampleTracker {
     // return output;
     return node_index_vector_[id];
   }
+  
+  /*! \brief Data indices for a given node */
+  std::vector<data_size_t>& NodeIndicesInternalIndex(int internal_category_id) {
+//    int32_t id = category_id_map_[category_id];
+    // std::vector<data_size_t>::iterator start = indices_.begin() + category_begin_[id];
+    // std::vector<data_size_t>::iterator end = indices_.begin() + category_begin_[id] + category_length_[id];
+    // std::vector<data_size_t> output(start, end);
+    // return output;
+    return node_index_vector_[internal_category_id];
+  }
 
   /*! \brief Returns label index map */
   std::map<int32_t, int32_t>& GetLabelMap() {return category_id_map_;}
 
+  std::vector<int32_t>& GetUniqueGroupIds() {return unique_category_ids_;}
+
  private:
   // Vectors tracking indices in each node
   std::vector<data_size_t> category_begin_;
   std::vector<data_size_t> category_length_;
   std::map<int32_t, int32_t> category_id_map_;
+  std::vector<int32_t> unique_category_ids_;
   std::vector<std::vector<data_size_t>> node_index_vector_;
   int32_t category_count_;
 };

include/stochtree/data.h

@@ -118,6 +118,7 @@ class RandomEffectsDataset {
     group_labels_ = group_labels;
     has_group_labels_ = true;
   }
+  inline data_size_t NumObservations() {return basis_.NumRows();}
   inline bool HasBasis() {return has_basis_;}
   inline bool HasVarWeights() {return has_var_weights_;}
   inline bool HasGroupLabels() {return has_group_labels_;}