tree.h 15.5 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
8
#ifndef LIGHTGBM_TREE_H_
#define LIGHTGBM_TREE_H_

#include <LightGBM/meta.h>
#include <LightGBM/dataset.h>

#include <string>
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
9
#include <memory>
Guolin Ke's avatar
Guolin Ke committed
10
11
12

namespace LightGBM {

13
#define kMaxTreeOutput (100)
Guolin Ke's avatar
Guolin Ke committed
14
15
#define kCategoricalMask (1)
#define kDefaultLeftMask (2)
Guolin Ke's avatar
Guolin Ke committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

/*!
* \brief Tree model
*/
class Tree {
public:
  /*!
  * \brief Constructor
  * \param max_leaves The number of max leaves
  */
  explicit Tree(int max_leaves);

  /*!
  * \brief Construtor, from a string
  * \param str Model string
  */
  explicit Tree(const std::string& str);

  ~Tree();

  /*!
Qiwei Ye's avatar
Qiwei Ye committed
37
38
39
  * \brief Performing a split on tree leaves.
  * \param leaf Index of leaf to be split
  * \param feature Index of feature; the converted index after removing useless features
Guolin Ke's avatar
Guolin Ke committed
40
  * \param real_feature Index of feature, the original index on data
41
  * \param threshold_bin Threshold(bin) of split
42
  * \param threshold_double Threshold on feature value
Guolin Ke's avatar
Guolin Ke committed
43
44
  * \param left_value Model Left child output
  * \param right_value Model Right child output
Guolin Ke's avatar
Guolin Ke committed
45
46
  * \param left_cnt Count of left child
  * \param right_cnt Count of right child
Guolin Ke's avatar
Guolin Ke committed
47
  * \param gain Split gain
Guolin Ke's avatar
Guolin Ke committed
48
49
  * \param missing_type missing type
  * \param default_left default direction for missing value
Guolin Ke's avatar
Guolin Ke committed
50
51
  * \return The index of new leaf.
  */
52
53
  int Split(int leaf, int feature, int real_feature, uint32_t threshold_bin,
            double threshold_double, double left_value, double right_value,
Guolin Ke's avatar
Guolin Ke committed
54
            data_size_t left_cnt, data_size_t right_cnt, double gain, MissingType missing_type, bool default_left);
Guolin Ke's avatar
Guolin Ke committed
55

56
57
58
59
60
61
62
  /*!
  * \brief Performing a split on tree leaves, with categorical feature
  * \param leaf Index of leaf to be split
  * \param feature Index of feature; the converted index after removing useless features
  * \param real_feature Index of feature, the original index on data
  * \param threshold_bin Threshold(bin) of split, use bitset to represent
  * \param num_threshold_bin size of threshold_bin
63
64
  * \param threshold Thresholds of real feature value, use bitset to represent
  * \param num_threshold size of threshold
65
66
67
68
69
70
71
  * \param left_value Model Left child output
  * \param right_value Model Right child output
  * \param left_cnt Count of left child
  * \param right_cnt Count of right child
  * \param gain Split gain
  * \return The index of new leaf.
  */
72
73
  int SplitCategorical(int leaf, int feature, int real_feature, const uint32_t* threshold_bin, int num_threshold_bin,
                       const uint32_t* threshold, int num_threshold, double left_value, double right_value,
74
75
                       data_size_t left_cnt, data_size_t right_cnt, double gain, MissingType missing_type);

Guolin Ke's avatar
Guolin Ke committed
76
  /*! \brief Get the output of one leaf */
77
  inline double LeafOutput(int leaf) const { return leaf_value_[leaf]; }
Guolin Ke's avatar
Guolin Ke committed
78

Guolin Ke's avatar
Guolin Ke committed
79
80
81
82
83
  /*! \brief Set the output of one leaf */
  inline void SetLeafOutput(int leaf, double output) {
    leaf_value_[leaf] = output;
  }

Guolin Ke's avatar
Guolin Ke committed
84
  /*!
Qiwei Ye's avatar
Qiwei Ye committed
85
  * \brief Adding prediction value of this tree model to scores
Guolin Ke's avatar
Guolin Ke committed
86
87
88
89
  * \param data The dataset
  * \param num_data Number of total data
  * \param score Will add prediction to score
  */
90
91
92
  void AddPredictionToScore(const Dataset* data,
                            data_size_t num_data,
                            double* score) const;
Guolin Ke's avatar
Guolin Ke committed
93
94

