c_api.cpp 72.9 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
Guolin Ke's avatar
Guolin Ke committed
5
#include <LightGBM/c_api.h>
Guolin Ke's avatar
Guolin Ke committed
6

Guolin Ke's avatar
Guolin Ke committed
7
8
#include <LightGBM/boosting.h>
#include <LightGBM/config.h>
9
10
11
#include <LightGBM/dataset.h>
#include <LightGBM/dataset_loader.h>
#include <LightGBM/metric.h>
12
#include <LightGBM/network.h>
13
14
15
16
17
18
19
#include <LightGBM/objective_function.h>
#include <LightGBM/prediction_early_stop.h>
#include <LightGBM/utils/common.h>
#include <LightGBM/utils/log.h>
#include <LightGBM/utils/openmp_wrapper.h>
#include <LightGBM/utils/random.h>
#include <LightGBM/utils/threading.h>
Guolin Ke's avatar
Guolin Ke committed
20
21

#include <string>
22
23
#include <cstdio>
#include <functional>
Guolin Ke's avatar
Guolin Ke committed
24
#include <memory>
wxchan's avatar
wxchan committed
25
#include <mutex>
26
27
#include <stdexcept>
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
28

29
#include "application/predictor.hpp"
Guolin Ke's avatar
Guolin Ke committed
30

Guolin Ke's avatar
Guolin Ke committed
31
32
namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
inline int LGBM_APIHandleException(const std::exception& ex) {
  LGBM_SetLastError(ex.what());
  return -1;
}
inline int LGBM_APIHandleException(const std::string& ex) {
  LGBM_SetLastError(ex.c_str());
  return -1;
}

#define API_BEGIN() try {
#define API_END() } \
catch(std::exception& ex) { return LGBM_APIHandleException(ex); } \
catch(std::string& ex) { return LGBM_APIHandleException(ex); } \
catch(...) { return LGBM_APIHandleException("unknown exception"); } \
return 0;

49
50
51
52
53
54
55
56
const int PREDICTOR_TYPES = 4;

// Single row predictor to abstract away caching logic
class SingleRowPredictor {
 public:
  PredictFunction predict_function;
  int64_t num_pred_in_one_row;

Guolin Ke's avatar
Guolin Ke committed
57
  SingleRowPredictor(int predict_type, Boosting* boosting, const Config& config, int iter) {
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    bool is_predict_leaf = false;
    bool is_raw_score = false;
    bool predict_contrib = false;
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
      is_raw_score = true;
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
      predict_contrib = true;
    } else {
      is_raw_score = false;
    }
    early_stop_ = config.pred_early_stop;
    early_stop_freq_ = config.pred_early_stop_freq;
    early_stop_margin_ = config.pred_early_stop_margin;
    iter_ = iter;
Guolin Ke's avatar
Guolin Ke committed
74
    predictor_.reset(new Predictor(boosting, iter_, is_raw_score, is_predict_leaf, predict_contrib,
75
                                   early_stop_, early_stop_freq_, early_stop_margin_));
Guolin Ke's avatar
Guolin Ke committed
76
    num_pred_in_one_row = boosting->NumPredictOneRow(iter_, is_predict_leaf, predict_contrib);
77
    predict_function = predictor_->GetPredictFunction();
Guolin Ke's avatar
Guolin Ke committed
78
    num_total_model_ = boosting->NumberOfTotalModel();
79
80
  }
  ~SingleRowPredictor() {}
Guolin Ke's avatar
Guolin Ke committed
81
  bool IsPredictorEqual(const Config& config, int iter, Boosting* boosting) {
82
83
84
85
    return early_stop_ != config.pred_early_stop ||
      early_stop_freq_ != config.pred_early_stop_freq ||
      early_stop_margin_ != config.pred_early_stop_margin ||
      iter_ != iter ||
Guolin Ke's avatar
Guolin Ke committed
86
      num_total_model_ != boosting->NumberOfTotalModel();
87
  }
Guolin Ke's avatar
Guolin Ke committed
88

89
90
91
92
93
94
95
96
97
 private:
  std::unique_ptr<Predictor> predictor_;
  bool early_stop_;
  int early_stop_freq_;
  double early_stop_margin_;
  int iter_;
  int num_total_model_;
};

Guolin Ke's avatar
Guolin Ke committed
98
class Booster {
Nikita Titov's avatar
Nikita Titov committed
99
 public:
Guolin Ke's avatar
Guolin Ke committed
100
  explicit Booster(const char* filename) {
101
    boosting_.reset(Boosting::CreateBoosting("gbdt", filename));
102
103
  }

Guolin Ke's avatar
Guolin Ke committed
104
  Booster(const Dataset* train_data,
105
          const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
106
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
107
    config_.Set(param);
108
109
110
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
111
    // create boosting
Guolin Ke's avatar
Guolin Ke committed
112
    if (config_.input_model.size() > 0) {
113
114
      Log::Warning("Continued train from model is not supported for c_api,\n"
                   "please use continued train with input score");
Guolin Ke's avatar
Guolin Ke committed
115
    }
Guolin Ke's avatar
Guolin Ke committed
116

Guolin Ke's avatar
Guolin Ke committed
117
    boosting_.reset(Boosting::CreateBoosting(config_.boosting, nullptr));
Guolin Ke's avatar
Guolin Ke committed
118

119
120
    train_data_ = train_data;
    CreateObjectiveAndMetrics();
Guolin Ke's avatar
Guolin Ke committed
121
    // initialize the boosting
Guolin Ke's avatar
Guolin Ke committed
122
    if (config_.tree_learner == std::string("feature")) {
123
      Log::Fatal("Do not support feature parallel in c api");
124
    }
Guolin Ke's avatar
Guolin Ke committed
125
    if (Network::num_machines() == 1 && config_.tree_learner != std::string("serial")) {
126
      Log::Warning("Only find one worker, will switch to serial tree learner");
Guolin Ke's avatar
Guolin Ke committed
127
      config_.tree_learner = "serial";
128
    }
Guolin Ke's avatar
Guolin Ke committed
129
    boosting_->Init(&config_, train_data_, objective_fun_.get(),
130
                    Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
131
132
133
134
135
  }

  void MergeFrom(const Booster* other) {
    std::lock_guard<std::mutex> lock(mutex_);
    boosting_->MergeFrom(other->boosting_.get());
Guolin Ke's avatar
Guolin Ke committed
136
137
138
139
  }

  ~Booster() {
  }
140

141
  void CreateObjectiveAndMetrics() {
Guolin Ke's avatar
Guolin Ke committed
142
    // create objective function
Guolin Ke's avatar
Guolin Ke committed
143
144
    objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                    config_));
Guolin Ke's avatar
Guolin Ke committed
145
146
147
148
149
150
151
152
153
154
    if (objective_fun_ == nullptr) {
      Log::Warning("Using self-defined objective function");
    }
    // initialize the objective function
    if (objective_fun_ != nullptr) {
      objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
    }

    // create training metric
    train_metric_.clear();
Guolin Ke's avatar
Guolin Ke committed
155
    for (auto metric_type : config_.metric) {
Guolin Ke's avatar
Guolin Ke committed
156
      auto metric = std::unique_ptr<Metric>(
Guolin Ke's avatar
Guolin Ke committed
157
        Metric::CreateMetric(metric_type, config_));
Guolin Ke's avatar
Guolin Ke committed
158
159
160
161
162
      if (metric == nullptr) { continue; }
      metric->Init(train_data_->metadata(), train_data_->num_data());
      train_metric_.push_back(std::move(metric));
    }
    train_metric_.shrink_to_fit();
163
164
165
166
167
168
169
170
171
172
173
  }

  void ResetTrainingData(const Dataset* train_data) {
    if (train_data != train_data_) {
      std::lock_guard<std::mutex> lock(mutex_);
      train_data_ = train_data;
      CreateObjectiveAndMetrics();
      // reset the boosting
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
    }
wxchan's avatar
wxchan committed
174
175
  }

176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
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
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
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
  static void CheckDatasetResetConfig(
      const Config& old_config,
      const std::unordered_map<std::string, std::string>& new_param) {
    Config new_config;
    new_config.Set(new_param);
    if (new_param.count("data_random_seed") &&
        new_config.data_random_seed != old_config.data_random_seed) {
      Log::Fatal("Cannot change data_random_seed after constructed Dataset handle.");
    }
    if (new_param.count("max_bin") &&
        new_config.max_bin != old_config.max_bin) {
      Log::Fatal("Cannot change max_bin after constructed Dataset handle.");
    }
    if (new_param.count("max_bin_by_feature") &&
        new_config.max_bin_by_feature != old_config.max_bin_by_feature) {
      Log::Fatal(
          "Cannot change max_bin_by_feature after constructed Dataset handle.");
    }
    if (new_param.count("bin_construct_sample_cnt") &&
        new_config.bin_construct_sample_cnt !=
            old_config.bin_construct_sample_cnt) {
      Log::Fatal(
          "Cannot change bin_construct_sample_cnt after constructed Dataset "
          "handle.");
    }
    if (new_param.count("min_data_in_bin") &&
        new_config.min_data_in_bin != old_config.min_data_in_bin) {
      Log::Fatal(
          "Cannot change min_data_in_bin after constructed Dataset handle.");
    }
    if (new_param.count("use_missing") &&
        new_config.use_missing != old_config.use_missing) {
      Log::Fatal("Cannot change use_missing after constructed Dataset handle.");
    }
    if (new_param.count("zero_as_missing") &&
        new_config.zero_as_missing != old_config.zero_as_missing) {
      Log::Fatal(
          "Cannot change zero_as_missing after constructed Dataset handle.");
    }
    if (new_param.count("categorical_feature") &&
        new_config.categorical_feature != old_config.categorical_feature) {
      Log::Fatal(
          "Cannot change categorical_feature after constructed Dataset "
          "handle.");
    }
    if (new_param.count("feature_pre_filter") &&
        new_config.feature_pre_filter != old_config.feature_pre_filter) {
      Log::Fatal(
          "Cannot change feature_pre_filter after constructed Dataset handle.");
    }
    if (new_param.count("is_enable_sparse") &&
        new_config.is_enable_sparse != old_config.is_enable_sparse) {
      Log::Fatal(
          "Cannot change is_enable_sparse after constructed Dataset handle.");
    }
    if (new_param.count("pre_partition") &&
        new_config.pre_partition != old_config.pre_partition) {
      Log::Fatal(
          "Cannot change pre_partition after constructed Dataset handle.");
    }
    if (new_param.count("enable_bundle") &&
        new_config.enable_bundle != old_config.enable_bundle) {
      Log::Fatal(
          "Cannot change enable_bundle after constructed Dataset handle.");
    }
    if (new_param.count("header") && new_config.header != old_config.header) {
      Log::Fatal("Cannot change header after constructed Dataset handle.");
    }
    if (new_param.count("two_round") &&
        new_config.two_round != old_config.two_round) {
      Log::Fatal("Cannot change two_round after constructed Dataset handle.");
    }
    if (new_param.count("label_column") &&
        new_config.label_column != old_config.label_column) {
      Log::Fatal(
          "Cannot change label_column after constructed Dataset handle.");
    }
    if (new_param.count("weight_column") &&
        new_config.weight_column != old_config.weight_column) {
      Log::Fatal(
          "Cannot change weight_column after constructed Dataset handle.");
    }
    if (new_param.count("group_column") &&
        new_config.group_column != old_config.group_column) {
      Log::Fatal(
          "Cannot change group_column after constructed Dataset handle.");
    }
    if (new_param.count("ignore_column") &&
        new_config.ignore_column != old_config.ignore_column) {
      Log::Fatal(
          "Cannot change ignore_column after constructed Dataset handle.");
    }
    if (new_param.count("forcedbins_filename")) {
      Log::Fatal("Cannot change forced bins after constructed Dataset handle.");
    }
    if (new_param.count("min_data_in_leaf") &&
        new_config.min_data_in_leaf < old_config.min_data_in_leaf &&
        old_config.feature_pre_filter) {
      Log::Fatal(
          "Reducing `min_data_in_leaf` with `feature_pre_filter=true` may "
          "cause unexpected behaviour "
          "for features that were pre-filtered by the larger "
          "`min_data_in_leaf`.\n"
          "You need to set `feature_pre_filter=false` to dynamically change "
          "the `min_data_in_leaf`.");
    }
  }

