tree.h

 
#ifndef STOCHTREE_TREE_H_
#define STOCHTREE_TREE_H_
 
#include <nlohmann/json.hpp>
#include <stochtree/data.h>
#include <stochtree/log.h>
#include <stochtree/meta.h>
#include <Eigen/Dense>
 
#include <cstdint>
#include <map>
#include <optional>
#include <set>
#include <stack>
#include <string>
 
using json = nlohmann::json;
 
namespace StochTree {
 
enum TreeNodeType {
  kLeafNode = 0,
  kNumericalSplitNode = 1,
  kCategoricalSplitNode = 2
};
 
// template<typename T>
// int enum_to_int(T& input_enum) {
//   return static_cast<int>(input_enum);
// }
 
// template<typename T>
// T json_to_enum(json& input_json) {
//   return static_cast<T>(input_json);
// }
 
std::string TreeNodeTypeToString(TreeNodeType type);
 
TreeNodeType TreeNodeTypeFromString(std::string const& name);
 
enum FeatureSplitType {
  kNumericSplit,
  kOrderedCategoricalSplit,
  kUnorderedCategoricalSplit
};
 
class TreeSplit;
 
class Tree {
 public:
  static constexpr std::int32_t kInvalidNodeId{-1};
  static constexpr std::int32_t kDeletedNodeMarker = std::numeric_limits<node_t>::max();
  static constexpr std::int32_t kRoot{0};
  
  Tree() = default;
  // ~Tree() = default;
  Tree(Tree const&) = delete;
  Tree& operator=(Tree const&) = delete;
  Tree(Tree&&) noexcept = default;
  Tree& operator=(Tree&&) noexcept = default;
  void CloneFromTree(Tree* tree);
 
  std::int32_t num_nodes{0};
  std::int32_t num_deleted_nodes{0};
 
  void Reset();
  void Init(int output_dimension = 1, bool is_log_scale = false);
  int AllocNode();
  void DeleteNode(std::int32_t nid);
  void ExpandNode(std::int32_t nid, int split_index, double split_value, double left_value, double right_value);
  void ExpandNode(std::int32_t nid, int split_index, std::vector<std::uint32_t> const& categorical_indices, double left_value, double right_value);
  void ExpandNode(std::int32_t nid, int split_index, double split_value, std::vector<double> left_value_vector, std::vector<double> right_value_vector);
  void ExpandNode(std::int32_t nid, int split_index, std::vector<std::uint32_t> const& categorical_indices, std::vector<double> left_value_vector, std::vector<double> right_value_vector);
  void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, double left_value, double right_value);
  void ExpandNode(std::int32_t nid, int split_index, TreeSplit& split, std::vector<double> left_value_vector, std::vector<double> right_value_vector);
 
  inline bool IsRoot() {return leaves_.size() == 1;}
  
  json to_json();
  void from_json(const json& tree_json);
 
