tree.h

Go to the documentation of this file.

 
#ifndef TREELITE_TREE_H_
#define TREELITE_TREE_H_
 
#include <treelite/contiguous_array.h>
#include <treelite/enum/operator.h>
#include <treelite/enum/task_type.h>
#include <treelite/enum/tree_node_type.h>
#include <treelite/enum/typeinfo.h>
#include <treelite/logging.h>
#include <treelite/pybuffer_frame.h>
#include <treelite/version.h>
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <limits>
#include <map>
#include <memory>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
 
/* Indicator that certain functions should be visible from a library (Windows only) */
#if defined(_MSC_VER) || defined(_WIN32)
#define TREELITE_DLL_EXPORT __declspec(dllexport)
#else
#define TREELITE_DLL_EXPORT
#endif
 
namespace treelite {
 
template <typename ThresholdType, typename LeafOutputType>
class ModelPreset;
 
}
 
namespace treelite::detail::serializer {
 
template <typename MixIn>
class Serializer;
template <typename MixIn>
class Deserializer;
 
}  // namespace treelite::detail::serializer
 
namespace treelite::detail::field_accessor {
 
template <typename ThresholdType, typename LeafOutputType>
PyBufferFrame GetTreeFieldImpl(
    ModelPreset<ThresholdType, LeafOutputType>&, std::uint64_t, std::string const&);
template <typename ThresholdType, typename LeafOutputType>
void SetTreeFieldImpl(
    ModelPreset<ThresholdType, LeafOutputType>&, std::uint64_t, std::string const&, PyBufferFrame);
 
}  // namespace treelite::detail::field_accessor
 
namespace treelite {
 
// Used for returning version triple from a Model object
struct Version {
  std::int32_t major_ver;
  std::int32_t minor_ver;
  std::int32_t patch_ver;
};
 
template <typename ThresholdType, typename LeafOutputType>
class Tree {
 public:
  static_assert(std::is_same_v<ThresholdType, float> || std::is_same_v<ThresholdType, double>,
      "ThresholdType must be either float32 or float64");
  static_assert(std::is_same_v<LeafOutputType, float> || std::is_same_v<LeafOutputType, double>,
      "LeafOutputType must be one of uint32_t, float32 or float64");
  static_assert(std::is_same_v<ThresholdType, LeafOutputType>,
      "Unsupported combination of ThresholdType and LeafOutputType");
 
  Tree() = default;
  ~Tree() = default;
  Tree(Tree const&) = delete;
  Tree& operator=(Tree const&) = delete;
  Tree(Tree&&) noexcept = default;
  Tree& operator=(Tree&&) noexcept = default;
 
  inline Tree<ThresholdType, LeafOutputType> Clone() const;
 
 private:
  ContiguousArray<TreeNodeType> node_type_;
  ContiguousArray<std::int32_t> cleft_;
  ContiguousArray<std::int32_t> cright_;
  ContiguousArray<std::int32_t> split_index_;
  ContiguousArray<bool> default_left_;
  ContiguousArray<LeafOutputType> leaf_value_;
  ContiguousArray<ThresholdType> threshold_;
  ContiguousArray<Operator> cmp_;
  ContiguousArray<bool> category_list_right_child_;
 
  // Leaf vector
  ContiguousArray<LeafOutputType> leaf_vector_;
  ContiguousArray<std::uint64_t> leaf_vector_begin_;
  ContiguousArray<std::uint64_t> leaf_vector_end_;
 
  // Category list
  ContiguousArray<std::uint32_t> category_list_;
  ContiguousArray<std::uint64_t> category_list_begin_;
  ContiguousArray<std::uint64_t> category_list_end_;
 
  // Node statistics
  ContiguousArray<std::uint64_t> data_count_;
  ContiguousArray<double> sum_hess_;
  ContiguousArray<double> gain_;
  ContiguousArray<bool> data_count_present_;
  ContiguousArray<bool> sum_hess_present_;
  ContiguousArray<bool> gain_present_;
 
  bool has_categorical_split_{false};
 
  /* Note: the following member fields shall be re-computed at serialization time */
 
  // Number of optional fields in the extension slots
  std::int32_t num_opt_field_per_tree_{0};
  std::int32_t num_opt_field_per_node_{0};
 
  template <typename WriterType, typename X, typename Y>
  friend void DumpTreeAsJSON(WriterType& writer, Tree<X, Y> const& tree);
 
  template <typename MixIn>
  friend class detail::serializer::Serializer;
  template <typename MixIn>
  friend class detail::serializer::Deserializer;
  template <typename X, typename Y>
  friend PyBufferFrame detail::field_accessor::GetTreeFieldImpl(
      ModelPreset<X, Y>&, std::uint64_t, std::string const&);
  template <typename X, typename Y>
  friend void detail::field_accessor::SetTreeFieldImpl(
      ModelPreset<X, Y>&, std::uint64_t, std::string const&, PyBufferFrame);
 
