partition_tracker.h

 
#ifndef STOCHTREE_PARTITION_TRACKER_H_
#define STOCHTREE_PARTITION_TRACKER_H_
 
#include <stochtree/data.h>
#include <stochtree/ensemble.h>
#include <stochtree/log.h>
#include <stochtree/openmp_utils.h>
#include <stochtree/tree.h>
 
#include <numeric>
 
namespace StochTree {
 
class SampleNodeMapper;
class SamplePredMapper;
class UnsortedNodeSampleTracker;
class SortedNodeSampleTracker;
class FeaturePresortRootContainer;
 
class ForestTracker {
 public:
  ForestTracker(Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types, int num_trees, int num_observations);
  ~ForestTracker() {}
  void ReconstituteFromForest(TreeEnsemble& forest, ForestDataset& dataset, ColumnVector& residual, bool is_mean_model);
  void AssignAllSamplesToRoot();
  void AssignAllSamplesToRoot(int32_t tree_num);
  void AssignAllSamplesToConstantPrediction(double value);
  void AssignAllSamplesToConstantPrediction(int32_t tree_num, double value);
  void UpdatePredictions(TreeEnsemble* ensemble, ForestDataset& dataset);
  void UpdateSampleTrackers(TreeEnsemble& forest, ForestDataset& dataset);
  void UpdateSampleTrackersResidual(TreeEnsemble& forest, ForestDataset& dataset, ColumnVector& residual, bool is_mean_model);
  void ResetRoot(Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types, int32_t tree_num);
  void AddSplit(Eigen::MatrixXd& covariates, TreeSplit& split, int32_t split_feature, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id, bool keep_sorted = false, int num_threads = -1);
  void RemoveSplit(Eigen::MatrixXd& covariates, Tree* tree, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id, bool keep_sorted = false);
  double GetSamplePrediction(data_size_t sample_id);
  double GetTreeSamplePrediction(data_size_t sample_id, int tree_id);
  void UpdateVarWeightsFromInternalPredictions(ForestDataset& dataset);
  void SetSamplePrediction(data_size_t sample_id, double value);
  void SetTreeSamplePrediction(data_size_t sample_id, int tree_id, double value);
  void SyncPredictions();
  data_size_t GetNodeId(int observation_num, int tree_num);
  data_size_t UnsortedNodeBegin(int tree_id, int node_id);
  data_size_t UnsortedNodeEnd(int tree_id, int node_id);
  data_size_t UnsortedNodeSize(int tree_id, int node_id);
  data_size_t SortedNodeBegin(int node_id, int feature_id);
  data_size_t SortedNodeEnd(int node_id, int feature_id);
  data_size_t SortedNodeSize(int node_id, int feature_id);
  std::vector<data_size_t>::iterator UnsortedNodeBeginIterator(int tree_id, int node_id);
  std::vector<data_size_t>::iterator UnsortedNodeEndIterator(int tree_id, int node_id);
  std::vector<data_size_t>::iterator SortedNodeBeginIterator(int node_id, int feature_id);
  std::vector<data_size_t>::iterator SortedNodeEndIterator(int node_id, int feature_id);
  SamplePredMapper* GetSamplePredMapper() {return sample_pred_mapper_.get();}
  SampleNodeMapper* GetSampleNodeMapper() {return sample_node_mapper_.get();}
  UnsortedNodeSampleTracker* GetUnsortedNodeSampleTracker() {return unsorted_node_sample_tracker_.get();}
  SortedNodeSampleTracker* GetSortedNodeSampleTracker() {return sorted_node_sample_tracker_.get();}
  int GetNumObservations() {return num_observations_;}
  int GetNumTrees() {return num_trees_;}
  int GetNumFeatures() {return num_features_;}
  bool Initialized() {return initialized_;}
 private:
  std::vector<double> sum_predictions_;
  std::unique_ptr<SamplePredMapper> sample_pred_mapper_;
  std::unique_ptr<SampleNodeMapper> sample_node_mapper_;
  std::unique_ptr<UnsortedNodeSampleTracker> unsorted_node_sample_tracker_;
  std::unique_ptr<FeaturePresortRootContainer> presort_container_;
  std::unique_ptr<SortedNodeSampleTracker> sorted_node_sample_tracker_;
  std::vector<FeatureType> feature_types_;
  int num_trees_;
  int num_observations_;
  int num_features_;
  bool initialized_{false};
 