  void ChangeToLeaf(std::int32_t nid, double value) {
    CHECK(this->IsLeaf(this->LeftChild(nid)));
    CHECK(this->IsLeaf(this->RightChild(nid)));
    this->DeleteNode(this->LeftChild(nid));
    this->DeleteNode(this->RightChild(nid));
    this->SetLeaf(nid, value);
 
    // Add nid to leaves and remove from internal nodes and leaf parents (if it was there)
    leaves_.push_back(nid);
    leaf_parents_.erase(std::remove(leaf_parents_.begin(), leaf_parents_.end(), nid), leaf_parents_.end());
    internal_nodes_.erase(std::remove(internal_nodes_.begin(), internal_nodes_.end(), nid), internal_nodes_.end());
 
    // Check if the other child of nid's parent node is also a leaf, if so, add parent back to leaf parents
    // TODO refactor and add this to the multivariate case as well
    if (!IsRoot(nid)) {
      int parent_id = Parent(nid);
      if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))){
        leaf_parents_.push_back(parent_id);
      }
    }
  }
 
  void CollapseToLeaf(std::int32_t nid, double value) {
    CHECK_EQ(output_dimension_, 1);
    if (this->IsLeaf(nid)) return;
    if (!this->IsLeaf(this->LeftChild(nid))) {
      CollapseToLeaf(this->LeftChild(nid), value);
    }
    if (!this->IsLeaf(this->RightChild(nid))) {
      CollapseToLeaf(this->RightChild(nid), value);
    }
    this->ChangeToLeaf(nid, value);
  }
 
  void ChangeToLeaf(std::int32_t nid, std::vector<double> value_vector) {
    CHECK(this->IsLeaf(this->LeftChild(nid)));
    CHECK(this->IsLeaf(this->RightChild(nid)));
    this->DeleteNode(this->LeftChild(nid));
    this->DeleteNode(this->RightChild(nid));
    this->SetLeafVector(nid, value_vector);
 
    // Add nid to leaves and remove from internal nodes and leaf parents (if it was there)
    leaves_.push_back(nid);
    leaf_parents_.erase(std::remove(leaf_parents_.begin(), leaf_parents_.end(), nid), leaf_parents_.end());
    internal_nodes_.erase(std::remove(internal_nodes_.begin(), internal_nodes_.end(), nid), internal_nodes_.end());
 
    // Check if the other child of nid's parent node is also a leaf, if so, add parent back to leaf parents
    // TODO refactor and add this to the multivariate case as well
    if (!IsRoot(nid)) {
      int parent_id = Parent(nid);
      if ((IsLeaf(LeftChild(parent_id))) && (IsLeaf(RightChild(parent_id)))){
        leaf_parents_.push_back(parent_id);
      }
    }
  }
  
  void CollapseToLeaf(std::int32_t nid, std::vector<double> value_vector) {
    CHECK_GT(output_dimension_, 1);
    CHECK_EQ(output_dimension_, value_vector.size());
    if (this->IsLeaf(nid)) return;
    if (!this->IsLeaf(this->LeftChild(nid))) {
      CollapseToLeaf(this->LeftChild(nid), value_vector);
    }
    if (!this->IsLeaf(this->RightChild(nid))) {
      CollapseToLeaf(this->RightChild(nid), value_vector);
    }
    this->ChangeToLeaf(nid, value_vector);
  }
 
  template <typename Func> void WalkTree(Func func) const {
    std::stack<std::int32_t> nodes;
    nodes.push(kRoot);
    auto &self = *this;
    while (!nodes.empty()) {
      auto nidx = nodes.top();
      nodes.pop();
      if (!func(nidx)) {
        return;
      }
      auto left = self.LeftChild(nidx);
      auto right = self.RightChild(nidx);
      if (left != Tree::kInvalidNodeId) {
        nodes.push(left);
      }
      if (right != Tree::kInvalidNodeId) {
        nodes.push(right);
      }
    }
  }
 
  std::vector<double> PredictFromNodes(std::vector<std::int32_t> node_indices);
  std::vector<double> PredictFromNodes(std::vector<std::int32_t> node_indices, Eigen::MatrixXd& basis);
  double PredictFromNode(std::int32_t node_id);
  double PredictFromNode(std::int32_t node_id, Eigen::MatrixXd& basis, int row_idx);
 
  bool HasVectorOutput() const {
    return output_dimension_ > 1;
  }
 
  std::int32_t OutputDimension() const {
    return output_dimension_;
  }
 
  bool IsLogScale() const {
    return is_log_scale_;
  }
  
  std::int32_t Parent(std::int32_t nid) const {
    return parent_[nid];
  }
  
  std::int32_t LeftChild(std::int32_t nid) const {
    return cleft_[nid];
  }
  
  std::int32_t RightChild(std::int32_t nid) const {
    return cright_[nid];
  }
  
  std::int32_t DefaultChild(std::int32_t nid) const {
    return cleft_[nid];
  }
  
  std::int32_t SplitIndex(std::int32_t nid) const {
    return split_index_[nid];
  }
  
  bool IsLeaf(std::int32_t nid) const {
    return cleft_[nid] == kInvalidNodeId;
  }
  
  bool IsRoot(std::int32_t nid) const {
    return parent_[nid] == kInvalidNodeId;
  }
 
  bool IsDeleted(std::int32_t nid) const {
    return node_deleted_[nid];
  }
 
  double LeafValue(std::int32_t nid) const {
    return leaf_value_[nid];
  }
  
  double LeafValue(std::int32_t nid, std::int32_t dim_id) const {
    CHECK_LT(dim_id, output_dimension_);
    if (output_dimension_ == 1 && dim_id == 0) {
      return leaf_value_[nid];
    } else {
      std::size_t const offset_begin = leaf_vector_begin_[nid];
      std::size_t const offset_end = leaf_vector_end_[nid];
      if (offset_begin >= leaf_vector_.size() || offset_end > leaf_vector_.size()) {
        Log::Fatal("No leaf vector set for node nid");
      }
      return leaf_vector_[offset_begin + dim_id];
    }
  }
 
  std::int32_t MaxLeafDepth() const {
    std::int32_t max_depth = 0;
    std::stack<std::int32_t> nodes;
    std::stack<std::int32_t> node_depths;
    nodes.push(kRoot);
    node_depths.push(0);
    auto &self = *this;
    while (!nodes.empty()) {
      auto nidx = nodes.top();
      nodes.pop();
      auto node_depth = node_depths.top();
      node_depths.pop();
      bool valid_node = !self.IsDeleted(nidx);
      if (valid_node) {
        if (node_depth > max_depth) max_depth = node_depth;
        auto left = self.LeftChild(nidx);
        auto right = self.RightChild(nidx);
        if (left != Tree::kInvalidNodeId) {
          nodes.push(left);
          node_depths.push(node_depth+1);
        }
        if (right != Tree::kInvalidNodeId) {
          nodes.push(right);
          node_depths.push(node_depth+1);
        }
      }
    }
    return max_depth;
  }
 
  std::vector<double> LeafVector(std::int32_t nid) const {
    std::size_t const offset_begin = leaf_vector_begin_[nid];
    std::size_t const offset_end = leaf_vector_end_[nid];
    if (offset_begin >= leaf_vector_.size() || offset_end > leaf_vector_.size()) {
      // Return empty vector, to indicate the lack of leaf vector
      return std::vector<double>();
    }
    return std::vector<double>(&leaf_vector_[offset_begin], &leaf_vector_[offset_end]);
    // Use unsafe access here, since we may need to take the address of one past the last
    // element, to follow with the range semantic of std::vector<>.
  }
 
  double SumSquaredNodeValues(std::int32_t nid) const {
    if (output_dimension_ == 1) {
      return std::pow(leaf_value_[nid], 2.0);
    } else {
      double result = 0.;
      std::size_t const offset_begin = leaf_vector_begin_[nid];
      std::size_t const offset_end = leaf_vector_end_[nid];
      if (offset_begin >= leaf_vector_.size() || offset_end > leaf_vector_.size()) {
        Log::Fatal("No leaf vector set for node nid");
      }
      for (std::size_t i = offset_begin; i < offset_end; i++) {
        result += std::pow(leaf_vector_[i], 2.0);
      }
      return result;
    }
  }
 
  double SumSquaredLeafValues() const {
    double result = 0.;
    for (auto& leaf : leaves_) {
      result += SumSquaredNodeValues(leaf);
    }
    return result;
  }
  
  bool HasLeafVector(std::int32_t nid) const {
    return leaf_vector_begin_[nid] != leaf_vector_end_[nid];
  }
 
  double Threshold(std::int32_t nid) const {
    return threshold_[nid];
  }
 
  std::vector<std::uint32_t> CategoryList(std::int32_t nid) const {
    std::size_t const offset_begin = category_list_begin_[nid];
    std::size_t const offset_end = category_list_end_[nid];
    if (offset_begin >= category_list_.size() || offset_end > category_list_.size()) {
      // Return empty vector, to indicate the lack of any category list
      // The node might be a numerical split
      return {};
    }
    return std::vector<std::uint32_t>(&category_list_[offset_begin], &category_list_[offset_end]);
    // Use unsafe access here, since we may need to take the address of one past the last
    // element, to follow with the range semantic of std::vector<>.
  }
 
  TreeNodeType NodeType(std::int32_t nid) const {
    return node_type_[nid];
  }
 
  bool IsNumericSplitNode(std::int32_t nid) const {
    return node_type_[nid] == TreeNodeType::kNumericalSplitNode;
  }
 
  bool IsCategoricalSplitNode(std::int32_t nid) const {
    return node_type_[nid] == TreeNodeType::kCategoricalSplitNode;
  }
 
  bool HasCategoricalSplit() const {
    return has_categorical_split_;
  }
 
  /* \brief Count number of leaves in tree. */
  [[nodiscard]] std::int32_t NumLeaves() const;
  [[nodiscard]] std::int32_t NumLeafParents() const;
  [[nodiscard]] std::int32_t NumSplitNodes() const;
 
  /* \brief Determine whether nid is leaf parent */
  [[nodiscard]] bool IsLeafParent(std::int32_t nid) const {
    // False until we deduce left and right node are
    // available and both are leaves
    bool is_left_leaf = false;
    bool is_right_leaf = false;
    // Check if node nidx is a leaf, if so, return false
    bool is_leaf = this->IsLeaf(nid);
    if (is_leaf){
      return false;
    } else {
      // If nidx is not a leaf, it must have left and right nodes
      // so we check if those are leaves
      std::int32_t left_node = LeftChild(nid);
      std::int32_t right_node = RightChild(nid);
      is_left_leaf = IsLeaf(left_node);
      is_right_leaf = IsLeaf(right_node);
    }
    return is_left_leaf && is_right_leaf;
  }
 
  [[nodiscard]] std::vector<std::int32_t> const& GetInternalNodes() const {
    return internal_nodes_;
  }
 
  [[nodiscard]] std::vector<std::int32_t> const& GetLeaves() const {
    return leaves_;
  }
 
  [[nodiscard]] std::vector<std::int32_t> const& GetLeafParents() const {
    return leaf_parents_;
  }
  
  [[nodiscard]] std::vector<std::int32_t> GetNodes() {
    std::vector<std::int32_t> output;
    auto const& self = *this;
    this->WalkTree([&output, &self](std::int32_t nidx) {
                    if (!self.IsDeleted(nidx)) {
                      output.push_back(nidx);
                    }
                    return true;
                  });
    return output;
  }
 
  [[nodiscard]] std::int32_t GetDepth(std::int32_t nid) const {
    int depth = 0;
    while (!IsRoot(nid)) {
      ++depth;
      nid = Parent(nid);
    }
    return depth;
  }
 
  [[nodiscard]] std::int32_t NumNodes() const noexcept { return num_nodes; }
  
  [[nodiscard]] std::int32_t NumDeletedNodes() const noexcept { return num_deleted_nodes; }
  
  [[nodiscard]] std::int32_t NumValidNodes() const noexcept {
    return num_nodes - num_deleted_nodes;
  }
 
  void SetLeftChild(std::int32_t nid, std::int32_t left_child) {
    cleft_[nid] = left_child;
  }
 
  void SetRightChild(std::int32_t nid, std::int32_t right_child) {
    cright_[nid] = right_child;
  }
 
  void SetChildren(std::int32_t nid, std::int32_t left_child, std::int32_t right_child) {
    SetLeftChild(nid, left_child);
    SetRightChild(nid, right_child);
  }
 
  void SetParent(std::int32_t child_node, std::int32_t parent_node) {
    parent_[child_node] = parent_node;
  }
 
  void SetParents(std::int32_t nid, std::int32_t left_child, std::int32_t right_child) {
    SetParent(left_child, nid);
    SetParent(right_child, nid);
  }
 
  void SetNumericSplit(
      std::int32_t nid, std::int32_t split_index, double threshold);
  
  void SetCategoricalSplit(std::int32_t nid, std::int32_t split_index,
      std::vector<std::uint32_t> const& category_list);
  
  void SetLeaf(std::int32_t nid, double value);
 
  void SetLeafVector(std::int32_t nid, std::vector<double> const& leaf_vector);
 
  void PredictLeafIndexInplace(ForestDataset* dataset, std::vector<int32_t>& output, int32_t offset, int32_t max_leaf);
 
  void PredictLeafIndexInplace(Eigen::MatrixXd& covariates, std::vector<int32_t>& output, int32_t offset, int32_t max_leaf);
 
  void PredictLeafIndexInplace(Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& covariates, std::vector<int32_t>& output, int32_t offset, int32_t max_leaf);
 
  void PredictLeafIndexInplace(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, int32_t offset, int32_t max_leaf);
 
  // Node info
  std::vector<TreeNodeType> node_type_;
  std::vector<std::int32_t> parent_;
  std::vector<std::int32_t> cleft_;
  std::vector<std::int32_t> cright_;
  std::vector<std::int32_t> split_index_;
  std::vector<double> leaf_value_;
  std::vector<double> threshold_;
  std::vector<bool> node_deleted_;
  std::vector<std::int32_t> internal_nodes_;
  std::vector<std::int32_t> leaves_;
  std::vector<std::int32_t> leaf_parents_;
  std::vector<std::int32_t> deleted_nodes_;
  
  // Leaf vector
  std::vector<double> leaf_vector_;
  std::vector<std::uint64_t> leaf_vector_begin_;
  std::vector<std::uint64_t> leaf_vector_end_;
 
  // Category list
  std::vector<std::uint32_t> category_list_;
  std::vector<std::uint64_t> category_list_begin_;
  std::vector<std::uint64_t> category_list_end_;
 
  bool has_categorical_split_{false};
  int output_dimension_{1};
  bool is_log_scale_{false};
};
 