 public:
  std::int32_t num_nodes{0};
  inline void Init();
  inline int AllocNode();
 
  inline int LeftChild(int nid) const {
    return cleft_[nid];
  }
  inline int RightChild(int nid) const {
    return cright_[nid];
  }
  inline int DefaultChild(int nid) const {
    return default_left_[nid] ? cleft_[nid] : cright_[nid];
  }
  inline std::int32_t SplitIndex(int nid) const {
    return split_index_[nid];
  }
  inline bool DefaultLeft(int nid) const {
    return default_left_[nid];
  }
  inline bool IsLeaf(int nid) const {
    return cleft_[nid] == -1;
  }
  inline LeafOutputType LeafValue(int nid) const {
    return leaf_value_[nid];
  }
  inline std::vector<LeafOutputType> LeafVector(int 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<LeafOutputType>();
    }
    return std::vector<LeafOutputType>(&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<>.
  }
  inline bool HasLeafVector(int nid) const {
    return leaf_vector_begin_[nid] != leaf_vector_end_[nid];
  }
  inline ThresholdType Threshold(int nid) const {
    return threshold_[nid];
  }
  inline Operator ComparisonOp(int nid) const {
    return cmp_[nid];
  }
  inline std::vector<std::uint32_t> CategoryList(int 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<>.
  }
  inline TreeNodeType NodeType(int nid) const {
    return node_type_[nid];
  }
  inline bool HasDataCount(int nid) const {
    return !data_count_present_.Empty() && data_count_present_[nid];
  }
  inline std::uint64_t DataCount(int nid) const {
    return data_count_[nid];
  }
 
  inline bool HasSumHess(int nid) const {
    return !sum_hess_present_.Empty() && sum_hess_present_[nid];
  }
  inline double SumHess(int nid) const {
    return sum_hess_[nid];
  }
  inline bool HasGain(int nid) const {
    return !gain_present_.Empty() && gain_present_[nid];
  }
  inline double Gain(int nid) const {
    return gain_[nid];
  }
  inline bool CategoryListRightChild(int nid) const {
    return category_list_right_child_[nid];
  }
 
  inline bool HasCategoricalSplit() const {
    return has_categorical_split_;
  }
 
  inline void SetChildren(int nid, int left_child, int right_child) {
    cleft_[nid] = left_child;
    cright_[nid] = right_child;
  }
  inline void SetNumericalTest(
      int nid, std::int32_t split_index, ThresholdType threshold, bool default_left, Operator cmp);
  inline void SetCategoricalTest(int nid, std::int32_t split_index, bool default_left,
      std::vector<std::uint32_t> const& category_list, bool category_list_right_child);
  inline void SetLeaf(int nid, LeafOutputType value);
  inline void SetLeafVector(int nid, std::vector<LeafOutputType> const& leaf_vector);
  inline void SetSumHess(int nid, double sum_hess);
  inline void SetDataCount(int nid, std::uint64_t data_count);
  inline void SetGain(int nid, double gain);
};
 
template <typename ThresholdT, typename LeafOutputT>
class ModelPreset {
 public:
  std::vector<Tree<ThresholdT, LeafOutputT>> trees;
 
  using threshold_type = ThresholdT;
  using leaf_output_type = LeafOutputT;
 
  ModelPreset() = default;
  ~ModelPreset() = default;
  ModelPreset(ModelPreset const&) = delete;
  ModelPreset& operator=(ModelPreset const&) = delete;
  ModelPreset(ModelPreset&&) noexcept = default;
  ModelPreset& operator=(ModelPreset&&) noexcept = default;
 
  inline TypeInfo GetThresholdType() const {
    return TypeInfoFromType<ThresholdT>();
  }
  inline TypeInfo GetLeafOutputType() const {
    return TypeInfoFromType<LeafOutputT>();
  }
  inline std::size_t GetNumTree() const {
    return trees.size();
  }
  void SetTreeLimit(std::size_t limit) {
    return trees.resize(limit);
  }
};
 
using ModelPresetVariant = std::variant<ModelPreset<float, float>, ModelPreset<double, double>>;
 
template <int variant_index>
ModelPresetVariant SetModelPresetVariant(int target_variant_index) {
  ModelPresetVariant result;
  if constexpr (variant_index != std::variant_size_v<ModelPresetVariant>) {
    if (variant_index == target_variant_index) {
      using ModelPresetT = std::variant_alternative_t<variant_index, ModelPresetVariant>;
      result = ModelPresetT();
    } else {
      result = SetModelPresetVariant<variant_index + 1>(target_variant_index);
    }
  }
  return result;
}
 
class Model {
 public:
  Model()
      : major_ver_(TREELITE_VER_MAJOR),
        minor_ver_(TREELITE_VER_MINOR),
        patch_ver_(TREELITE_VER_PATCH) {}
  virtual ~Model() = default;
  Model(Model const&) = delete;
  Model& operator=(Model const&) = delete;
  Model(Model&&) = default;
  Model& operator=(Model&&) = default;
 