  void UpdatePredictionsInternal(TreeEnsemble* ensemble, Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis);
  void UpdatePredictionsInternal(TreeEnsemble* ensemble, Eigen::MatrixXd& covariates);
  void UpdateSampleTrackersInternal(TreeEnsemble& forest, Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis);
  void UpdateSampleTrackersInternal(TreeEnsemble& forest, Eigen::MatrixXd& covariates);
  void UpdateSampleTrackersResidualInternalBasis(TreeEnsemble& forest, ForestDataset& dataset, ColumnVector& residual, bool is_mean_model);
  void UpdateSampleTrackersResidualInternalNoBasis(TreeEnsemble& forest, ForestDataset& dataset, ColumnVector& residual, bool is_mean_model);
};
 
class SamplePredMapper {
 public:
  SamplePredMapper(int num_trees, data_size_t num_observations) {
    num_trees_ = num_trees;
    num_observations_ = num_observations;
    // Initialize the vector of vectors of leaf indices for each tree
    tree_preds_.resize(num_trees_);
    for (int j = 0; j < num_trees_; j++) {
      tree_preds_[j].resize(num_observations_);
    }
  }
 
  inline double GetPred(data_size_t sample_id, int tree_id) {
    CHECK_LT(sample_id, num_observations_);
    CHECK_LT(tree_id, num_trees_);
    return tree_preds_[tree_id][sample_id];
  }
 
  inline void SetPred(data_size_t sample_id, int tree_id, double value) {
    CHECK_LT(sample_id, num_observations_);
    CHECK_LT(tree_id, num_trees_);
    tree_preds_[tree_id][sample_id] = value;
  }
  
  inline int NumTrees() {return num_trees_;}
  
  inline int NumObservations() {return num_observations_;}
 
  inline void AssignAllSamplesToConstantPrediction(int tree_id, double value) {
    for (data_size_t i = 0; i < num_observations_; i++) {
      tree_preds_[tree_id][i] = value;
    }
  }
 
 private:
  std::vector<std::vector<double>> tree_preds_;
  int num_trees_;
  data_size_t num_observations_;
};
 
class SampleNodeMapper {
 public:
  SampleNodeMapper(int num_trees, data_size_t num_observations) {
    num_trees_ = num_trees;
    num_observations_ = num_observations;
    // Initialize the vector of vectors of leaf indices for each tree
    tree_observation_indices_.resize(num_trees_);
    for (int j = 0; j < num_trees_; j++) {
      tree_observation_indices_[j].resize(num_observations_);
    }
  }
  
  SampleNodeMapper(SampleNodeMapper& other){
    num_trees_ = other.NumTrees();
    num_observations_ = other.NumObservations();
    // Initialize the vector of vectors of leaf indices for each tree
    tree_observation_indices_.resize(num_trees_);
    for (int j = 0; j < num_trees_; j++) {
      tree_observation_indices_[j].resize(num_observations_);
      for (int i = 0; i < num_observations_; i++) {
        tree_observation_indices_[j][i] = other.GetNodeId(i, j);
      }
    }
  }
 