  /*!
Qiwei Ye's avatar
Qiwei Ye committed
95
  * \brief Adding prediction value of this tree model to scorese
Guolin Ke's avatar
Guolin Ke committed
96
97
98
99
100
101
  * \param data The dataset
  * \param used_data_indices Indices of used data
  * \param num_data Number of total data
  * \param score Will add prediction to score
  */
  void AddPredictionToScore(const Dataset* data,
Qiwei Ye's avatar
Qiwei Ye committed
102
                            const data_size_t* used_data_indices,
103
                            data_size_t num_data, double* score) const;
Guolin Ke's avatar
Guolin Ke committed
104
105

  /*!
106
  * \brief Prediction on one record
Guolin Ke's avatar
Guolin Ke committed
107
108
109
  * \param feature_values Feature value of this record
  * \return Prediction result
  */
110
  inline double Predict(const double* feature_values) const;
111

112
  inline int PredictLeafIndex(const double* feature_values) const;
Guolin Ke's avatar
Guolin Ke committed
113

114
115
  inline void PredictContrib(const double* feature_values, int num_features, double* output) const;

Guolin Ke's avatar
Guolin Ke committed
116
117
118
  /*! \brief Get Number of leaves*/
  inline int num_leaves() const { return num_leaves_; }

Guolin Ke's avatar
Guolin Ke committed
119
120
121
  /*! \brief Get depth of specific leaf*/
  inline int leaf_depth(int leaf_idx) const { return leaf_depth_[leaf_idx]; }

wxchan's avatar
wxchan committed
122
  /*! \brief Get feature of specific split*/
Guolin Ke's avatar
Guolin Ke committed
123
  inline int split_feature(int split_idx) const { return split_feature_[split_idx]; }
wxchan's avatar
wxchan committed
124

Guolin Ke's avatar
Guolin Ke committed
125
126
  inline double split_gain(int split_idx) const { return split_gain_[split_idx]; }

127
  /*! \brief Get the number of data points that fall at or below this node*/
Guolin Ke's avatar
Guolin Ke committed
128
  inline int data_count(int node) const { return node >= 0 ? internal_count_[node] : leaf_count_[~node]; }
129

Guolin Ke's avatar
Guolin Ke committed
130
131
  /*!
  * \brief Shrinkage for the tree's output
Qiwei Ye's avatar
Qiwei Ye committed
132
  *        shrinkage rate (a.k.a learning rate) is used to tune the traning process
Guolin Ke's avatar
Guolin Ke committed
133
134
  * \param rate The factor of shrinkage
  */
135
  inline void Shrinkage(double rate) {
136
    #pragma omp parallel for schedule(static, 1024) if (num_leaves_ >= 2048)
Guolin Ke's avatar
Guolin Ke committed
137
    for (int i = 0; i < num_leaves_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
138
      leaf_value_[i] *= rate;
Guolin Ke's avatar
Guolin Ke committed
139
      if (leaf_value_[i] > kMaxTreeOutput) { leaf_value_[i] = kMaxTreeOutput; } else if (leaf_value_[i] < -kMaxTreeOutput) { leaf_value_[i] = -kMaxTreeOutput; }
Guolin Ke's avatar
Guolin Ke committed
140
    }
Guolin Ke's avatar
Guolin Ke committed
141
    shrinkage_ *= rate;
Guolin Ke's avatar
Guolin Ke committed
142
143
  }

Guolin Ke's avatar
Guolin Ke committed
144
145
146
147
148
149
150
151
152
  inline void AddBias(double val) {
    #pragma omp parallel for schedule(static, 1024) if (num_leaves_ >= 2048)
    for (int i = 0; i < num_leaves_; ++i) {
      leaf_value_[i] = val + leaf_value_[i];
    }
    // force to 1.0
    shrinkage_ = 1.0f;
  }

153
154
155
156
157
158
  inline void AsConstantTree(double val) {
    num_leaves_ = 1;
    shrinkage_ = 1.0f;
    leaf_value_[0] = val;
  }

wxchan's avatar
wxchan committed
159
  /*! \brief Serialize this object to string*/
Guolin Ke's avatar
Guolin Ke committed
160
  std::string ToString() const;
Guolin Ke's avatar
Guolin Ke committed
161

wxchan's avatar
wxchan committed
162
  /*! \brief Serialize this object to json*/
Guolin Ke's avatar
Guolin Ke committed
163
  std::string ToJSON() const;
wxchan's avatar
wxchan committed
164

165
  /*! \brief Serialize this object to if-else statement*/
Guolin Ke's avatar
Guolin Ke committed
166
  std::string ToIfElse(int index, bool is_predict_leaf_index) const;
167

Guolin Ke's avatar
Guolin Ke committed
168
169
170
  inline static bool IsZero(double fval) {
    if (fval > -kZeroAsMissingValueRange && fval <= kZeroAsMissingValueRange) {
      return true;
Guolin Ke's avatar
Guolin Ke committed
171
    } else {
Guolin Ke's avatar
Guolin Ke committed
172
      return false;
Guolin Ke's avatar
Guolin Ke committed
173
174
175
    }
  }

Guolin Ke's avatar
Guolin Ke committed
176
177
178
179
180
181
182
  inline static bool GetDecisionType(int8_t decision_type, int8_t mask) {
    return (decision_type & mask) > 0;
  }