  ModelPresetVariant variant_;
 
  template <typename ThresholdType, typename LeafOutputType>
  inline static std::unique_ptr<Model> Create();
  inline static std::unique_ptr<Model> Create(TypeInfo threshold_type, TypeInfo leaf_output_type);
  inline TypeInfo GetThresholdType() const {
    return std::visit([](auto&& inner) { return inner.GetThresholdType(); }, variant_);
  }
  inline TypeInfo GetLeafOutputType() const {
    return std::visit([](auto&& inner) { return inner.GetLeafOutputType(); }, variant_);
  }
 
  inline std::size_t GetNumTree() const {
    return std::visit([](auto&& inner) { return inner.GetNumTree(); }, variant_);
  }
  inline void SetTreeLimit(std::size_t limit) {
    std::visit([=](auto&& inner) { return inner.SetTreeLimit(limit); }, variant_);
  }
  void DumpAsJSON(std::ostream& fo, bool pretty_print) const;
 
  inline std::string DumpAsJSON(bool pretty_print) const {
    std::ostringstream oss;
    DumpAsJSON(oss, pretty_print);
    return oss.str();
  }
 
  /* Compatibility Matrix:
     +------------------+----------+----------+----------------+-----------+
     |                  | To: =3.9 | To: =4.0 | To: >=4.1,<5.0 | To: >=5.0 |
     +------------------+----------+----------+----------------+-----------+
     | From: =3.9       | Yes      | Yes      | Yes            | No        |
     | From: =4.0       | No       | Yes      | Yes            | Yes       |
     | From: >=4.1,<5.0 | No       | Yes      | Yes            | Yes       |
     | From: >=5.0      | No       | No       | No             | Yes       |
     +------------------+----------+----------+----------------+-----------+ */
 
  /* In-memory serialization, zero-copy */
  TREELITE_DLL_EXPORT std::vector<PyBufferFrame> SerializeToPyBuffer();
  TREELITE_DLL_EXPORT static std::unique_ptr<Model> DeserializeFromPyBuffer(
      std::vector<PyBufferFrame> const& frames);
 
  /* Serialization to a file stream */
  void SerializeToStream(std::ostream& os);
  static std::unique_ptr<Model> DeserializeFromStream(std::istream& is);
  inline Version GetVersion() const {
    return {major_ver_, minor_ver_, patch_ver_};
  }
 
  PyBufferFrame GetHeaderField(std::string const& name);
  PyBufferFrame GetTreeField(std::uint64_t tree_id, std::string const& name);
  void SetHeaderField(std::string const& name, PyBufferFrame frame);
  void SetTreeField(std::uint64_t tree_id, std::string const& name, PyBufferFrame frame);
 
  std::int32_t num_feature{0};
  TaskType task_type;
  bool average_tree_output{false};
 
  /* Task parameters */
  std::int32_t num_target;
  ContiguousArray<std::int32_t> num_class;
  ContiguousArray<std::int32_t> leaf_vector_shape;
  /* Per-tree metadata */
  ContiguousArray<std::int32_t> target_id;
  ContiguousArray<std::int32_t> class_id;
  /* Other model parameters */
  std::string postprocessor;
  float sigmoid_alpha{1.0f};
  float ratio_c{1.0f};
  ContiguousArray<double> base_scores;
  std::string attributes;
 
 private:
  /* Note: the following member fields shall be re-computed at serialization time */
  // Number of trees
  std::uint64_t num_tree_{0};
  // Number of optional fields in the extension slot
  std::int32_t num_opt_field_per_model_{0};
  // Which Treelite version produced this model
  std::int32_t major_ver_;
  std::int32_t minor_ver_;
  std::int32_t patch_ver_;
  // Type parameters
  TypeInfo threshold_type_{TypeInfo::kInvalid};
  TypeInfo leaf_output_type_{TypeInfo::kInvalid};
 
  template <typename MixIn>
  friend class detail::serializer::Serializer;
  template <typename MixIn>
  friend class detail::serializer::Deserializer;
};
 
std::unique_ptr<Model> ConcatenateModelObjects(std::vector<Model const*> const& objs);
 
}  // namespace treelite
 
#include <treelite/detail/tree.h>
 
#endif  // TREELITE_TREE_H_