  void AddSplit(Eigen::MatrixXd& covariates, TreeSplit& split, int32_t split_feature, int32_t tree_id, int32_t split_node_id, int32_t left_node_id, int32_t right_node_id) {
    CHECK_EQ(num_observations_, covariates.rows());
    // Eigen::MatrixXd X = covariates.GetData();
    for (int i = 0; i < num_observations_; i++) {
      if (tree_observation_indices_[tree_id][i] == split_node_id) {
        auto fvalue = covariates(i, split_feature);
        if (split.SplitTrue(fvalue)) {
          tree_observation_indices_[tree_id][i] = left_node_id;
        } else {
          tree_observation_indices_[tree_id][i] = right_node_id;
        }
      }
    }
  }
 
  inline data_size_t GetNodeId(data_size_t sample_id, int tree_id) {
    CHECK_LT(sample_id, num_observations_);
    CHECK_LT(tree_id, num_trees_);
    return tree_observation_indices_[tree_id][sample_id];
  }
 
  inline void SetNodeId(data_size_t sample_id, int tree_id, int node_id) {
    CHECK_LT(sample_id, num_observations_);
    CHECK_LT(tree_id, num_trees_);
    tree_observation_indices_[tree_id][sample_id] = node_id;
  }
  
  inline int NumTrees() {return num_trees_;}
  
  inline int NumObservations() {return num_observations_;}
 
  inline void AssignAllSamplesToRoot(int tree_id) {
    for (data_size_t i = 0; i < num_observations_; i++) {
      tree_observation_indices_[tree_id][i] = 0;
    }
  }
 
 private:
  std::vector<std::vector<int>> tree_observation_indices_;
  int num_trees_;
  data_size_t num_observations_;
};
 
class FeatureUnsortedPartition {
 public:
  FeatureUnsortedPartition(data_size_t n);
 
  void ReconstituteFromTree(Tree& tree, ForestDataset& dataset);
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int left_node_id, int right_node_id, int feature_split, TreeSplit& split);
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int left_node_id, int right_node_id, int feature_split, double split_value);
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int left_node_id, int right_node_id, int feature_split, std::vector<std::uint32_t> const& category_list);
 
  void PruneNodeToLeaf(int node_id);
 
  bool IsLeaf(int node_id);
 
  bool IsValidNode(int node_id);
 
  bool LeftNodeIsLeaf(int node_id);
 
  bool RightNodeIsLeaf(int node_id);
 
  data_size_t NodeBegin(int node_id);
 
  data_size_t NodeEnd(int node_id);
 
  data_size_t NodeSize(int node_id);
 
  int Parent(int node_id);
 
  int LeftNode(int node_id);
 
  int RightNode(int node_id);
 
  std::vector<data_size_t> indices_;
 
  std::vector<data_size_t> NodeIndices(int node_id);
 
  void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper);
 
 private:
  // Vectors tracking indices in each node
  std::vector<data_size_t> node_begin_;
  std::vector<data_size_t> node_length_;
  std::vector<int32_t> parent_nodes_;
  std::vector<int32_t> left_nodes_;
  std::vector<int32_t> right_nodes_;
  int num_nodes_, num_deleted_nodes_;
  std::vector<int> deleted_nodes_;
 
  // Private helper functions
  void ExpandNodeTrackingVectors(int node_id, int left_node_id, int right_node_id, data_size_t node_start_idx, data_size_t num_left, data_size_t num_right);
  void ConvertLeafParentToLeaf(int node_id);
};
 
class UnsortedNodeSampleTracker {
 public:
  UnsortedNodeSampleTracker(data_size_t n, int num_trees) {
    feature_partitions_.resize(num_trees);
    num_trees_ = num_trees;
    for (int i = 0; i < num_trees; i++) {
      feature_partitions_[i].reset(new FeatureUnsortedPartition(n));
    }
  }
 
  void ReconstituteFromForest(TreeEnsemble& forest, ForestDataset& dataset);
 
  void PartitionTreeNode(Eigen::MatrixXd& covariates, int tree_id, int node_id, int left_node_id, int right_node_id, int feature_split, TreeSplit& split) {
    return feature_partitions_[tree_id]->PartitionNode(covariates, node_id, left_node_id, right_node_id, feature_split, split);
  }
 