  inline static void SetDecisionType(int8_t* decision_type, bool input, int8_t mask) {
    if (input) {
      (*decision_type) |= mask;
Guolin Ke's avatar
Guolin Ke committed
183
    } else {
Guolin Ke's avatar
Guolin Ke committed
184
      (*decision_type) &= (127 - mask);
Guolin Ke's avatar
Guolin Ke committed
185
186
187
    }
  }

Guolin Ke's avatar
Guolin Ke committed
188
189
190
191
192
193
194
195
196
  inline static int8_t GetMissingType(int8_t decision_type) {
    return (decision_type >> 2) & 3;
  }

  inline static void SetMissingType(int8_t* decision_type, int8_t input) {
    (*decision_type) &= 3;
    (*decision_type) |= (input << 2);
  }

197
198
private:

Guolin Ke's avatar
Guolin Ke committed
199
  std::string NumericalDecisionIfElse(int node) const;
Guolin Ke's avatar
Guolin Ke committed
200

Guolin Ke's avatar
Guolin Ke committed
201
  std::string CategoricalDecisionIfElse(int node) const;
202
203
204

  inline int NumericalDecision(double fval, int node) const {
    uint8_t missing_type = GetMissingType(decision_type_[node]);
Guolin Ke's avatar
Guolin Ke committed
205
206
207
208
209
210
211
    if (std::isnan(fval)) {
      if (missing_type != 2) {
        fval = 0.0f;
      }
    }
    if ((missing_type == 1 && IsZero(fval))
        || (missing_type == 2 && std::isnan(fval))) {
212
213
      if (GetDecisionType(decision_type_[node], kDefaultLeftMask)) {
        return left_child_[node];
Guolin Ke's avatar
Guolin Ke committed
214
      } else {
215
        return right_child_[node];
Guolin Ke's avatar
Guolin Ke committed
216
217
      }
    }
218
219
220
221
222
    if (fval <= threshold_[node]) {
      return left_child_[node];
    } else {
      return right_child_[node];
    }
Guolin Ke's avatar
Guolin Ke committed
223
  }
Guolin Ke's avatar
Guolin Ke committed
224

225
226
227
228
229
230
231
232
233
234
235
236
  inline int NumericalDecisionInner(uint32_t fval, int node, uint32_t default_bin, uint32_t max_bin) const {
    uint8_t missing_type = GetMissingType(decision_type_[node]);
    if ((missing_type == 1 && fval == default_bin)
        || (missing_type == 2 && fval == max_bin)) {
      if (GetDecisionType(decision_type_[node], kDefaultLeftMask)) {
        return left_child_[node];
      } else {
        return right_child_[node];
      }
    }
    if (fval <= threshold_in_bin_[node]) {
      return left_child_[node];
237
    } else {
238
      return right_child_[node];
239
240
    }
  }
Guolin Ke's avatar
Guolin Ke committed
241

242
243
244
245
246
247
248
249
250
251
252
253
  inline int CategoricalDecision(double fval, int node) const {
    uint8_t missing_type = GetMissingType(decision_type_[node]);
    int int_fval = static_cast<int>(fval);
    if (int_fval < 0) {
      return right_child_[node];;
    } else if (std::isnan(fval)) {
      // NaN is always in the right
      if (missing_type == 2) {
        return right_child_[node];
      }
      int_fval = 0;
    }
254
255
256
    int cat_idx = int(threshold_[node]);
    if (Common::FindInBitset(cat_threshold_.data() + cat_boundaries_[cat_idx],
                             cat_boundaries_[cat_idx + 1] - cat_boundaries_[cat_idx], int_fval)) {
257
258
259
260
      return left_child_[node];
    }
    return right_child_[node];
  }
Guolin Ke's avatar
Guolin Ke committed
261

262
  inline int CategoricalDecisionInner(uint32_t fval, int node) const {
263
264
265
    int cat_idx = int(threshold_in_bin_[node]);
    if (Common::FindInBitset(cat_threshold_inner_.data() + cat_boundaries_inner_[cat_idx],
                             cat_boundaries_inner_[cat_idx + 1] - cat_boundaries_inner_[cat_idx], fval)) {
266
267
268
269
      return left_child_[node];
    }
    return right_child_[node];
  }
Guolin Ke's avatar
Guolin Ke committed
270

271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
  inline int Decision(double fval, int node) const {
    if (GetDecisionType(decision_type_[node], kCategoricalMask)) {
      return CategoricalDecision(fval, node);
    } else {
      return NumericalDecision(fval, node);
    }
  }

