split_info.hpp 8.38 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
#ifndef LIGHTGBM_TREELEARNER_SPLIT_INFO_HPP_
#define LIGHTGBM_TREELEARNER_SPLIT_INFO_HPP_

#include <LightGBM/meta.h>

#include <cmath>
#include <cstdint>
#include <cstring>

#include <functional>

namespace LightGBM {

/*!
* \brief Used to store some information for gain split point
*/
struct SplitInfo {
public:
  /*! \brief Feature index */
20
  int feature = -1;
Guolin Ke's avatar
Guolin Ke committed
21
  /*! \brief Split threshold */
22
23
24
25
26
  uint32_t threshold = 0;
  /*! \brief Left number of data after split */
  data_size_t left_count = 0;
  /*! \brief Right number of data after split */
  data_size_t right_count = 0;
27
  int num_cat_threshold = 0;
Guolin Ke's avatar
Guolin Ke committed
28
  /*! \brief Left output after split */
29
  double left_output = 0.0;
Guolin Ke's avatar
Guolin Ke committed
30
  /*! \brief Right output after split */
31
  double right_output = 0.0;
Guolin Ke's avatar
Guolin Ke committed
32
  /*! \brief Split gain */
33
  double gain = kMinScore;
Guolin Ke's avatar
Guolin Ke committed
34
  /*! \brief Left sum gradient after split */
35
  double left_sum_gradient = 0;
Guolin Ke's avatar
Guolin Ke committed
36
  /*! \brief Left sum hessian after split */
37
  double left_sum_hessian = 0;
Guolin Ke's avatar
Guolin Ke committed
38
  /*! \brief Right sum gradient after split */
39
  double right_sum_gradient = 0;
Guolin Ke's avatar
Guolin Ke committed
40
  /*! \brief Right sum hessian after split */
41
  double right_sum_hessian = 0;
42
  std::vector<uint32_t> cat_threshold;
43
44
  /*! \brief True if default split is left */
  bool default_left = true;
Guolin Ke's avatar
Guolin Ke committed
45

46
47
  inline static int Size(int max_cat_threshold) {
    return 2 * sizeof(int) + sizeof(uint32_t) + sizeof(bool) + sizeof(double) * 7 + sizeof(data_size_t) * 2 + max_cat_threshold * sizeof(uint32_t);
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
  }

  inline void CopyTo(char* buffer) const {
    std::memcpy(buffer, &feature, sizeof(feature));
    buffer += sizeof(feature);
    std::memcpy(buffer, &left_count, sizeof(left_count));
    buffer += sizeof(left_count);
    std::memcpy(buffer, &right_count, sizeof(right_count));
    buffer += sizeof(right_count);
    std::memcpy(buffer, &gain, sizeof(gain));
    buffer += sizeof(gain);
    std::memcpy(buffer, &threshold, sizeof(threshold));
    buffer += sizeof(threshold);
    std::memcpy(buffer, &left_output, sizeof(left_output));
    buffer += sizeof(left_output);
    std::memcpy(buffer, &right_output, sizeof(right_output));
    buffer += sizeof(right_output);
    std::memcpy(buffer, &left_sum_gradient, sizeof(left_sum_gradient));
    buffer += sizeof(left_sum_gradient);
    std::memcpy(buffer, &left_sum_hessian, sizeof(left_sum_hessian));
    buffer += sizeof(left_sum_hessian);
    std::memcpy(buffer, &right_sum_gradient, sizeof(right_sum_gradient));
    buffer += sizeof(right_sum_gradient);
    std::memcpy(buffer, &right_sum_hessian, sizeof(right_sum_hessian));
    buffer += sizeof(right_sum_hessian);
    std::memcpy(buffer, &default_left, sizeof(default_left));
    buffer += sizeof(default_left);
Guolin Ke's avatar
Guolin Ke committed
75
76
77
    std::memcpy(buffer, &num_cat_threshold, sizeof(num_cat_threshold));
    buffer += sizeof(num_cat_threshold);
    std::memcpy(buffer, cat_threshold.data(), sizeof(uint32_t) * num_cat_threshold);
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
  }

  void CopyFrom(const char* buffer) {
    std::memcpy(&feature, buffer, sizeof(feature));
    buffer += sizeof(feature);
    std::memcpy(&left_count, buffer, sizeof(left_count));
    buffer += sizeof(left_count);
    std::memcpy(&right_count, buffer, sizeof(right_count));
    buffer += sizeof(right_count);
    std::memcpy(&gain, buffer, sizeof(gain));
    buffer += sizeof(gain);
    std::memcpy(&threshold, buffer, sizeof(threshold));
    buffer += sizeof(threshold);
    std::memcpy(&left_output, buffer, sizeof(left_output));
    buffer += sizeof(left_output);
    std::memcpy(&right_output, buffer, sizeof(right_output));
    buffer += sizeof(right_output);
    std::memcpy(&left_sum_gradient, buffer, sizeof(left_sum_gradient));
    buffer += sizeof(left_sum_gradient);
    std::memcpy(&left_sum_hessian, buffer, sizeof(left_sum_hessian));
    buffer += sizeof(left_sum_hessian);
    std::memcpy(&right_sum_gradient, buffer, sizeof(right_sum_gradient));
    buffer += sizeof(right_sum_gradient);
    std::memcpy(&right_sum_hessian, buffer, sizeof(right_sum_hessian));
    buffer += sizeof(right_sum_hessian);
    std::memcpy(&default_left, buffer, sizeof(default_left));
    buffer += sizeof(default_left);
Guolin Ke's avatar
Guolin Ke committed
105
106
    std::memcpy(&num_cat_threshold, buffer, sizeof(num_cat_threshold));
    buffer += sizeof(num_cat_threshold);
107
    cat_threshold.resize(num_cat_threshold);
Guolin Ke's avatar
Guolin Ke committed
108
    std::memcpy(cat_threshold.data(), buffer, sizeof(uint32_t) * num_cat_threshold);
Guolin Ke's avatar
Guolin Ke committed
109
110
111
  }