wxchan's avatar
wxchan committed
284
  void ResetConfig(const char* parameters) {
Guolin Ke's avatar
Guolin Ke committed
285
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
286
    auto param = Config::Str2Map(parameters);
wxchan's avatar
wxchan committed
287
    if (param.count("num_class")) {
288
      Log::Fatal("Cannot change num_class during training");
wxchan's avatar
wxchan committed
289
    }
Guolin Ke's avatar
Guolin Ke committed
290
291
    if (param.count("boosting")) {
      Log::Fatal("Cannot change boosting during training");
wxchan's avatar
wxchan committed
292
    }
Guolin Ke's avatar
Guolin Ke committed
293
    if (param.count("metric")) {
294
      Log::Fatal("Cannot change metric during training");
Guolin Ke's avatar
Guolin Ke committed
295
    }
Guolin Ke's avatar
Guolin Ke committed
296

297
298
    CheckDatasetResetConfig(config_, param);

Guolin Ke's avatar
Guolin Ke committed
299
    config_.Set(param);
300

301
302
303
    if (config_.num_threads > 0) {
      omp_set_num_threads(config_.num_threads);
    }
Guolin Ke's avatar
Guolin Ke committed
304
305
306

    if (param.count("objective")) {
      // create objective function
Guolin Ke's avatar
Guolin Ke committed
307
308
      objective_fun_.reset(ObjectiveFunction::CreateObjectiveFunction(config_.objective,
                                                                      config_));
Guolin Ke's avatar
Guolin Ke committed
309
310
311
312
313
314
315
      if (objective_fun_ == nullptr) {
        Log::Warning("Using self-defined objective function");
      }
      // initialize the objective function
      if (objective_fun_ != nullptr) {
        objective_fun_->Init(train_data_->metadata(), train_data_->num_data());
      }
316
317
      boosting_->ResetTrainingData(train_data_,
                                   objective_fun_.get(), Common::ConstPtrInVectorWrapper<Metric>(train_metric_));
wxchan's avatar
wxchan committed
318
    }
Guolin Ke's avatar
Guolin Ke committed
319

Guolin Ke's avatar
Guolin Ke committed
320
    boosting_->ResetConfig(&config_);
wxchan's avatar
wxchan committed
321
322
323
324
325
  }

  void AddValidData(const Dataset* valid_data) {
    std::lock_guard<std::mutex> lock(mutex_);
    valid_metrics_.emplace_back();
Guolin Ke's avatar
Guolin Ke committed
326
327
    for (auto metric_type : config_.metric) {
      auto metric = std::unique_ptr<Metric>(Metric::CreateMetric(metric_type, config_));
wxchan's avatar
wxchan committed
328
329
330
331
332
333
      if (metric == nullptr) { continue; }
      metric->Init(valid_data->metadata(), valid_data->num_data());
      valid_metrics_.back().push_back(std::move(metric));
    }
    valid_metrics_.back().shrink_to_fit();
    boosting_->AddValidDataset(valid_data,
334
                               Common::ConstPtrInVectorWrapper<Metric>(valid_metrics_.back()));
wxchan's avatar
wxchan committed
335
  }
Guolin Ke's avatar
Guolin Ke committed
336

337
  bool TrainOneIter() {
wxchan's avatar
wxchan committed
338
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
339
    return boosting_->TrainOneIter(nullptr, nullptr);
340
341
  }

Guolin Ke's avatar
Guolin Ke committed
342
343
344
345
346
347
348
349
350
351
352
  void Refit(const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<std::vector<int32_t>> v_leaf_preds(nrow, std::vector<int32_t>(ncol, 0));
    for (int i = 0; i < nrow; ++i) {
      for (int j = 0; j < ncol; ++j) {
        v_leaf_preds[i][j] = leaf_preds[i * ncol + j];
      }
    }
    boosting_->RefitTree(v_leaf_preds);
  }

353
  bool TrainOneIter(const score_t* gradients, const score_t* hessians) {
wxchan's avatar
wxchan committed
354
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
355
    return boosting_->TrainOneIter(gradients, hessians);
356
357
  }

wxchan's avatar
wxchan committed
358
359
360
361
362
  void RollbackOneIter() {
    std::lock_guard<std::mutex> lock(mutex_);
    boosting_->RollbackOneIter();
  }

363
  void PredictSingleRow(int num_iteration, int predict_type, int ncol,
364
365
366
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
               const Config& config,
               double* out_result, int64_t* out_len) {
367
368
369
    if (!config.predict_disable_shape_check && ncol != boosting_->MaxFeatureIdx() + 1) {
      Log::Fatal("The number of features in data (%d) is not the same as it was in training data (%d).\n"\
                 "You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.", ncol, boosting_->MaxFeatureIdx() + 1);
370
    }
371
    std::lock_guard<std::mutex> lock(mutex_);
372
    if (single_row_predictor_[predict_type].get() == nullptr ||
Guolin Ke's avatar
Guolin Ke committed
373
374
        !single_row_predictor_[predict_type]->IsPredictorEqual(config, num_iteration, boosting_.get())) {
      single_row_predictor_[predict_type].reset(new SingleRowPredictor(predict_type, boosting_.get(),
375
                                                                       config, num_iteration));
376
377
378
    }
    auto one_row = get_row_fun(0);
    auto pred_wrt_ptr = out_result;
379
    single_row_predictor_[predict_type]->predict_function(one_row, pred_wrt_ptr);
380

381
    *out_len = single_row_predictor_[predict_type]->num_pred_in_one_row;
382
383
384
  }


385
  void Predict(int num_iteration, int predict_type, int nrow, int ncol,
Guolin Ke's avatar
Guolin Ke committed
386
               std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
387
               const Config& config,
Guolin Ke's avatar
Guolin Ke committed
388
               double* out_result, int64_t* out_len) {
389
390
391
    if (!config.predict_disable_shape_check && ncol != boosting_->MaxFeatureIdx() + 1) {
      Log::Fatal("The number of features in data (%d) is not the same as it was in training data (%d).\n" \
                 "You can set ``predict_disable_shape_check=true`` to discard this error, but please be aware what you are doing.", ncol, boosting_->MaxFeatureIdx() + 1);
392
    }
wxchan's avatar
wxchan committed
393
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
394
395
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
396
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
397
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
Guolin Ke's avatar
Guolin Ke committed
398
      is_predict_leaf = true;
Guolin Ke's avatar
Guolin Ke committed
399
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
Guolin Ke's avatar
Guolin Ke committed
400
      is_raw_score = true;
401
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
402
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
403
404
    } else {
      is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
405
    }
Guolin Ke's avatar
Guolin Ke committed
406

Guolin Ke's avatar
Guolin Ke committed
407
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
408
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
409
    int64_t num_pred_in_one_row = boosting_->NumPredictOneRow(num_iteration, is_predict_leaf, predict_contrib);
Guolin Ke's avatar
Guolin Ke committed
410
    auto pred_fun = predictor.GetPredictFunction();
411
412
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
413
    for (int i = 0; i < nrow; ++i) {
414
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
415
      auto one_row = get_row_fun(i);
Tony-Y's avatar
Tony-Y committed
416
      auto pred_wrt_ptr = out_result + static_cast<size_t>(num_pred_in_one_row) * i;
Guolin Ke's avatar
Guolin Ke committed
417
      pred_fun(one_row, pred_wrt_ptr);
418
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
419
    }
420
    OMP_THROW_EX();
421
    *out_len = num_pred_in_one_row * nrow;
Guolin Ke's avatar
Guolin Ke committed
422
423
424
  }

  void Predict(int num_iteration, int predict_type, const char* data_filename,
Guolin Ke's avatar
Guolin Ke committed
425
               int data_has_header, const Config& config,
cbecker's avatar
cbecker committed
426
               const char* result_filename) {
Guolin Ke's avatar
Guolin Ke committed
427
428
429
    std::lock_guard<std::mutex> lock(mutex_);
    bool is_predict_leaf = false;
    bool is_raw_score = false;
Guolin Ke's avatar
Guolin Ke committed
430
    bool predict_contrib = false;
Guolin Ke's avatar
Guolin Ke committed
431
432
433
434
    if (predict_type == C_API_PREDICT_LEAF_INDEX) {
      is_predict_leaf = true;
    } else if (predict_type == C_API_PREDICT_RAW_SCORE) {
      is_raw_score = true;
435
    } else if (predict_type == C_API_PREDICT_CONTRIB) {
Guolin Ke's avatar
Guolin Ke committed
436
      predict_contrib = true;
Guolin Ke's avatar
Guolin Ke committed
437
438
439
    } else {
      is_raw_score = false;
    }
Guolin Ke's avatar
Guolin Ke committed
440
    Predictor predictor(boosting_.get(), num_iteration, is_raw_score, is_predict_leaf, predict_contrib,
441
                        config.pred_early_stop, config.pred_early_stop_freq, config.pred_early_stop_margin);
Guolin Ke's avatar
Guolin Ke committed
442
    bool bool_data_has_header = data_has_header > 0 ? true : false;
443
    predictor.Predict(data_filename, result_filename, bool_data_has_header, config.predict_disable_shape_check);
Guolin Ke's avatar
Guolin Ke committed
444
445
  }