  inline int DecisionInner(uint32_t fval, int node, uint32_t default_bin, uint32_t max_bin) const {
    if (GetDecisionType(decision_type_[node], kCategoricalMask)) {
      return CategoricalDecisionInner(fval, node);
    } else {
      return NumericalDecisionInner(fval, node, default_bin, max_bin);
    }
  }

  inline void Split(int leaf, int feature, int real_feature,
                    double left_value, double right_value, data_size_t left_cnt, data_size_t right_cnt, double gain);
Guolin Ke's avatar
Guolin Ke committed
289
  /*!
Qiwei Ye's avatar
Qiwei Ye committed
290
  * \brief Find leaf index of which record belongs by features
Guolin Ke's avatar
Guolin Ke committed
291
292
293
  * \param feature_values Feature value of this record
  * \return Leaf index
  */
294
  inline int GetLeaf(const double* feature_values) const;
Guolin Ke's avatar
Guolin Ke committed
295

wxchan's avatar
wxchan committed
296
  /*! \brief Serialize one node to json*/
Guolin Ke's avatar
Guolin Ke committed
297
  std::string NodeToJSON(int index) const;
wxchan's avatar
wxchan committed
298

299
  /*! \brief Serialize one node to if-else statement*/
Guolin Ke's avatar
Guolin Ke committed
300
301
  std::string NodeToIfElse(int index, bool is_predict_leaf_index) const;

302
  double ExpectedValue(int node) const;
Guolin Ke's avatar
Guolin Ke committed
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326

  int MaxDepth() const;

  /*!
  * \brief Used by TreeSHAP for data we keep about our decision path
  */
  struct PathElement {
    int feature_index;
    double zero_fraction;
    double one_fraction;

    // note that pweight is included for convenience and is not tied with the other attributes,
    // the pweight of the i'th path element is the permuation weight of paths with i-1 ones in them
    double pweight;

    PathElement() {}
    PathElement(int i, double z, double o, double w) : feature_index(i), zero_fraction(z), one_fraction(o), pweight(w) {}
  };

