Added convenience method to random effects and category tracker classes

Author: andrewherren

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);
@@ -163,11 +164,14 @@ class CategorySampleTracker {
   /*! \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/random_effects.h

@@ -43,6 +43,7 @@ class RandomEffectsTracker {
   std::vector<data_size_t>::iterator UnsortedNodeBeginIterator(int category_id);
   std::vector<data_size_t>::iterator UnsortedNodeEndIterator(int category_id);
   std::map<int32_t, int32_t>& GetLabelMap() {return category_sample_tracker_->GetLabelMap();}
+  std::vector<int32_t>& GetUniqueGroupIds() {return category_sample_tracker_->GetUniqueGroupIds();}
   std::vector<data_size_t>& NodeIndices(int category_id) {return category_sample_tracker_->NodeIndices(category_id);}
   std::vector<data_size_t>& NodeIndicesInternalIndex(int internal_category_id) {return category_sample_tracker_->NodeIndicesInternalIndex(internal_category_id);}
   double GetPrediction(data_size_t observation_num) {return rfx_predictions_.at(observation_num);}
@@ -248,6 +249,9 @@ class RandomEffectsContainer {
   ~RandomEffectsContainer() {}
   void AddSample(MultivariateRegressionRandomEffectsModel& model);
   void Predict(RandomEffectsDataset& dataset, LabelMapper& label_mapper, std::vector<double>& output);
+  int NumSamples() {return num_samples_;}
+  int NumComponents() {return num_components_;}
+  int NumGroups() {return num_groups_;}
   std::vector<double> GetBeta() {return beta_;}
   std::vector<double> GetAlpha() {return alpha_;}
   std::vector<double> GetXi() {return xi_;}