basic.py 92 KB
Newer Older
wxchan's avatar
wxchan committed
1
# coding: utf-8
2
3
# pylint: disable = invalid-name, C0111, C0301
# pylint: disable = R0912, R0913, R0914, W0105, W0201, W0212
wxchan's avatar
wxchan committed
4
5
6
"""Wrapper c_api of LightGBM"""
from __future__ import absolute_import

7
import copy
wxchan's avatar
wxchan committed
8
import ctypes
9
import os
wxchan's avatar
wxchan committed
10
import warnings
wxchan's avatar
wxchan committed
11
from tempfile import NamedTemporaryFile
wxchan's avatar
wxchan committed
12
13
14
15

import numpy as np
import scipy.sparse

16
17
18
from .compat import (DataFrame, LGBMDeprecationWarning, Series,
                     decode_string, integer_types,
                     json, json_default_with_numpy,
19
                     numeric_types, range_, zip_, string_type)
wxchan's avatar
wxchan committed
20
21
from .libpath import find_lib_path

wxchan's avatar
wxchan committed
22

wxchan's avatar
wxchan committed
23
24
25
26
def _load_lib():
    """Load LightGBM Library."""
    lib_path = find_lib_path()
    if len(lib_path) == 0:
27
        return None
wxchan's avatar
wxchan committed
28
29
30
31
    lib = ctypes.cdll.LoadLibrary(lib_path[0])
    lib.LGBM_GetLastError.restype = ctypes.c_char_p
    return lib

wxchan's avatar
wxchan committed
32

wxchan's avatar
wxchan committed
33
34
_LIB = _load_lib()

wxchan's avatar
wxchan committed
35

wxchan's avatar
wxchan committed
36
37
38
39
40
41
42
43
def _safe_call(ret):
    """Check the return value of C API call
    Parameters
    ----------
    ret : int
        return value from API calls
    """
    if ret != 0:
44
        raise LightGBMError(decode_string(_LIB.LGBM_GetLastError()))
wxchan's avatar
wxchan committed
45

wxchan's avatar
wxchan committed
46

wxchan's avatar
wxchan committed
47
48
49
50
51
def is_numeric(obj):
    """Check is a number or not, include numpy number etc."""
    try:
        float(obj)
        return True
wxchan's avatar
wxchan committed
52
53
54
    except (TypeError, ValueError):
        # TypeError: obj is not a string or a number
        # ValueError: invalid literal
wxchan's avatar
wxchan committed
55
56
        return False

wxchan's avatar
wxchan committed
57

wxchan's avatar
wxchan committed
58
def is_numpy_1d_array(data):
Guolin Ke's avatar
Guolin Ke committed
59
    """Check is 1d numpy array"""
60
    return isinstance(data, np.ndarray) and len(data.shape) == 1
wxchan's avatar
wxchan committed
61

wxchan's avatar
wxchan committed
62

wxchan's avatar
wxchan committed
63
def is_1d_list(data):
Guolin Ke's avatar
Guolin Ke committed
64
    """Check is 1d list"""
65
    return isinstance(data, list) and \
66
        (not data or is_numeric(data[0]))
wxchan's avatar
wxchan committed
67

wxchan's avatar
wxchan committed
68

69
def list_to_1d_numpy(data, dtype=np.float32, name='list'):
Guolin Ke's avatar
Guolin Ke committed
70
    """convert to 1d numpy array"""
wxchan's avatar
wxchan committed
71
72
73
74
75
76
77
    if is_numpy_1d_array(data):
        if data.dtype == dtype:
            return data
        else:
            return data.astype(dtype=dtype, copy=False)
    elif is_1d_list(data):
        return np.array(data, dtype=dtype, copy=False)
78
79
    elif isinstance(data, Series):
        return data.values.astype(dtype)
wxchan's avatar
wxchan committed
80
    else:
81
82
        raise TypeError("Wrong type({0}) for {1}.\n"
                        "It should be list, numpy 1-D array or pandas Series".format(type(data).__name__, name))
wxchan's avatar
wxchan committed
83

wxchan's avatar
wxchan committed
84

wxchan's avatar
wxchan committed
85
86
87
88
def cfloat32_array_to_numpy(cptr, length):
    """Convert a ctypes float pointer array to a numpy array.
    """
    if isinstance(cptr, ctypes.POINTER(ctypes.c_float)):
89
        return np.fromiter(cptr, dtype=np.float32, count=length)
wxchan's avatar
wxchan committed
90
    else:
91
        raise RuntimeError('Expected float pointer')
wxchan's avatar
wxchan committed
92

Guolin Ke's avatar
Guolin Ke committed
93

Guolin Ke's avatar
Guolin Ke committed
94
95
96
97
98
99
100
101
def cfloat64_array_to_numpy(cptr, length):
    """Convert a ctypes double pointer array to a numpy array.
    """
    if isinstance(cptr, ctypes.POINTER(ctypes.c_double)):
        return np.fromiter(cptr, dtype=np.float64, count=length)
    else:
        raise RuntimeError('Expected double pointer')

wxchan's avatar
wxchan committed
102

wxchan's avatar
wxchan committed
103
104
105
106
def cint32_array_to_numpy(cptr, length):
    """Convert a ctypes float pointer array to a numpy array.
    """
    if isinstance(cptr, ctypes.POINTER(ctypes.c_int32)):
107
        return np.fromiter(cptr, dtype=np.int32, count=length)
wxchan's avatar
wxchan committed
108
    else:
109
        raise RuntimeError('Expected int pointer')
wxchan's avatar
wxchan committed
110

wxchan's avatar
wxchan committed
111

wxchan's avatar
wxchan committed
112
113
114
115
def c_str(string):
    """Convert a python string to cstring."""
    return ctypes.c_char_p(string.encode('utf-8'))

wxchan's avatar
wxchan committed
116

wxchan's avatar
wxchan committed
117
118
119
120
def c_array(ctype, values):
    """Convert a python array to c array."""
    return (ctype * len(values))(*values)

wxchan's avatar
wxchan committed
121

wxchan's avatar
wxchan committed
122
def param_dict_to_str(data):
123
    if data is None or not data:
wxchan's avatar
wxchan committed
124
125
126
        return ""
    pairs = []
    for key, val in data.items():
127
        if isinstance(val, (list, tuple, set)) or is_numpy_1d_array(val):
wxchan's avatar
wxchan committed
128
            pairs.append(str(key) + '=' + ','.join(map(str, val)))
wxchan's avatar
wxchan committed
129
        elif isinstance(val, string_type) or isinstance(val, numeric_types) or is_numeric(val):
wxchan's avatar
wxchan committed
130
            pairs.append(str(key) + '=' + str(val))
131
        elif val is not None:
132
            raise TypeError('Unknown type of parameter:%s, got:%s'
wxchan's avatar
wxchan committed
133
134
                            % (key, type(val).__name__))
    return ' '.join(pairs)
135

wxchan's avatar
wxchan committed
136

137
class _TempFile(object):
138
139
140
141
    def __enter__(self):
        with NamedTemporaryFile(prefix="lightgbm_tmp_", delete=True) as f:
            self.name = f.name
        return self
wxchan's avatar
wxchan committed
142

143
144
145
    def __exit__(self, exc_type, exc_val, exc_tb):
        if os.path.isfile(self.name):
            os.remove(self.name)
wxchan's avatar
wxchan committed
146

147
148
149
150
    def readlines(self):
        with open(self.name, "r+") as f:
            ret = f.readlines()
        return ret
wxchan's avatar
wxchan committed
151

152
153
    def writelines(self, lines):
        with open(self.name, "w+") as f:
154
            f.writelines(lines)
155

wxchan's avatar
wxchan committed
156

157
158
159
160
161
162
163
class LightGBMError(Exception):
    """Error throwed by LightGBM"""
    pass


MAX_INT32 = (1 << 31) - 1

wxchan's avatar
wxchan committed
164
165
166
167
168
"""marco definition of data type in c_api of LightGBM"""
C_API_DTYPE_FLOAT32 = 0
C_API_DTYPE_FLOAT64 = 1
C_API_DTYPE_INT32 = 2
C_API_DTYPE_INT64 = 3
Guolin Ke's avatar
Guolin Ke committed
169

wxchan's avatar
wxchan committed
170
171
172
"""Matric is row major in python"""
C_API_IS_ROW_MAJOR = 1

Guolin Ke's avatar
Guolin Ke committed
173
"""marco definition of prediction type in c_api of LightGBM"""
wxchan's avatar
wxchan committed
174
175
176
C_API_PREDICT_NORMAL = 0
C_API_PREDICT_RAW_SCORE = 1
C_API_PREDICT_LEAF_INDEX = 2
177
C_API_PREDICT_CONTRIB = 3
wxchan's avatar
wxchan committed
178

Guolin Ke's avatar
Guolin Ke committed
179
"""data type of data field"""
wxchan's avatar
wxchan committed
180
181
FIELD_TYPE_MAPPER = {"label": C_API_DTYPE_FLOAT32,
                     "weight": C_API_DTYPE_FLOAT32,
Guolin Ke's avatar
Guolin Ke committed
182
                     "init_score": C_API_DTYPE_FLOAT64,
wxchan's avatar
wxchan committed
183
184
                     "group": C_API_DTYPE_INT32}

185
186
187
188
189
PANDAS_DTYPE_MAPPER = {'int8': 'int', 'int16': 'int', 'int32': 'int',
                       'int64': 'int', 'uint8': 'int', 'uint16': 'int',
                       'uint32': 'int', 'uint64': 'int', 'float16': 'float',
                       'float32': 'float', 'float64': 'float', 'bool': 'int'}

wxchan's avatar
wxchan committed
190

191
192
193
194
def convert_from_sliced_object(data):
    """fix the memory of multi-dimensional sliced object"""
    if data.base is not None and isinstance(data, np.ndarray) and isinstance(data.base, np.ndarray):
        if not data.flags.c_contiguous:
195
196
            warnings.warn("Usage of np.ndarray subset (sliced data) is not recommended "
                          "due to it will double the peak memory cost in LightGBM.")
197
198
199
200
            return np.copy(data)
    return data


wxchan's avatar
wxchan committed
201
def c_float_array(data):
Guolin Ke's avatar
Guolin Ke committed
202
    """get pointer of float numpy array / list"""
wxchan's avatar
wxchan committed
203
204
205
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
206
207
        data = convert_from_sliced_object(data)
        assert data.flags.c_contiguous
wxchan's avatar
wxchan committed
208
209
210
211
212
213
214
        if data.dtype == np.float32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
            type_data = C_API_DTYPE_FLOAT32
        elif data.dtype == np.float64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
            type_data = C_API_DTYPE_FLOAT64
        else:
215
            raise TypeError("Expected np.float32 or np.float64, met type({})"
wxchan's avatar
wxchan committed
216
217
                            .format(data.dtype))
    else:
218
        raise TypeError("Unknown type({})".format(type(data).__name__))
219
    return (ptr_data, type_data, data)  # return `data` to avoid the temporary copy is freed
wxchan's avatar
wxchan committed
220

wxchan's avatar
wxchan committed
221

wxchan's avatar
wxchan committed
222
def c_int_array(data):
Guolin Ke's avatar
Guolin Ke committed
223
    """get pointer of int numpy array / list"""
wxchan's avatar
wxchan committed
224
225
226
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
227
228
        data = convert_from_sliced_object(data)
        assert data.flags.c_contiguous
wxchan's avatar
wxchan committed
229
230
231
232
233
234
235
        if data.dtype == np.int32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
            type_data = C_API_DTYPE_INT32
        elif data.dtype == np.int64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int64))
            type_data = C_API_DTYPE_INT64
        else:
236
            raise TypeError("Expected np.int32 or np.int64, met type({})"
wxchan's avatar
wxchan committed
237
238
                            .format(data.dtype))
    else:
239
        raise TypeError("Unknown type({})".format(type(data).__name__))
240
    return (ptr_data, type_data, data)  # return `data` to avoid the temporary copy is freed
wxchan's avatar
wxchan committed
241

wxchan's avatar
wxchan committed
242

243
def _data_from_pandas(data, feature_name, categorical_feature, pandas_categorical):
244
    if isinstance(data, DataFrame):
245
246
        if len(data.shape) != 2 or data.shape[0] < 1:
            raise ValueError('Input data must be 2 dimensional and non empty.')
247
248
249
250
251
252
253
254
        if feature_name == 'auto' or feature_name is None:
            data = data.rename(columns=str)
        cat_cols = data.select_dtypes(include=['category']).columns
        if pandas_categorical is None:  # train dataset
            pandas_categorical = [list(data[col].cat.categories) for col in cat_cols]
        else:
            if len(cat_cols) != len(pandas_categorical):
                raise ValueError('train and valid dataset categorical_feature do not match.')
255
            for col, category in zip_(cat_cols, pandas_categorical):
256
257
258
259
                if list(data[col].cat.categories) != list(category):
                    data[col] = data[col].cat.set_categories(category)
        if len(cat_cols):  # cat_cols is pandas Index object
            data = data.copy()  # not alter origin DataFrame
260
            data[cat_cols] = data[cat_cols].apply(lambda x: x.cat.codes).replace({-1: np.nan})
