feature_histogram.hpp 21.9 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
#ifndef LIGHTGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_
#define LIGHTGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_

#include "split_info.hpp"
Guolin Ke's avatar
Guolin Ke committed
5
6
7

#include <LightGBM/utils/array_args.h>
#include <LightGBM/dataset.h>
Guolin Ke's avatar
Guolin Ke committed
8
9
10

#include <cstring>

Guolin Ke's avatar
Guolin Ke committed
11
namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
12
{
Guolin Ke's avatar
Guolin Ke committed
13

Guolin Ke's avatar
Guolin Ke committed
14
15
16
class FeatureMetainfo {
public:
  int num_bin;
Guolin Ke's avatar
Guolin Ke committed
17
  MissingType missing_type;
Guolin Ke's avatar
Guolin Ke committed
18
  int bias = 0;
Guolin Ke's avatar
Guolin Ke committed
19
  uint32_t default_bin;
Guolin Ke's avatar
Guolin Ke committed
20
21
22
  /*! \brief pointer of tree config */
  const TreeConfig* tree_config;
};
Guolin Ke's avatar
Guolin Ke committed
23
24
25
26
27
/*!
* \brief FeatureHistogram is used to construct and store a histogram for a feature.
*/
class FeatureHistogram {
public:
Guolin Ke's avatar
Guolin Ke committed
28
  FeatureHistogram() {
Guolin Ke's avatar
Guolin Ke committed
29
    data_ = nullptr;
Guolin Ke's avatar
Guolin Ke committed
30
  }
Guolin Ke's avatar
Guolin Ke committed
31

Guolin Ke's avatar
Guolin Ke committed
32
33
34
  ~FeatureHistogram() {
  }

Guolin Ke's avatar
Guolin Ke committed
35
36
37
38
39
  /*! \brief Disable copy */
  FeatureHistogram& operator=(const FeatureHistogram&) = delete;
  /*! \brief Disable copy */
  FeatureHistogram(const FeatureHistogram&) = delete;

Guolin Ke's avatar
Guolin Ke committed
40
41
42
43
44
  /*!
  * \brief Init the feature histogram
  * \param feature the feature data for this histogram
  * \param min_num_data_one_leaf minimal number of data in one leaf
  */
45
  void Init(HistogramBinEntry* data, const FeatureMetainfo* meta, BinType bin_type) {
Guolin Ke's avatar
Guolin Ke committed
46
47
    meta_ = meta;
    data_ = data;
48
49
    if (bin_type == BinType::NumericalBin) {
      find_best_threshold_fun_ = std::bind(&FeatureHistogram::FindBestThresholdNumerical, this, std::placeholders::_1
Guolin Ke's avatar
Guolin Ke committed
50
                                           , std::placeholders::_2, std::placeholders::_3, std::placeholders::_4);
51
52
    } else {
      find_best_threshold_fun_ = std::bind(&FeatureHistogram::FindBestThresholdCategorical, this, std::placeholders::_1
Guolin Ke's avatar
Guolin Ke committed
53
                                           , std::placeholders::_2, std::placeholders::_3, std::placeholders::_4);
54
    }
Guolin Ke's avatar
Guolin Ke committed
55
56
  }

Guolin Ke's avatar
Guolin Ke committed
57
58
  HistogramBinEntry* RawData() {
    return data_;
Guolin Ke's avatar
Guolin Ke committed
59
60
61
62
63
64
  }
  /*!
  * \brief Subtract current histograms with other
  * \param other The histogram that want to subtract
  */
  void Subtract(const FeatureHistogram& other) {
Guolin Ke's avatar
Guolin Ke committed
65
    for (int i = 0; i < meta_->num_bin - meta_->bias; ++i) {
Guolin Ke's avatar
Guolin Ke committed
66
67
68
69
70
      data_[i].cnt -= other.data_[i].cnt;
      data_[i].sum_gradients -= other.data_[i].sum_gradients;
      data_[i].sum_hessians -= other.data_[i].sum_hessians;
    }
  }
71

Guolin Ke's avatar
Guolin Ke committed
72
  void FindBestThreshold(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
73
                         SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
74
    output->default_left = true;
Guolin Ke's avatar
Guolin Ke committed
75
    output->gain = kMinScore;
76
77
78
79
    find_best_threshold_fun_(sum_gradient, sum_hessian + 2 * kEpsilon, num_data, output);
  }

  void FindBestThresholdNumerical(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
80
                                  SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
81
82

    is_splittable_ = false;
Guolin Ke's avatar
Guolin Ke committed
83
84
    double gain_shift = GetLeafSplitGain(sum_gradient, sum_hessian,
                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
Guolin Ke's avatar
Guolin Ke committed
85
    double min_gain_shift = gain_shift + meta_->tree_config->min_gain_to_split;
Guolin Ke's avatar
Guolin Ke committed
86
87
88
89
90
91
92
93
    if (meta_->num_bin > 2 && meta_->missing_type != MissingType::None) {
      if (meta_->missing_type == MissingType::Zero) {
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, true, false);
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, 1, true, false);
      } else {
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, false, true);
        FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, 1, false, true);
      }
94
    } else {
Guolin Ke's avatar
Guolin Ke committed
95
96
97
98
99
      FindBestThresholdSequence(sum_gradient, sum_hessian, num_data, min_gain_shift, output, -1, false, false);
      // fix the direction error when only have 2 bins
      if (meta_->missing_type == MissingType::NaN) {
        output->default_left = false;
      }
Guolin Ke's avatar
Guolin Ke committed
100
    }
