ensemble.h

 
#ifndef STOCHTREE_ENSEMBLE_H_
#define STOCHTREE_ENSEMBLE_H_
 
#include <stochtree/data.h>
#include <stochtree/tree.h>
#include <nlohmann/json.hpp>
 
#include <algorithm>
#include <deque>
#include <optional>
#include <random>
#include <unordered_map>
 
using json = nlohmann::json;
 
namespace StochTree {
 
class TreeEnsemble {
 public:
  TreeEnsemble(int num_trees, int output_dimension = 1, bool is_leaf_constant = true, bool is_exponentiated = false) {
    // Initialize trees in the ensemble
    trees_ = std::vector<std::unique_ptr<Tree>>(num_trees);
    for (int i = 0; i < num_trees; i++) {
      trees_[i].reset(new Tree());
      trees_[i]->Init(output_dimension, is_exponentiated);
    }
    // Store ensemble configurations
    num_trees_ = num_trees;
    output_dimension_ = output_dimension;
    is_leaf_constant_ = is_leaf_constant;
    is_exponentiated_ = is_exponentiated;
  }
  
  TreeEnsemble(TreeEnsemble& ensemble) {
    // Unpack ensemble configurations
    num_trees_ = ensemble.num_trees_;
    output_dimension_ = ensemble.output_dimension_;
    is_leaf_constant_ = ensemble.is_leaf_constant_;
    is_exponentiated_ = ensemble.is_exponentiated_;
    // Initialize trees in the ensemble
    trees_ = std::vector<std::unique_ptr<Tree>>(num_trees_);
    for (int i = 0; i < num_trees_; i++) {
      trees_[i].reset(new Tree());
    }
    // Clone trees in the ensemble
    for (int j = 0; j < num_trees_; j++) {
      Tree* tree = ensemble.GetTree(j);
      this->CloneFromExistingTree(j, tree);
    }
  }
  
  ~TreeEnsemble() {}
 
  inline Tree* GetTree(int i) {
    return trees_[i].get();
  }
 
  inline void ResetRoot() {
    for (int i = 0; i < num_trees_; i++) {
      ResetInitTree(i);
    }
  }
 
  inline void ResetTree(int i) {
    trees_[i].reset(new Tree());
  }
 
  inline void ResetInitTree(int i) {
    trees_[i].reset(new Tree());
    trees_[i]->Init(output_dimension_, is_exponentiated_);
  }
 
  inline void CloneFromExistingTree(int i, Tree* tree) {
    return trees_[i]->CloneFromTree(tree);
  }
 
  inline void ReconstituteFromForest(TreeEnsemble& ensemble) {
    // Delete old tree pointers
    trees_.clear();
    // Unpack ensemble configurations
    num_trees_ = ensemble.num_trees_;
    output_dimension_ = ensemble.output_dimension_;
    is_leaf_constant_ = ensemble.is_leaf_constant_;
    is_exponentiated_ = ensemble.is_exponentiated_;
    // Initialize trees in the ensemble
    trees_ = std::vector<std::unique_ptr<Tree>>(num_trees_);
    for (int i = 0; i < num_trees_; i++) {
      trees_[i].reset(new Tree());
    }
    // Clone trees in the ensemble
    for (int j = 0; j < num_trees_; j++) {
      Tree* tree = ensemble.GetTree(j);
      this->CloneFromExistingTree(j, tree);
    }
  }
 
  std::vector<double> Predict(ForestDataset& dataset) {
    data_size_t n = dataset.NumObservations();
    std::vector<double> output(n);
    PredictInplace(dataset, output, 0);
    return output;
  }
 
  std::vector<double> PredictRaw(ForestDataset& dataset) {
    data_size_t n = dataset.NumObservations();
    data_size_t total_output_size = n * output_dimension_;
    std::vector<double> output(total_output_size);
    PredictRawInplace(dataset, output, 0);
    return output;
  }
  
  inline void PredictInplace(ForestDataset& dataset, std::vector<double> &output, data_size_t offset = 0) {
    PredictInplace(dataset, output, 0, trees_.size(), offset);
  }
 
  inline void PredictInplace(ForestDataset& dataset, std::vector<double> &output, 
                             int tree_begin, int tree_end, data_size_t offset = 0) {
    if (is_leaf_constant_) {
      PredictInplace(dataset.GetCovariates(), output, tree_begin, tree_end, offset);
    } else {
      CHECK(dataset.HasBasis());
      PredictInplace(dataset.GetCovariates(), dataset.GetBasis(), output, tree_begin, tree_end, offset);
    }
  }
 
  inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector<double> &output, data_size_t offset = 0) {
    PredictInplace(covariates, basis, output, 0, trees_.size(), offset);
  }
 