261
262
263
264
265
266
267
268
269
270
271
272
273
274
        if categorical_feature is not None:
            if feature_name is None:
                feature_name = list(data.columns)
            if categorical_feature == 'auto':
                categorical_feature = list(cat_cols)
            else:
                categorical_feature = list(categorical_feature) + list(cat_cols)
        if feature_name == 'auto':
            feature_name = list(data.columns)
        data_dtypes = data.dtypes
        if not all(dtype.name in PANDAS_DTYPE_MAPPER for dtype in data_dtypes):
            bad_fields = [data.columns[i] for i, dtype in
                          enumerate(data_dtypes) if dtype.name not in PANDAS_DTYPE_MAPPER]

275
276
            msg = ("DataFrame.dtypes for data must be int, float or bool.\n"
                   "Did not expect the data types in fields ")
277
            raise ValueError(msg + ', '.join(bad_fields))
278
        data = data.values.astype('float')
279
280
281
282
283
284
    else:
        if feature_name == 'auto':
            feature_name = None
        if categorical_feature == 'auto':
            categorical_feature = None
    return data, feature_name, categorical_feature, pandas_categorical
285
286
287
288
289
290
291
292
293


def _label_from_pandas(label):
    if isinstance(label, DataFrame):
        if len(label.columns) > 1:
            raise ValueError('DataFrame for label cannot have multiple columns')
        label_dtypes = label.dtypes
        if not all(dtype.name in PANDAS_DTYPE_MAPPER for dtype in label_dtypes):
            raise ValueError('DataFrame.dtypes for label must be int, float or bool')
294
        label = label.values.astype('float').flatten()
295
296
297
    return label


298
299
def _save_pandas_categorical(file_name, pandas_categorical):
    with open(file_name, 'a') as f:
300
301
        f.write('\npandas_categorical:'
                + json.dumps(pandas_categorical, default=json_default_with_numpy) + '\n')
302
303
304
305


def _load_pandas_categorical(file_name):
    with open(file_name, 'r') as f:
Guolin Ke's avatar
Guolin Ke committed
306
307
308
309
        lines = f.readlines()
        last_line = lines[-1]
        if last_line.strip() == "":
            last_line = lines[-2]
310
311
312
313
314
        if last_line.startswith('pandas_categorical:'):
            return json.loads(last_line[len('pandas_categorical:'):])
    return None


Guolin Ke's avatar
Guolin Ke committed
315
316
class _InnerPredictor(object):
    """
317
318
    A _InnerPredictor of LightGBM.
    Only used for prediction, usually used for continued-train
Guolin Ke's avatar
Guolin Ke committed
319
    Note: Can convert from Booster, but cannot convert to Booster
wxchan's avatar
wxchan committed
320
    """
321
    def __init__(self, model_file=None, booster_handle=None, pred_parameter=None):
Guolin Ke's avatar
Guolin Ke committed
322
        """Initialize the _InnerPredictor. Not expose to user
wxchan's avatar
wxchan committed
323
324
325
326
327

        Parameters
        ----------
        model_file : string
            Path to the model file.
Guolin Ke's avatar
Guolin Ke committed
328
329
        booster_handle : Handle of Booster
            use handle to init
330
331
        pred_parameter: dict
            Other parameters for the prediciton
wxchan's avatar
wxchan committed
332
333
334
335
336
        """
        self.handle = ctypes.c_void_p()
        self.__is_manage_handle = True
        if model_file is not None:
            """Prediction task"""
Guolin Ke's avatar
Guolin Ke committed
337
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
338
339
340
341
            _safe_call(_LIB.LGBM_BoosterCreateFromModelfile(
                c_str(model_file),
                ctypes.byref(out_num_iterations),
                ctypes.byref(self.handle)))
Guolin Ke's avatar
Guolin Ke committed
342
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
343
344
345
346
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
347
            self.num_total_iteration = out_num_iterations.value
348
            self.pandas_categorical = _load_pandas_categorical(model_file)
wxchan's avatar
wxchan committed
349
        elif booster_handle is not None:
Guolin Ke's avatar
Guolin Ke committed
350
            self.__is_manage_handle = False
wxchan's avatar
wxchan committed
351
            self.handle = booster_handle
Guolin Ke's avatar
Guolin Ke committed
352
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
353
354
355
356
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
Guolin Ke's avatar
Guolin Ke committed
357
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
358
359
360
            _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
                self.handle,
                ctypes.byref(out_num_iterations)))
361
            self.num_total_iteration = out_num_iterations.value
362
            self.pandas_categorical = None
wxchan's avatar
wxchan committed
363
        else:
Guolin Ke's avatar
Guolin Ke committed
364
            raise TypeError('Need Model file or Booster handle to create a predictor')
wxchan's avatar
wxchan committed
365

366
367
        pred_parameter = {} if pred_parameter is None else pred_parameter
        self.pred_parameter = param_dict_to_str(pred_parameter)
cbecker's avatar
cbecker committed
368

wxchan's avatar
wxchan committed
369
    def __del__(self):
370
371
372
373
374
        try:
            if self.__is_manage_handle:
                _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
375

376
377
378
379
380
    def __getstate__(self):
        this = self.__dict__.copy()
        this.pop('handle', None)
        return this

wxchan's avatar
wxchan committed
381
    def predict(self, data, num_iteration=-1,
382
                raw_score=False, pred_leaf=False, pred_contrib=False, data_has_header=False,
wxchan's avatar
wxchan committed
383
384
385
386
387
388
                is_reshape=True):
        """
        Predict logic

        Parameters
        ----------
389
        data : string, numpy array, pandas DataFrame or scipy.sparse
wxchan's avatar
wxchan committed
390
            Data source for prediction
391
            When data type is string, it represents the path of txt file
wxchan's avatar
wxchan committed
392
        num_iteration : int
393
            Used iteration for prediction
wxchan's avatar
wxchan committed
394
395
396
397
        raw_score : bool
            True for predict raw score
        pred_leaf : bool
            True for predict leaf index
398
399
        pred_contrib : bool
            True for predict feature contributions
wxchan's avatar
wxchan committed
400
        data_has_header : bool
Guolin Ke's avatar
Guolin Ke committed
401
            Used for txt data, True if txt data has header
wxchan's avatar
wxchan committed
402
        is_reshape : bool
403
            Reshape to (nrow, ncol) if true
wxchan's avatar
wxchan committed
404
405
406
407
408

        Returns
        -------
        Prediction result
        """
wxchan's avatar
wxchan committed
409
        if isinstance(data, Dataset):
410
            raise TypeError("Cannot use Dataset instance for prediction, please use raw data instead")
411
        data = _data_from_pandas(data, None, None, self.pandas_categorical)[0]
wxchan's avatar
wxchan committed
412
413
414
415
416
        predict_type = C_API_PREDICT_NORMAL
        if raw_score:
            predict_type = C_API_PREDICT_RAW_SCORE
        if pred_leaf:
            predict_type = C_API_PREDICT_LEAF_INDEX
417
418
        if pred_contrib:
            predict_type = C_API_PREDICT_CONTRIB
wxchan's avatar
wxchan committed
419
        int_data_has_header = 1 if data_has_header else 0
420
421
        if num_iteration > self.num_total_iteration:
            num_iteration = self.num_total_iteration
cbecker's avatar
cbecker committed
422

wxchan's avatar
wxchan committed
423
        if isinstance(data, string_type):
424
            with _TempFile() as f:
wxchan's avatar
wxchan committed
425
426
427
                _safe_call(_LIB.LGBM_BoosterPredictForFile(
                    self.handle,
                    c_str(data),
Guolin Ke's avatar
Guolin Ke committed
428
429
430
                    ctypes.c_int(int_data_has_header),
                    ctypes.c_int(predict_type),
                    ctypes.c_int(num_iteration),
431
                    c_str(self.pred_parameter),
wxchan's avatar
wxchan committed
432
433
                    c_str(f.name)))
                lines = f.readlines()
434
435
                nrow = len(lines)
                preds = [float(token) for line in lines for token in line.split('\t')]
Guolin Ke's avatar
Guolin Ke committed
436
                preds = np.array(preds, dtype=np.float64, copy=False)
wxchan's avatar
wxchan committed
437
        elif isinstance(data, scipy.sparse.csr_matrix):
438
            preds, nrow = self.__pred_for_csr(data, num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
439
        elif isinstance(data, scipy.sparse.csc_matrix):
440
            preds, nrow = self.__pred_for_csc(data, num_iteration, predict_type)
wxchan's avatar
wxchan committed
441
        elif isinstance(data, np.ndarray):
442
            preds, nrow = self.__pred_for_np2d(data, num_iteration, predict_type)
443
444
445
        elif isinstance(data, list):
            try:
                data = np.array(data)
446
            except BaseException:
447
                raise ValueError('Cannot convert data list to numpy array.')
448
            preds, nrow = self.__pred_for_np2d(data, num_iteration, predict_type)
wxchan's avatar
wxchan committed
449
450
        else:
            try:
451
                warnings.warn('Converting data to scipy sparse matrix.')
wxchan's avatar
wxchan committed
452
                csr = scipy.sparse.csr_matrix(data)
453
            except BaseException:
454
                raise TypeError('Cannot predict data for type {}'.format(type(data).__name__))
455
            preds, nrow = self.__pred_for_csr(csr, num_iteration, predict_type)
wxchan's avatar
wxchan committed
456
457
        if pred_leaf:
            preds = preds.astype(np.int32)
458
        if is_reshape and preds.size != nrow:
wxchan's avatar
wxchan committed
459
            if preds.size % nrow == 0:
460
                preds = preds.reshape(nrow, -1)
wxchan's avatar
wxchan committed
461
            else:
462
                raise ValueError('Length of predict result (%d) cannot be divide nrow (%d)'
wxchan's avatar
wxchan committed
463
464
465
466
                                 % (preds.size, nrow))
        return preds

    def __get_num_preds(self, num_iteration, nrow, predict_type):
Guolin Ke's avatar
Guolin Ke committed
467
468
469
        """
        Get size of prediction result
        """
470
471
472
473
474
        if nrow > MAX_INT32:
            raise LightGBMError('LightGBM cannot perform prediction for data'
                                'with number of rows greater than MAX_INT32 (%d).\n'
                                'You can split your data into chunks'
                                'and then concatenate predictions for them' % MAX_INT32)
Guolin Ke's avatar
Guolin Ke committed
475
476
477
        n_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterCalcNumPredict(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
478
479
480
            ctypes.c_int(nrow),
            ctypes.c_int(predict_type),
            ctypes.c_int(num_iteration),
Guolin Ke's avatar
Guolin Ke committed
481
482
            ctypes.byref(n_preds)))
        return n_preds.value
wxchan's avatar
wxchan committed
483
484
485
486
487
488

    def __pred_for_np2d(self, mat, num_iteration, predict_type):
        """
        Predict for a 2-D numpy matrix.
        """
        if len(mat.shape) != 2:
489
            raise ValueError('Input numpy.ndarray or list must be 2 dimensional')
wxchan's avatar
wxchan committed
490