  inline void Reset() {
112
    // initialize with -1 and -inf gain
Guolin Ke's avatar
Guolin Ke committed
113
114
115
116
    feature = -1;
    gain = kMinScore;
  }

117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
  inline bool operator > (const SplitInfo& si) const {
    double local_gain = this->gain;
    double other_gain = si.gain;
    // replace nan with -inf
    if (local_gain == NAN) {
      local_gain = kMinScore;
    }
    // replace nan with -inf
    if (other_gain == NAN) {
      other_gain = kMinScore;
    }
    int local_feature = this->feature;
    int other_feature = si.feature;
    // replace -1 with max int
    if (local_feature == -1) {
      local_feature = INT32_MAX;
    }
    // replace -1 with max int
    if (other_feature == -1) {
      other_feature = INT32_MAX;
    }
    if (local_gain != other_gain) {
      return local_gain > other_gain;
    } else {
      // if same gain, use smaller feature
      return local_feature < other_feature;
Guolin Ke's avatar
Guolin Ke committed
143
144
145
    }
  }

146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
  inline bool operator == (const SplitInfo& si) const {
    double local_gain = this->gain;
    double other_gain = si.gain;
    // replace nan with -inf
    if (local_gain == NAN) {
      local_gain = kMinScore;
    }
    // replace nan with -inf
    if (other_gain == NAN) {
      other_gain = kMinScore;
    }
    int local_feature = this->feature;
    int other_feature = si.feature;
    // replace -1 with max int
    if (local_feature == -1) {
      local_feature = INT32_MAX;
    }
    // replace -1 with max int
    if (other_feature == -1) {
      other_feature = INT32_MAX;
    }
    if (local_gain != other_gain) {
      return local_gain == other_gain;
    } else {
      // if same gain, use smaller feature
      return local_feature == other_feature;
    }
  }
Guolin Ke's avatar
Guolin Ke committed
174

175
};
Guolin Ke's avatar
Guolin Ke committed
176

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
struct LightSplitInfo {
public:
  /*! \brief Feature index */
  int feature = -1;
  /*! \brief Split gain */
  double gain = kMinScore;
  /*! \brief Left number of data after split */
  data_size_t left_count = 0;
  /*! \brief Right number of data after split */
  data_size_t right_count = 0;

  inline void Reset() {
    // initialize with -1 and -inf gain
    feature = -1;
    gain = kMinScore;
Guolin Ke's avatar
Guolin Ke committed
192
  }
193
194
195
196
197
198

  void CopyFrom(const SplitInfo& other) {
    feature = other.feature;
    gain = other.gain;
    left_count = other.left_count;
    right_count = other.right_count;
Guolin Ke's avatar
Guolin Ke committed
199
  }
200
201
202
203
204
205
206
207
208
209

  void CopyFrom(const char* buffer) {
    std::memcpy(&feature, buffer, sizeof(feature));
    buffer += sizeof(feature);
    std::memcpy(&left_count, buffer, sizeof(left_count));
    buffer += sizeof(left_count);
    std::memcpy(&right_count, buffer, sizeof(right_count));
    buffer += sizeof(right_count);
    std::memcpy(&gain, buffer, sizeof(gain));
    buffer += sizeof(gain);
Guolin Ke's avatar
Guolin Ke committed
210
  }
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238

  inline bool operator > (const LightSplitInfo& si) const {
    double local_gain = this->gain;
    double other_gain = si.gain;
    // replace nan with -inf
    if (local_gain == NAN) {
      local_gain = kMinScore;
    }
    // replace nan with -inf
    if (other_gain == NAN) {
      other_gain = kMinScore;
    }
    int local_feature = this->feature;
    int other_feature = si.feature;
    // replace -1 with max int
    if (local_feature == -1) {
      local_feature = INT32_MAX;
    }
    // replace -1 with max int
    if (other_feature == -1) {
      other_feature = INT32_MAX;
    }
    if (local_gain != other_gain) {
      return local_gain > other_gain;
    } else {
      // if same gain, use smaller feature
      return local_feature < other_feature;
    }
Guolin Ke's avatar
Guolin Ke committed
239
  }
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267

  inline bool operator == (const LightSplitInfo& si) const {
    double local_gain = this->gain;
    double other_gain = si.gain;
    // replace nan with -inf
    if (local_gain == NAN) {
      local_gain = kMinScore;
    }
    // replace nan with -inf
    if (other_gain == NAN) {
      other_gain = kMinScore;
    }
    int local_feature = this->feature;
    int other_feature = si.feature;
    // replace -1 with max int
    if (local_feature == -1) {
      local_feature = INT32_MAX;
    }
    // replace -1 with max int
    if (other_feature == -1) {
      other_feature = INT32_MAX;
    }
    if (local_gain != other_gain) {
      return local_gain == other_gain;
    } else {
      // if same gain, use smaller feature
      return local_feature == other_feature;
    }
Guolin Ke's avatar
Guolin Ke committed
268
  }
269
270

};
Guolin Ke's avatar
Guolin Ke committed
271

Guolin Ke's avatar
Guolin Ke committed
272
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
273
#endif   // LightGBM_TREELEARNER_SPLIT_INFO_HPP_