Guolin Ke's avatar
Guolin Ke committed
446
  void GetPredictAt(int data_idx, double* out_result, int64_t* out_len) {
wxchan's avatar
wxchan committed
447
448
449
    boosting_->GetPredictAt(data_idx, out_result, out_len);
  }

450
451
  void SaveModelToFile(int start_iteration, int num_iteration, const char* filename) {
    boosting_->SaveModelToFile(start_iteration, num_iteration, filename);
Guolin Ke's avatar
Guolin Ke committed
452
  }
453

454
  void LoadModelFromString(const char* model_str) {
455
456
    size_t len = std::strlen(model_str);
    boosting_->LoadModelFromString(model_str, len);
457
458
  }

459
460
  std::string SaveModelToString(int start_iteration, int num_iteration) {
    return boosting_->SaveModelToString(start_iteration, num_iteration);
461
462
  }

463
  std::string DumpModel(int start_iteration, int num_iteration) {
464
    return boosting_->DumpModel(start_iteration, num_iteration);
wxchan's avatar
wxchan committed
465
  }
466

467
468
469
470
  std::vector<double> FeatureImportance(int num_iteration, int importance_type) {
    return boosting_->FeatureImportance(num_iteration, importance_type);
  }

471
472
473
474
475
476
477
478
479
480
  double UpperBoundValue() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return boosting_->GetUpperBoundValue();
  }

  double LowerBoundValue() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return boosting_->GetLowerBoundValue();
  }

Guolin Ke's avatar
Guolin Ke committed
481
  double GetLeafValue(int tree_idx, int leaf_idx) const {
Guolin Ke's avatar
Guolin Ke committed
482
    return dynamic_cast<GBDTBase*>(boosting_.get())->GetLeafValue(tree_idx, leaf_idx);
Guolin Ke's avatar
Guolin Ke committed
483
484
485
486
  }

  void SetLeafValue(int tree_idx, int leaf_idx, double val) {
    std::lock_guard<std::mutex> lock(mutex_);
Guolin Ke's avatar
Guolin Ke committed
487
    dynamic_cast<GBDTBase*>(boosting_.get())->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
488
489
  }

490
  void ShuffleModels(int start_iter, int end_iter) {
491
    std::lock_guard<std::mutex> lock(mutex_);
492
    boosting_->ShuffleModels(start_iter, end_iter);
493
494
  }

wxchan's avatar
wxchan committed
495
496
497
498
499
500
501
  int GetEvalCounts() const {
    int ret = 0;
    for (const auto& metric : train_metric_) {
      ret += static_cast<int>(metric->GetName().size());
    }
    return ret;
  }
502

wxchan's avatar
wxchan committed
503
504
505
506
  int GetEvalNames(char** out_strs) const {
    int idx = 0;
    for (const auto& metric : train_metric_) {
      for (const auto& name : metric->GetName()) {
Guolin Ke's avatar
Guolin Ke committed
507
        std::memcpy(out_strs[idx], name.c_str(), name.size() + 1);
wxchan's avatar
wxchan committed
508
509
510
511
512
513
        ++idx;
      }
    }
    return idx;
  }

wxchan's avatar
wxchan committed
514
515
516
  int GetFeatureNames(char** out_strs) const {
    int idx = 0;
    for (const auto& name : boosting_->FeatureNames()) {
Guolin Ke's avatar
Guolin Ke committed
517
      std::memcpy(out_strs[idx], name.c_str(), name.size() + 1);
wxchan's avatar
wxchan committed
518
519
520
521
522
      ++idx;
    }
    return idx;
  }

wxchan's avatar
wxchan committed
523
  const Boosting* GetBoosting() const { return boosting_.get(); }
Guolin Ke's avatar
Guolin Ke committed
524

Nikita Titov's avatar
Nikita Titov committed
525
 private:
wxchan's avatar
wxchan committed
526
  const Dataset* train_data_;
Guolin Ke's avatar
Guolin Ke committed
527
  std::unique_ptr<Boosting> boosting_;
528
  std::unique_ptr<SingleRowPredictor> single_row_predictor_[PREDICTOR_TYPES];
529

Guolin Ke's avatar
Guolin Ke committed
530
  /*! \brief All configs */
Guolin Ke's avatar
Guolin Ke committed
531
  Config config_;
Guolin Ke's avatar
Guolin Ke committed
532
  /*! \brief Metric for training data */
Guolin Ke's avatar
Guolin Ke committed
533
  std::vector<std::unique_ptr<Metric>> train_metric_;
Guolin Ke's avatar
Guolin Ke committed
534
  /*! \brief Metrics for validation data */
Guolin Ke's avatar
Guolin Ke committed
535
  std::vector<std::vector<std::unique_ptr<Metric>>> valid_metrics_;
Guolin Ke's avatar
Guolin Ke committed
536
  /*! \brief Training objective function */
Guolin Ke's avatar
Guolin Ke committed
537
  std::unique_ptr<ObjectiveFunction> objective_fun_;
wxchan's avatar
wxchan committed
538
  /*! \brief mutex for threading safe call */
539
  mutable std::mutex mutex_;
Guolin Ke's avatar
Guolin Ke committed
540
541
};

542
}  // namespace LightGBM
Guolin Ke's avatar
Guolin Ke committed
543
544
545

using namespace LightGBM;

Guolin Ke's avatar
Guolin Ke committed
546
547
548
549
550
551
552
553
// some help functions used to convert data

std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major);

554
555
556
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type);

Guolin Ke's avatar
Guolin Ke committed
557
558
std::function<std::vector<std::pair<int, double>>(int idx)>
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices,
559
                   const void* data, int data_type, int64_t nindptr, int64_t nelem);
Guolin Ke's avatar
Guolin Ke committed
560
561
562

// Row iterator of on column for CSC matrix
class CSC_RowIterator {
Nikita Titov's avatar
Nikita Titov committed
563
 public:
Guolin Ke's avatar
Guolin Ke committed
564
  CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
565
                  const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx);
Guolin Ke's avatar
Guolin Ke committed
566
567
568
569
570
  ~CSC_RowIterator() {}
  // return value at idx, only can access by ascent order
  double Get(int idx);
  // return next non-zero pair, if index < 0, means no more data
  std::pair<int, double> NextNonZero();
Nikita Titov's avatar
Nikita Titov committed
571
572

 private:
Guolin Ke's avatar
Guolin Ke committed
573
574
575
576
577
578
579
580
581
  int nonzero_idx_ = 0;
  int cur_idx_ = -1;
  double cur_val_ = 0.0f;
  bool is_end_ = false;
  std::function<std::pair<int, double>(int idx)> iter_fun_;
};

// start of c_api functions

Guolin Ke's avatar
Guolin Ke committed
582
const char* LGBM_GetLastError() {
wxchan's avatar
wxchan committed
583
  return LastErrorMsg();
Guolin Ke's avatar
Guolin Ke committed
584
585
}