inline bool operator==(const Tree& lhs, const Tree& rhs) {
  return (
    (lhs.has_categorical_split_ == rhs.has_categorical_split_) && 
    (lhs.output_dimension_ == rhs.output_dimension_) && 
    (lhs.is_log_scale_ == rhs.is_log_scale_) && 
    (lhs.node_type_ == rhs.node_type_) && 
    (lhs.parent_ == rhs.parent_) && 
    (lhs.cleft_ == rhs.cleft_) && 
    (lhs.cright_ == rhs.cright_) && 
    (lhs.split_index_ == rhs.split_index_) && 
    (lhs.leaf_value_ == rhs.leaf_value_) && 
    (lhs.threshold_ == rhs.threshold_) && 
    (lhs.internal_nodes_ == rhs.internal_nodes_) && 
    (lhs.leaves_ == rhs.leaves_) && 
    (lhs.leaf_parents_ == rhs.leaf_parents_) && 
    (lhs.deleted_nodes_ == rhs.deleted_nodes_) && 
    (lhs.leaf_vector_ == rhs.leaf_vector_) && 
    (lhs.leaf_vector_begin_ == rhs.leaf_vector_begin_) && 
    (lhs.leaf_vector_end_ == rhs.leaf_vector_end_) && 
    (lhs.category_list_ == rhs.category_list_) && 
    (lhs.category_list_begin_ == rhs.category_list_begin_) && 
    (lhs.category_list_end_ == rhs.category_list_end_)
  );
}
 