  void PartitionTreeNode(Eigen::MatrixXd& covariates, int tree_id, int node_id, int left_node_id, int right_node_id, int feature_split, double split_value) {
    return feature_partitions_[tree_id]->PartitionNode(covariates, node_id, left_node_id, right_node_id, feature_split, split_value);
  }
 
  void PartitionTreeNode(Eigen::MatrixXd& covariates, int tree_id, int node_id, int left_node_id, int right_node_id, int feature_split, std::vector<std::uint32_t> const& category_list) {
    return feature_partitions_[tree_id]->PartitionNode(covariates, node_id, left_node_id, right_node_id, feature_split, category_list);
  }
  
  void ResetTreeToRoot(int tree_id, data_size_t n) {
    feature_partitions_[tree_id].reset(new FeatureUnsortedPartition(n));;
  }
 
  void PruneTreeNodeToLeaf(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->PruneNodeToLeaf(node_id);
  }
 
  bool IsLeaf(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->IsLeaf(node_id);
  }
 
  bool IsValidNode(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->IsValidNode(node_id);
  }
 
  bool LeftNodeIsLeaf(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->LeftNodeIsLeaf(node_id);
  }
 
  bool RightNodeIsLeaf(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->RightNodeIsLeaf(node_id);
  }
 
  data_size_t NodeBegin(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->NodeBegin(node_id);
  }
 
  data_size_t NodeEnd(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->NodeEnd(node_id);
  }
 
  std::vector<data_size_t>::iterator NodeBeginIterator(int tree_id, int node_id) {
    data_size_t node_begin = feature_partitions_[tree_id]->NodeBegin(node_id);
    auto begin_iter = feature_partitions_[tree_id]->indices_.begin();
    return begin_iter + node_begin;
  }
 
  std::vector<data_size_t>::iterator NodeEndIterator(int tree_id, int node_id) {
    int node_end = feature_partitions_[tree_id]->NodeEnd(node_id);
    auto begin_iter = feature_partitions_[tree_id]->indices_.begin();
    return begin_iter + node_end;
  }
 
  data_size_t NodeSize(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->NodeSize(node_id);
  }
 
  int Parent(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->Parent(node_id);
  }
 
  int LeftNode(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->LeftNode(node_id);
  }
 
  int RightNode(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->RightNode(node_id);
  }
 
  std::vector<data_size_t> TreeNodeIndices(int tree_id, int node_id) {
    return feature_partitions_[tree_id]->NodeIndices(node_id);
  }
 
  void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper) {
    feature_partitions_[tree_id]->UpdateObservationMapping(node_id, tree_id, sample_node_mapper);
  }
 
  void UpdateObservationMapping(Tree* tree, int tree_id, SampleNodeMapper* sample_node_mapper) {
    std::vector<int> leaves = tree->GetLeaves();
    int leaf;
    for (int i = 0; i < leaves.size(); i++) {
      leaf = leaves[i];
      UpdateObservationMapping(leaf, tree_id, sample_node_mapper);
    }
  }
 
  int NumTrees() { return num_trees_; }
 
  FeatureUnsortedPartition* GetFeaturePartition(int i) { return feature_partitions_[i].get(); }
 
 private:
  // Vectors of feature partitions
  std::vector<std::unique_ptr<FeatureUnsortedPartition>> feature_partitions_;
  int num_trees_;
};
 
class NodeOffsetSize {
 public:
  NodeOffsetSize(data_size_t node_offset, data_size_t node_size) : node_begin_{node_offset}, node_size_{node_size}, presorted_{false} {
    node_end_ = node_begin_ + node_size_;
  }
 
  ~NodeOffsetSize() {}
 
  void SetSorted() {presorted_ = true;}
 
  bool IsSorted() {return presorted_;}
 