  inline void PredictInplace(Eigen::MatrixXd& covariates, Eigen::MatrixXd& basis, std::vector<double> &output, 
                             int tree_begin, int tree_end, data_size_t offset = 0) {
    double pred;
    CHECK_EQ(covariates.rows(), basis.rows());
    CHECK_EQ(output_dimension_, trees_[0]->OutputDimension());
    CHECK_EQ(output_dimension_, basis.cols());
    data_size_t n = covariates.rows();
    data_size_t total_output_size = n;
    if (output.size() < total_output_size + offset) {
      Log::Fatal("Mismatched size of prediction vector and training data");
    }
    for (data_size_t i = 0; i < n; i++) {
      pred = 0.0;
      for (size_t j = tree_begin; j < tree_end; j++) {
        auto &tree = *trees_[j];
        std::int32_t nidx = EvaluateTree(tree, covariates, i);
        for (int32_t k = 0; k < output_dimension_; k++) {
          pred += tree.LeafValue(nidx, k) * basis(i, k);
        }
      }
      if (is_exponentiated_) output[i + offset] = std::exp(pred);
      else output[i + offset] = pred;
    }
  }
 
  inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector<double> &output, data_size_t offset = 0) {
    PredictInplace(covariates, output, 0, trees_.size(), offset);
  }
 
  inline void PredictInplace(Eigen::MatrixXd& covariates, std::vector<double> &output, int tree_begin, int tree_end, data_size_t offset = 0) {
    double pred;
    data_size_t n = covariates.rows();
    data_size_t total_output_size = n;
    if (output.size() < total_output_size + offset) {
      Log::Fatal("Mismatched size of prediction vector and training data");
    }
    for (data_size_t i = 0; i < n; i++) {
      pred = 0.0;
      for (size_t j = tree_begin; j < tree_end; j++) {
        auto &tree = *trees_[j];
        std::int32_t nidx = EvaluateTree(tree, covariates, i);
        pred += tree.LeafValue(nidx, 0);
      }
      if (is_exponentiated_) output[i + offset] = std::exp(pred);
      else output[i + offset] = pred;
    }
  }
 
  inline void PredictRawInplace(ForestDataset& dataset, std::vector<double> &output, data_size_t offset = 0) {
    PredictRawInplace(dataset, output, 0, trees_.size(), offset);
  }
 
  inline void PredictRawInplace(ForestDataset& dataset, std::vector<double> &output, 
                             int tree_begin, int tree_end, data_size_t offset = 0) {
    double pred;
    Eigen::MatrixXd covariates = dataset.GetCovariates();
    CHECK_EQ(output_dimension_, trees_[0]->OutputDimension());
    data_size_t n = covariates.rows();
    data_size_t total_output_size = n * output_dimension_;
    if (output.size() < total_output_size + offset) {
      Log::Fatal("Mismatched size of raw prediction vector and training data");
    }
    for (data_size_t i = 0; i < n; i++) {
      for (int32_t k = 0; k < output_dimension_; k++) {
        pred = 0.0;
        for (size_t j = tree_begin; j < tree_end; j++) {
          auto &tree = *trees_[j];
          int32_t nidx = EvaluateTree(tree, covariates, i);
          pred += tree.LeafValue(nidx, k);
        }
        output[i*output_dimension_ + k + offset] = pred;
      }
    }
  }
 
  inline int32_t NumTrees() {
    return num_trees_;
  }
 
  inline int32_t NumLeaves() {
    int32_t result = 0;
    for (int i = 0; i < num_trees_; i++) {
      result += trees_[i]->NumLeaves();
    }
    return result;
  }
 
  inline double SumLeafSquared() {
    double result = 0.;
    for (int i = 0; i < num_trees_; i++) {
      result += trees_[i]->SumSquaredLeafValues();
    }
    return result;
  }
 
  inline int32_t OutputDimension() {
    return output_dimension_;
  }
 
  inline bool IsLeafConstant() {
    return is_leaf_constant_;
  }
 
  inline bool IsExponentiated() {
    return is_exponentiated_;
  }
 
  inline int32_t TreeMaxDepth(int tree_num) {
    return trees_[tree_num]->MaxLeafDepth();
  }
 
  inline double AverageMaxDepth() {
    double numerator = 0.;
    double denominator = 0.;
    for (int i = 0; i < num_trees_; i++) {
      numerator += static_cast<double>(TreeMaxDepth(i));
      denominator += 1.;
    }
    return numerator / denominator;
  }
 
  inline bool AllRoots() {
    for (int i = 0; i < num_trees_; i++) {
      if (!trees_[i]->IsRoot()) {
        return false;
      }
    }
    return true;
  }
 
  inline void SetLeafValue(double leaf_value) {
    CHECK_EQ(output_dimension_, 1);
    for (int i = 0; i < num_trees_; i++) {
      CHECK(trees_[i]->IsRoot());
      trees_[i]->SetLeaf(0, leaf_value);
    }
  }
 
  inline void SetLeafVector(std::vector<double>& leaf_vector) {
    CHECK_EQ(output_dimension_, leaf_vector.size());
    for (int i = 0; i < num_trees_; i++) {
      CHECK(trees_[i]->IsRoot());
      trees_[i]->SetLeafVector(0, leaf_vector);
    }
  }
 
  int GetMaxLeafIndex() {
    int max_leaf = 0;
    for (int j = 0; j < num_trees_; j++) {
      auto &tree = *trees_[j];
      max_leaf += tree.NumLeaves();
    }
    return max_leaf;
  }
 
  void PredictLeafIndicesInplace(ForestDataset* dataset, std::vector<int32_t>& output, int num_trees, data_size_t n) {
    PredictLeafIndicesInplace(dataset->GetCovariates(), output, num_trees, n);
  }
 
  void PredictLeafIndicesInplace(Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& covariates, std::vector<int32_t>& output, int num_trees, data_size_t n) {
    CHECK_GE(output.size(), num_trees*n);
    int offset = 0;
    int max_leaf = 0;
    for (int j = 0; j < num_trees; j++) {
      auto &tree = *trees_[j];
      int num_leaves = tree.NumLeaves();
      tree.PredictLeafIndexInplace(covariates, output, offset, max_leaf);
      offset += n;
      max_leaf += num_leaves;
    }
  }
 
  void PredictLeafIndicesInplace(Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& covariates, 
                                 Eigen::Map<Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& output, 
                                 int column_ind, int num_trees, data_size_t n) {
    CHECK_GE(output.size(), num_trees*n);
    int offset = 0;
    int max_leaf = 0;
    for (int j = 0; j < num_trees; j++) {
      auto &tree = *trees_[j];
      int num_leaves = tree.NumLeaves();
      tree.PredictLeafIndexInplace(covariates, output, column_ind, offset, max_leaf);
      offset += n;
      max_leaf += num_leaves;
    }
  }
 
  void PredictLeafIndicesInplace(Eigen::MatrixXd& covariates, std::vector<int32_t>& output, int num_trees, data_size_t n) {
    CHECK_GE(output.size(), num_trees*n);
    int offset = 0;
    int max_leaf = 0;
    for (int j = 0; j < num_trees; j++) {
      auto &tree = *trees_[j];
      int num_leaves = tree.NumLeaves();
      tree.PredictLeafIndexInplace(covariates, output, offset, max_leaf);
      offset += n;
      max_leaf += num_leaves;
    }
  }
 
  std::vector<int32_t> PredictLeafIndices(ForestDataset* dataset) {
    int num_trees = num_trees_;
    data_size_t n = dataset->NumObservations();
    std::vector<int32_t> output(n*num_trees);
    PredictLeafIndicesInplace(dataset, output, num_trees, n);
    return output;
  }
 
  json to_json() {
    json result_obj;
    result_obj.emplace("num_trees", this->num_trees_);
    result_obj.emplace("output_dimension", this->output_dimension_);
    result_obj.emplace("is_leaf_constant", this->is_leaf_constant_);
    result_obj.emplace("is_exponentiated", this->is_exponentiated_);
 
    std::string tree_label;
    for (int i = 0; i < trees_.size(); i++) {
      tree_label = "tree_" + std::to_string(i);
      result_obj.emplace(tree_label, trees_[i]->to_json());
    }
    
    return result_obj;
  }
  
  void from_json(const json& ensemble_json) {
    this->num_trees_ = ensemble_json.at("num_trees");
    this->output_dimension_ = ensemble_json.at("output_dimension");
    this->is_leaf_constant_ = ensemble_json.at("is_leaf_constant");
    this->is_exponentiated_ = ensemble_json.at("is_exponentiated");
 
    std::string tree_label;
    trees_.clear();
    trees_.resize(this->num_trees_);
    for (int i = 0; i < this->num_trees_; i++) {
      tree_label = "tree_" + std::to_string(i);
      trees_[i] = std::make_unique<Tree>();
      trees_[i]->from_json(ensemble_json.at(tree_label));
    }
  }
 
 private:
  std::vector<std::unique_ptr<Tree>> trees_;
  int num_trees_;
  int output_dimension_;
  bool is_leaf_constant_;
  bool is_exponentiated_;
};
 
 
} // namespace StochTree
 
#endif // STOCHTREE_ENSEMBLE_H_