Guolin Ke's avatar
Guolin Ke committed
101
    output->gain -= min_gain_shift;
Guolin Ke's avatar
Guolin Ke committed
102
  }
103
104

  void FindBestThresholdCategorical(double sum_gradient, double sum_hessian, data_size_t num_data,
Guolin Ke's avatar
Guolin Ke committed
105
                                    SplitInfo* output) {
Guolin Ke's avatar
Guolin Ke committed
106
    output->default_left = false;
107
108
    double best_gain = kMinScore;
    uint32_t best_threshold = static_cast<uint32_t>(meta_->num_bin);
109
110
111
    data_size_t best_left_count = 0;
    double best_sum_left_gradient = 0.0f;
    double best_sum_left_hessian = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
112
113
    double gain_shift = GetLeafSplitGain(sum_gradient, sum_hessian,
                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
114
115
116
117
118
119
120
121
122
    double min_gain_shift = gain_shift + meta_->tree_config->min_gain_to_split;
    is_splittable_ = false;
    const int bias = meta_->bias;
    int t = meta_->num_bin - 1 - bias;
    const int t_end = 0;
    // from right to left, and we don't need data in bin0
    for (; t >= t_end; --t) {
      // if data not enough, or sum hessian too small
      if (data_[t].cnt < meta_->tree_config->min_data_in_leaf
Guolin Ke's avatar
Guolin Ke committed
123
          || data_[t].sum_hessians < meta_->tree_config->min_sum_hessian_in_leaf) continue;
124
125
126
127
128
129
130
131
132
133
      data_size_t other_count = num_data - data_[t].cnt;
      // if data not enough
      if (other_count < meta_->tree_config->min_data_in_leaf) continue;

      double sum_other_hessian = sum_hessian - data_[t].sum_hessians - kEpsilon;
      // if sum hessian too small
      if (sum_other_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;

      double sum_other_gradient = sum_gradient - data_[t].sum_gradients;
      // current split gain
Guolin Ke's avatar
Guolin Ke committed
134
135
136
137
      double current_gain = GetLeafSplitGain(sum_other_gradient, sum_other_hessian,
                                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
        + GetLeafSplitGain(data_[t].sum_gradients, data_[t].sum_hessians + kEpsilon,
                           meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
138
139
140
141
142
143
144
145
      // gain with split is worse than without split
      if (current_gain <= min_gain_shift) continue;

      // mark to is splittable
      is_splittable_ = true;
      // better split point
      if (current_gain > best_gain) {
        best_threshold = static_cast<uint32_t>(t + bias);
146
147
148
        best_sum_left_gradient = data_[t].sum_gradients;
        best_sum_left_hessian = data_[t].sum_hessians + kEpsilon;
        best_left_count = data_[t].cnt;
149
150
151
152
153
154
155
        best_gain = current_gain;
      }
    }
    // need restore zero bin
    if (bias == 1) {
      t = meta_->num_bin - 1 - bias;
      double sum_bin0_gradient = sum_gradient;
156
      double sum_bin0_hessian = sum_hessian - 2 * kEpsilon;
157
158
159
160
161
162
163
164
165
      data_size_t cnt_bin0 = num_data;
      for (; t >= 0; --t) {
        sum_bin0_gradient -= data_[t].sum_gradients;
        sum_bin0_hessian -= data_[t].sum_hessians;
        cnt_bin0 -= data_[t].cnt;
      }
      data_size_t other_count = num_data - cnt_bin0;
      double sum_other_hessian = sum_hessian - sum_bin0_hessian - kEpsilon;
      if (cnt_bin0 >= meta_->tree_config->min_data_in_leaf
Guolin Ke's avatar
Guolin Ke committed
166
167
168
          && sum_bin0_hessian >= meta_->tree_config->min_sum_hessian_in_leaf
          && other_count >= meta_->tree_config->min_data_in_leaf
          && sum_other_hessian >= meta_->tree_config->min_sum_hessian_in_leaf) {
169
        double sum_other_gradient = sum_gradient - sum_bin0_gradient;
Guolin Ke's avatar
Guolin Ke committed
170
171
172
173
        double current_gain = GetLeafSplitGain(sum_other_gradient, sum_other_hessian,
                                               meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
          + GetLeafSplitGain(sum_bin0_gradient, sum_bin0_hessian + kEpsilon,
                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
174
175
176
177
178
        if (current_gain > min_gain_shift) {
          is_splittable_ = true;
          // better split point
          if (current_gain > best_gain) {
            best_threshold = static_cast<uint32_t>(0);
179
180
181
            best_sum_left_gradient = sum_bin0_gradient;
            best_sum_left_hessian = sum_bin0_hessian + kEpsilon;
            best_left_count = cnt_bin0;
182
183
184
185
186
187
188
189
            best_gain = current_gain;
          }
        }
      }
    }
    if (is_splittable_) {
      // update split information
      output->threshold = best_threshold;
Guolin Ke's avatar
Guolin Ke committed
190
191
      output->left_output = CalculateSplittedLeafOutput(best_sum_left_gradient, best_sum_left_hessian,
                                                        meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
192
193
194
195
      output->left_count = best_left_count;
      output->left_sum_gradient = best_sum_left_gradient;
      output->left_sum_hessian = best_sum_left_hessian - kEpsilon;
      output->right_output = CalculateSplittedLeafOutput(sum_gradient - best_sum_left_gradient,
Guolin Ke's avatar
Guolin Ke committed
196
197
                                                         sum_hessian - best_sum_left_hessian,
                                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
198
199
200
      output->right_count = num_data - best_left_count;
      output->right_sum_gradient = sum_gradient - best_sum_left_gradient;
      output->right_sum_hessian = sum_hessian - best_sum_left_hessian - kEpsilon;
Guolin Ke's avatar
Guolin Ke committed
201
202
      output->gain = best_gain - min_gain_shift;
    } 
203
204
  }

Guolin Ke's avatar
Guolin Ke committed
205
206
207
208
  /*!
  * \brief Binary size of this histogram
  */
  int SizeOfHistgram() const {
Guolin Ke's avatar
Guolin Ke committed
209
    return (meta_->num_bin - meta_->bias) * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
210
211
212
213
214
  }

  /*!
  * \brief Restore histogram from memory
  */
Guolin Ke's avatar
Guolin Ke committed
215
216
  void FromMemory(char* memory_data) {
    std::memcpy(data_, memory_data, (meta_->num_bin - meta_->bias) * sizeof(HistogramBinEntry));
Guolin Ke's avatar
Guolin Ke committed
217
218
219
220
221
222
223
224
225
226
227
228
229
  }

  /*!
  * \brief True if this histogram can be splitted
  */
  bool is_splittable() { return is_splittable_; }

  /*!
  * \brief Set splittable to this histogram
  */
  void set_is_splittable(bool val) { is_splittable_ = val; }

  /*!
230
  * \brief Calculate the split gain based on regularized sum_gradients and sum_hessians
Guolin Ke's avatar
Guolin Ke committed
231
232
233
234
  * \param sum_gradients
  * \param sum_hessians
  * \return split gain
  */
Guolin Ke's avatar
Guolin Ke committed
235
  static double GetLeafSplitGain(double sum_gradients, double sum_hessians, double l1, double l2) {
236
    double abs_sum_gradients = std::fabs(sum_gradients);
Guolin Ke's avatar
Guolin Ke committed
237
    double reg_abs_sum_gradients = std::max(0.0, abs_sum_gradients - l1);
Guolin Ke's avatar
Guolin Ke committed
238
    return (reg_abs_sum_gradients * reg_abs_sum_gradients)
Guolin Ke's avatar
Guolin Ke committed
239
      / (sum_hessians + l2);
Guolin Ke's avatar
Guolin Ke committed
240

Guolin Ke's avatar
Guolin Ke committed
241
242
243
  }

  /*!
244
  * \brief Calculate the output of a leaf based on regularized sum_gradients and sum_hessians
Guolin Ke's avatar
Guolin Ke committed
245
246
247
248
  * \param sum_gradients
  * \param sum_hessians
  * \return leaf output
  */
Guolin Ke's avatar
Guolin Ke committed
249
  static double CalculateSplittedLeafOutput(double sum_gradients, double sum_hessians, double l1, double l2) {
250
    double abs_sum_gradients = std::fabs(sum_gradients);
Guolin Ke's avatar
Guolin Ke committed
251
    double reg_abs_sum_gradients = std::max(0.0, abs_sum_gradients - l1);
Guolin Ke's avatar
Guolin Ke committed
252
    return -std::copysign(reg_abs_sum_gradients, sum_gradients)
Guolin Ke's avatar
Guolin Ke committed
253
      / (sum_hessians + l2);
Guolin Ke's avatar
Guolin Ke committed
254
  }
Guolin Ke's avatar
Guolin Ke committed
255
256
257

private:

Guolin Ke's avatar
Guolin Ke committed
258
  void FindBestThresholdSequence(double sum_gradient, double sum_hessian, data_size_t num_data, double min_gain_shift,
Guolin Ke's avatar
Guolin Ke committed
259
                                 SplitInfo* output, int dir, bool skip_default_bin, bool use_na_as_missing) {
Guolin Ke's avatar
Guolin Ke committed
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274

    const int bias = meta_->bias;

    double best_sum_left_gradient = NAN;
    double best_sum_left_hessian = NAN;
    double best_gain = kMinScore;
    data_size_t best_left_count = 0;
    uint32_t best_threshold = static_cast<uint32_t>(meta_->num_bin);

    if (dir == -1) {

      double sum_right_gradient = 0.0f;
      double sum_right_hessian = kEpsilon;
      data_size_t right_count = 0;

Guolin Ke's avatar
Guolin Ke committed
275
      int t = meta_->num_bin - 1 - bias - use_na_as_missing;
Guolin Ke's avatar
Guolin Ke committed
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
      const int t_end = 1 - bias;

      // from right to left, and we don't need data in bin0
      for (; t >= t_end; --t) {

        // need to skip default bin
        if (skip_default_bin && (t + bias) == static_cast<int>(meta_->default_bin)) { continue; }

        sum_right_gradient += data_[t].sum_gradients;
        sum_right_hessian += data_[t].sum_hessians;
        right_count += data_[t].cnt;
        // if data not enough, or sum hessian too small
        if (right_count < meta_->tree_config->min_data_in_leaf
            || sum_right_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;
        data_size_t left_count = num_data - right_count;
        // if data not enough
        if (left_count < meta_->tree_config->min_data_in_leaf) break;

        double sum_left_hessian = sum_hessian - sum_right_hessian;
        // if sum hessian too small
        if (sum_left_hessian < meta_->tree_config->min_sum_hessian_in_leaf) break;

        double sum_left_gradient = sum_gradient - sum_right_gradient;
        // current split gain
        double current_gain = GetLeafSplitGain(sum_left_gradient, sum_left_hessian,
                                               meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
          + GetLeafSplitGain(sum_right_gradient, sum_right_hessian,
                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
        // gain with split is worse than without split
        if (current_gain <= min_gain_shift) continue;

        // mark to is splittable
        is_splittable_ = true;
        // better split point
        if (current_gain > best_gain) {
          best_left_count = left_count;
          best_sum_left_gradient = sum_left_gradient;
          best_sum_left_hessian = sum_left_hessian;
          // left is <= threshold, right is > threshold.  so this is t-1
          best_threshold = static_cast<uint32_t>(t - 1 + bias);
          best_gain = current_gain;
        }
      }
    } else{
      double sum_left_gradient = 0.0f;
      double sum_left_hessian = kEpsilon;
      data_size_t left_count = 0;

      int t = 0;
      const int t_end = meta_->num_bin - 2 - bias;

Guolin Ke's avatar
Guolin Ke committed
327
328
329
330
331
332
333
334
335
336
337
338
      if (use_na_as_missing && bias == 1) {
        sum_left_gradient = sum_gradient;
        sum_left_hessian = sum_hessian - kEpsilon;
        left_count = num_data;
        for (int i = 0; i < meta_->num_bin - bias; ++i) {
          sum_left_gradient -= data_[i].sum_gradients;
          sum_left_hessian -= data_[i].sum_hessians;
          left_count -= data_[i].cnt;
        }
        t = -1;
      }

Guolin Ke's avatar
Guolin Ke committed
339
340
341
342
      for (; t <= t_end; ++t) {

        // need to skip default bin
        if (skip_default_bin && (t + bias) == static_cast<int>(meta_->default_bin)) { continue; }
Guolin Ke's avatar
Guolin Ke committed
343
344
345
346
347
        if (t >= 0) {
          sum_left_gradient += data_[t].sum_gradients;
          sum_left_hessian += data_[t].sum_hessians;
          left_count += data_[t].cnt;
        }
Guolin Ke's avatar
Guolin Ke committed
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
        // if data not enough, or sum hessian too small
        if (left_count < meta_->tree_config->min_data_in_leaf
            || sum_left_hessian < meta_->tree_config->min_sum_hessian_in_leaf) continue;
        data_size_t right_count = num_data - left_count;
        // if data not enough
        if (right_count < meta_->tree_config->min_data_in_leaf) break;

        double sum_right_hessian = sum_hessian - sum_left_hessian;
        // if sum hessian too small
        if (sum_right_hessian < meta_->tree_config->min_sum_hessian_in_leaf) break;

        double sum_right_gradient = sum_gradient - sum_left_gradient;
        // current split gain
        double current_gain = GetLeafSplitGain(sum_left_gradient, sum_left_hessian,
                                               meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2)
          + GetLeafSplitGain(sum_right_gradient, sum_right_hessian,
                             meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
        // gain with split is worse than without split
        if (current_gain <= min_gain_shift) continue;

        // mark to is splittable
        is_splittable_ = true;
        // better split point
        if (current_gain > best_gain) {
          best_left_count = left_count;
          best_sum_left_gradient = sum_left_gradient;
          best_sum_left_hessian = sum_left_hessian;
          best_threshold = static_cast<uint32_t>(t + bias);
          best_gain = current_gain;
        }
      }
    }

    if (is_splittable_ && best_gain > output->gain) {
      // update split information
      output->threshold = best_threshold;
      output->left_output = CalculateSplittedLeafOutput(best_sum_left_gradient, best_sum_left_hessian,
                                                        meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
      output->left_count = best_left_count;
      output->left_sum_gradient = best_sum_left_gradient;
      output->left_sum_hessian = best_sum_left_hessian - kEpsilon;
      output->right_output = CalculateSplittedLeafOutput(sum_gradient - best_sum_left_gradient,
                                                         sum_hessian - best_sum_left_hessian,
                                                         meta_->tree_config->lambda_l1, meta_->tree_config->lambda_l2);
      output->right_count = num_data - best_left_count;
      output->right_sum_gradient = sum_gradient - best_sum_left_gradient;
      output->right_sum_hessian = sum_hessian - best_sum_left_hessian - kEpsilon;
      output->gain = best_gain;
Guolin Ke's avatar
Guolin Ke committed
396
      output->default_left = dir == -1;
Guolin Ke's avatar
Guolin Ke committed
397
398
399
    }
  }

Guolin Ke's avatar
Guolin Ke committed
400
  const FeatureMetainfo* meta_;
Guolin Ke's avatar
Guolin Ke committed
401
  /*! \brief sum of gradient of each bin */
Guolin Ke's avatar
Guolin Ke committed
402
403
  HistogramBinEntry* data_;
  //std::vector<HistogramBinEntry> data_;
Guolin Ke's avatar
Guolin Ke committed
404
405
  /*! \brief False if this histogram cannot split */
  bool is_splittable_ = true;
406
407

  std::function<void(double, double, data_size_t, SplitInfo*)> find_best_threshold_fun_;
Guolin Ke's avatar
Guolin Ke committed
408
};
Guolin Ke's avatar
Guolin Ke committed
409
410
411
412
413
414
class HistogramPool {
public:
  /*!
  * \brief Constructor
  */
  HistogramPool() {
Guolin Ke's avatar
Guolin Ke committed
415
416
    cache_size_ = 0;
    total_size_ = 0;
Guolin Ke's avatar
Guolin Ke committed
417
418
419
420
421
422
423
424
425
426
427
  }
  /*!
  * \brief Destructor
  */
  ~HistogramPool() {
  }
  /*!
  * \brief Reset pool size
  * \param cache_size Max cache size
  * \param total_size Total size will be used
  */
Guolin Ke's avatar
Guolin Ke committed
428
  void Reset(int cache_size, int total_size) {
Guolin Ke's avatar
Guolin Ke committed
429
430
431
432
433
434
435
436
437
    cache_size_ = cache_size;
    // at least need 2 bucket to store smaller leaf and larger leaf
    CHECK(cache_size_ >= 2);
    total_size_ = total_size;
    if (cache_size_ > total_size_) {
      cache_size_ = total_size_;
    }
    is_enough_ = (cache_size_ == total_size_);
    if (!is_enough_) {
438
439
440
      mapper_.resize(total_size_);
      inverse_mapper_.resize(cache_size_);
      last_used_time_.resize(cache_size_);
Guolin Ke's avatar
Guolin Ke committed
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
      ResetMap();
    }
  }
  /*!
  * \brief Reset mapper
  */
  void ResetMap() {
    if (!is_enough_) {
      cur_time_ = 0;
      std::fill(mapper_.begin(), mapper_.end(), -1);
      std::fill(inverse_mapper_.begin(), inverse_mapper_.end(), -1);
      std::fill(last_used_time_.begin(), last_used_time_.end(), 0);
    }
  }

Guolin Ke's avatar
Guolin Ke committed
456
457
  void DynamicChangeSize(const Dataset* train_data, const TreeConfig* tree_config, int cache_size, int total_size) {
    if (feature_metas_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
458
459
460
461
      int num_feature = train_data->num_features();
      feature_metas_.resize(num_feature);
      #pragma omp parallel for schedule(static, 512) if(num_feature >= 1024)
      for (int i = 0; i < num_feature; ++i) {
Guolin Ke's avatar
Guolin Ke committed
462
        feature_metas_[i].num_bin = train_data->FeatureNumBin(i);
Guolin Ke's avatar
Guolin Ke committed
463
        feature_metas_[i].default_bin = train_data->FeatureBinMapper(i)->GetDefaultBin();
Guolin Ke's avatar
Guolin Ke committed
464
        feature_metas_[i].missing_type = train_data->FeatureBinMapper(i)->missing_type();
Guolin Ke's avatar
Guolin Ke committed
465
466
467
468
469
470
471
        if (train_data->FeatureBinMapper(i)->GetDefaultBin() == 0) {
          feature_metas_[i].bias = 1;
        } else {
          feature_metas_[i].bias = 0;
        }
        feature_metas_[i].tree_config = tree_config;
      }
Guolin Ke's avatar
Guolin Ke committed
472
    }
Guolin Ke's avatar
Guolin Ke committed
473
474
    uint64_t num_total_bin = train_data->NumTotalBin();
    Log::Info("Total Bins %d", num_total_bin);
Guolin Ke's avatar
Guolin Ke committed
475
    int old_cache_size = static_cast<int>(pool_.size());
Guolin Ke's avatar
Guolin Ke committed
476
    Reset(cache_size, total_size);
Guolin Ke's avatar
Guolin Ke committed
477
478
479
480
481
482

    if (cache_size > old_cache_size) {
      pool_.resize(cache_size);
      data_.resize(cache_size);
    }

483
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
484
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
485
    for (int i = old_cache_size; i < cache_size; ++i) {
486
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
487
488
489
490
491
      pool_[i].reset(new FeatureHistogram[train_data->num_features()]);
      data_[i].resize(num_total_bin);
      uint64_t offset = 0;
      for (int j = 0; j < train_data->num_features(); ++j) {
        offset += static_cast<uint64_t>(train_data->SubFeatureBinOffset(j));
492
        pool_[i][j].Init(data_[i].data() + offset, &feature_metas_[j], train_data->FeatureBinMapper(j)->bin_type());
Guolin Ke's avatar
Guolin Ke committed
493
494
495
496
497
498
499
        auto num_bin = train_data->FeatureNumBin(j);
        if (train_data->FeatureBinMapper(j)->GetDefaultBin() == 0) {
          num_bin -= 1;
        }
        offset += static_cast<uint64_t>(num_bin);
      }
      CHECK(offset == num_total_bin);
500
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
501
    }
502
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
503
504
  }

Guolin Ke's avatar
Guolin Ke committed
505
  void ResetConfig(const TreeConfig* tree_config) {
Guolin Ke's avatar
Guolin Ke committed
506
507
508
    int size = static_cast<int>(feature_metas_.size());
    #pragma omp parallel for schedule(static, 512) if(size >= 1024)
    for (int i = 0; i < size; ++i) {
Guolin Ke's avatar
Guolin Ke committed
509
      feature_metas_[i].tree_config = tree_config;
Guolin Ke's avatar
Guolin Ke committed
510
511
    }
  }
Guolin Ke's avatar
Guolin Ke committed
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
  /*!
  * \brief Get data for the specific index
  * \param idx which index want to get
  * \param out output data will store into this
  * \return True if this index is in the pool, False if this index is not in the pool
  */
  bool Get(int idx, FeatureHistogram** out) {
    if (is_enough_) {
      *out = pool_[idx].get();
      return true;
    } else if (mapper_[idx] >= 0) {
      int slot = mapper_[idx];
      *out = pool_[slot].get();
      last_used_time_[slot] = ++cur_time_;
      return true;
    } else {
      // choose the least used slot 
      int slot = static_cast<int>(ArrayArgs<int>::ArgMin(last_used_time_));
      *out = pool_[slot].get();
      last_used_time_[slot] = ++cur_time_;

      // reset previous mapper
      if (inverse_mapper_[slot] >= 0) mapper_[inverse_mapper_[slot]] = -1;

      // update current mapper
      mapper_[idx] = slot;
      inverse_mapper_[slot] = idx;
      return false;
    }
  }

  /*!
  * \brief Move data from one index to another index
  * \param src_idx
  * \param dst_idx
  */
  void Move(int src_idx, int dst_idx) {
    if (is_enough_) {
      std::swap(pool_[src_idx], pool_[dst_idx]);
      return;
    }
    if (mapper_[src_idx] < 0) {
      return;
    }
    // get slot of src idx
    int slot = mapper_[src_idx];
    // reset src_idx
    mapper_[src_idx] = -1;

    // move to dst idx
    mapper_[dst_idx] = slot;
    last_used_time_[slot] = ++cur_time_;
    inverse_mapper_[slot] = dst_idx;
  }
private:
  std::vector<std::unique_ptr<FeatureHistogram[]>> pool_;
Guolin Ke's avatar
Guolin Ke committed
568
569
  std::vector<std::vector<HistogramBinEntry>> data_;
  std::vector<FeatureMetainfo> feature_metas_;
Guolin Ke's avatar
Guolin Ke committed
570
571
572
573
574
575
576
577
578
  int cache_size_;
  int total_size_;
  bool is_enough_ = false;
  std::vector<int> mapper_;
  std::vector<int> inverse_mapper_;
  std::vector<int> last_used_time_;
  int cur_time_ = 0;
};

Guolin Ke's avatar
Guolin Ke committed
579
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
580
#endif   // LightGBM_TREELEARNER_FEATURE_HISTOGRAM_HPP_