  data_size_t Begin() {return node_begin_;}
 
  data_size_t End() {return node_end_;}
 
  data_size_t Size() {return node_size_;}
 
 private:
  data_size_t node_begin_;
  data_size_t node_size_;
  data_size_t node_end_;
  bool presorted_;
};
 
class FeaturePresortPartition;
 
class FeaturePresortRoot {
 friend FeaturePresortPartition; 
 public:
  FeaturePresortRoot(Eigen::MatrixXd& covariates, int32_t feature_index, FeatureType feature_type) {
    feature_index_ = feature_index;
    ArgsortRoot(covariates);
  }
 
  ~FeaturePresortRoot() {}
 
  void ArgsortRoot(Eigen::MatrixXd& covariates) {
    data_size_t num_obs = covariates.rows();
    
    // Make a vector of indices from 0 to num_obs - 1
    if (feature_sort_indices_.size() != num_obs){
      feature_sort_indices_.resize(num_obs, 0);
    }
    std::iota(feature_sort_indices_.begin(), feature_sort_indices_.end(), 0);
 
    // Define a custom comparator to be used with stable_sort:
    // For every two indices l and r store as elements of `data_sort_indices_`, 
    // compare them for sorting purposes by indexing the covariate's raw data with both l and r
    auto comp_op = [&](size_t const &l, size_t const &r) { return std::less<double>{}(covariates(l, feature_index_), covariates(r, feature_index_)); };
    std::stable_sort(feature_sort_indices_.begin(), feature_sort_indices_.end(), comp_op);
  }
 
 private:
  std::vector<data_size_t> feature_sort_indices_;
  int32_t feature_index_;
};
 
class FeaturePresortRootContainer {
 public:
  FeaturePresortRootContainer(Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types) {
    num_features_ = covariates.cols();
    feature_presort_.resize(num_features_);
    for (int i = 0; i < num_features_; i++) {
      feature_presort_[i].reset(new FeaturePresortRoot(covariates, i, feature_types[i]));
    }
  }
 
  ~FeaturePresortRootContainer() {}
 
  FeaturePresortRoot* GetFeaturePresort(int feature_num) {return feature_presort_[feature_num].get(); }
 
 private:
  std::vector<std::unique_ptr<FeaturePresortRoot>> feature_presort_;
  int num_features_;
};
 
class FeaturePresortPartition {
 public:
  FeaturePresortPartition(FeaturePresortRoot* feature_presort_root, Eigen::MatrixXd& covariates, int32_t feature_index, FeatureType feature_type) {
    // Unpack all feature details
    feature_index_ = feature_index;
    feature_type_ = feature_type;
    num_obs_ = covariates.rows();
    feature_sort_indices_ = feature_presort_root->feature_sort_indices_;
 
    // Initialize new tree to root
    data_size_t node_offset = 0;
    node_offset_sizes_.emplace_back(node_offset, num_obs_);
  }
 
  ~FeaturePresortPartition() {}
 
  void SplitFeature(Eigen::MatrixXd& covariates, int32_t node_id, int32_t feature_index, TreeSplit& split);
 
  void SplitFeatureNumeric(Eigen::MatrixXd& covariates, int32_t node_id, int32_t feature_index, double split_value);
 
  void SplitFeatureCategorical(Eigen::MatrixXd& covariates, int32_t node_id, int32_t feature_index, std::vector<std::uint32_t> const& category_list);
 
  data_size_t NodeBegin(int32_t node_id) {return node_offset_sizes_[node_id].Begin();}
 
  data_size_t NodeEnd(int32_t node_id) {return node_offset_sizes_[node_id].End();}
 
  data_size_t NodeSize(int32_t node_id) {return node_offset_sizes_[node_id].Size();}
 
  std::vector<data_size_t> NodeIndices(int node_id);
 
  data_size_t SortIndex(data_size_t j) {return feature_sort_indices_[j];}
 