  /*! \brief Polynomial time algorithm for SHAP values (https://arxiv.org/abs/1706.06060) */
  void TreeSHAP(const double *feature_values, double *phi,
                int node, int unique_depth,
                PathElement *parent_unique_path, double parent_zero_fraction,
                double parent_one_fraction, int parent_feature_index) const;
327

328
  /*! \brief Extend our decision path with a fraction of one and zero extensions for TreeSHAP*/
Guolin Ke's avatar
Guolin Ke committed
329
330
  static void ExtendPath(PathElement *unique_path, int unique_depth,
                         double zero_fraction, double one_fraction, int feature_index);
331
332

  /*! \brief Undo a previous extension of the decision path for TreeSHAP*/
Guolin Ke's avatar
Guolin Ke committed
333
  static void UnwindPath(PathElement *unique_path, int unique_depth, int path_index);
334
335

  /*! determine what the total permuation weight would be if we unwound a previous extension in the decision path*/
Guolin Ke's avatar
Guolin Ke committed
336
  static double UnwoundPathSum(const PathElement *unique_path, int unique_depth, int path_index);
337

Guolin Ke's avatar
Guolin Ke committed
338
339
340
341
342
343
  /*! \brief Number of max leaves*/
  int max_leaves_;
  /*! \brief Number of current levas*/
  int num_leaves_;
  // following values used for non-leaf node
  /*! \brief A non-leaf node's left child */
Guolin Ke's avatar
Guolin Ke committed
344
  std::vector<int> left_child_;
Guolin Ke's avatar
Guolin Ke committed
345
  /*! \brief A non-leaf node's right child */
Guolin Ke's avatar
Guolin Ke committed
346
  std::vector<int> right_child_;
Guolin Ke's avatar
Guolin Ke committed
347
  /*! \brief A non-leaf node's split feature */
Guolin Ke's avatar
Guolin Ke committed
348
  std::vector<int> split_feature_inner_;
Guolin Ke's avatar
Guolin Ke committed
349
  /*! \brief A non-leaf node's split feature, the original index */
Guolin Ke's avatar
Guolin Ke committed
350
  std::vector<int> split_feature_;
Guolin Ke's avatar
Guolin Ke committed
351
  /*! \brief A non-leaf node's split threshold in bin */
Guolin Ke's avatar
Guolin Ke committed
352
  std::vector<uint32_t> threshold_in_bin_;
Guolin Ke's avatar
Guolin Ke committed
353
  /*! \brief A non-leaf node's split threshold in feature value */
Guolin Ke's avatar
Guolin Ke committed
354
  std::vector<double> threshold_;
355
  int num_cat_;
356
357
358
359
  std::vector<int> cat_boundaries_inner_;
  std::vector<uint32_t> cat_threshold_inner_;
  std::vector<int> cat_boundaries_;
  std::vector<uint32_t> cat_threshold_;
Guolin Ke's avatar
Guolin Ke committed
360
  /*! \brief Store the information for categorical feature handle and mising value handle. */
361
  std::vector<int8_t> decision_type_;
Guolin Ke's avatar
Guolin Ke committed
362
  /*! \brief A non-leaf node's split gain */
Guolin Ke's avatar
Guolin Ke committed
363
  std::vector<double> split_gain_;
Guolin Ke's avatar
Guolin Ke committed
364
365
  // used for leaf node
  /*! \brief The parent of leaf */
Guolin Ke's avatar
Guolin Ke committed
366
  std::vector<int> leaf_parent_;
Guolin Ke's avatar
Guolin Ke committed
367
  /*! \brief Output of leaves */
Guolin Ke's avatar
Guolin Ke committed
368
  std::vector<double> leaf_value_;
Guolin Ke's avatar
Guolin Ke committed
369
370
371
372
373
374
  /*! \brief DataCount of leaves */
  std::vector<data_size_t> leaf_count_;
  /*! \brief Output of non-leaf nodes */
  std::vector<double> internal_value_;
  /*! \brief DataCount of non-leaf nodes */
  std::vector<data_size_t> internal_count_;
Guolin Ke's avatar
Guolin Ke committed
375
  /*! \brief Depth for leaves */
Guolin Ke's avatar
Guolin Ke committed
376
  std::vector<int> leaf_depth_;
Guolin Ke's avatar
Guolin Ke committed
377
  double shrinkage_;
Guolin Ke's avatar
Guolin Ke committed
378
379
};

380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
inline void Tree::Split(int leaf, int feature, int real_feature,
                        double left_value, double right_value, data_size_t left_cnt, data_size_t right_cnt, double gain) {
  int new_node_idx = num_leaves_ - 1;
  // update parent info
  int parent = leaf_parent_[leaf];
  if (parent >= 0) {
    // if cur node is left child
    if (left_child_[parent] == ~leaf) {
      left_child_[parent] = new_node_idx;
    } else {
      right_child_[parent] = new_node_idx;
    }
  }
  // add new node
  split_feature_inner_[new_node_idx] = feature;
  split_feature_[new_node_idx] = real_feature;

  split_gain_[new_node_idx] = Common::AvoidInf(gain);
  // add two new leaves
  left_child_[new_node_idx] = ~leaf;
  right_child_[new_node_idx] = ~num_leaves_;
  // update new leaves
  leaf_parent_[leaf] = new_node_idx;
  leaf_parent_[num_leaves_] = new_node_idx;
  // save current leaf value to internal node before change
  internal_value_[new_node_idx] = leaf_value_[leaf];
  internal_count_[new_node_idx] = left_cnt + right_cnt;
  leaf_value_[leaf] = std::isnan(left_value) ? 0.0f : left_value;
  leaf_count_[leaf] = left_cnt;
  leaf_value_[num_leaves_] = std::isnan(right_value) ? 0.0f : right_value;
  leaf_count_[num_leaves_] = right_cnt;
  // update leaf depth
  leaf_depth_[num_leaves_] = leaf_depth_[leaf] + 1;
  leaf_depth_[leaf]++;
}

416
inline double Tree::Predict(const double* feature_values) const {
Guolin Ke's avatar
Guolin Ke committed
417
418
419
420
  if (num_leaves_ > 1) {
    int leaf = GetLeaf(feature_values);
    return LeafOutput(leaf);
  } else {
421
    return leaf_value_[0];
Guolin Ke's avatar
Guolin Ke committed
422
  }
Guolin Ke's avatar
Guolin Ke committed
423
424
}

425
inline int Tree::PredictLeafIndex(const double* feature_values) const {
Guolin Ke's avatar
Guolin Ke committed
426
427
428
429
430
431
  if (num_leaves_ > 1) {
    int leaf = GetLeaf(feature_values);
    return leaf;
  } else {
    return 0;
  }
wxchan's avatar
wxchan committed
432
433
}

Guolin Ke's avatar
Guolin Ke committed
434
inline void Tree::PredictContrib(const double* feature_values, int num_features, double* output) const {
435
  output[num_features] += ExpectedValue(0);
436
  // Run the recursion with preallocated space for the unique path data
Guolin Ke's avatar
Guolin Ke committed
437
438
439
  const int max_path_len = MaxDepth() + 1;
  std::vector<PathElement> unique_path_data((max_path_len*(max_path_len + 1)) / 2);
  TreeSHAP(feature_values, output, 0, 0, unique_path_data.data(), 1, 1, -1);
440
441
}

442
inline int Tree::GetLeaf(const double* feature_values) const {
Guolin Ke's avatar
Guolin Ke committed
443
  int node = 0;
444
  if (num_cat_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
445
    while (node >= 0) {
446
      node = Decision(feature_values[split_feature_[node]], node);
Guolin Ke's avatar
Guolin Ke committed
447
448
449
    }
  } else {
    while (node >= 0) {
450
      node = NumericalDecision(feature_values[split_feature_[node]], node);
Guolin Ke's avatar
Guolin Ke committed
451
452
453
454
455
456
457
    }
  }
  return ~node;
}

}  // namespace LightGBM

Guolin Ke's avatar
Guolin Ke committed
458
#endif   // LightGBM_TREE_H_