491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
        def inner_predict(mat, num_iteration, predict_type, preds=None):
            if mat.dtype == np.float32 or mat.dtype == np.float64:
                data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
            else:
                """change non-float data to float data, need to copy"""
                data = np.array(mat.reshape(mat.size), dtype=np.float32)
            ptr_data, type_ptr_data, _ = c_float_array(data)
            n_preds = self.__get_num_preds(num_iteration, mat.shape[0], predict_type)
            if preds is None:
                preds = np.zeros(n_preds, dtype=np.float64)
            elif len(preds.shape) != 1 or len(preds) != n_preds:
                raise ValueError("Wrong length of pre-allocated predict array")
            out_num_preds = ctypes.c_int64(0)
            _safe_call(_LIB.LGBM_BoosterPredictForMat(
                self.handle,
                ptr_data,
                ctypes.c_int(type_ptr_data),
                ctypes.c_int(mat.shape[0]),
                ctypes.c_int(mat.shape[1]),
                ctypes.c_int(C_API_IS_ROW_MAJOR),
                ctypes.c_int(predict_type),
                ctypes.c_int(num_iteration),
                c_str(self.pred_parameter),
                ctypes.byref(out_num_preds),
                preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
            if n_preds != out_num_preds.value:
                raise ValueError("Wrong length for predict results")
            return preds, mat.shape[0]

        nrow = mat.shape[0]
        if nrow > MAX_INT32:
            sections = np.arange(start=MAX_INT32, stop=nrow, step=MAX_INT32)
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
            n_preds = [self.__get_num_preds(num_iteration, i, predict_type) for i in np.diff([0] + list(sections) + [nrow])]
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
            preds = np.zeros(sum(n_preds), dtype=np.float64)
527
528
            for chunk, (start_idx_pred, end_idx_pred) in zip_(np.array_split(mat, sections),
                                                              zip_(n_preds_sections, n_preds_sections[1:])):
529
530
531
                # avoid memory consumption by arrays concatenation operations
                inner_predict(chunk, num_iteration, predict_type, preds[start_idx_pred:end_idx_pred])
            return preds, nrow
wxchan's avatar
wxchan committed
532
        else:
533
            return inner_predict(mat, num_iteration, predict_type)
wxchan's avatar
wxchan committed
534
535
536
537
538

    def __pred_for_csr(self, csr, num_iteration, predict_type):
        """
        Predict for a csr data
        """
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
        def inner_predict(csr, num_iteration, predict_type, preds=None):
            nrow = len(csr.indptr) - 1
            n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
            if preds is None:
                preds = np.zeros(n_preds, dtype=np.float64)
            elif len(preds.shape) != 1 or len(preds) != n_preds:
                raise ValueError("Wrong length of pre-allocated predict array")
            out_num_preds = ctypes.c_int64(0)

            ptr_indptr, type_ptr_indptr, __ = c_int_array(csr.indptr)
            ptr_data, type_ptr_data, _ = c_float_array(csr.data)

            _safe_call(_LIB.LGBM_BoosterPredictForCSR(
                self.handle,
                ptr_indptr,
                ctypes.c_int32(type_ptr_indptr),
                csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
                ptr_data,
                ctypes.c_int(type_ptr_data),
                ctypes.c_int64(len(csr.indptr)),
                ctypes.c_int64(len(csr.data)),
                ctypes.c_int64(csr.shape[1]),
                ctypes.c_int(predict_type),
                ctypes.c_int(num_iteration),
                c_str(self.pred_parameter),
                ctypes.byref(out_num_preds),
                preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
            if n_preds != out_num_preds.value:
                raise ValueError("Wrong length for predict results")
            return preds, nrow
wxchan's avatar
wxchan committed
569

570
571
572
573
574
575
576
577
578
579
580
581
582
583
        nrow = len(csr.indptr) - 1
        if nrow > MAX_INT32:
            sections = [0] + list(np.arange(start=MAX_INT32, stop=nrow, step=MAX_INT32)) + [nrow]
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
            n_preds = [self.__get_num_preds(num_iteration, i, predict_type) for i in np.diff(sections)]
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
            preds = np.zeros(sum(n_preds), dtype=np.float64)
            for (start_idx, end_idx), (start_idx_pred, end_idx_pred) in zip_(zip_(sections, sections[1:]),
                                                                             zip_(n_preds_sections, n_preds_sections[1:])):
                # avoid memory consumption by arrays concatenation operations
                inner_predict(csr[start_idx:end_idx], num_iteration, predict_type, preds[start_idx_pred:end_idx_pred])
            return preds, nrow
        else:
            return inner_predict(csr, num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
584
585
586
587
588
589

    def __pred_for_csc(self, csc, num_iteration, predict_type):
        """
        Predict for a csc data
        """
        nrow = csc.shape[0]
590
591
        if nrow > MAX_INT32:
            return self.__pred_for_csr(csc.tocsr(), num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
592
593
594
595
        n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
        preds = np.zeros(n_preds, dtype=np.float64)
        out_num_preds = ctypes.c_int64(0)

596
597
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csc.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csc.data)
Guolin Ke's avatar
Guolin Ke committed
598
599
600
601

        _safe_call(_LIB.LGBM_BoosterPredictForCSC(
            self.handle,
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
602
            ctypes.c_int32(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
603
604
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
605
606
607
608
609
610
            ctypes.c_int(type_ptr_data),
            ctypes.c_int64(len(csc.indptr)),
            ctypes.c_int64(len(csc.data)),
            ctypes.c_int64(csc.shape[0]),
            ctypes.c_int(predict_type),
            ctypes.c_int(num_iteration),
611
            c_str(self.pred_parameter),
Guolin Ke's avatar
Guolin Ke committed
612
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
613
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
614
        if n_preds != out_num_preds.value:
615
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
616
617
        return preds, nrow

wxchan's avatar
wxchan committed
618

wxchan's avatar
wxchan committed
619
620
class Dataset(object):
    """Dataset in LightGBM."""
621
    def __init__(self, data, label=None, reference=None,
622
                 weight=None, group=None, init_score=None, silent=False,
623
                 feature_name='auto', categorical_feature='auto', params=None,
wxchan's avatar
wxchan committed
624
                 free_raw_data=True):
625
626
        """Constract Dataset.

wxchan's avatar
wxchan committed
627
628
        Parameters
        ----------
629
        data : string, numpy array, pandas DataFrame, scipy.sparse or list of numpy arrays
wxchan's avatar
wxchan committed
630
            Data source of Dataset.
631
            If string, it represents the path to txt file.
632
        label : list, numpy 1-D array, pandas one-column DataFrame/Series or None, optional (default=None)
633
634
635
            Label of the data.
        reference : Dataset or None, optional (default=None)
            If this is Dataset for validation, training data should be used as reference.
636
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
wxchan's avatar
wxchan committed
637
            Weight for each instance.
638
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
639
            Group/query size for Dataset.
640
        init_score : list, numpy 1-D array, pandas Series or None, optional (default=None)
641
            Init score for Dataset.
642
643
644
645
646
647
648
649
650
651
        silent : bool, optional (default=False)
            Whether to print messages during construction.
        feature_name : list of strings or 'auto', optional (default="auto")
            Feature names.
            If 'auto' and data is pandas DataFrame, data columns names are used.
        categorical_feature : list of strings or int, or 'auto', optional (default="auto")
            Categorical features.
            If list of int, interpreted as indices.
            If list of strings, interpreted as feature names (need to specify ``feature_name`` as well).
            If 'auto' and data is pandas DataFrame, pandas categorical columns are used.
652
            All values in categorical features should be less than int32 max value (2147483647).
653
            Large values could be memory consuming. Consider to use consecutive integers started from zero.
654
            All negative values in categorical features will be treated as missing values.
Nikita Titov's avatar
Nikita Titov committed
655
        params : dict or None, optional (default=None)
656
            Other parameters.
Nikita Titov's avatar
Nikita Titov committed
657
        free_raw_data : bool, optional (default=True)
658
            If True, raw data is freed after constructing inner Dataset.
wxchan's avatar
wxchan committed
659
        """
wxchan's avatar
wxchan committed
660
661
662
663
664
665
        self.handle = None
        self.data = data
        self.label = label
        self.reference = reference
        self.weight = weight
        self.group = group
666
        self.init_score = init_score
wxchan's avatar
wxchan committed
667
668
        self.silent = silent
        self.feature_name = feature_name
669
        self.categorical_feature = categorical_feature
670
        self.params = copy.deepcopy(params)
wxchan's avatar
wxchan committed
671
672
673
        self.free_raw_data = free_raw_data
        self.used_indices = None
        self._predictor = None
674
        self.pandas_categorical = None
675
        self.params_back_up = None
wxchan's avatar
wxchan committed
676
677

    def __del__(self):
678
679
680
681
        try:
            self._free_handle()
        except AttributeError:
            pass
682
683

    def _free_handle(self):
684
        if self.handle is not None:
685
            _safe_call(_LIB.LGBM_DatasetFree(self.handle))
686
            self.handle = None
Nikita Titov's avatar
Nikita Titov committed
687
        return self
wxchan's avatar
wxchan committed
688

689
    def _lazy_init(self, data, label=None, reference=None,
690
                   weight=None, group=None, init_score=None, predictor=None,
wxchan's avatar
wxchan committed
691
                   silent=False, feature_name='auto',
692
                   categorical_feature='auto', params=None):
wxchan's avatar
wxchan committed
693
694
        if data is None:
            self.handle = None
Nikita Titov's avatar
Nikita Titov committed
695
            return self
Guolin Ke's avatar
Guolin Ke committed
696
697
698
        if reference is not None:
            self.pandas_categorical = reference.pandas_categorical
            categorical_feature = reference.categorical_feature
699
700
701
702
        data, feature_name, categorical_feature, self.pandas_categorical = _data_from_pandas(data,
                                                                                             feature_name,
                                                                                             categorical_feature,
                                                                                             self.pandas_categorical)
wxchan's avatar
wxchan committed
703
704
        label = _label_from_pandas(label)
        self.data_has_header = False
705
        # process for args
wxchan's avatar
wxchan committed
706
        params = {} if params is None else params
707
708
709
        args_names = (getattr(self.__class__, '_lazy_init')
                      .__code__
                      .co_varnames[:getattr(self.__class__, '_lazy_init').__code__.co_argcount])
710
711
        for key, _ in params.items():
            if key in args_names:
712
713
714
                warnings.warn('{0} keyword has been found in `params` and will be ignored.\n'
                              'Please use {0} argument of the Dataset constructor to pass this parameter.'
                              .format(key))
wxchan's avatar
wxchan committed
715
        self.predictor = predictor
716
717
718
        # user can set verbose with params, it has higher priority
        if not any(verbose_alias in params for verbose_alias in ('verbose', 'verbosity')) and silent:
            params["verbose"] = -1
719
        # get categorical features
720
721
722
723
724
725
726
727
728
729
730
731
732
        if categorical_feature is not None:
            categorical_indices = set()
            feature_dict = {}
            if feature_name is not None:
                feature_dict = {name: i for i, name in enumerate(feature_name)}
            for name in categorical_feature:
                if isinstance(name, string_type) and name in feature_dict:
                    categorical_indices.add(feature_dict[name])
                elif isinstance(name, integer_types):
                    categorical_indices.add(name)
                else:
                    raise TypeError("Wrong type({}) or unknown name({}) in categorical_feature"
                                    .format(type(name).__name__, name))
733
            if categorical_indices:
734
                if "categorical_feature" in params or "categorical_column" in params:
735
                    warnings.warn('categorical_feature in param dict is overridden.')
736
737
                    params.pop("categorical_feature", None)
                    params.pop("categorical_column", None)
738
                params['categorical_column'] = sorted(categorical_indices)
739

wxchan's avatar
wxchan committed
740
        params_str = param_dict_to_str(params)
741
        # process for reference dataset
wxchan's avatar
wxchan committed
742
        ref_dataset = None
wxchan's avatar
wxchan committed
743
        if isinstance(reference, Dataset):
744
            ref_dataset = reference.construct().handle
wxchan's avatar
wxchan committed
745
746
        elif reference is not None:
            raise TypeError('Reference dataset should be None or dataset instance')
747
        # start construct data
wxchan's avatar
wxchan committed
748
        if isinstance(data, string_type):
749
            # check data has header or not
Guolin Ke's avatar
Guolin Ke committed
750
            if str(params.get("has_header", "")).lower() == "true" \
wxchan's avatar
wxchan committed
751
                    or str(params.get("header", "")).lower() == "true":
752
                self.data_has_header = True
wxchan's avatar
wxchan committed
753
754
755
756
757
758
759
760
            self.handle = ctypes.c_void_p()
            _safe_call(_LIB.LGBM_DatasetCreateFromFile(
                c_str(data),
                c_str(params_str),
                ref_dataset,
                ctypes.byref(self.handle)))
        elif isinstance(data, scipy.sparse.csr_matrix):
            self.__init_from_csr(data, params_str, ref_dataset)
Guolin Ke's avatar
Guolin Ke committed
761
762
        elif isinstance(data, scipy.sparse.csc_matrix):
            self.__init_from_csc(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
763
764
        elif isinstance(data, np.ndarray):
            self.__init_from_np2d(data, params_str, ref_dataset)
765
766
        elif isinstance(data, list) and len(data) > 0 and all(isinstance(x, np.ndarray) for x in data):
            self.__init_from_list_np2d(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
767
768
769
770
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                self.__init_from_csr(csr, params_str, ref_dataset)
771
            except BaseException:
wxchan's avatar
wxchan committed
772
                raise TypeError('Cannot initialize Dataset from {}'.format(type(data).__name__))
wxchan's avatar
wxchan committed
773
774
775
        if label is not None:
            self.set_label(label)
        if self.get_label() is None:
776
            raise ValueError("Label should not be None")
wxchan's avatar
wxchan committed
777
778
779
780
781
        if weight is not None:
            self.set_weight(weight)
        if group is not None:
            self.set_group(group)
        # load init score
782
783
784
        if init_score is not None:
            self.set_init_score(init_score)
            if self.predictor is not None:
785
                warnings.warn("The prediction of init_model will be overridden by init_score.")
786
        elif isinstance(self.predictor, _InnerPredictor):
wxchan's avatar
wxchan committed
787
788
789
790
791
792
            init_score = self.predictor.predict(data,
                                                raw_score=True,
                                                data_has_header=self.data_has_header,
                                                is_reshape=False)
            if self.predictor.num_class > 1:
                # need re group init score
wxchan's avatar
wxchan committed
793
                new_init_score = np.zeros(init_score.size, dtype=np.float32)
wxchan's avatar
wxchan committed
794
                num_data = self.num_data()
wxchan's avatar
wxchan committed
795
796
                for i in range_(num_data):
                    for j in range_(self.predictor.num_class):
wxchan's avatar
wxchan committed
797
798
799
                        new_init_score[j * num_data + i] = init_score[i * self.predictor.num_class + j]
                init_score = new_init_score
            self.set_init_score(init_score)
Guolin Ke's avatar
Guolin Ke committed
800
801
        elif self.predictor is not None:
            raise TypeError('wrong predictor type {}'.format(type(self.predictor).__name__))
Guolin Ke's avatar
Guolin Ke committed
802
        # set feature names
Nikita Titov's avatar
Nikita Titov committed
803
        return self.set_feature_name(feature_name)
wxchan's avatar
wxchan committed
804
805
806
807
808
809
810
811
812
813
814
815

    def __init_from_np2d(self, mat, params_str, ref_dataset):
        """
        Initialize data from a 2-D numpy matrix.
        """
        if len(mat.shape) != 2:
            raise ValueError('Input numpy.ndarray must be 2 dimensional')

        self.handle = ctypes.c_void_p()
        if mat.dtype == np.float32 or mat.dtype == np.float64:
            data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
        else:
816
            # change non-float data to float data, need to copy
wxchan's avatar
wxchan committed
817
818
            data = np.array(mat.reshape(mat.size), dtype=np.float32)

819
        ptr_data, type_ptr_data, _ = c_float_array(data)
wxchan's avatar
wxchan committed
820
821
        _safe_call(_LIB.LGBM_DatasetCreateFromMat(
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
822
823
824
825
            ctypes.c_int(type_ptr_data),
            ctypes.c_int(mat.shape[0]),
            ctypes.c_int(mat.shape[1]),
            ctypes.c_int(C_API_IS_ROW_MAJOR),
wxchan's avatar
wxchan committed
826
827
828
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
829
        return self
wxchan's avatar
wxchan committed
830

831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
    def __init_from_list_np2d(self, mats, params_str, ref_dataset):
        """
        Initialize data from list of 2-D numpy matrices.
        """
        ncol = mats[0].shape[1]
        nrow = np.zeros((len(mats),), np.int32)
        if mats[0].dtype == np.float64:
            ptr_data = (ctypes.POINTER(ctypes.c_double) * len(mats))()
        else:
            ptr_data = (ctypes.POINTER(ctypes.c_float) * len(mats))()

        holders = []
        type_ptr_data = None

        for i, mat in enumerate(mats):
            if len(mat.shape) != 2:
                raise ValueError('Input numpy.ndarray must be 2 dimensional')

            if mat.shape[1] != ncol:
                raise ValueError('Input arrays must have same number of columns')

            nrow[i] = mat.shape[0]

            if mat.dtype == np.float32 or mat.dtype == np.float64:
                mats[i] = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
            else:
                # change non-float data to float data, need to copy
                mats[i] = np.array(mat.reshape(mat.size), dtype=np.float32)

            chunk_ptr_data, chunk_type_ptr_data, holder = c_float_array(mats[i])
            if type_ptr_data is not None and chunk_type_ptr_data != type_ptr_data:
                raise ValueError('Input chunks must have same type')
            ptr_data[i] = chunk_ptr_data
            type_ptr_data = chunk_type_ptr_data
            holders.append(holder)

        self.handle = ctypes.c_void_p()
        _safe_call(_LIB.LGBM_DatasetCreateFromMats(
            ctypes.c_int(len(mats)),
            ctypes.cast(ptr_data, ctypes.POINTER(ctypes.POINTER(ctypes.c_double))),
            ctypes.c_int(type_ptr_data),
            nrow.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ctypes.c_int(ncol),
            ctypes.c_int(C_API_IS_ROW_MAJOR),
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
878
        return self
879

wxchan's avatar
wxchan committed
880
881
882
883
884
    def __init_from_csr(self, csr, params_str, ref_dataset):
        """
        Initialize data from a CSR matrix.
        """
        if len(csr.indices) != len(csr.data):
885
            raise ValueError('Length mismatch: {} vs {}'.format(len(csr.indices), len(csr.data)))
wxchan's avatar
wxchan committed
886
887
        self.handle = ctypes.c_void_p()

888
889
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csr.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csr.data)
wxchan's avatar
wxchan committed
890
891
892

        _safe_call(_LIB.LGBM_DatasetCreateFromCSR(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
893
            ctypes.c_int(type_ptr_indptr),
wxchan's avatar
wxchan committed
894
895
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
896
897
898
899
            ctypes.c_int(type_ptr_data),
            ctypes.c_int64(len(csr.indptr)),
            ctypes.c_int64(len(csr.data)),
            ctypes.c_int64(csr.shape[1]),
wxchan's avatar
wxchan committed
900
901
902
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
903
        return self
wxchan's avatar
wxchan committed
904

Guolin Ke's avatar
Guolin Ke committed
905
906
907
908
909
910
911
912
    def __init_from_csc(self, csc, params_str, ref_dataset):
        """
        Initialize data from a csc matrix.
        """
        if len(csc.indices) != len(csc.data):
            raise ValueError('Length mismatch: {} vs {}'.format(len(csc.indices), len(csc.data)))
        self.handle = ctypes.c_void_p()

913
914
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csc.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csc.data)
Guolin Ke's avatar
Guolin Ke committed
915
916
917

        _safe_call(_LIB.LGBM_DatasetCreateFromCSC(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
918
            ctypes.c_int(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
919
920
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
921
922
923
924
            ctypes.c_int(type_ptr_data),
            ctypes.c_int64(len(csc.indptr)),
            ctypes.c_int64(len(csc.data)),
            ctypes.c_int64(csc.shape[0]),
Guolin Ke's avatar
Guolin Ke committed
925
926
927
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
928
        return self
Guolin Ke's avatar
Guolin Ke committed
929

wxchan's avatar
wxchan committed
930
    def construct(self):
931
932
933
934
935
        """Lazy init.

        Returns
        -------
        self : Dataset
Nikita Titov's avatar
Nikita Titov committed
936
            Constructed Dataset object.
937
        """
938
        if self.handle is None:
wxchan's avatar
wxchan committed
939
940
            if self.reference is not None:
                if self.used_indices is None:
941
                    # create valid
942
                    self._lazy_init(self.data, label=self.label, reference=self.reference,
943
944
                                    weight=self.weight, group=self.group,
                                    init_score=self.init_score, predictor=self._predictor,
945
                                    silent=self.silent, feature_name=self.feature_name, params=self.params)
wxchan's avatar
wxchan committed
946
                else:
947
                    # construct subset
wxchan's avatar
wxchan committed
948
                    used_indices = list_to_1d_numpy(self.used_indices, np.int32, name='used_indices')
949
                    assert used_indices.flags.c_contiguous
Guolin Ke's avatar
Guolin Ke committed
950
951
                    if self.reference.group is not None:
                        group_info = np.array(self.reference.group).astype(int)
952
953
                        _, self.group = np.unique(np.repeat(range_(len(group_info)), repeats=group_info)[self.used_indices],
                                                  return_counts=True)
954
                    self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
955
956
                    params_str = param_dict_to_str(self.params)
                    _safe_call(_LIB.LGBM_DatasetGetSubset(
957
                        self.reference.construct().handle,
wxchan's avatar
wxchan committed
958
                        used_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
959
                        ctypes.c_int(used_indices.shape[0]),
wxchan's avatar
wxchan committed
960
961
                        c_str(params_str),
                        ctypes.byref(self.handle)))
Guolin Ke's avatar
Guolin Ke committed
962
963
                    if self.group is not None:
                        self.set_group(self.group)
wxchan's avatar
wxchan committed
964
965
966
                    if self.get_label() is None:
                        raise ValueError("Label should not be None.")
            else:
967
                # create train
968
                self._lazy_init(self.data, label=self.label,
969
970
971
                                weight=self.weight, group=self.group,
                                init_score=self.init_score, predictor=self._predictor,
                                silent=self.silent, feature_name=self.feature_name,
972
                                categorical_feature=self.categorical_feature, params=self.params)
wxchan's avatar
wxchan committed
973
974
975
            if self.free_raw_data:
                self.data = None
        return self
wxchan's avatar
wxchan committed
976

wxchan's avatar
wxchan committed
977
    def create_valid(self, data, label=None, weight=None, group=None,
978
                     init_score=None, silent=False, params=None):
979
        """Create validation data align with current Dataset.
wxchan's avatar
wxchan committed
980
981
982

        Parameters
        ----------
983
        data : string, numpy array, pandas DataFrame, scipy.sparse or list of numpy arrays
wxchan's avatar
wxchan committed
984
            Data source of Dataset.
985
            If string, it represents the path to txt file.
986
987
988
        label : list, numpy 1-D array, pandas one-column DataFrame/Series or None, optional (default=None)
            Label of the data.
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
wxchan's avatar
wxchan committed
989
            Weight for each instance.
990
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
991
            Group/query size for Dataset.
992
        init_score : list, numpy 1-D array, pandas Series or None, optional (default=None)
993
            Init score for Dataset.
994
995
        silent : bool, optional (default=False)
            Whether to print messages during construction.
Nikita Titov's avatar
Nikita Titov committed
996
        params : dict or None, optional (default=None)
997
998
999
1000
            Other parameters.

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1001
1002
        valid : Dataset
            Validation Dataset with reference to self.
wxchan's avatar
wxchan committed
1003
        """
1004
        ret = Dataset(data, label=label, reference=self,
1005
1006
                      weight=weight, group=group, init_score=init_score,
                      silent=silent, params=params, free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
1007
        ret._predictor = self._predictor
1008
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1009
        return ret
wxchan's avatar
wxchan committed
1010

wxchan's avatar
wxchan committed
1011
    def subset(self, used_indices, params=None):
1012
        """Get subset of current Dataset.
wxchan's avatar
wxchan committed
1013
1014
1015
1016

        Parameters
        ----------
        used_indices : list of int
1017
            Indices used to create the subset.
Nikita Titov's avatar
Nikita Titov committed
1018
        params : dict or None, optional (default=None)
1019
1020
1021
1022
1023
1024
            Other parameters.

        Returns
        -------
        subset : Dataset
            Subset of the current Dataset.
wxchan's avatar
wxchan committed
1025
        """
wxchan's avatar
wxchan committed
1026
1027
        if params is None:
            params = self.params
wxchan's avatar
wxchan committed
1028
        ret = Dataset(None, reference=self, feature_name=self.feature_name,
1029
                      categorical_feature=self.categorical_feature, params=params)
wxchan's avatar
wxchan committed
1030
        ret._predictor = self._predictor
1031
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1032
1033
1034
1035
        ret.used_indices = used_indices
        return ret

    def save_binary(self, filename):
1036
        """Save Dataset to binary file.
wxchan's avatar
wxchan committed
1037
1038
1039
1040
1041

        Parameters
        ----------
        filename : string
            Name of the output file.
Nikita Titov's avatar
Nikita Titov committed
1042
1043
1044
1045
1046

        Returns
        -------
        self : Dataset
            Returns self.
wxchan's avatar
wxchan committed
1047
1048
1049
1050
        """
        _safe_call(_LIB.LGBM_DatasetSaveBinary(
            self.construct().handle,
            c_str(filename)))
Nikita Titov's avatar
Nikita Titov committed
1051
        return self
wxchan's avatar
wxchan committed
1052
1053
1054
1055

    def _update_params(self, params):
        if not self.params:
            self.params = params
wxchan's avatar
wxchan committed
1056
        else:
1057
            self.params_back_up = copy.deepcopy(self.params)
wxchan's avatar
wxchan committed
1058
            self.params.update(params)
Nikita Titov's avatar
Nikita Titov committed
1059
        return self
wxchan's avatar
wxchan committed
1060

1061
1062
1063
    def _reverse_update_params(self):
        self.params = copy.deepcopy(self.params_back_up)
        self.params_back_up = None
Nikita Titov's avatar
Nikita Titov committed
1064
        return self
1065

wxchan's avatar
wxchan committed
1066
    def set_field(self, field_name, data):
wxchan's avatar
wxchan committed
1067
        """Set property into the Dataset.
wxchan's avatar
wxchan committed
1068
1069
1070

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1071
        field_name : string
1072
            The field name of the information.
1073
        data : list, numpy 1-D array, pandas Series or None
1074
            The array of data to be set.
Nikita Titov's avatar
Nikita Titov committed
1075
1076
1077
1078
1079

        Returns
        -------
        self : Dataset
            Dataset with set property.
wxchan's avatar
wxchan committed
1080
        """
1081
1082
        if self.handle is None:
            raise Exception("Cannot set %s before construct dataset" % field_name)
wxchan's avatar
wxchan committed
1083
        if data is None:
1084
            # set to None
wxchan's avatar
wxchan committed
1085
1086
1087
1088
            _safe_call(_LIB.LGBM_DatasetSetField(
                self.handle,
                c_str(field_name),
                None,
Guolin Ke's avatar
Guolin Ke committed
1089
1090
                ctypes.c_int(0),
                ctypes.c_int(FIELD_TYPE_MAPPER[field_name])))
Nikita Titov's avatar
Nikita Titov committed
1091
            return self
Guolin Ke's avatar
Guolin Ke committed
1092
1093
1094
1095
1096
        dtype = np.float32
        if field_name == 'group':
            dtype = np.int32
        elif field_name == 'init_score':
            dtype = np.float64
1097
        data = list_to_1d_numpy(data, dtype, name=field_name)
1098
1099
        if data.dtype == np.float32 or data.dtype == np.float64:
            ptr_data, type_data, _ = c_float_array(data)
wxchan's avatar
wxchan committed
1100
        elif data.dtype == np.int32:
1101
            ptr_data, type_data, _ = c_int_array(data)
wxchan's avatar
wxchan committed
1102
        else:
Guolin Ke's avatar
Guolin Ke committed
1103
            raise TypeError("Excepted np.float32/64 or np.int32, meet type({})".format(data.dtype))
wxchan's avatar
wxchan committed
1104
        if type_data != FIELD_TYPE_MAPPER[field_name]:
1105
            raise TypeError("Input type error for set_field")
wxchan's avatar
wxchan committed
1106
1107
1108
1109
        _safe_call(_LIB.LGBM_DatasetSetField(
            self.handle,
            c_str(field_name),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
1110
1111
            ctypes.c_int(len(data)),
            ctypes.c_int(type_data)))
Nikita Titov's avatar
Nikita Titov committed
1112
        return self
wxchan's avatar
wxchan committed
1113

wxchan's avatar
wxchan committed
1114
1115
    def get_field(self, field_name):
        """Get property from the Dataset.
wxchan's avatar
wxchan committed
1116
1117
1118

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1119
        field_name : string
1120
            The field name of the information.
wxchan's avatar
wxchan committed
1121
1122
1123

        Returns
        -------
1124
1125
        info : numpy array
            A numpy array with information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1126
        """
1127
        if self.handle is None:
1128
            raise Exception("Cannot get %s before construct Dataset" % field_name)
Guolin Ke's avatar
Guolin Ke committed
1129
1130
        tmp_out_len = ctypes.c_int()
        out_type = ctypes.c_int()
wxchan's avatar
wxchan committed
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
        ret = ctypes.POINTER(ctypes.c_void_p)()
        _safe_call(_LIB.LGBM_DatasetGetField(
            self.handle,
            c_str(field_name),
            ctypes.byref(tmp_out_len),
            ctypes.byref(ret),
            ctypes.byref(out_type)))
        if out_type.value != FIELD_TYPE_MAPPER[field_name]:
            raise TypeError("Return type error for get_field")
        if tmp_out_len.value == 0:
            return None
        if out_type.value == C_API_DTYPE_INT32:
            return cint32_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_int32)), tmp_out_len.value)
        elif out_type.value == C_API_DTYPE_FLOAT32:
            return cfloat32_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_float)), tmp_out_len.value)
Guolin Ke's avatar
Guolin Ke committed
1146
1147
        elif out_type.value == C_API_DTYPE_FLOAT64:
            return cfloat64_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_double)), tmp_out_len.value)
1148
        else:
wxchan's avatar
wxchan committed
1149
            raise TypeError("Unknown type")
Guolin Ke's avatar
Guolin Ke committed
1150

1151
    def set_categorical_feature(self, categorical_feature):
1152
        """Set categorical features.
1153
1154
1155

        Parameters
        ----------
1156
1157
        categorical_feature : list of int or strings
            Names or indices of categorical features.
Nikita Titov's avatar
Nikita Titov committed
1158
1159
1160
1161
1162

        Returns
        -------
        self : Dataset
            Dataset with set categorical features.
1163
1164
        """
        if self.categorical_feature == categorical_feature:
Nikita Titov's avatar
Nikita Titov committed
1165
            return self
1166
        if self.data is not None:
1167
1168
            if self.categorical_feature is None:
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
1169
                return self._free_handle()
1170
1171
            elif categorical_feature == 'auto':
                warnings.warn('Using categorical_feature in Dataset.')
Nikita Titov's avatar
Nikita Titov committed
1172
                return self
1173
            else:
1174
1175
                warnings.warn('categorical_feature in Dataset is overridden.\n'
                              'New categorical_feature is {}'.format(sorted(list(categorical_feature))))
1176
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
1177
                return self._free_handle()
1178
        else:
1179
1180
            raise LightGBMError("Cannot set categorical feature after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
1181

Guolin Ke's avatar
Guolin Ke committed
1182
1183
    def _set_predictor(self, predictor):
        """
1184
        Set predictor for continued training, not recommended for user to call this function.
Guolin Ke's avatar
Guolin Ke committed
1185
1186
1187
        Please set init_model in engine.train or engine.cv
        """
        if predictor is self._predictor:
Nikita Titov's avatar
Nikita Titov committed
1188
            return self
Guolin Ke's avatar
Guolin Ke committed
1189
1190
        if self.data is not None:
            self._predictor = predictor
Nikita Titov's avatar
Nikita Titov committed
1191
            return self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1192
        else:
1193
1194
            raise LightGBMError("Cannot set predictor after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
Guolin Ke's avatar
Guolin Ke committed
1195
1196

    def set_reference(self, reference):
1197
        """Set reference Dataset.
Guolin Ke's avatar
Guolin Ke committed
1198
1199
1200
1201

        Parameters
        ----------
        reference : Dataset
1202
            Reference that is used as a template to consturct the current Dataset.
Nikita Titov's avatar
Nikita Titov committed
1203
1204
1205
1206
1207

        Returns
        -------
        self : Dataset
            Dataset with set reference.
Guolin Ke's avatar
Guolin Ke committed
1208
        """
1209
1210
1211
        self.set_categorical_feature(reference.categorical_feature) \
            .set_feature_name(reference.feature_name) \
            ._set_predictor(reference._predictor)
1212
1213
        # we're done if self and reference share a common upstrem reference
        if self.get_ref_chain().intersection(reference.get_ref_chain()):
Nikita Titov's avatar
Nikita Titov committed
1214
            return self
Guolin Ke's avatar
Guolin Ke committed
1215
1216
        if self.data is not None:
            self.reference = reference
Nikita Titov's avatar
Nikita Titov committed
1217
            return self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1218
        else:
1219
1220
            raise LightGBMError("Cannot set reference after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
Guolin Ke's avatar
Guolin Ke committed
1221
1222

    def set_feature_name(self, feature_name):
1223
        """Set feature name.
Guolin Ke's avatar
Guolin Ke committed
1224
1225
1226

        Parameters
        ----------
1227
1228
        feature_name : list of strings
            Feature names.
Nikita Titov's avatar
Nikita Titov committed
1229
1230
1231
1232
1233

        Returns
        -------
        self : Dataset
            Dataset with set feature name.
Guolin Ke's avatar
Guolin Ke committed
1234
        """
1235
1236
        if feature_name != 'auto':
            self.feature_name = feature_name
1237
        if self.handle is not None and feature_name is not None and feature_name != 'auto':
wxchan's avatar
wxchan committed
1238
            if len(feature_name) != self.num_feature():
1239
1240
                raise ValueError("Length of feature_name({}) and num_feature({}) don't match"
                                 .format(len(feature_name), self.num_feature()))
1241
            c_feature_name = [c_str(name) for name in feature_name]
wxchan's avatar
wxchan committed
1242
1243
1244
            _safe_call(_LIB.LGBM_DatasetSetFeatureNames(
                self.handle,
                c_array(ctypes.c_char_p, c_feature_name),
Guolin Ke's avatar
Guolin Ke committed
1245
                ctypes.c_int(len(feature_name))))
Nikita Titov's avatar
Nikita Titov committed
1246
        return self
Guolin Ke's avatar
Guolin Ke committed
1247
1248

    def set_label(self, label):
1249
        """Set label of Dataset
Guolin Ke's avatar
Guolin Ke committed
1250
1251
1252

        Parameters
        ----------
1253
        label : list, numpy 1-D array, pandas one-column DataFrame/Series or None
1254
            The label information to be set into Dataset.
Nikita Titov's avatar
Nikita Titov committed
1255
1256
1257
1258
1259

        Returns
        -------
        self : Dataset
            Dataset with set label.
Guolin Ke's avatar
Guolin Ke committed
1260
1261
        """
        self.label = label
1262
        if self.handle is not None:
1263
            label = list_to_1d_numpy(_label_from_pandas(label), name='label')
wxchan's avatar
wxchan committed
1264
            self.set_field('label', label)
Nikita Titov's avatar
Nikita Titov committed
1265
        return self
Guolin Ke's avatar
Guolin Ke committed
1266
1267

    def set_weight(self, weight):
1268
        """Set weight of each instance.
Guolin Ke's avatar
Guolin Ke committed
1269
1270
1271

        Parameters
        ----------
1272
        weight : list, numpy 1-D array, pandas Series or None
1273
            Weight to be set for each data point.
Nikita Titov's avatar
Nikita Titov committed
1274
1275
1276
1277
1278

        Returns
        -------
        self : Dataset
            Dataset with set weight.
Guolin Ke's avatar
Guolin Ke committed
1279
        """
1280
1281
        if weight is not None and np.all(weight == 1):
            weight = None
Guolin Ke's avatar
Guolin Ke committed
1282
        self.weight = weight
1283
        if self.handle is not None and weight is not None:
wxchan's avatar
wxchan committed
1284
1285
            weight = list_to_1d_numpy(weight, name='weight')
            self.set_field('weight', weight)
Nikita Titov's avatar
Nikita Titov committed
1286
        return self
Guolin Ke's avatar
Guolin Ke committed
1287
1288

    def set_init_score(self, init_score):
1289
        """Set init score of Booster to start from.
Guolin Ke's avatar
Guolin Ke committed
1290
1291
1292

        Parameters
        ----------
1293
        init_score : list, numpy 1-D array, pandas Series or None
1294
            Init score for Booster.
Nikita Titov's avatar
Nikita Titov committed
1295
1296
1297
1298
1299

        Returns
        -------
        self : Dataset
            Dataset with set init score.
Guolin Ke's avatar
Guolin Ke committed
1300
1301
        """
        self.init_score = init_score
1302
        if self.handle is not None and init_score is not None:
Guolin Ke's avatar
Guolin Ke committed
1303
            init_score = list_to_1d_numpy(init_score, np.float64, name='init_score')
wxchan's avatar
wxchan committed
1304
            self.set_field('init_score', init_score)
Nikita Titov's avatar
Nikita Titov committed
1305
        return self
Guolin Ke's avatar
Guolin Ke committed
1306
1307

    def set_group(self, group):
1308
        """Set group size of Dataset (used for ranking).
Guolin Ke's avatar
Guolin Ke committed
1309
1310
1311

        Parameters
        ----------
1312
        group : list, numpy 1-D array, pandas Series or None
1313
            Group size of each group.
Nikita Titov's avatar
Nikita Titov committed
1314
1315
1316
1317
1318

        Returns
        -------
        self : Dataset
            Dataset with set group.
Guolin Ke's avatar
Guolin Ke committed
1319
1320
        """
        self.group = group
1321
        if self.handle is not None and group is not None:
wxchan's avatar
wxchan committed
1322
1323
            group = list_to_1d_numpy(group, np.int32, name='group')
            self.set_field('group', group)
Nikita Titov's avatar
Nikita Titov committed
1324
        return self
Guolin Ke's avatar
Guolin Ke committed
1325
1326

    def get_label(self):
1327
        """Get the label of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1328
1329
1330

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1331
        label : numpy array or None
1332
            The label information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1333
        """
1334
        if self.label is None:
wxchan's avatar
wxchan committed
1335
            self.label = self.get_field('label')
Guolin Ke's avatar
Guolin Ke committed
1336
1337
1338
        return self.label

    def get_weight(self):
1339
        """Get the weight of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1340
1341
1342

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1343
        weight : numpy array or None
1344
            Weight for each data point from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1345
        """
1346
        if self.weight is None:
wxchan's avatar
wxchan committed
1347
            self.weight = self.get_field('weight')
Guolin Ke's avatar
Guolin Ke committed
1348
1349
1350
        return self.weight

    def get_init_score(self):
1351
        """Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1352
1353
1354

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1355
        init_score : numpy array or None
1356
            Init score of Booster.
Guolin Ke's avatar
Guolin Ke committed
1357
        """
1358
        if self.init_score is None:
wxchan's avatar
wxchan committed
1359
            self.init_score = self.get_field('init_score')
Guolin Ke's avatar
Guolin Ke committed
1360
1361
1362
        return self.init_score

    def get_group(self):
1363
        """Get the group of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1364
1365
1366

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1367
        group : numpy array or None
1368
            Group size of each group.
Guolin Ke's avatar
Guolin Ke committed
1369
        """
1370
        if self.group is None:
wxchan's avatar
wxchan committed
1371
            self.group = self.get_field('group')
Guolin Ke's avatar
Guolin Ke committed
1372
1373
            if self.group is not None:
                # group data from LightGBM is boundaries data, need to convert to group size
Nikita Titov's avatar
Nikita Titov committed
1374
                self.group = np.diff(self.group)
Guolin Ke's avatar
Guolin Ke committed
1375
1376
1377
        return self.group

    def num_data(self):
1378
        """Get the number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1379
1380
1381

        Returns
        -------
1382
1383
        number_of_rows : int
            The number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1384
        """
1385
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1386
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1387
1388
1389
            _safe_call(_LIB.LGBM_DatasetGetNumData(self.handle,
                                                   ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1390
        else:
1391
            raise LightGBMError("Cannot get num_data before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1392
1393

    def num_feature(self):
1394
        """Get the number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1395
1396
1397

        Returns
        -------
1398
1399
        number_of_columns : int
            The number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1400
        """
1401
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1402
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1403
1404
1405
            _safe_call(_LIB.LGBM_DatasetGetNumFeature(self.handle,
                                                      ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1406
        else:
1407
            raise LightGBMError("Cannot get num_feature before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1408

1409
    def get_ref_chain(self, ref_limit=100):
1410
1411
1412
1413
1414
1415
1416
        """Get a chain of Dataset objects, starting with r, then going to r.reference if exists,
        then to r.reference.reference, etc. until we hit ``ref_limit`` or a reference loop.

        Parameters
        ----------
        ref_limit : int, optional (default=100)
            The limit number of references.
1417
1418
1419

        Returns
        -------
1420
1421
1422
        ref_chain : set of Dataset
            Chain of references of the Datasets.
        """
1423
        head = self
1424
        ref_chain = set()
1425
1426
        while len(ref_chain) < ref_limit:
            if isinstance(head, Dataset):
1427
                ref_chain.add(head)
1428
1429
1430
1431
1432
1433
                if (head.reference is not None) and (head.reference not in ref_chain):
                    head = head.reference
                else:
                    break
            else:
                break
Nikita Titov's avatar
Nikita Titov committed
1434
        return ref_chain
1435

wxchan's avatar
wxchan committed
1436

wxchan's avatar
wxchan committed
1437
class Booster(object):
1438
    """Booster in LightGBM."""
wxchan's avatar
wxchan committed
1439
    def __init__(self, params=None, train_set=None, model_file=None, silent=False):
1440
        """Initialize the Booster.
wxchan's avatar
wxchan committed
1441
1442
1443

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1444
        params : dict or None, optional (default=None)
1445
1446
1447
1448
            Parameters for Booster.
        train_set : Dataset or None, optional (default=None)
            Training dataset.
        model_file : string or None, optional (default=None)
wxchan's avatar
wxchan committed
1449
            Path to the model file.
1450
1451
        silent : bool, optional (default=False)
            Whether to print messages during construction.
wxchan's avatar
wxchan committed
1452
        """
1453
        self.handle = None
1454
        self.network = False
wxchan's avatar
wxchan committed
1455
1456
1457
        self.__need_reload_eval_info = True
        self.__train_data_name = "training"
        self.__attr = {}
1458
        self.__set_objective_to_none = False
wxchan's avatar
wxchan committed
1459
        self.best_iteration = -1
wxchan's avatar
wxchan committed
1460
        self.best_score = {}
wxchan's avatar
wxchan committed
1461
        params = {} if params is None else params
1462
1463
1464
        # user can set verbose with params, it has higher priority
        if not any(verbose_alias in params for verbose_alias in ('verbose', 'verbosity')) and silent:
            params["verbose"] = -1
wxchan's avatar
wxchan committed
1465
        if train_set is not None:
1466
            # Training task
wxchan's avatar
wxchan committed
1467
            if not isinstance(train_set, Dataset):
1468
1469
                raise TypeError('Training data should be Dataset instance, met {}'
                                .format(type(train_set).__name__))
wxchan's avatar
wxchan committed
1470
            params_str = param_dict_to_str(params)
1471
            # construct booster object
1472
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1473
            _safe_call(_LIB.LGBM_BoosterCreate(
wxchan's avatar
wxchan committed
1474
                train_set.construct().handle,
wxchan's avatar
wxchan committed
1475
1476
                c_str(params_str),
                ctypes.byref(self.handle)))
1477
            # save reference to data
wxchan's avatar
wxchan committed
1478
1479
1480
1481
            self.train_set = train_set
            self.valid_sets = []
            self.name_valid_sets = []
            self.__num_dataset = 1
Guolin Ke's avatar
Guolin Ke committed
1482
1483
            self.__init_predictor = train_set._predictor
            if self.__init_predictor is not None:
wxchan's avatar
wxchan committed
1484
1485
                _safe_call(_LIB.LGBM_BoosterMerge(
                    self.handle,
Guolin Ke's avatar
Guolin Ke committed
1486
                    self.__init_predictor.handle))
Guolin Ke's avatar
Guolin Ke committed
1487
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1488
1489
1490
1491
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1492
            # buffer for inner predict
wxchan's avatar
wxchan committed
1493
1494
1495
            self.__inner_predict_buffer = [None]
            self.__is_predicted_cur_iter = [False]
            self.__get_eval_info()
1496
            self.pandas_categorical = train_set.pandas_categorical
1497
            # set network if necessary
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
            for alias in ["machines", "workers", "nodes"]:
                if alias in params:
                    machines = params[alias]
                    if isinstance(machines, string_type):
                        num_machines = len(machines.split(','))
                    elif isinstance(machines, (list, set)):
                        num_machines = len(machines)
                        machines = ','.join(machines)
                    else:
                        raise ValueError("Invalid machines in params.")
                    self.set_network(machines,
                                     local_listen_port=params.get("local_listen_port", 12400),
                                     listen_time_out=params.get("listen_time_out", 120),
                                     num_machines=params.get("num_machines", num_machines))
                    break
wxchan's avatar
wxchan committed
1513
        elif model_file is not None:
1514
            # Prediction task
Guolin Ke's avatar
Guolin Ke committed
1515
            out_num_iterations = ctypes.c_int(0)
1516
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1517
1518
1519
1520
            _safe_call(_LIB.LGBM_BoosterCreateFromModelfile(
                c_str(model_file),
                ctypes.byref(out_num_iterations),
                ctypes.byref(self.handle)))
Guolin Ke's avatar
Guolin Ke committed
1521
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1522
1523
1524
1525
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1526
            self.pandas_categorical = _load_pandas_categorical(model_file)
1527
        elif 'model_str' in params:
1528
            self.model_from_string(params['model_str'])
wxchan's avatar
wxchan committed
1529
        else:
1530
            raise TypeError('Need at least one training dataset or model file to create booster instance')
Guolin Ke's avatar
Guolin Ke committed
1531
        self.params = params.copy()
wxchan's avatar
wxchan committed
1532
1533

    def __del__(self):
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
        try:
            if self.network:
                self.free_network()
        except AttributeError:
            pass
        try:
            if self.handle is not None:
                _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
1544

wxchan's avatar
wxchan committed
1545
1546
1547
1548
    def __copy__(self):
        return self.__deepcopy__(None)

    def __deepcopy__(self, _):
1549
        model_str = self.model_to_string(num_iteration=-1)
1550
        booster = Booster({'model_str': model_str})
1551
        booster.pandas_categorical = self.pandas_categorical
1552
        return booster
wxchan's avatar
wxchan committed
1553
1554
1555
1556
1557
1558
1559

    def __getstate__(self):
        this = self.__dict__.copy()
        handle = this['handle']
        this.pop('train_set', None)
        this.pop('valid_sets', None)
        if handle is not None:
1560
            this["handle"] = self.model_to_string(num_iteration=-1)
wxchan's avatar
wxchan committed
1561
1562
1563
        return this

    def __setstate__(self, state):
1564
1565
        model_str = state.get('handle', None)
        if model_str is not None:
wxchan's avatar
wxchan committed
1566
            handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
1567
            out_num_iterations = ctypes.c_int(0)
1568
1569
1570
1571
            _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
                c_str(model_str),
                ctypes.byref(out_num_iterations),
                ctypes.byref(handle)))
wxchan's avatar
wxchan committed
1572
1573
1574
            state['handle'] = handle
        self.__dict__.update(state)

wxchan's avatar
wxchan committed
1575
    def free_dataset(self):
Nikita Titov's avatar
Nikita Titov committed
1576
1577
1578
1579
1580
1581
1582
        """Free Booster's Datasets.

        Returns
        -------
        self : Booster
            Booster without Datasets.
        """
wxchan's avatar
wxchan committed
1583
1584
        self.__dict__.pop('train_set', None)
        self.__dict__.pop('valid_sets', None)
1585
        self.__num_dataset = 0
Nikita Titov's avatar
Nikita Titov committed
1586
        return self
wxchan's avatar
wxchan committed
1587

1588
1589
1590
    def _free_buffer(self):
        self.__inner_predict_buffer = []
        self.__is_predicted_cur_iter = []
Nikita Titov's avatar
Nikita Titov committed
1591
        return self
1592

1593
1594
1595
1596
1597
1598
    def set_network(self, machines, local_listen_port=12400,
                    listen_time_out=120, num_machines=1):
        """Set the network configuration.

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1599
        machines : list, set or string
1600
            Names of machines.
Nikita Titov's avatar
Nikita Titov committed
1601
        local_listen_port : int, optional (default=12400)
1602
            TCP listen port for local machines.
Nikita Titov's avatar
Nikita Titov committed
1603
        listen_time_out : int, optional (default=120)
1604
            Socket time-out in minutes.
Nikita Titov's avatar
Nikita Titov committed
1605
        num_machines : int, optional (default=1)
1606
            The number of machines for parallel learning application.
Nikita Titov's avatar
Nikita Titov committed
1607
1608
1609
1610
1611

        Returns
        -------
        self : Booster
            Booster with set network.
1612
1613
1614
1615
1616
1617
        """
        _safe_call(_LIB.LGBM_NetworkInit(c_str(machines),
                                         ctypes.c_int(local_listen_port),
                                         ctypes.c_int(listen_time_out),
                                         ctypes.c_int(num_machines)))
        self.network = True
Nikita Titov's avatar
Nikita Titov committed
1618
        return self
1619
1620

    def free_network(self):
Nikita Titov's avatar
Nikita Titov committed
1621
1622
1623
1624
1625
1626
1627
        """Free Booster's network.

        Returns
        -------
        self : Booster
            Booster with freed network.
        """
1628
1629
        _safe_call(_LIB.LGBM_NetworkFree())
        self.network = False
Nikita Titov's avatar
Nikita Titov committed
1630
        return self
1631

wxchan's avatar
wxchan committed
1632
    def set_train_data_name(self, name):
1633
1634
1635
1636
        """Set the name to the training Dataset.

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1637
1638
1639
1640
1641
1642
1643
        name : string
            Name for the training Dataset.

        Returns
        -------
        self : Booster
            Booster with set training Dataset name.
1644
        """
wxchan's avatar
wxchan committed
1645
        self.__train_data_name = name
Nikita Titov's avatar
Nikita Titov committed
1646
        return self
wxchan's avatar
wxchan committed
1647
1648

    def add_valid(self, data, name):
1649
        """Add validation data.
wxchan's avatar
wxchan committed
1650
1651
1652
1653

        Parameters
        ----------
        data : Dataset
1654
1655
1656
            Validation data.
        name : string
            Name of validation data.
Nikita Titov's avatar
Nikita Titov committed
1657
1658
1659
1660
1661

        Returns
        -------
        self : Booster
            Booster with set validation data.
wxchan's avatar
wxchan committed
1662
        """
Guolin Ke's avatar
Guolin Ke committed
1663
        if not isinstance(data, Dataset):
1664
1665
            raise TypeError('Validation data should be Dataset instance, met {}'
                            .format(type(data).__name__))
Guolin Ke's avatar
Guolin Ke committed
1666
        if data._predictor is not self.__init_predictor:
1667
1668
            raise LightGBMError("Add validation data failed, "
                                "you should use same predictor for these data")
wxchan's avatar
wxchan committed
1669
1670
        _safe_call(_LIB.LGBM_BoosterAddValidData(
            self.handle,
wxchan's avatar
wxchan committed
1671
            data.construct().handle))
wxchan's avatar
wxchan committed
1672
1673
1674
1675
1676
        self.valid_sets.append(data)
        self.name_valid_sets.append(name)
        self.__num_dataset += 1
        self.__inner_predict_buffer.append(None)
        self.__is_predicted_cur_iter.append(False)
Nikita Titov's avatar
Nikita Titov committed
1677
        return self
wxchan's avatar
wxchan committed
1678
1679

    def reset_parameter(self, params):
1680
        """Reset parameters of Booster.
wxchan's avatar
wxchan committed
1681
1682
1683
1684

        Parameters
        ----------
        params : dict
1685
            New parameters for Booster.
Nikita Titov's avatar
Nikita Titov committed
1686
1687
1688
1689
1690

        Returns
        -------
        self : Booster
            Booster with new parameters.
wxchan's avatar
wxchan committed
1691
        """
1692
        if any(metric_alias in params for metric_alias in ('metric', 'metrics', 'metric_types')):
wxchan's avatar
wxchan committed
1693
1694
1695
1696
1697
1698
            self.__need_reload_eval_info = True
        params_str = param_dict_to_str(params)
        if params_str:
            _safe_call(_LIB.LGBM_BoosterResetParameter(
                self.handle,
                c_str(params_str)))
Guolin Ke's avatar
Guolin Ke committed
1699
        self.params.update(params)
Nikita Titov's avatar
Nikita Titov committed
1700
        return self
wxchan's avatar
wxchan committed
1701
1702

    def update(self, train_set=None, fobj=None):
Nikita Titov's avatar
Nikita Titov committed
1703
        """Update Booster for one iteration.
1704

wxchan's avatar
wxchan committed
1705
1706
        Parameters
        ----------
1707
1708
1709
1710
        train_set : Dataset or None, optional (default=None)
            Training data.
            If None, last training data is used.
        fobj : callable or None, optional (default=None)
wxchan's avatar
wxchan committed
1711
1712
            Customized objective function.

1713
1714
1715
1716
            For multi-class task, the score is group by class_id first, then group by row_id.
            If you want to get i-th row score in j-th class, the access way is score[j * num_data + i]
            and you should group grad and hess in this way as well.

wxchan's avatar
wxchan committed
1717
1718
        Returns
        -------
1719
1720
        is_finished : bool
            Whether the update was successfully finished.
wxchan's avatar
wxchan committed
1721
1722
        """

1723
        # need reset training data
wxchan's avatar
wxchan committed
1724
        if train_set is not None and train_set is not self.train_set:
Guolin Ke's avatar
Guolin Ke committed
1725
            if not isinstance(train_set, Dataset):
1726
1727
                raise TypeError('Training data should be Dataset instance, met {}'
                                .format(type(train_set).__name__))
Guolin Ke's avatar
Guolin Ke committed
1728
            if train_set._predictor is not self.__init_predictor:
1729
1730
                raise LightGBMError("Replace training data failed, "
                                    "you should use same predictor for these data")
wxchan's avatar
wxchan committed
1731
1732
1733
            self.train_set = train_set
            _safe_call(_LIB.LGBM_BoosterResetTrainingData(
                self.handle,
wxchan's avatar
wxchan committed
1734
                self.train_set.construct().handle))
wxchan's avatar
wxchan committed
1735
1736
1737
            self.__inner_predict_buffer[0] = None
        is_finished = ctypes.c_int(0)
        if fobj is None:
1738
1739
            if self.__set_objective_to_none:
                raise ValueError('Cannot update due to null objective function.')
wxchan's avatar
wxchan committed
1740
1741
1742
            _safe_call(_LIB.LGBM_BoosterUpdateOneIter(
                self.handle,
                ctypes.byref(is_finished)))
wxchan's avatar
wxchan committed
1743
            self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1744
1745
            return is_finished.value == 1
        else:
1746
            if not self.__set_objective_to_none:
Nikita Titov's avatar
Nikita Titov committed
1747
                self.reset_parameter({"objective": "none"}).__set_objective_to_none = True
wxchan's avatar
wxchan committed
1748
1749
1750
1751
1752
            grad, hess = fobj(self.__inner_predict(0), self.train_set)
            return self.__boost(grad, hess)

    def __boost(self, grad, hess):
        """
Nikita Titov's avatar
Nikita Titov committed
1753
1754
1755
1756
1757
        Boost Booster for one iteration with customized gradient statistics.

        Note: For multi-class task, the score is group by class_id first, then group by row_id.
              If you want to get i-th row score in j-th class, the access way is score[j * num_data + i]
              and you should group grad and hess in this way as well.
1758

wxchan's avatar
wxchan committed
1759
1760
        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1761
1762
        grad : 1d numpy array or list
            The first order derivative (gradient).
wxchan's avatar
wxchan committed
1763
        hess : 1d numpy or 1d list
Nikita Titov's avatar
Nikita Titov committed
1764
            The second order derivative (Hessian).
wxchan's avatar
wxchan committed
1765
1766
1767

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1768
1769
        is_finished : bool
            Whether the boost was successfully finished.
wxchan's avatar
wxchan committed
1770
        """
1771
1772
        grad = list_to_1d_numpy(grad, name='gradient')
        hess = list_to_1d_numpy(hess, name='hessian')
1773
1774
        assert grad.flags.c_contiguous
        assert hess.flags.c_contiguous
wxchan's avatar
wxchan committed
1775
        if len(grad) != len(hess):
1776
1777
            raise ValueError("Lengths of gradient({}) and hessian({}) don't match"
                             .format(len(grad), len(hess)))
wxchan's avatar
wxchan committed
1778
1779
1780
1781
1782
1783
        is_finished = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterUpdateOneIterCustom(
            self.handle,
            grad.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            hess.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            ctypes.byref(is_finished)))
wxchan's avatar
wxchan committed
1784
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1785
1786
1787
        return is_finished.value == 1

    def rollback_one_iter(self):
Nikita Titov's avatar
Nikita Titov committed
1788
1789
1790
1791
1792
1793
1794
        """Rollback one iteration.

        Returns
        -------
        self : Booster
            Booster with rolled back one iteration.
        """
wxchan's avatar
wxchan committed
1795
1796
        _safe_call(_LIB.LGBM_BoosterRollbackOneIter(
            self.handle))
wxchan's avatar
wxchan committed
1797
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
Nikita Titov's avatar
Nikita Titov committed
1798
        return self
wxchan's avatar
wxchan committed
1799
1800

    def current_iteration(self):
1801
1802
1803
1804
1805
1806
1807
        """Get the index of the current iteration.

        Returns
        -------
        cur_iter : int
            The index of the current iteration.
        """
Guolin Ke's avatar
Guolin Ke committed
1808
        out_cur_iter = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1809
1810
1811
1812
1813
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
            self.handle,
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
    def num_model_per_iteration(self):
        """Get number of models per iteration.

        Returns
        -------
        model_per_iter : int
            The number of models per iteration.
        """
        model_per_iter = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterNumModelPerIteration(
            self.handle,
            ctypes.byref(model_per_iter)))
        return model_per_iter.value

    def num_trees(self):
        """Get number of weak sub-models.

        Returns
        -------
        num_trees : int
            The number of weak sub-models.
        """
        num_trees = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterNumberOfTotalModel(
            self.handle,
            ctypes.byref(num_trees)))
        return num_trees.value

wxchan's avatar
wxchan committed
1842
    def eval(self, data, name, feval=None):
1843
        """Evaluate for data.
wxchan's avatar
wxchan committed
1844
1845
1846

        Parameters
        ----------
1847
1848
1849
1850
1851
        data : Dataset
            Data for the evaluating.
        name : string
            Name of the data.
        feval : callable or None, optional (default=None)
1852
1853
1854
1855
1856
            Customized evaluation function.
            Should accept two parameters: preds, train_data.
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
            Note: should return (eval_name, eval_result, is_higher_better) or list of such tuples.
1857

wxchan's avatar
wxchan committed
1858
1859
        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1860
        result : list
1861
            List with evaluation results.
wxchan's avatar
wxchan committed
1862
        """
Guolin Ke's avatar
Guolin Ke committed
1863
1864
        if not isinstance(data, Dataset):
            raise TypeError("Can only eval for Dataset instance")
wxchan's avatar
wxchan committed
1865
1866
1867
1868
        data_idx = -1
        if data is self.train_set:
            data_idx = 0
        else:
wxchan's avatar
wxchan committed
1869
            for i in range_(len(self.valid_sets)):
wxchan's avatar
wxchan committed
1870
1871
1872
                if data is self.valid_sets[i]:
                    data_idx = i + 1
                    break
1873
        # need to push new valid data
wxchan's avatar
wxchan committed
1874
1875
1876
1877
1878
1879
1880
        if data_idx == -1:
            self.add_valid(data, name)
            data_idx = self.__num_dataset - 1

        return self.__inner_eval(name, data_idx, feval)

    def eval_train(self, feval=None):
1881
        """Evaluate for training data.
wxchan's avatar
wxchan committed
1882
1883
1884

        Parameters
        ----------
1885
        feval : callable or None, optional (default=None)
1886
1887
1888
1889
1890
            Customized evaluation function.
            Should accept two parameters: preds, train_data.
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
            Note: should return (eval_name, eval_result, is_higher_better) or list of such tuples.
wxchan's avatar
wxchan committed
1891
1892
1893

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1894
        result : list
1895
            List with evaluation results.
wxchan's avatar
wxchan committed
1896
1897
1898
1899
        """
        return self.__inner_eval(self.__train_data_name, 0, feval)

    def eval_valid(self, feval=None):
1900
        """Evaluate for validation data.
wxchan's avatar
wxchan committed
1901
1902
1903

        Parameters
        ----------
1904
        feval : callable or None, optional (default=None)
1905
1906
1907
1908
1909
            Customized evaluation function.
            Should accept two parameters: preds, train_data.
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
            Note: should return (eval_name, eval_result, is_higher_better) or list of such tuples.
wxchan's avatar
wxchan committed
1910
1911
1912

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1913
        result : list
1914
            List with evaluation results.
wxchan's avatar
wxchan committed
1915
        """
wxchan's avatar
wxchan committed
1916
        return [item for i in range_(1, self.__num_dataset)
wxchan's avatar
wxchan committed
1917
                for item in self.__inner_eval(self.name_valid_sets[i - 1], i, feval)]
wxchan's avatar
wxchan committed
1918

1919
    def save_model(self, filename, num_iteration=None, start_iteration=0):
1920
        """Save Booster to file.
wxchan's avatar
wxchan committed
1921
1922
1923

        Parameters
        ----------
1924
1925
        filename : string
            Filename to save Booster.
1926
1927
1928
1929
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
Nikita Titov's avatar
Nikita Titov committed
1930
        start_iteration : int, optional (default=0)
1931
            Start index of the iteration that should be saved.
Nikita Titov's avatar
Nikita Titov committed
1932
1933
1934
1935
1936

        Returns
        -------
        self : Booster
            Returns self.
wxchan's avatar
wxchan committed
1937
        """
1938
        if num_iteration is None:
1939
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1940
1941
        _safe_call(_LIB.LGBM_BoosterSaveModel(
            self.handle,
1942
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
1943
            ctypes.c_int(num_iteration),
wxchan's avatar
wxchan committed
1944
            c_str(filename)))
1945
        _save_pandas_categorical(filename, self.pandas_categorical)
Nikita Titov's avatar
Nikita Titov committed
1946
        return self
wxchan's avatar
wxchan committed
1947

1948
    def shuffle_models(self, start_iteration=0, end_iteration=-1):
1949
        """Shuffle models.
Nikita Titov's avatar
Nikita Titov committed
1950

1951
1952
1953
1954
1955
1956
1957
1958
        Parameters
        ----------
        start_iteration : int, optional (default=0)
            Index of the iteration that will start to shuffle.
        end_iteration : int, optional (default=-1)
            The last iteration that will be shuffled.
            If <= 0, means the last iteration.

Nikita Titov's avatar
Nikita Titov committed
1959
1960
1961
1962
        Returns
        -------
        self : Booster
            Booster with shuffled models.
1963
        """
1964
1965
1966
1967
        _safe_call(_LIB.LGBM_BoosterShuffleModels(
            self.handle,
            ctypes.c_int(start_iter),
            ctypes.c_int(end_iter)))
Nikita Titov's avatar
Nikita Titov committed
1968
        return self
1969
1970
1971
1972
1973
1974

    def model_from_string(self, model_str, verbose=True):
        """Load Booster from a string.

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1975
        model_str : string
1976
            Model will be loaded from this string.
Nikita Titov's avatar
Nikita Titov committed
1977
1978
        verbose : bool, optional (default=True)
            Whether to print messages while loading model.
1979
1980
1981

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1982
        self : Booster
1983
1984
            Loaded Booster object.
        """
1985
1986
1987
1988
        if self.handle is not None:
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        self._free_buffer()
        self.handle = ctypes.c_void_p()
1989
1990
1991
1992
1993
1994
1995
1996
1997
        out_num_iterations = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
            c_str(model_str),
            ctypes.byref(out_num_iterations),
            ctypes.byref(self.handle)))
        out_num_class = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetNumClasses(
            self.handle,
            ctypes.byref(out_num_class)))
1998
        if verbose:
Nikita Titov's avatar
Nikita Titov committed
1999
            print('Finished loading model, total used %d iterations' % int(out_num_iterations.value))
2000
        self.__num_class = out_num_class.value
2001
2002
2003
2004
        return self

    def model_to_string(self, num_iteration=None, start_iteration=0):
        """Save Booster to string.
2005

2006
2007
2008
2009
2010
2011
        Parameters
        ----------
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
Nikita Titov's avatar
Nikita Titov committed
2012
        start_iteration : int, optional (default=0)
2013
2014
2015
2016
            Start index of the iteration that should be saved.

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2017
        str_repr : string
2018
2019
            String representation of Booster.
        """
2020
        if num_iteration is None:
2021
2022
            num_iteration = self.best_iteration
        buffer_len = 1 << 20
2023
        tmp_out_len = ctypes.c_int64(0)
2024
2025
2026
2027
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterSaveModelToString(
            self.handle,
2028
            ctypes.c_int(start_iteration),
2029
            ctypes.c_int(num_iteration),
2030
            ctypes.c_int64(buffer_len),
2031
2032
2033
2034
2035
2036
2037
2038
2039
            ctypes.byref(tmp_out_len),
            ptr_string_buffer))
        actual_len = tmp_out_len.value
        '''if buffer length is not long enough, re-allocate a buffer'''
        if actual_len > buffer_len:
            string_buffer = ctypes.create_string_buffer(actual_len)
            ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
            _safe_call(_LIB.LGBM_BoosterSaveModelToString(
                self.handle,
2040
                ctypes.c_int(start_iteration),
2041
                ctypes.c_int(num_iteration),
2042
                ctypes.c_int64(actual_len),
2043
2044
2045
2046
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
        return string_buffer.value.decode()

2047
    def dump_model(self, num_iteration=None, start_iteration=0):
Nikita Titov's avatar
Nikita Titov committed
2048
        """Dump Booster to JSON format.
wxchan's avatar
wxchan committed
2049

2050
2051
        Parameters
        ----------
2052
2053
2054
2055
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be dumped.
            If None, if the best iteration exists, it is dumped; otherwise, all iterations are dumped.
            If <= 0, all iterations are dumped.
Nikita Titov's avatar
Nikita Titov committed
2056
        start_iteration : int, optional (default=0)
2057
            Start index of the iteration that should be dumped.
2058

wxchan's avatar
wxchan committed
2059
2060
        Returns
        -------
2061
        json_repr : dict
Nikita Titov's avatar
Nikita Titov committed
2062
            JSON format of Booster.
wxchan's avatar
wxchan committed
2063
        """
2064
        if num_iteration is None:
2065
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
2066
        buffer_len = 1 << 20
2067
        tmp_out_len = ctypes.c_int64(0)
wxchan's avatar
wxchan committed
2068
2069
2070
2071
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterDumpModel(
            self.handle,
2072
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
2073
            ctypes.c_int(num_iteration),
2074
            ctypes.c_int64(buffer_len),
wxchan's avatar
wxchan committed
2075
            ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2076
            ptr_string_buffer))
wxchan's avatar
wxchan committed
2077
        actual_len = tmp_out_len.value
2078
        '''if buffer length is not long enough, reallocate a buffer'''
wxchan's avatar
wxchan committed
2079
2080
2081
2082
2083
        if actual_len > buffer_len:
            string_buffer = ctypes.create_string_buffer(actual_len)
            ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
            _safe_call(_LIB.LGBM_BoosterDumpModel(
                self.handle,
2084
                ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
2085
                ctypes.c_int(num_iteration),
2086
                ctypes.c_int64(actual_len),
wxchan's avatar
wxchan committed
2087
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2088
                ptr_string_buffer))
wxchan's avatar
wxchan committed
2089
2090
        return json.loads(string_buffer.value.decode())

2091
2092
    def predict(self, data, num_iteration=None,
                raw_score=False, pred_leaf=False, pred_contrib=False,
2093
                data_has_header=False, is_reshape=True, pred_parameter=None, **kwargs):
2094
        """Make a prediction.
wxchan's avatar
wxchan committed
2095
2096
2097

        Parameters
        ----------
2098
        data : string, numpy array, pandas DataFrame or scipy.sparse
2099
2100
            Data source for prediction.
            If string, it represents the path to txt file.
2101
2102
2103
2104
        num_iteration : int or None, optional (default=None)
            Limit number of iterations in the prediction.
            If None, if the best iteration exists, it is used; otherwise, all iterations are used.
            If <= 0, all iterations are used (no limits).
2105
2106
2107
2108
        raw_score : bool, optional (default=False)
            Whether to predict raw scores.
        pred_leaf : bool, optional (default=False)
            Whether to predict leaf index.
2109
2110
        pred_contrib : bool, optional (default=False)
            Whether to predict feature contributions.
2111
2112
2113
2114
2115
2116
2117

            Note
            ----
            If you want to get more explanation for your model's predictions using SHAP values
            like SHAP interaction values,
            you can install shap package (https://github.com/slundberg/shap).

2118
2119
2120
2121
2122
        data_has_header : bool, optional (default=False)
            Whether the data has header.
            Used only if data is string.
        is_reshape : bool, optional (default=True)
            If True, result is reshaped to [nrow, ncol].
2123
        **kwargs : other parameters for the prediction
wxchan's avatar
wxchan committed
2124
2125
2126

        Returns
        -------
2127
2128
        result : numpy array
            Prediction result.
wxchan's avatar
wxchan committed
2129
        """
Nikita Titov's avatar
Nikita Titov committed
2130
        predictor = self._to_predictor(kwargs)
2131
        if num_iteration is None:
2132
            num_iteration = self.best_iteration
2133
2134
2135
        return predictor.predict(data, num_iteration,
                                 raw_score, pred_leaf, pred_contrib,
                                 data_has_header, is_reshape)
wxchan's avatar
wxchan committed
2136

2137
    def refit(self, data, label, decay_rate=0.9, **kwargs):
Guolin Ke's avatar
Guolin Ke committed
2138
2139
2140
2141
        """Refit the existing Booster by new data.

        Parameters
        ----------
2142
        data : string, numpy array, pandas DataFrame or scipy.sparse
Guolin Ke's avatar
Guolin Ke committed
2143
2144
            Data source for refit.
            If string, it represents the path to txt file.
2145
        label : list, numpy 1-D array or pandas one-column DataFrame/Series
Guolin Ke's avatar
Guolin Ke committed
2146
2147
            Label for refit.
        decay_rate : float, optional (default=0.9)
2148
2149
            Decay rate of refit,
            will use ``leaf_output = decay_rate * old_leaf_output + (1.0 - decay_rate) * new_leaf_output`` to refit trees.
2150
2151
        **kwargs : other parameters for refit
            These parameters will be passed to ``predict`` method.
Guolin Ke's avatar
Guolin Ke committed
2152
2153
2154
2155
2156
2157

        Returns
        -------
        result : Booster
            Refitted Booster.
        """
2158
2159
        predictor = self._to_predictor(kwargs)
        leaf_preds = predictor.predict(data, -1, pred_leaf=True)
2160
        nrow, ncol = leaf_preds.shape
2161
        train_set = Dataset(data, label, silent=True)
Guolin Ke's avatar
Guolin Ke committed
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
        new_booster = Booster(self.params, train_set, silent=True)
        # Copy models
        _safe_call(_LIB.LGBM_BoosterMerge(
            new_booster.handle,
            predictor.handle))
        leaf_preds = leaf_preds.reshape(-1)
        ptr_data, type_ptr_data, _ = c_int_array(leaf_preds)
        _safe_call(_LIB.LGBM_BoosterRefit(
            new_booster.handle,
            ptr_data,
            ctypes.c_int(nrow),
            ctypes.c_int(ncol)))
        return new_booster

2176
    def get_leaf_output(self, tree_id, leaf_id):
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
        """Get the output of a leaf.

        Parameters
        ----------
        tree_id : int
            The index of the tree.
        leaf_id : int
            The index of the leaf in the tree.

        Returns
        -------
        result : float
            The output of the leaf.
        """
2191
2192
2193
2194
2195
2196
2197
2198
        ret = ctypes.c_double(0)
        _safe_call(_LIB.LGBM_BoosterGetLeafValue(
            self.handle,
            ctypes.c_int(tree_id),
            ctypes.c_int(leaf_id),
            ctypes.byref(ret)))
        return ret.value

2199
    def _to_predictor(self, pred_parameter=None):
wxchan's avatar
wxchan committed
2200
        """Convert to predictor"""
2201
        predictor = _InnerPredictor(booster_handle=self.handle, pred_parameter=pred_parameter)
2202
        predictor.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
2203
2204
        return predictor

2205
    def num_feature(self):
2206
2207
2208
2209
2210
2211
2212
        """Get number of features.

        Returns
        -------
        num_feature : int
            The number of features.
        """
2213
2214
2215
2216
2217
2218
        out_num_feature = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetNumFeature(
            self.handle,
            ctypes.byref(out_num_feature)))
        return out_num_feature.value

wxchan's avatar
wxchan committed
2219
    def feature_name(self):
2220
        """Get names of features.
wxchan's avatar
wxchan committed
2221
2222
2223

        Returns
        -------
2224
2225
        result : list
            List with names of features.
wxchan's avatar
wxchan committed
2226
        """
2227
        num_feature = self.num_feature()
2228
        # Get name of features
wxchan's avatar
wxchan committed
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
        tmp_out_len = ctypes.c_int(0)
        string_buffers = [ctypes.create_string_buffer(255) for i in range_(num_feature)]
        ptr_string_buffers = (ctypes.c_char_p * num_feature)(*map(ctypes.addressof, string_buffers))
        _safe_call(_LIB.LGBM_BoosterGetFeatureNames(
            self.handle,
            ctypes.byref(tmp_out_len),
            ptr_string_buffers))
        if num_feature != tmp_out_len.value:
            raise ValueError("Length of feature names doesn't equal with num_feature")
        return [string_buffers[i].value.decode() for i in range_(num_feature)]

2240
    def feature_importance(self, importance_type='split', iteration=-1):
2241
        """Get feature importances.
2242

2243
2244
        Parameters
        ----------
2245
2246
2247
2248
        importance_type : string, optional (default="split")
            How the importance is calculated.
            If "split", result contains numbers of times the feature is used in a model.
            If "gain", result contains total gains of splits which use the feature.
2249

2250
2251
        Returns
        -------
2252
2253
        result : numpy array
            Array with feature importances.
2254
        """
2255
2256
2257
2258
2259
2260
        if importance_type == "split":
            importance_type_int = 0
        elif importance_type == "gain":
            importance_type_int = 1
        else:
            importance_type_int = -1
Nikita Titov's avatar
Nikita Titov committed
2261
        result = np.zeros(self.num_feature(), dtype=np.float64)
2262
2263
2264
2265
2266
2267
2268
2269
2270
        _safe_call(_LIB.LGBM_BoosterFeatureImportance(
            self.handle,
            ctypes.c_int(iteration),
            ctypes.c_int(importance_type_int),
            result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
        if importance_type_int == 0:
            return result.astype(int)
        else:
            return result
2271

wxchan's avatar
wxchan committed
2272
2273
    def __inner_eval(self, data_name, data_idx, feval=None):
        """
Nikita Titov's avatar
Nikita Titov committed
2274
        Evaluate training or validation data
wxchan's avatar
wxchan committed
2275
2276
        """
        if data_idx >= self.__num_dataset:
2277
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
2278
2279
2280
        self.__get_eval_info()
        ret = []
        if self.__num_inner_eval > 0:
2281
            result = np.zeros(self.__num_inner_eval, dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
2282
            tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
2283
2284
            _safe_call(_LIB.LGBM_BoosterGetEval(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
2285
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
2286
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2287
                result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
2288
            if tmp_out_len.value != self.__num_inner_eval:
2289
                raise ValueError("Wrong length of eval results")
wxchan's avatar
wxchan committed
2290
            for i in range_(self.__num_inner_eval):
2291
2292
                ret.append((data_name, self.__name_inner_eval[i],
                            result[i], self.__higher_better_inner_eval[i]))
wxchan's avatar
wxchan committed
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
        if feval is not None:
            if data_idx == 0:
                cur_data = self.train_set
            else:
                cur_data = self.valid_sets[data_idx - 1]
            feval_ret = feval(self.__inner_predict(data_idx), cur_data)
            if isinstance(feval_ret, list):
                for eval_name, val, is_higher_better in feval_ret:
                    ret.append((data_name, eval_name, val, is_higher_better))
            else:
                eval_name, val, is_higher_better = feval_ret
                ret.append((data_name, eval_name, val, is_higher_better))
        return ret

    def __inner_predict(self, data_idx):
        """
        Predict for training and validation dataset
        """
        if data_idx >= self.__num_dataset:
2312
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
2313
2314
2315
2316
2317
        if self.__inner_predict_buffer[data_idx] is None:
            if data_idx == 0:
                n_preds = self.train_set.num_data() * self.__num_class
            else:
                n_preds = self.valid_sets[data_idx - 1].num_data() * self.__num_class
2318
            self.__inner_predict_buffer[data_idx] = np.zeros(n_preds, dtype=np.float64)
2319
        # avoid to predict many time in one iteration
wxchan's avatar
wxchan committed
2320
2321
        if not self.__is_predicted_cur_iter[data_idx]:
            tmp_out_len = ctypes.c_int64(0)
Guolin Ke's avatar
Guolin Ke committed
2322
            data_ptr = self.__inner_predict_buffer[data_idx].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
wxchan's avatar
wxchan committed
2323
2324
            _safe_call(_LIB.LGBM_BoosterGetPredict(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
2325
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
2326
2327
2328
                ctypes.byref(tmp_out_len),
                data_ptr))
            if tmp_out_len.value != len(self.__inner_predict_buffer[data_idx]):
2329
                raise ValueError("Wrong length of predict results for data %d" % (data_idx))
wxchan's avatar
wxchan committed
2330
2331
2332
2333
2334
2335
2336
2337
2338
            self.__is_predicted_cur_iter[data_idx] = True
        return self.__inner_predict_buffer[data_idx]

    def __get_eval_info(self):
        """
        Get inner evaluation count and names
        """
        if self.__need_reload_eval_info:
            self.__need_reload_eval_info = False
Guolin Ke's avatar
Guolin Ke committed
2339
            out_num_eval = ctypes.c_int(0)
2340
            # Get num of inner evals
wxchan's avatar
wxchan committed
2341
2342
2343
2344
2345
            _safe_call(_LIB.LGBM_BoosterGetEvalCounts(
                self.handle,
                ctypes.byref(out_num_eval)))
            self.__num_inner_eval = out_num_eval.value
            if self.__num_inner_eval > 0:
2346
                # Get name of evals
Guolin Ke's avatar
Guolin Ke committed
2347
                tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
2348
                string_buffers = [ctypes.create_string_buffer(255) for i in range_(self.__num_inner_eval)]
wxchan's avatar
wxchan committed
2349
                ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
wxchan's avatar
wxchan committed
2350
2351
2352
2353
2354
                _safe_call(_LIB.LGBM_BoosterGetEvalNames(
                    self.handle,
                    ctypes.byref(tmp_out_len),
                    ptr_string_buffers))
                if self.__num_inner_eval != tmp_out_len.value:
2355
                    raise ValueError("Length of eval names doesn't equal with num_evals")
2356
                self.__name_inner_eval = \
wxchan's avatar
wxchan committed
2357
                    [string_buffers[i].value.decode() for i in range_(self.__num_inner_eval)]
2358
                self.__higher_better_inner_eval = \
2359
                    [name.startswith(('auc', 'ndcg@', 'map@')) for name in self.__name_inner_eval]
2360

wxchan's avatar
wxchan committed
2361
    def attr(self, key):
2362
        """Get attribute string from the Booster.
wxchan's avatar
wxchan committed
2363
2364
2365

        Parameters
        ----------
2366
2367
        key : string
            The name of the attribute.
wxchan's avatar
wxchan committed
2368
2369
2370

        Returns
        -------
2371
2372
        value : string or None
            The attribute value.
Nikita Titov's avatar
Nikita Titov committed
2373
            Returns None if attribute does not exist.
wxchan's avatar
wxchan committed
2374
        """
2375
        return self.__attr.get(key, None)
wxchan's avatar
wxchan committed
2376
2377

    def set_attr(self, **kwargs):
2378
        """Set the attribute of the Booster.
wxchan's avatar
wxchan committed
2379
2380
2381
2382

        Parameters
        ----------
        **kwargs
2383
2384
            The attributes to set.
            Setting a value to None deletes an attribute.
Nikita Titov's avatar
Nikita Titov committed
2385
2386
2387
2388
2389

        Returns
        -------
        self : Booster
            Booster with set attribute.
wxchan's avatar
wxchan committed
2390
2391
2392
        """
        for key, value in kwargs.items():
            if value is not None:
wxchan's avatar
wxchan committed
2393
                if not isinstance(value, string_type):
Nikita Titov's avatar
Nikita Titov committed
2394
                    raise ValueError("Only string values are accepted")
wxchan's avatar
wxchan committed
2395
2396
2397
                self.__attr[key] = value
            else:
                self.__attr.pop(key, None)
Nikita Titov's avatar
Nikita Titov committed
2398
        return self