Guolin Ke's avatar
Guolin Ke committed
586
int LGBM_DatasetCreateFromFile(const char* filename,
587
588
589
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
590
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
591
592
  auto param = Config::Str2Map(parameters);
  Config config;
593
594
595
596
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
597
  DatasetLoader loader(config, nullptr, 1, filename);
Guolin Ke's avatar
Guolin Ke committed
598
  if (reference == nullptr) {
599
    if (Network::num_machines() == 1) {
600
      *out = loader.LoadFromFile(filename);
601
    } else {
602
      *out = loader.LoadFromFile(filename, Network::rank(), Network::num_machines());
603
    }
Guolin Ke's avatar
Guolin Ke committed
604
  } else {
605
    *out = loader.LoadFromFileAlignWithOtherDataset(filename,
606
                                                    reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
607
  }
608
  API_END();
Guolin Ke's avatar
Guolin Ke committed
609
610
}

611

Guolin Ke's avatar
Guolin Ke committed
612
int LGBM_DatasetCreateFromSampledColumn(double** sample_data,
613
614
615
616
617
618
619
                                        int** sample_indices,
                                        int32_t ncol,
                                        const int* num_per_col,
                                        int32_t num_sample_row,
                                        int32_t num_total_row,
                                        const char* parameters,
                                        DatasetHandle* out) {
620
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
621
622
  auto param = Config::Str2Map(parameters);
  Config config;
623
624
625
626
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
627
  DatasetLoader loader(config, nullptr, 1, nullptr);
628
629
630
631
  *out = loader.CostructFromSampleData(sample_data, sample_indices, ncol, num_per_col,
                                       num_sample_row,
                                       static_cast<data_size_t>(num_total_row));
  API_END();
Guolin Ke's avatar
Guolin Ke committed
632
633
}

634

Guolin Ke's avatar
Guolin Ke committed
635
int LGBM_DatasetCreateByReference(const DatasetHandle reference,
636
637
                                  int64_t num_total_row,
                                  DatasetHandle* out) {
Guolin Ke's avatar
Guolin Ke committed
638
639
640
641
642
643
644
645
  API_BEGIN();
  std::unique_ptr<Dataset> ret;
  ret.reset(new Dataset(static_cast<data_size_t>(num_total_row)));
  ret->CreateValid(reinterpret_cast<const Dataset*>(reference));
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
646
int LGBM_DatasetPushRows(DatasetHandle dataset,
647
648
649
650
651
                         const void* data,
                         int data_type,
                         int32_t nrow,
                         int32_t ncol,
                         int32_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
652
653
654
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromDenseMatric(data, nrow, ncol, data_type, 1);
655
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
656
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
657
  for (int i = 0; i < nrow; ++i) {
658
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
659
660
661
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid, start_row + i, one_row);
662
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
663
  }
664
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
665
666
667
668
669
670
  if (start_row + nrow == p_dataset->num_data()) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
671
int LGBM_DatasetPushRowsByCSR(DatasetHandle dataset,
672
673
674
675
676
677
678
679
680
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
                              int64_t,
                              int64_t start_row) {
Guolin Ke's avatar
Guolin Ke committed
681
682
683
684
  API_BEGIN();
  auto p_dataset = reinterpret_cast<Dataset*>(dataset);
  auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
685
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
686
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
687
  for (int i = 0; i < nrow; ++i) {
688
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
689
690
691
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    p_dataset->PushOneRow(tid,
692
                          static_cast<data_size_t>(start_row + i), one_row);
693
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
694
  }
695
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
696
697
698
699
700
701
  if (start_row + nrow == static_cast<int64_t>(p_dataset->num_data())) {
    p_dataset->FinishLoad();
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
702
int LGBM_DatasetCreateFromMat(const void* data,
703
704
705
706
707
708
709
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
  return LGBM_DatasetCreateFromMats(1,
                                    &data,
                                    data_type,
                                    &nrow,
                                    ncol,
                                    is_row_major,
                                    parameters,
                                    reference,
                                    out);
}


int LGBM_DatasetCreateFromMats(int32_t nmat,
                               const void** data,
                               int data_type,
                               int32_t* nrow,
                               int32_t ncol,
                               int is_row_major,
                               const char* parameters,
                               const DatasetHandle reference,
                               DatasetHandle* out) {
731
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
732
733
  auto param = Config::Str2Map(parameters);
  Config config;
734
735
736
737
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
738
  std::unique_ptr<Dataset> ret;
739
740
741
742
743
744
745
746
747
  int32_t total_nrow = 0;
  for (int j = 0; j < nmat; ++j) {
    total_nrow += nrow[j];
  }

  std::vector<std::function<std::vector<double>(int row_idx)>> get_row_fun;
  for (int j = 0; j < nmat; ++j) {
    get_row_fun.push_back(RowFunctionFromDenseMatric(data[j], nrow[j], ncol, data_type, is_row_major));
  }
748

Guolin Ke's avatar
Guolin Ke committed
749
750
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
751
    Random rand(config.data_random_seed);
752
753
    int sample_cnt = static_cast<int>(total_nrow < config.bin_construct_sample_cnt ? total_nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(total_nrow, sample_cnt);
754
    sample_cnt = static_cast<int>(sample_indices.size());
755
    std::vector<std::vector<double>> sample_values(ncol);
Guolin Ke's avatar
Guolin Ke committed
756
    std::vector<std::vector<int>> sample_idx(ncol);
757
758
759

    int offset = 0;
    int j = 0;
Guolin Ke's avatar
Guolin Ke committed
760
    for (size_t i = 0; i < sample_indices.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
761
      auto idx = sample_indices[i];
762
763
764
765
      while ((idx - offset) >= nrow[j]) {
        offset += nrow[j];
        ++j;
      }
766

767
768
769
770
771
      auto row = get_row_fun[j](static_cast<int>(idx - offset));
      for (size_t k = 0; k < row.size(); ++k) {
        if (std::fabs(row[k]) > kZeroThreshold || std::isnan(row[k])) {
          sample_values[k].emplace_back(row[k]);
          sample_idx[k].emplace_back(static_cast<int>(i));
Guolin Ke's avatar
Guolin Ke committed
772
        }
Guolin Ke's avatar
Guolin Ke committed
773
774
      }
    }
Guolin Ke's avatar
Guolin Ke committed
775
    DatasetLoader loader(config, nullptr, 1, nullptr);
Guolin Ke's avatar
Guolin Ke committed
776
777
    ret.reset(loader.CostructFromSampleData(Common::Vector2Ptr<double>(&sample_values).data(),
                                            Common::Vector2Ptr<int>(&sample_idx).data(),
778
                                            ncol,
779
                                            Common::VectorSize<double>(sample_values).data(),
780
                                            sample_cnt, total_nrow));
Guolin Ke's avatar
Guolin Ke committed
781
  } else {
782
    ret.reset(new Dataset(total_nrow));
Guolin Ke's avatar
Guolin Ke committed
783
    ret->CreateValid(
784
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
785
  }
786
787
788
789
790
791
792
793
794
795
796
797
798
799
  int32_t start_row = 0;
  for (int j = 0; j < nmat; ++j) {
    OMP_INIT_EX();
    #pragma omp parallel for schedule(static)
    for (int i = 0; i < nrow[j]; ++i) {
      OMP_LOOP_EX_BEGIN();
      const int tid = omp_get_thread_num();
      auto one_row = get_row_fun[j](i);
      ret->PushOneRow(tid, start_row + i, one_row);
      OMP_LOOP_EX_END();
    }
    OMP_THROW_EX();

    start_row += nrow[j];
Guolin Ke's avatar
Guolin Ke committed
800
801
  }
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
802
  *out = ret.release();
803
  API_END();
804
805
}

Guolin Ke's avatar
Guolin Ke committed
806
int LGBM_DatasetCreateFromCSR(const void* indptr,
807
808
809
810
811
812
813
814
815
816
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
                              int64_t num_col,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
817
  API_BEGIN();
818
819
820
821
822
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
Guolin Ke's avatar
Guolin Ke committed
823
824
  auto param = Config::Str2Map(parameters);
  Config config;
825
826
827
828
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
829
  std::unique_ptr<Dataset> ret;
830
  auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
831
832
833
  int32_t nrow = static_cast<int32_t>(nindptr - 1);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
834
835
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
836
    auto sample_indices = rand.Sample(nrow, sample_cnt);
837
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
838
839
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
840
841
842
843
    for (size_t i = 0; i < sample_indices.size(); ++i) {
      auto idx = sample_indices[i];
      auto row = get_row_fun(static_cast<int>(idx));
      for (std::pair<int, double>& inner_data : row) {
Guolin Ke's avatar
Guolin Ke committed
844
        CHECK(inner_data.first < num_col);
Guolin Ke's avatar
Guolin Ke committed
845
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
Guolin Ke's avatar
Guolin Ke committed
846
847
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
848
849
850
        }
      }
    }
Guolin Ke's avatar
Guolin Ke committed
851
    DatasetLoader loader(config, nullptr, 1, nullptr);
Guolin Ke's avatar
Guolin Ke committed
852
853
    ret.reset(loader.CostructFromSampleData(Common::Vector2Ptr<double>(&sample_values).data(),
                                            Common::Vector2Ptr<int>(&sample_idx).data(),
854
                                            static_cast<int>(num_col),
855
856
                                            Common::VectorSize<double>(sample_values).data(),
                                            sample_cnt, nrow));
857
  } else {
858
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
859
    ret->CreateValid(
860
      reinterpret_cast<const Dataset*>(reference));
861
  }
862
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
863
  #pragma omp parallel for schedule(static)
864
  for (int i = 0; i < nindptr - 1; ++i) {
865
    OMP_LOOP_EX_BEGIN();
866
867
868
    const int tid = omp_get_thread_num();
    auto one_row = get_row_fun(i);
    ret->PushOneRow(tid, i, one_row);
869
    OMP_LOOP_EX_END();
870
  }
871
  OMP_THROW_EX();
872
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
873
  *out = ret.release();
874
  API_END();
875
876
}

877
int LGBM_DatasetCreateFromCSRFunc(void* get_row_funptr,
878
879
880
881
882
                                  int num_rows,
                                  int64_t num_col,
                                  const char* parameters,
                                  const DatasetHandle reference,
                                  DatasetHandle* out) {
883
  API_BEGIN();
884
885
886
887
888
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
  auto get_row_fun = *static_cast<std::function<void(int idx, std::vector<std::pair<int, double>>&)>*>(get_row_funptr);
  auto param = Config::Str2Map(parameters);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  std::unique_ptr<Dataset> ret;
  int32_t nrow = num_rows;
  if (reference == nullptr) {
    // sample data first
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
    auto sample_indices = rand.Sample(nrow, sample_cnt);
    sample_cnt = static_cast<int>(sample_indices.size());
    std::vector<std::vector<double>> sample_values(num_col);
    std::vector<std::vector<int>> sample_idx(num_col);
    // local buffer to re-use memory
    std::vector<std::pair<int, double>> buffer;
    for (size_t i = 0; i < sample_indices.size(); ++i) {
      auto idx = sample_indices[i];
      get_row_fun(static_cast<int>(idx), buffer);
      for (std::pair<int, double>& inner_data : buffer) {
        CHECK(inner_data.first < num_col);
        if (std::fabs(inner_data.second) > kZeroThreshold || std::isnan(inner_data.second)) {
          sample_values[inner_data.first].emplace_back(inner_data.second);
          sample_idx[inner_data.first].emplace_back(static_cast<int>(i));
        }
      }
    }
    DatasetLoader loader(config, nullptr, 1, nullptr);
Guolin Ke's avatar
Guolin Ke committed
920
921
    ret.reset(loader.CostructFromSampleData(Common::Vector2Ptr<double>(&sample_values).data(),
                                            Common::Vector2Ptr<int>(&sample_idx).data(),
922
                                            static_cast<int>(num_col),
923
924
925
926
927
928
929
                                            Common::VectorSize<double>(sample_values).data(),
                                            sample_cnt, nrow));
  } else {
    ret.reset(new Dataset(nrow));
    ret->CreateValid(
      reinterpret_cast<const Dataset*>(reference));
  }
930

931
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
932
933
  std::vector<std::pair<int, double>> thread_buffer;
  #pragma omp parallel for schedule(static) private(thread_buffer)
934
935
936
  for (int i = 0; i < num_rows; ++i) {
    OMP_LOOP_EX_BEGIN();
    {
937
      const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
938
939
      get_row_fun(i, thread_buffer);
      ret->PushOneRow(tid, i, thread_buffer);
940
941
942
943
944
945
946
947
948
    }
    OMP_LOOP_EX_END();
  }
  OMP_THROW_EX();
  ret->FinishLoad();
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
949
int LGBM_DatasetCreateFromCSC(const void* col_ptr,
950
951
952
953
954
955
956
957
958
959
                              int col_ptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t ncol_ptr,
                              int64_t nelem,
                              int64_t num_row,
                              const char* parameters,
                              const DatasetHandle reference,
                              DatasetHandle* out) {
960
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
961
962
  auto param = Config::Str2Map(parameters);
  Config config;
963
964
965
966
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
967
  std::unique_ptr<Dataset> ret;
Guolin Ke's avatar
Guolin Ke committed
968
969
970
  int32_t nrow = static_cast<int32_t>(num_row);
  if (reference == nullptr) {
    // sample data first
Guolin Ke's avatar
Guolin Ke committed
971
972
    Random rand(config.data_random_seed);
    int sample_cnt = static_cast<int>(nrow < config.bin_construct_sample_cnt ? nrow : config.bin_construct_sample_cnt);
Guolin Ke's avatar
Guolin Ke committed
973
    auto sample_indices = rand.Sample(nrow, sample_cnt);
974
    sample_cnt = static_cast<int>(sample_indices.size());
Guolin Ke's avatar
Guolin Ke committed
975
    std::vector<std::vector<double>> sample_values(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
976
    std::vector<std::vector<int>> sample_idx(ncol_ptr - 1);
977
    OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
978
    #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
979
    for (int i = 0; i < static_cast<int>(sample_values.size()); ++i) {
980
      OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
981
982
983
      CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
      for (int j = 0; j < sample_cnt; j++) {
        auto val = col_it.Get(sample_indices[j]);
Guolin Ke's avatar
Guolin Ke committed
984
        if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
Guolin Ke's avatar
Guolin Ke committed
985
986
          sample_values[i].emplace_back(val);
          sample_idx[i].emplace_back(j);
Guolin Ke's avatar
Guolin Ke committed
987
988
        }
      }
989
      OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
990
    }
991
    OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
992
    DatasetLoader loader(config, nullptr, 1, nullptr);
Guolin Ke's avatar
Guolin Ke committed
993
994
    ret.reset(loader.CostructFromSampleData(Common::Vector2Ptr<double>(&sample_values).data(),
                                            Common::Vector2Ptr<int>(&sample_idx).data(),
995
996
997
                                            static_cast<int>(sample_values.size()),
                                            Common::VectorSize<double>(sample_values).data(),
                                            sample_cnt, nrow));
Guolin Ke's avatar
Guolin Ke committed
998
  } else {
999
    ret.reset(new Dataset(nrow));
Guolin Ke's avatar
Guolin Ke committed
1000
    ret->CreateValid(
1001
      reinterpret_cast<const Dataset*>(reference));
Guolin Ke's avatar
Guolin Ke committed
1002
  }
1003
  OMP_INIT_EX();
Guolin Ke's avatar
Guolin Ke committed
1004
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
1005
  for (int i = 0; i < ncol_ptr - 1; ++i) {
1006
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1007
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
1008
    int feature_idx = ret->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
1009
    if (feature_idx < 0) { continue; }
Guolin Ke's avatar
Guolin Ke committed
1010
1011
    int group = ret->Feature2Group(feature_idx);
    int sub_feature = ret->Feture2SubFeature(feature_idx);
Guolin Ke's avatar
Guolin Ke committed
1012
    CSC_RowIterator col_it(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, i);
Guolin Ke's avatar
Guolin Ke committed
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
    auto bin_mapper = ret->FeatureBinMapper(feature_idx);
    if (bin_mapper->GetDefaultBin() == bin_mapper->GetMostFreqBin()) {
      int row_idx = 0;
      while (row_idx < nrow) {
        auto pair = col_it.NextNonZero();
        row_idx = pair.first;
        // no more data
        if (row_idx < 0) { break; }
        ret->PushOneData(tid, row_idx, group, sub_feature, pair.second);
      }
    } else {
      for (int row_idx = 0; row_idx < nrow; ++row_idx) {
        auto val = col_it.Get(row_idx);
        ret->PushOneData(tid, row_idx, group, sub_feature, val);
      }
Guolin Ke's avatar
Guolin Ke committed
1028
    }
1029
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
1030
  }
1031
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
1032
  ret->FinishLoad();
Guolin Ke's avatar
Guolin Ke committed
1033
  *out = ret.release();
1034
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1035
1036
}

Guolin Ke's avatar
Guolin Ke committed
1037
int LGBM_DatasetGetSubset(
1038
  const DatasetHandle handle,
wxchan's avatar
wxchan committed
1039
1040
1041
  const int32_t* used_row_indices,
  int32_t num_used_row_indices,
  const char* parameters,
Guolin Ke's avatar
typo  
Guolin Ke committed
1042
  DatasetHandle* out) {
wxchan's avatar
wxchan committed
1043
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1044
1045
  auto param = Config::Str2Map(parameters);
  Config config;
1046
1047
1048
1049
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
1050
  auto full_dataset = reinterpret_cast<const Dataset*>(handle);
Guolin Ke's avatar
Guolin Ke committed
1051
  CHECK(num_used_row_indices > 0);
1052
1053
1054
  const int32_t lower = 0;
  const int32_t upper = full_dataset->num_data() - 1;
  Common::CheckElementsIntervalClosed(used_row_indices, lower, upper, num_used_row_indices, "Used indices of subset");
1055
1056
1057
  if (!std::is_sorted(used_row_indices, used_row_indices + num_used_row_indices)) {
    Log::Fatal("used_row_indices should be sorted in Subset");
  }
Guolin Ke's avatar
Guolin Ke committed
1058
  auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
1059
  ret->CopyFeatureMapperFrom(full_dataset);
Guolin Ke's avatar
Guolin Ke committed
1060
  ret->CopySubset(full_dataset, used_row_indices, num_used_row_indices, true);
wxchan's avatar
wxchan committed
1061
1062
1063
1064
  *out = ret.release();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1065
int LGBM_DatasetSetFeatureNames(
Guolin Ke's avatar
typo  
Guolin Ke committed
1066
  DatasetHandle handle,
Guolin Ke's avatar
Guolin Ke committed
1067
  const char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1068
  int num_feature_names) {
Guolin Ke's avatar
Guolin Ke committed
1069
1070
1071
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  std::vector<std::string> feature_names_str;
Guolin Ke's avatar
Guolin Ke committed
1072
  for (int i = 0; i < num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1073
1074
1075
1076
1077
1078
    feature_names_str.emplace_back(feature_names[i]);
  }
  dataset->set_feature_names(feature_names_str);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1079
int LGBM_DatasetGetFeatureNames(
1080
1081
  DatasetHandle handle,
  char** feature_names,
Guolin Ke's avatar
Guolin Ke committed
1082
  int* num_feature_names) {
1083
1084
1085
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  auto inside_feature_name = dataset->feature_names();
Guolin Ke's avatar
Guolin Ke committed
1086
1087
  *num_feature_names = static_cast<int>(inside_feature_name.size());
  for (int i = 0; i < *num_feature_names; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1088
    std::memcpy(feature_names[i], inside_feature_name[i].c_str(), inside_feature_name[i].size() + 1);
1089
1090
1091
1092
  }
  API_END();
}

1093
#pragma warning(disable : 4702)
Guolin Ke's avatar
Guolin Ke committed
1094
int LGBM_DatasetFree(DatasetHandle handle) {
1095
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1096
  delete reinterpret_cast<Dataset*>(handle);
1097
  API_END();
1098
1099
}

Guolin Ke's avatar
Guolin Ke committed
1100
int LGBM_DatasetSaveBinary(DatasetHandle handle,
1101
                           const char* filename) {
1102
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1103
1104
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->SaveBinaryFile(filename);
1105
  API_END();
1106
1107
}

1108
1109
1110
1111
1112
1113
1114
1115
int LGBM_DatasetDumpText(DatasetHandle handle,
                         const char* filename) {
  API_BEGIN();
  auto dataset = reinterpret_cast<Dataset*>(handle);
  dataset->DumpTextFile(filename);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1116
int LGBM_DatasetSetField(DatasetHandle handle,
1117
1118
1119
1120
                         const char* field_name,
                         const void* field_data,
                         int num_element,
                         int type) {
1121
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1122
  auto dataset = reinterpret_cast<Dataset*>(handle);
1123
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1124
  if (type == C_API_DTYPE_FLOAT32) {
Guolin Ke's avatar
Guolin Ke committed
1125
    is_success = dataset->SetFloatField(field_name, reinterpret_cast<const float*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1126
  } else if (type == C_API_DTYPE_INT32) {
Guolin Ke's avatar
Guolin Ke committed
1127
    is_success = dataset->SetIntField(field_name, reinterpret_cast<const int*>(field_data), static_cast<int32_t>(num_element));
Guolin Ke's avatar
Guolin Ke committed
1128
1129
  } else if (type == C_API_DTYPE_FLOAT64) {
    is_success = dataset->SetDoubleField(field_name, reinterpret_cast<const double*>(field_data), static_cast<int32_t>(num_element));
1130
  }
1131
  if (!is_success) { throw std::runtime_error("Input data type error or field not found"); }
1132
  API_END();
1133
1134
}

Guolin Ke's avatar
Guolin Ke committed
1135
int LGBM_DatasetGetField(DatasetHandle handle,
1136
1137
1138
1139
                         const char* field_name,
                         int* out_len,
                         const void** out_ptr,
                         int* out_type) {
1140
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1141
  auto dataset = reinterpret_cast<Dataset*>(handle);
1142
  bool is_success = false;
Guolin Ke's avatar
Guolin Ke committed
1143
  if (dataset->GetFloatField(field_name, out_len, reinterpret_cast<const float**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1144
    *out_type = C_API_DTYPE_FLOAT32;
1145
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1146
  } else if (dataset->GetIntField(field_name, out_len, reinterpret_cast<const int**>(out_ptr))) {
Guolin Ke's avatar
Guolin Ke committed
1147
    *out_type = C_API_DTYPE_INT32;
1148
    is_success = true;
Guolin Ke's avatar
Guolin Ke committed
1149
1150
1151
  } else if (dataset->GetDoubleField(field_name, out_len, reinterpret_cast<const double**>(out_ptr))) {
    *out_type = C_API_DTYPE_FLOAT64;
    is_success = true;
1152
  } 
1153
  if (!is_success) { throw std::runtime_error("Field not found"); }
wxchan's avatar
wxchan committed
1154
  if (*out_ptr == nullptr) { *out_len = 0; }
1155
  API_END();
1156
1157
}

1158
int LGBM_DatasetUpdateParamChecking(const char* old_parameters, const char* new_parameters) {
1159
  API_BEGIN();
1160
1161
1162
1163
1164
  auto old_param = Config::Str2Map(old_parameters);
  Config old_config;
  old_config.Set(old_param);
  auto new_param = Config::Str2Map(new_parameters);
  Booster::CheckDatasetResetConfig(old_config, new_param);
1165
1166
1167
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1168
int LGBM_DatasetGetNumData(DatasetHandle handle,
1169
                           int* out) {
1170
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1171
1172
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_data();
1173
  API_END();
1174
1175
}

Guolin Ke's avatar
Guolin Ke committed
1176
int LGBM_DatasetGetNumFeature(DatasetHandle handle,
1177
                              int* out) {
1178
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1179
1180
  auto dataset = reinterpret_cast<Dataset*>(handle);
  *out = dataset->num_total_features();
1181
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1182
}
1183

1184
1185
1186
1187
1188
int LGBM_DatasetAddFeaturesFrom(DatasetHandle target,
                                DatasetHandle source) {
  API_BEGIN();
  auto target_d = reinterpret_cast<Dataset*>(target);
  auto source_d = reinterpret_cast<Dataset*>(source);
1189
  target_d->AddFeaturesFrom(source_d);
1190
1191
1192
  API_END();
}

1193
1194
// ---- start of booster

Guolin Ke's avatar
Guolin Ke committed
1195
int LGBM_BoosterCreate(const DatasetHandle train_data,
1196
1197
                       const char* parameters,
                       BoosterHandle* out) {
1198
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1199
  const Dataset* p_train_data = reinterpret_cast<const Dataset*>(train_data);
wxchan's avatar
wxchan committed
1200
1201
  auto ret = std::unique_ptr<Booster>(new Booster(p_train_data, parameters));
  *out = ret.release();
1202
  API_END();
1203
1204
}

Guolin Ke's avatar
Guolin Ke committed
1205
int LGBM_BoosterCreateFromModelfile(
1206
  const char* filename,
Guolin Ke's avatar
Guolin Ke committed
1207
  int* out_num_iterations,
1208
  BoosterHandle* out) {
1209
  API_BEGIN();
wxchan's avatar
wxchan committed
1210
  auto ret = std::unique_ptr<Booster>(new Booster(filename));
Guolin Ke's avatar
Guolin Ke committed
1211
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
wxchan's avatar
wxchan committed
1212
  *out = ret.release();
1213
  API_END();
1214
1215
}

Guolin Ke's avatar
Guolin Ke committed
1216
int LGBM_BoosterLoadModelFromString(
1217
1218
1219
1220
  const char* model_str,
  int* out_num_iterations,
  BoosterHandle* out) {
  API_BEGIN();
wxchan's avatar
wxchan committed
1221
  auto ret = std::unique_ptr<Booster>(new Booster(nullptr));
1222
1223
1224
1225
1226
1227
  ret->LoadModelFromString(model_str);
  *out_num_iterations = ret->GetBoosting()->GetCurrentIteration();
  *out = ret.release();
  API_END();
}

1228
#pragma warning(disable : 4702)
Guolin Ke's avatar
Guolin Ke committed
1229
int LGBM_BoosterFree(BoosterHandle handle) {
1230
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1231
  delete reinterpret_cast<Booster*>(handle);
1232
  API_END();
1233
1234
}

1235
int LGBM_BoosterShuffleModels(BoosterHandle handle, int start_iter, int end_iter) {
1236
1237
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1238
  ref_booster->ShuffleModels(start_iter, end_iter);
1239
1240
1241
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1242
int LGBM_BoosterMerge(BoosterHandle handle,
1243
                      BoosterHandle other_handle) {
wxchan's avatar
wxchan committed
1244
1245
1246
1247
1248
1249
1250
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  Booster* ref_other_booster = reinterpret_cast<Booster*>(other_handle);
  ref_booster->MergeFrom(ref_other_booster);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1251
int LGBM_BoosterAddValidData(BoosterHandle handle,
1252
                             const DatasetHandle valid_data) {
wxchan's avatar
wxchan committed
1253
1254
1255
1256
1257
1258
1259
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(valid_data);
  ref_booster->AddValidData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1260
int LGBM_BoosterResetTrainingData(BoosterHandle handle,
1261
                                  const DatasetHandle train_data) {
wxchan's avatar
wxchan committed
1262
1263
1264
1265
1266
1267
1268
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  const Dataset* p_dataset = reinterpret_cast<const Dataset*>(train_data);
  ref_booster->ResetTrainingData(p_dataset);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1269
int LGBM_BoosterResetParameter(BoosterHandle handle, const char* parameters) {
wxchan's avatar
wxchan committed
1270
1271
1272
1273
1274
1275
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->ResetConfig(parameters);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1276
int LGBM_BoosterGetNumClasses(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1277
1278
1279
1280
1281
1282
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->NumberOfClasses();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1283
1284
1285
1286
1287
1288
1289
int LGBM_BoosterRefit(BoosterHandle handle, const int32_t* leaf_preds, int32_t nrow, int32_t ncol) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->Refit(leaf_preds, nrow, ncol);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1290
int LGBM_BoosterUpdateOneIter(BoosterHandle handle, int* is_finished) {
1291
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1292
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1293
1294
1295
1296
1297
  if (ref_booster->TrainOneIter()) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1298
  API_END();
1299
1300
}

Guolin Ke's avatar
Guolin Ke committed
1301
int LGBM_BoosterUpdateOneIterCustom(BoosterHandle handle,
1302
1303
1304
                                    const float* grad,
                                    const float* hess,
                                    int* is_finished) {
1305
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1306
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1307
  #ifdef SCORE_T_USE_DOUBLE
1308
  Log::Fatal("Don't support custom loss function when SCORE_T_USE_DOUBLE is enabled");
1309
  #else
1310
1311
1312
1313
1314
  if (ref_booster->TrainOneIter(grad, hess)) {
    *is_finished = 1;
  } else {
    *is_finished = 0;
  }
1315
  #endif
1316
  API_END();
1317
1318
}

Guolin Ke's avatar
Guolin Ke committed
1319
int LGBM_BoosterRollbackOneIter(BoosterHandle handle) {
wxchan's avatar
wxchan committed
1320
1321
1322
1323
1324
1325
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  ref_booster->RollbackOneIter();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1326
int LGBM_BoosterGetCurrentIteration(BoosterHandle handle, int* out_iteration) {
wxchan's avatar
wxchan committed
1327
1328
1329
1330
1331
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_iteration = ref_booster->GetBoosting()->GetCurrentIteration();
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1332

1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
int LGBM_BoosterNumModelPerIteration(BoosterHandle handle, int* out_tree_per_iteration) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_tree_per_iteration = ref_booster->GetBoosting()->NumModelPerIteration();
  API_END();
}

int LGBM_BoosterNumberOfTotalModel(BoosterHandle handle, int* out_models) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_models = ref_booster->GetBoosting()->NumberOfTotalModel();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1347
int LGBM_BoosterGetEvalCounts(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1348
1349
1350
1351
1352
1353
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalCounts();
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1354
int LGBM_BoosterGetEvalNames(BoosterHandle handle, int* out_len, char** out_strs) {
wxchan's avatar
wxchan committed
1355
1356
1357
1358
1359
1360
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetEvalNames(out_strs);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1361
int LGBM_BoosterGetFeatureNames(BoosterHandle handle, int* out_len, char** out_strs) {
wxchan's avatar
wxchan committed
1362
1363
1364
1365
1366
1367
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetFeatureNames(out_strs);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1368
int LGBM_BoosterGetNumFeature(BoosterHandle handle, int* out_len) {
wxchan's avatar
wxchan committed
1369
1370
1371
1372
1373
1374
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  *out_len = ref_booster->GetBoosting()->MaxFeatureIdx() + 1;
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1375
int LGBM_BoosterGetEval(BoosterHandle handle,
1376
1377
1378
                        int data_idx,
                        int* out_len,
                        double* out_results) {
1379
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1380
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1381
  auto boosting = ref_booster->GetBoosting();
wxchan's avatar
wxchan committed
1382
  auto result_buf = boosting->GetEvalAt(data_idx);
Guolin Ke's avatar
Guolin Ke committed
1383
  *out_len = static_cast<int>(result_buf.size());
1384
  for (size_t i = 0; i < result_buf.size(); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1385
    (out_results)[i] = static_cast<double>(result_buf[i]);
1386
  }
1387
  API_END();
1388
1389
}

Guolin Ke's avatar
Guolin Ke committed
1390
int LGBM_BoosterGetNumPredict(BoosterHandle handle,
1391
1392
                              int data_idx,
                              int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1393
1394
1395
1396
1397
1398
  API_BEGIN();
  auto boosting = reinterpret_cast<Booster*>(handle)->GetBoosting();
  *out_len = boosting->GetNumPredictAt(data_idx);
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1399
int LGBM_BoosterGetPredict(BoosterHandle handle,
1400
1401
1402
                           int data_idx,
                           int64_t* out_len,
                           double* out_result) {
1403
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1404
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1405
  ref_booster->GetPredictAt(data_idx, out_result, out_len);
1406
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1407
1408
}

Guolin Ke's avatar
Guolin Ke committed
1409
int LGBM_BoosterPredictForFile(BoosterHandle handle,
1410
1411
1412
1413
                               const char* data_filename,
                               int data_has_header,
                               int predict_type,
                               int num_iteration,
1414
                               const char* parameter,
1415
                               const char* result_filename) {
1416
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1417
1418
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1419
1420
1421
1422
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1423
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
cbecker's avatar
cbecker committed
1424
  ref_booster->Predict(num_iteration, predict_type, data_filename, data_has_header,
Guolin Ke's avatar
Guolin Ke committed
1425
                       config, result_filename);
1426
  API_END();
1427
1428
}

Guolin Ke's avatar
Guolin Ke committed
1429
int LGBM_BoosterCalcNumPredict(BoosterHandle handle,
1430
1431
1432
1433
                               int num_row,
                               int predict_type,
                               int num_iteration,
                               int64_t* out_len) {
Guolin Ke's avatar
Guolin Ke committed
1434
1435
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1436
1437
  *out_len = static_cast<int64_t>(num_row) * ref_booster->GetBoosting()->NumPredictOneRow(
    num_iteration, predict_type == C_API_PREDICT_LEAF_INDEX, predict_type == C_API_PREDICT_CONTRIB);
Guolin Ke's avatar
Guolin Ke committed
1438
1439
1440
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1441
int LGBM_BoosterPredictForCSR(BoosterHandle handle,
1442
1443
1444
1445
1446
1447
1448
                              const void* indptr,
                              int indptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t nindptr,
                              int64_t nelem,
1449
                              int64_t num_col,
1450
1451
                              int predict_type,
                              int num_iteration,
1452
                              const char* parameter,
1453
1454
                              int64_t* out_len,
                              double* out_result) {
1455
  API_BEGIN();
1456
1457
1458
1459
1460
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
Guolin Ke's avatar
Guolin Ke committed
1461
1462
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1463
1464
1465
1466
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1467
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1468
  auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
Guolin Ke's avatar
Guolin Ke committed
1469
  int nrow = static_cast<int>(nindptr - 1);
1470
  ref_booster->Predict(num_iteration, predict_type, nrow, static_cast<int>(num_col), get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
1471
                       config, out_result, out_len);
1472
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1473
}
1474

1475
int LGBM_BoosterPredictForCSRSingleRow(BoosterHandle handle,
1476
1477
1478
1479
1480
1481
1482
                                       const void* indptr,
                                       int indptr_type,
                                       const int32_t* indices,
                                       const void* data,
                                       int data_type,
                                       int64_t nindptr,
                                       int64_t nelem,
1483
                                       int64_t num_col,
1484
1485
1486
1487
1488
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1489
  API_BEGIN();
1490
1491
1492
1493
1494
  if (num_col <= 0) {
    Log::Fatal("The number of columns should be greater than zero.");
  } else if (num_col >= INT32_MAX) {
    Log::Fatal("The number of columns should be smaller than INT32_MAX.");
  }
1495
1496
1497
1498
1499
1500
1501
1502
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowFunctionFromCSR(indptr, indptr_type, indices, data, data_type, nindptr, nelem);
1503
  ref_booster->PredictSingleRow(num_iteration, predict_type, static_cast<int32_t>(num_col), get_row_fun, config, out_result, out_len);
1504
1505
1506
1507
  API_END();
}


Guolin Ke's avatar
Guolin Ke committed
1508
int LGBM_BoosterPredictForCSC(BoosterHandle handle,
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
                              const void* col_ptr,
                              int col_ptr_type,
                              const int32_t* indices,
                              const void* data,
                              int data_type,
                              int64_t ncol_ptr,
                              int64_t nelem,
                              int64_t num_row,
                              int predict_type,
                              int num_iteration,
1519
                              const char* parameter,
1520
1521
                              int64_t* out_len,
                              double* out_result) {
Guolin Ke's avatar
Guolin Ke committed
1522
1523
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
1524
1525
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  int num_threads = 1;
  #pragma omp parallel
  #pragma omp master
  {
    num_threads = omp_get_num_threads();
  }
Guolin Ke's avatar
Guolin Ke committed
1536
  int ncol = static_cast<int>(ncol_ptr - 1);
Guolin Ke's avatar
Guolin Ke committed
1537
1538
1539
1540
1541
  std::vector<std::vector<CSC_RowIterator>> iterators(num_threads, std::vector<CSC_RowIterator>());
  for (int i = 0; i < num_threads; ++i) {
    for (int j = 0; j < ncol; ++j) {
      iterators[i].emplace_back(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, j);
    }
Guolin Ke's avatar
Guolin Ke committed
1542
1543
  }
  std::function<std::vector<std::pair<int, double>>(int row_idx)> get_row_fun =
Guolin Ke's avatar
Guolin Ke committed
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
      [&iterators, ncol](int i) {
        std::vector<std::pair<int, double>> one_row;
        one_row.reserve(ncol);
        const int tid = omp_get_thread_num();
        for (int j = 0; j < ncol; ++j) {
          auto val = iterators[tid][j].Get(i);
          if (std::fabs(val) > kZeroThreshold || std::isnan(val)) {
            one_row.emplace_back(j, val);
          }
        }
        return one_row;
      };
1556
  ref_booster->Predict(num_iteration, predict_type, static_cast<int>(num_row), ncol, get_row_fun, config,
cbecker's avatar
cbecker committed
1557
                       out_result, out_len);
Guolin Ke's avatar
Guolin Ke committed
1558
1559
1560
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1561
int LGBM_BoosterPredictForMat(BoosterHandle handle,
1562
1563
1564
1565
1566
1567
1568
                              const void* data,
                              int data_type,
                              int32_t nrow,
                              int32_t ncol,
                              int is_row_major,
                              int predict_type,
                              int num_iteration,
1569
                              const char* parameter,
1570
1571
                              int64_t* out_len,
                              double* out_result) {
1572
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1573
1574
  auto param = Config::Str2Map(parameter);
  Config config;
Guolin Ke's avatar
Guolin Ke committed
1575
1576
1577
1578
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
Guolin Ke's avatar
Guolin Ke committed
1579
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1580
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, nrow, ncol, data_type, is_row_major);
1581
  ref_booster->Predict(num_iteration, predict_type, nrow, ncol, get_row_fun,
Guolin Ke's avatar
Guolin Ke committed
1582
                       config, out_result, out_len);
1583
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1584
}
1585

1586
int LGBM_BoosterPredictForMatSingleRow(BoosterHandle handle,
1587
1588
1589
1590
1591
1592
1593
1594
1595
                                       const void* data,
                                       int data_type,
                                       int32_t ncol,
                                       int is_row_major,
                                       int predict_type,
                                       int num_iteration,
                                       const char* parameter,
                                       int64_t* out_len,
                                       double* out_result) {
1596
1597
1598
1599
1600
1601
1602
1603
1604
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseMatric(data, 1, ncol, data_type, is_row_major);
1605
  ref_booster->PredictSingleRow(num_iteration, predict_type, ncol, get_row_fun, config, out_result, out_len);
1606
1607
1608
1609
  API_END();
}


1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
int LGBM_BoosterPredictForMats(BoosterHandle handle,
                               const void** data,
                               int data_type,
                               int32_t nrow,
                               int32_t ncol,
                               int predict_type,
                               int num_iteration,
                               const char* parameter,
                               int64_t* out_len,
                               double* out_result) {
  API_BEGIN();
  auto param = Config::Str2Map(parameter);
  Config config;
  config.Set(param);
  if (config.num_threads > 0) {
    omp_set_num_threads(config.num_threads);
  }
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  auto get_row_fun = RowPairFunctionFromDenseRows(data, ncol, data_type);
1629
  ref_booster->Predict(num_iteration, predict_type, nrow, ncol, get_row_fun, config, out_result, out_len);
1630
1631
1632
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1633
int LGBM_BoosterSaveModel(BoosterHandle handle,
1634
                          int start_iteration,
1635
1636
                          int num_iteration,
                          const char* filename) {
1637
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1638
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1639
  ref_booster->SaveModelToFile(start_iteration, num_iteration, filename);
wxchan's avatar
wxchan committed
1640
1641
1642
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1643
int LGBM_BoosterSaveModelToString(BoosterHandle handle,
1644
                                  int start_iteration,
1645
                                  int num_iteration,
1646
                                  int64_t buffer_len,
1647
                                  int64_t* out_len,
1648
                                  char* out_str) {
1649
1650
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1651
  std::string model = ref_booster->SaveModelToString(start_iteration, num_iteration);
1652
  *out_len = static_cast<int64_t>(model.size()) + 1;
1653
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
1654
    std::memcpy(out_str, model.c_str(), *out_len);
1655
1656
1657
1658
  }
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1659
int LGBM_BoosterDumpModel(BoosterHandle handle,
1660
                          int start_iteration,
1661
                          int num_iteration,
1662
1663
                          int64_t buffer_len,
                          int64_t* out_len,
1664
                          char* out_str) {
wxchan's avatar
wxchan committed
1665
1666
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
1667
  std::string model = ref_booster->DumpModel(start_iteration, num_iteration);
1668
  *out_len = static_cast<int64_t>(model.size()) + 1;
wxchan's avatar
wxchan committed
1669
  if (*out_len <= buffer_len) {
Guolin Ke's avatar
Guolin Ke committed
1670
    std::memcpy(out_str, model.c_str(), *out_len);
wxchan's avatar
wxchan committed
1671
  }
1672
  API_END();
Guolin Ke's avatar
Guolin Ke committed
1673
}
1674

Guolin Ke's avatar
Guolin Ke committed
1675
int LGBM_BoosterGetLeafValue(BoosterHandle handle,
1676
1677
1678
                             int tree_idx,
                             int leaf_idx,
                             double* out_val) {
Guolin Ke's avatar
Guolin Ke committed
1679
1680
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
1681
  *out_val = static_cast<double>(ref_booster->GetLeafValue(tree_idx, leaf_idx));
Guolin Ke's avatar
Guolin Ke committed
1682
1683
1684
  API_END();
}

Guolin Ke's avatar
Guolin Ke committed
1685
int LGBM_BoosterSetLeafValue(BoosterHandle handle,
1686
1687
1688
                             int tree_idx,
                             int leaf_idx,
                             double val) {
Guolin Ke's avatar
Guolin Ke committed
1689
1690
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
Guolin Ke's avatar
Guolin Ke committed
1691
  ref_booster->SetLeafValue(tree_idx, leaf_idx, val);
Guolin Ke's avatar
Guolin Ke committed
1692
1693
1694
  API_END();
}

1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
int LGBM_BoosterFeatureImportance(BoosterHandle handle,
                                  int num_iteration,
                                  int importance_type,
                                  double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  std::vector<double> feature_importances = ref_booster->FeatureImportance(num_iteration, importance_type);
  for (size_t i = 0; i < feature_importances.size(); ++i) {
    (out_results)[i] = feature_importances[i];
  }
  API_END();
}

1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
int LGBM_BoosterGetUpperBoundValue(BoosterHandle handle,
                                   double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  double max_value = ref_booster->UpperBoundValue();
  *out_results = max_value;
  API_END();
}

int LGBM_BoosterGetLowerBoundValue(BoosterHandle handle,
                                   double* out_results) {
  API_BEGIN();
  Booster* ref_booster = reinterpret_cast<Booster*>(handle);
  double min_value = ref_booster->LowerBoundValue();
  *out_results = min_value;
  API_END();
}

1726
1727
1728
1729
1730
int LGBM_NetworkInit(const char* machines,
                     int local_listen_port,
                     int listen_time_out,
                     int num_machines) {
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1731
  Config config;
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
  config.machines = Common::RemoveQuotationSymbol(std::string(machines));
  config.local_listen_port = local_listen_port;
  config.num_machines = num_machines;
  config.time_out = listen_time_out;
  if (num_machines > 1) {
    Network::Init(config);
  }
  API_END();
}

int LGBM_NetworkFree() {
  API_BEGIN();
  Network::Dispose();
  API_END();
}

1748
1749
1750
int LGBM_NetworkInitWithFunctions(int num_machines, int rank,
                                  void* reduce_scatter_ext_fun,
                                  void* allgather_ext_fun) {
ww's avatar
ww committed
1751
  API_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
1752
  if (num_machines > 1) {
1753
    Network::Init(num_machines, rank, (ReduceScatterFunction)reduce_scatter_ext_fun, (AllgatherFunction)allgather_ext_fun);
ww's avatar
ww committed
1754
1755
1756
  }
  API_END();
}
Guolin Ke's avatar
Guolin Ke committed
1757

Guolin Ke's avatar
Guolin Ke committed
1758
// ---- start of some help functions
1759
1760
1761

std::function<std::vector<double>(int row_idx)>
RowFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major) {
Guolin Ke's avatar
Guolin Ke committed
1762
  if (data_type == C_API_DTYPE_FLOAT32) {
1763
1764
    const float* data_ptr = reinterpret_cast<const float*>(data);
    if (is_row_major) {
1765
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
1766
        std::vector<double> ret(num_col);
1767
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
1768
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1769
          ret[i] = static_cast<double>(*(tmp_ptr + i));
1770
1771
1772
1773
        }
        return ret;
      };
    } else {
1774
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
1775
        std::vector<double> ret(num_col);
1776
        for (int i = 0; i < num_col; ++i) {
1777
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
1778
1779
1780
1781
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
1782
  } else if (data_type == C_API_DTYPE_FLOAT64) {
1783
1784
    const double* data_ptr = reinterpret_cast<const double*>(data);
    if (is_row_major) {
1785
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
1786
        std::vector<double> ret(num_col);
1787
        auto tmp_ptr = data_ptr + static_cast<size_t>(num_col) * row_idx;
1788
        for (int i = 0; i < num_col; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1789
          ret[i] = static_cast<double>(*(tmp_ptr + i));
1790
1791
1792
1793
        }
        return ret;
      };
    } else {
1794
      return [=] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
1795
        std::vector<double> ret(num_col);
1796
        for (int i = 0; i < num_col; ++i) {
1797
          ret[i] = static_cast<double>(*(data_ptr + static_cast<size_t>(num_row) * i + row_idx));
1798
1799
1800
1801
1802
        }
        return ret;
      };
    }
  }
1803
  throw std::runtime_error("Unknown data type in RowFunctionFromDenseMatric");
1804
1805
1806
1807
}

std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseMatric(const void* data, int num_row, int num_col, int data_type, int is_row_major) {
Guolin Ke's avatar
Guolin Ke committed
1808
1809
  auto inner_function = RowFunctionFromDenseMatric(data, num_row, num_col, data_type, is_row_major);
  if (inner_function != nullptr) {
1810
    return [inner_function] (int row_idx) {
Guolin Ke's avatar
Guolin Ke committed
1811
1812
      auto raw_values = inner_function(row_idx);
      std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
1813
      ret.reserve(raw_values.size());
Guolin Ke's avatar
Guolin Ke committed
1814
      for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
Guolin Ke's avatar
Guolin Ke committed
1815
        if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
Guolin Ke's avatar
Guolin Ke committed
1816
          ret.emplace_back(i, raw_values[i]);
1817
        }
Guolin Ke's avatar
Guolin Ke committed
1818
1819
1820
      }
      return ret;
    };
1821
  }
Guolin Ke's avatar
Guolin Ke committed
1822
  return nullptr;
1823
1824
}

1825
1826
1827
1828
1829
1830
1831
// data is array of pointers to individual rows
std::function<std::vector<std::pair<int, double>>(int row_idx)>
RowPairFunctionFromDenseRows(const void** data, int num_col, int data_type) {
  return [=](int row_idx) {
    auto inner_function = RowFunctionFromDenseMatric(data[row_idx], 1, num_col, data_type, /* is_row_major */ true);
    auto raw_values = inner_function(0);
    std::vector<std::pair<int, double>> ret;
Guolin Ke's avatar
Guolin Ke committed
1832
    ret.reserve(raw_values.size());
1833
1834
1835
1836
1837
1838
1839
1840
1841
    for (int i = 0; i < static_cast<int>(raw_values.size()); ++i) {
      if (std::fabs(raw_values[i]) > kZeroThreshold || std::isnan(raw_values[i])) {
        ret.emplace_back(i, raw_values[i]);
      }
    }
    return ret;
  };
}

1842
std::function<std::vector<std::pair<int, double>>(int idx)>
1843
RowFunctionFromCSR(const void* indptr, int indptr_type, const int32_t* indices, const void* data, int data_type, int64_t , int64_t ) {
Guolin Ke's avatar
Guolin Ke committed
1844
  if (data_type == C_API_DTYPE_FLOAT32) {
1845
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
1846
    if (indptr_type == C_API_DTYPE_INT32) {
1847
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
1848
      return [=] (int idx) {
1849
1850
1851
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
1852
1853
1854
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
1855
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1856
          ret.emplace_back(indices[i], data_ptr[i]);
1857
1858
1859
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
1860
    } else if (indptr_type == C_API_DTYPE_INT64) {
1861
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
1862
      return [=] (int idx) {
1863
1864
1865
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
1866
1867
1868
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
1869
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1870
          ret.emplace_back(indices[i], data_ptr[i]);
1871
1872
1873
1874
        }
        return ret;
      };
    }
Guolin Ke's avatar
Guolin Ke committed
1875
  } else if (data_type == C_API_DTYPE_FLOAT64) {
1876
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
1877
    if (indptr_type == C_API_DTYPE_INT32) {
1878
      const int32_t* ptr_indptr = reinterpret_cast<const int32_t*>(indptr);
1879
      return [=] (int idx) {
1880
1881
1882
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
1883
1884
1885
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
1886
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1887
          ret.emplace_back(indices[i], data_ptr[i]);
1888
1889
1890
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
1891
    } else if (indptr_type == C_API_DTYPE_INT64) {
1892
      const int64_t* ptr_indptr = reinterpret_cast<const int64_t*>(indptr);
1893
      return [=] (int idx) {
1894
1895
1896
        std::vector<std::pair<int, double>> ret;
        int64_t start = ptr_indptr[idx];
        int64_t end = ptr_indptr[idx + 1];
1897
1898
1899
        if (end - start > 0)  {
          ret.reserve(end - start);
        }
Guolin Ke's avatar
Guolin Ke committed
1900
        for (int64_t i = start; i < end; ++i) {
Guolin Ke's avatar
Guolin Ke committed
1901
          ret.emplace_back(indices[i], data_ptr[i]);
1902
1903
1904
        }
        return ret;
      };
Guolin Ke's avatar
Guolin Ke committed
1905
1906
    }
  }
1907
  throw std::runtime_error("Unknown data type in RowFunctionFromCSR");
1908
1909
}

Guolin Ke's avatar
Guolin Ke committed
1910
std::function<std::pair<int, double>(int idx)>
1911
IterateFunctionFromCSC(const void* col_ptr, int col_ptr_type, const int32_t* indices, const void* data, int data_type, int64_t ncol_ptr, int64_t , int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
1912
  CHECK(col_idx < ncol_ptr && col_idx >= 0);
Guolin Ke's avatar
Guolin Ke committed
1913
  if (data_type == C_API_DTYPE_FLOAT32) {
1914
    const float* data_ptr = reinterpret_cast<const float*>(data);
Guolin Ke's avatar
Guolin Ke committed
1915
    if (col_ptr_type == C_API_DTYPE_INT32) {
1916
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
1917
1918
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
1919
1920
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
1921
1922
        if (i >= end) {
          return std::make_pair(-1, 0.0);
1923
        }
Guolin Ke's avatar
Guolin Ke committed
1924
1925
1926
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
1927
      };
Guolin Ke's avatar
Guolin Ke committed
1928
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
1929
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
1930
1931
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
1932
1933
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
1934
1935
        if (i >= end) {
          return std::make_pair(-1, 0.0);
1936
        }
Guolin Ke's avatar
Guolin Ke committed
1937
1938
1939
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
1940
      };
Guolin Ke's avatar
Guolin Ke committed
1941
    }
Guolin Ke's avatar
Guolin Ke committed
1942
  } else if (data_type == C_API_DTYPE_FLOAT64) {
1943
    const double* data_ptr = reinterpret_cast<const double*>(data);
Guolin Ke's avatar
Guolin Ke committed
1944
    if (col_ptr_type == C_API_DTYPE_INT32) {
1945
      const int32_t* ptr_col_ptr = reinterpret_cast<const int32_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
1946
1947
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
1948
1949
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
1950
1951
        if (i >= end) {
          return std::make_pair(-1, 0.0);
1952
        }
Guolin Ke's avatar
Guolin Ke committed
1953
1954
1955
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
1956
      };
Guolin Ke's avatar
Guolin Ke committed
1957
    } else if (col_ptr_type == C_API_DTYPE_INT64) {
1958
      const int64_t* ptr_col_ptr = reinterpret_cast<const int64_t*>(col_ptr);
Guolin Ke's avatar
Guolin Ke committed
1959
1960
      int64_t start = ptr_col_ptr[col_idx];
      int64_t end = ptr_col_ptr[col_idx + 1];
1961
1962
      return [=] (int offset) {
        int64_t i = static_cast<int64_t>(start + offset);
Guolin Ke's avatar
Guolin Ke committed
1963
1964
        if (i >= end) {
          return std::make_pair(-1, 0.0);
1965
        }
Guolin Ke's avatar
Guolin Ke committed
1966
1967
1968
        int idx = static_cast<int>(indices[i]);
        double val = static_cast<double>(data_ptr[i]);
        return std::make_pair(idx, val);
1969
      };
Guolin Ke's avatar
Guolin Ke committed
1970
1971
    }
  }
1972
  throw std::runtime_error("Unknown data type in CSC matrix");
1973
1974
}

Guolin Ke's avatar
Guolin Ke committed
1975
CSC_RowIterator::CSC_RowIterator(const void* col_ptr, int col_ptr_type, const int32_t* indices,
1976
                                 const void* data, int data_type, int64_t ncol_ptr, int64_t nelem, int col_idx) {
Guolin Ke's avatar
Guolin Ke committed
1977
1978
1979
1980
1981
1982
1983
1984
1985
  iter_fun_ = IterateFunctionFromCSC(col_ptr, col_ptr_type, indices, data, data_type, ncol_ptr, nelem, col_idx);
}

double CSC_RowIterator::Get(int idx) {
  while (idx > cur_idx_ && !is_end_) {
    auto ret = iter_fun_(nonzero_idx_);
    if (ret.first < 0) {
      is_end_ = true;
      break;
1986
    }
Guolin Ke's avatar
Guolin Ke committed
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
    cur_idx_ = ret.first;
    cur_val_ = ret.second;
    ++nonzero_idx_;
  }
  if (idx == cur_idx_) {
    return cur_val_;
  } else {
    return 0.0f;
  }
}

std::pair<int, double> CSC_RowIterator::NextNonZero() {
  if (!is_end_) {
    auto ret = iter_fun_(nonzero_idx_);
    ++nonzero_idx_;
    if (ret.first < 0) {
      is_end_ = true;
2004
    }
Guolin Ke's avatar
Guolin Ke committed
2005
2006
2007
    return ret;
  } else {
    return std::make_pair(-1, 0.0);
2008
  }
Guolin Ke's avatar
Guolin Ke committed
2009
}