inline bool SplitTrueNumeric(double fvalue, double threshold) {
  return (fvalue <= threshold);
}
 
inline bool SplitTrueCategorical(double fvalue, std::vector<std::uint32_t> const& category_list) {
  bool category_matched;
  // A valid (integer) category must satisfy two criteria:
  // 1) it must be exactly representable as double
  // 2) it must fit into uint32_t
  auto max_representable_int
      = std::min(static_cast<double>(std::numeric_limits<std::uint32_t>::max()),
          static_cast<double>(std::uint64_t(1) << std::numeric_limits<double>::digits));
  if (fvalue < 0 || std::fabs(fvalue) > max_representable_int) {
    category_matched = false;
  } else {
    auto const category_value = static_cast<std::uint32_t>(fvalue);
    category_matched = (std::find(category_list.begin(), category_list.end(), category_value)
                        != category_list.end());
  }
  return category_matched;
}
 
inline int NextNodeNumeric(double fvalue, double threshold, int left_child, int right_child) {
  return (SplitTrueNumeric(fvalue, threshold) ? left_child : right_child);
}
 
inline int NextNodeCategorical(double fvalue, std::vector<std::uint32_t> const& category_list, int left_child, int right_child) {
  return SplitTrueCategorical(fvalue, category_list) ? left_child : right_child;
}
 
