leaf_splits.hpp

#ifndef LIGHTGBM_TREELEARNER_LEAF_SPLITS_HPP_
#define LIGHTGBM_TREELEARNER_LEAF_SPLITS_HPP_
 
#include <limits>
 
#include <LightGBM/meta.h>
#include "data_partition.hpp"
 
#include <vector>
 
namespace LightGBM {
 
class LeafSplits {
public:
  LeafSplits(data_size_t num_data)
    :num_data_in_leaf_(num_data), num_data_(num_data),
    data_indices_(nullptr) {
  }
  void ResetNumData(data_size_t num_data) {
    num_data_ = num_data;
    num_data_in_leaf_ = num_data;
  }
  ~LeafSplits() {
  }
 
  void Init(int leaf, const DataPartition* data_partition, double sum_gradients, double sum_hessians) {
    leaf_index_ = leaf;
    data_indices_ = data_partition->GetIndexOnLeaf(leaf, &num_data_in_leaf_);
    sum_gradients_ = sum_gradients;
    sum_hessians_ = sum_hessians;
    min_val_ = -std::numeric_limits<double>::max();
    max_val_ = std::numeric_limits<double>::max();
  }
 
  void SetValueConstraint(double min, double max) {
    min_val_ = min;
    max_val_ = max;
  }
 
 
  void Init(const score_t* gradients, const score_t* hessians) {
    num_data_in_leaf_ = num_data_;
    leaf_index_ = 0;
    data_indices_ = nullptr;
    double tmp_sum_gradients = 0.0f;
    double tmp_sum_hessians = 0.0f;
#pragma omp parallel for schedule(static) reduction(+:tmp_sum_gradients, tmp_sum_hessians)
    for (data_size_t i = 0; i < num_data_in_leaf_; ++i) {
      tmp_sum_gradients += gradients[i];
      tmp_sum_hessians += hessians[i];
    }
    sum_gradients_ = tmp_sum_gradients;
    sum_hessians_ = tmp_sum_hessians;
    min_val_ = -std::numeric_limits<double>::max();
    max_val_ = std::numeric_limits<double>::max();
  }
 
  void Init(int leaf, const DataPartition* data_partition, const score_t* gradients, const score_t* hessians) {
    leaf_index_ = leaf;
    data_indices_ = data_partition->GetIndexOnLeaf(leaf, &num_data_in_leaf_);
    double tmp_sum_gradients = 0.0f;
    double tmp_sum_hessians = 0.0f;
#pragma omp parallel for schedule(static) reduction(+:tmp_sum_gradients, tmp_sum_hessians)
    for (data_size_t i = 0; i < num_data_in_leaf_; ++i) {
      data_size_t idx = data_indices_[i];
      tmp_sum_gradients += gradients[idx];
      tmp_sum_hessians += hessians[idx];
    }
    sum_gradients_ = tmp_sum_gradients;
    sum_hessians_ = tmp_sum_hessians;
    min_val_ = -std::numeric_limits<double>::max();
    max_val_ = std::numeric_limits<double>::max();
  }
 
 
  void Init(double sum_gradients, double sum_hessians) {
    leaf_index_ = 0;
    sum_gradients_ = sum_gradients;
    sum_hessians_ = sum_hessians;
    min_val_ = -std::numeric_limits<double>::max();
    max_val_ = std::numeric_limits<double>::max();
  }
 
  void Init() {
    leaf_index_ = -1;
    data_indices_ = nullptr;
    num_data_in_leaf_ = 0;
    min_val_ = -std::numeric_limits<double>::max();
    max_val_ = std::numeric_limits<double>::max();
  }
 
 
  int LeafIndex() const { return leaf_index_; }
 
  data_size_t num_data_in_leaf() const { return num_data_in_leaf_; }
 
  double sum_gradients() const { return sum_gradients_; }
 
  double sum_hessians() const { return sum_hessians_; }
 
  double max_constraint() const { return max_val_; }
 
  double min_constraint() const { return min_val_; }
 
  const data_size_t* data_indices() const { return data_indices_; }
 
 
private:
  int leaf_index_;
  data_size_t num_data_in_leaf_;
  data_size_t num_data_;
  double sum_gradients_;
  double sum_hessians_;
  const data_size_t* data_indices_;
  double min_val_;
  double max_val_;
};
 
}  // namespace LightGBM
#endif   // LightGBM_TREELEARNER_LEAF_SPLITS_HPP_