  FeatureType GetFeatureType() {return feature_type_;}
 
  void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper);
 
  std::vector<data_size_t> feature_sort_indices_;
 private:
  void AddLeftRightNodes(data_size_t left_node_begin, data_size_t left_node_size, data_size_t right_node_begin, data_size_t right_node_size);
 
  std::vector<NodeOffsetSize> node_offset_sizes_;
  int32_t feature_index_;
  FeatureType feature_type_;
  data_size_t num_obs_;
};
 
class SortedNodeSampleTracker {
 public:
  SortedNodeSampleTracker(FeaturePresortRootContainer* feature_presort_root_container, Eigen::MatrixXd& covariates, std::vector<FeatureType>& feature_types) {
    num_features_ = covariates.cols();
    feature_partitions_.resize(num_features_);
    FeaturePresortRoot* feature_presort_root;
    for (int i = 0; i < num_features_; i++) {
      feature_presort_root = feature_presort_root_container->GetFeaturePresort(i);
      feature_partitions_[i].reset(new FeaturePresortPartition(feature_presort_root, covariates, i, feature_types[i]));
    }
  }
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int feature_split, TreeSplit& split, int num_threads = -1) {
    StochTree::ParallelFor(0, num_features_, num_threads, [&](int i) {
      feature_partitions_[i]->SplitFeature(covariates, node_id, feature_split, split);
    });
  }
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int feature_split, double split_value, int num_threads = -1) {
    StochTree::ParallelFor(0, num_features_, num_threads, [&](int i) {
      feature_partitions_[i]->SplitFeatureNumeric(covariates, node_id, feature_split, split_value);
    });
  }
 
  void PartitionNode(Eigen::MatrixXd& covariates, int node_id, int feature_split, std::vector<std::uint32_t> const& category_list, int num_threads = -1) {
    StochTree::ParallelFor(0, num_features_, num_threads, [&](int i) {
      feature_partitions_[i]->SplitFeatureCategorical(covariates, node_id, feature_split, category_list);
    });
  }
 
  data_size_t NodeBegin(int node_id, int feature_index) {
    return feature_partitions_[feature_index]->NodeBegin(node_id);
  }
 
  data_size_t NodeEnd(int node_id, int feature_index) {
    return feature_partitions_[feature_index]->NodeEnd(node_id);
  }
 
  data_size_t NodeSize(int node_id, int feature_index) {
    return feature_partitions_[feature_index]->NodeSize(node_id);
  }
 
  std::vector<data_size_t>::iterator NodeBeginIterator(int node_id, int feature_index) {
    data_size_t node_begin = NodeBegin(node_id, feature_index);
    auto begin_iter = feature_partitions_[feature_index]->feature_sort_indices_.begin();
    return begin_iter + node_begin;
  }
 
  std::vector<data_size_t>::iterator NodeEndIterator(int node_id, int feature_index) {
    data_size_t node_end = NodeEnd(node_id, feature_index);
    auto begin_iter = feature_partitions_[feature_index]->feature_sort_indices_.begin();
    return begin_iter + node_end;
  }
 
  std::vector<data_size_t> NodeIndices(int node_id, int feature_index) {
    return feature_partitions_[feature_index]->NodeIndices(node_id);
  }
 
  data_size_t SortIndex(data_size_t j, int feature_index) {return feature_partitions_[feature_index]->SortIndex(j); }
 
  void UpdateObservationMapping(int node_id, int tree_id, SampleNodeMapper* sample_node_mapper, int feature_index = 0) {
    feature_partitions_[feature_index]->UpdateObservationMapping(node_id, tree_id, sample_node_mapper);
  }
 
 private:
  std::vector<std::unique_ptr<FeaturePresortPartition>> feature_partitions_;
  int num_features_;
};
 
} // namespace StochTree
 
#endif // STOCHTREE_PARTITION_TRACKER_H_