inline int EvaluateTree(Tree const& tree, Eigen::MatrixXd& data, int row) {
  int node_id = 0;
  while (!tree.IsLeaf(node_id)) {
    auto const split_index = tree.SplitIndex(node_id);
    double const fvalue = data(row, split_index);
    if (std::isnan(fvalue)) {
      node_id = tree.DefaultChild(node_id);
    } else {
      if (tree.NodeType(node_id) == StochTree::TreeNodeType::kCategoricalSplitNode) {
        node_id = NextNodeCategorical(fvalue, tree.CategoryList(node_id),
            tree.LeftChild(node_id), tree.RightChild(node_id));
      } else {
        node_id = NextNodeNumeric(fvalue, tree.Threshold(node_id), tree.LeftChild(node_id), tree.RightChild(node_id));
      }
    }
  }
  return node_id;
}
 
inline int EvaluateTree(Tree const& tree, Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>& data, int row) {
  int node_id = 0;
  while (!tree.IsLeaf(node_id)) {
    auto const split_index = tree.SplitIndex(node_id);
    double const fvalue = data(row, split_index);
    if (std::isnan(fvalue)) {
      node_id = tree.DefaultChild(node_id);
    } else {
      if (tree.NodeType(node_id) == StochTree::TreeNodeType::kCategoricalSplitNode) {
        node_id = NextNodeCategorical(fvalue, tree.CategoryList(node_id),
            tree.LeftChild(node_id), tree.RightChild(node_id));
      } else {
        node_id = NextNodeNumeric(fvalue, tree.Threshold(node_id), tree.LeftChild(node_id), tree.RightChild(node_id));
      }
    }
  }
  return node_id;
}
 
inline bool RowSplitLeft(Eigen::MatrixXd& covariates, int row, int split_index, double split_value) {
  double const fvalue = covariates(row, split_index);
  return SplitTrueNumeric(fvalue, split_value);
}
 
inline bool RowSplitLeft(Eigen::MatrixXd& covariates, int row, int split_index, std::vector<std::uint32_t> const& category_list) {
  double const fvalue = covariates(row, split_index);
  return SplitTrueCategorical(fvalue, category_list);
}
 
class TreeSplit {
 public:
  TreeSplit() {}
  TreeSplit(double split_value) {
    numeric_ = true;
    split_value_ = split_value;
    split_set_ = true;
  }
  TreeSplit(std::vector<std::uint32_t>& split_categories) {
    numeric_ = false;
    split_categories_ = split_categories;
    split_set_ = true;
  }
  ~TreeSplit() {}
  bool SplitSet() {return split_set_;}
  bool NumericSplit() {return numeric_;}
  bool SplitTrue(double fvalue) {
    if (numeric_) return SplitTrueNumeric(fvalue, split_value_);
    else return SplitTrueCategorical(fvalue, split_categories_);
  }
  double SplitValue() {return split_value_;}
  std::vector<std::uint32_t> SplitCategories() {return split_categories_;}
 private:
  bool split_set_{false};
  bool numeric_;
  double split_value_;
  std::vector<std::uint32_t> split_categories_;
};
 
 
} // namespace StochTree
 
#endif // STOCHTREE_TREE_H_