basic.py 75.7 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
from .compat import (DataFrame, Series, integer_types, json,
                     json_default_with_numpy, numeric_types, range_,
18
                     string_type, LGBMDeprecationWarning)
wxchan's avatar
wxchan committed
19
20
from .libpath import find_lib_path

wxchan's avatar
wxchan committed
21

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

wxchan's avatar
wxchan committed
31

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

wxchan's avatar
wxchan committed
34

wxchan's avatar
wxchan committed
35
36
37
38
class LightGBMError(Exception):
    """Error throwed by LightGBM"""
    pass

wxchan's avatar
wxchan committed
39

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

wxchan's avatar
wxchan committed
50

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

wxchan's avatar
wxchan committed
61

wxchan's avatar
wxchan committed
62
def is_numpy_1d_array(data):
Guolin Ke's avatar
Guolin Ke committed
63
    """Check is 1d numpy array"""
64
    return isinstance(data, np.ndarray) and len(data.shape) == 1
wxchan's avatar
wxchan committed
65

wxchan's avatar
wxchan committed
66

wxchan's avatar
wxchan committed
67
def is_1d_list(data):
Guolin Ke's avatar
Guolin Ke committed
68
    """Check is 1d list"""
69
    return isinstance(data, list) and \
70
        (not data or is_numeric(data[0]))
wxchan's avatar
wxchan committed
71

wxchan's avatar
wxchan committed
72

73
def list_to_1d_numpy(data, dtype=np.float32, name='list'):
Guolin Ke's avatar
Guolin Ke committed
74
    """convert to 1d numpy array"""
wxchan's avatar
wxchan committed
75
76
77
78
79
80
81
    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)
82
83
    elif isinstance(data, Series):
        return data.values.astype(dtype)
wxchan's avatar
wxchan committed
84
    else:
85
        raise TypeError("Wrong type({}) for {}, should be list or numpy array".format(type(data).__name__, name))
wxchan's avatar
wxchan committed
86

wxchan's avatar
wxchan committed
87

wxchan's avatar
wxchan committed
88
89
90
91
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)):
92
        return np.fromiter(cptr, dtype=np.float32, count=length)
wxchan's avatar
wxchan committed
93
    else:
94
        raise RuntimeError('Expected float pointer')
wxchan's avatar
wxchan committed
95

Guolin Ke's avatar
Guolin Ke committed
96

Guolin Ke's avatar
Guolin Ke committed
97
98
99
100
101
102
103
104
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
105

wxchan's avatar
wxchan committed
106
107
108
109
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)):
110
        return np.fromiter(cptr, dtype=np.int32, count=length)
wxchan's avatar
wxchan committed
111
    else:
112
        raise RuntimeError('Expected int pointer')
wxchan's avatar
wxchan committed
113

wxchan's avatar
wxchan committed
114

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

wxchan's avatar
wxchan committed
119

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

wxchan's avatar
wxchan committed
124

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

wxchan's avatar
wxchan committed
139

140
class _temp_file(object):
141
142
143
144
    def __enter__(self):
        with NamedTemporaryFile(prefix="lightgbm_tmp_", delete=True) as f:
            self.name = f.name
        return self
wxchan's avatar
wxchan committed
145

146
147
148
    def __exit__(self, exc_type, exc_val, exc_tb):
        if os.path.isfile(self.name):
            os.remove(self.name)
wxchan's avatar
wxchan committed
149

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

155
156
    def writelines(self, lines):
        with open(self.name, "w+") as f:
157
            f.writelines(lines)
158

wxchan's avatar
wxchan committed
159

wxchan's avatar
wxchan committed
160
161
162
163
164
"""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
165

wxchan's avatar
wxchan committed
166
167
168
"""Matric is row major in python"""
C_API_IS_ROW_MAJOR = 1

Guolin Ke's avatar
Guolin Ke committed
169
"""marco definition of prediction type in c_api of LightGBM"""
wxchan's avatar
wxchan committed
170
171
172
C_API_PREDICT_NORMAL = 0
C_API_PREDICT_RAW_SCORE = 1
C_API_PREDICT_LEAF_INDEX = 2
173
C_API_PREDICT_CONTRIB = 3
wxchan's avatar
wxchan committed
174

Guolin Ke's avatar
Guolin Ke committed
175
"""data type of data field"""
wxchan's avatar
wxchan committed
176
177
FIELD_TYPE_MAPPER = {"label": C_API_DTYPE_FLOAT32,
                     "weight": C_API_DTYPE_FLOAT32,
Guolin Ke's avatar
Guolin Ke committed
178
                     "init_score": C_API_DTYPE_FLOAT64,
wxchan's avatar
wxchan committed
179
180
                     "group": C_API_DTYPE_INT32}

wxchan's avatar
wxchan committed
181

wxchan's avatar
wxchan committed
182
def c_float_array(data):
Guolin Ke's avatar
Guolin Ke committed
183
    """get pointer of float numpy array / list"""
wxchan's avatar
wxchan committed
184
185
186
187
188
189
190
191
192
193
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
        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:
194
            raise TypeError("Expected np.float32 or np.float64, met type({})"
wxchan's avatar
wxchan committed
195
196
                            .format(data.dtype))
    else:
197
        raise TypeError("Unknown type({})".format(type(data).__name__))
wxchan's avatar
wxchan committed
198
199
    return (ptr_data, type_data)

wxchan's avatar
wxchan committed
200

wxchan's avatar
wxchan committed
201
def c_int_array(data):
Guolin Ke's avatar
Guolin Ke committed
202
    """get pointer of int numpy array / list"""
wxchan's avatar
wxchan committed
203
204
205
206
207
208
209
210
211
212
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
        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:
213
            raise TypeError("Expected np.int32 or np.int64, met type({})"
wxchan's avatar
wxchan committed
214
215
                            .format(data.dtype))
    else:
216
        raise TypeError("Unknown type({})".format(type(data).__name__))
wxchan's avatar
wxchan committed
217
218
    return (ptr_data, type_data)

wxchan's avatar
wxchan committed
219

220
221
222
223
224
225
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'}


226
def _data_from_pandas(data, feature_name, categorical_feature, pandas_categorical):
227
    if isinstance(data, DataFrame):
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
        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.')
            for col, category in zip(cat_cols, pandas_categorical):
                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
            data[cat_cols] = data[cat_cols].apply(lambda x: x.cat.codes)
        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]

            msg = """DataFrame.dtypes for data must be int, float or bool. Did not expect the data types in fields """
            raise ValueError(msg + ', '.join(bad_fields))
258
        data = data.values.astype('float')
259
260
261
262
263
264
    else:
        if feature_name == 'auto':
            feature_name = None
        if categorical_feature == 'auto':
            categorical_feature = None
    return data, feature_name, categorical_feature, pandas_categorical
265
266
267
268
269
270
271
272
273
274
275
276
277


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')
        label = label.values.astype('float')
    return label


278
279
def _save_pandas_categorical(file_name, pandas_categorical):
    with open(file_name, 'a') as f:
Guolin Ke's avatar
Guolin Ke committed
280
        f.write('\npandas_categorical:' + json.dumps(pandas_categorical, default=json_default_with_numpy) + '\n')
281
282
283
284


def _load_pandas_categorical(file_name):
    with open(file_name, 'r') as f:
Guolin Ke's avatar
Guolin Ke committed
285
286
287
288
        lines = f.readlines()
        last_line = lines[-1]
        if last_line.strip() == "":
            last_line = lines[-2]
289
290
291
292
293
        if last_line.startswith('pandas_categorical:'):
            return json.loads(last_line[len('pandas_categorical:'):])
    return None


Guolin Ke's avatar
Guolin Ke committed
294
295
class _InnerPredictor(object):
    """
296
297
    A _InnerPredictor of LightGBM.
    Only used for prediction, usually used for continued-train
Guolin Ke's avatar
Guolin Ke committed
298
    Note: Can convert from Booster, but cannot convert to Booster
wxchan's avatar
wxchan committed
299
    """
300
    def __init__(self, model_file=None, booster_handle=None, pred_parameter=None):
Guolin Ke's avatar
Guolin Ke committed
301
        """Initialize the _InnerPredictor. Not expose to user
wxchan's avatar
wxchan committed
302
303
304
305
306

        Parameters
        ----------
        model_file : string
            Path to the model file.
Guolin Ke's avatar
Guolin Ke committed
307
308
        booster_handle : Handle of Booster
            use handle to init
309
310
        pred_parameter: dict
            Other parameters for the prediciton
wxchan's avatar
wxchan committed
311
312
313
314
315
        """
        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
316
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
317
318
319
320
            _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
321
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
322
323
324
325
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
326
            self.num_total_iteration = out_num_iterations.value
327
            self.pandas_categorical = _load_pandas_categorical(model_file)
wxchan's avatar
wxchan committed
328
        elif booster_handle is not None:
Guolin Ke's avatar
Guolin Ke committed
329
            self.__is_manage_handle = False
wxchan's avatar
wxchan committed
330
            self.handle = booster_handle
Guolin Ke's avatar
Guolin Ke committed
331
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
332
333
334
335
            _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
336
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
337
338
339
            _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
                self.handle,
                ctypes.byref(out_num_iterations)))
340
            self.num_total_iteration = out_num_iterations.value
341
            self.pandas_categorical = None
wxchan's avatar
wxchan committed
342
        else:
Guolin Ke's avatar
Guolin Ke committed
343
            raise TypeError('Need Model file or Booster handle to create a predictor')
wxchan's avatar
wxchan committed
344

345
346
        pred_parameter = {} if pred_parameter is None else pred_parameter
        self.pred_parameter = param_dict_to_str(pred_parameter)
cbecker's avatar
cbecker committed
347

wxchan's avatar
wxchan committed
348
349
350
351
    def __del__(self):
        if self.__is_manage_handle:
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))

352
353
354
355
356
    def __getstate__(self):
        this = self.__dict__.copy()
        this.pop('handle', None)
        return this

wxchan's avatar
wxchan committed
357
    def predict(self, data, num_iteration=-1,
358
                raw_score=False, pred_leaf=False, pred_contrib=False, data_has_header=False,
wxchan's avatar
wxchan committed
359
360
361
362
363
364
365
366
                is_reshape=True):
        """
        Predict logic

        Parameters
        ----------
        data : string/numpy array/scipy.sparse
            Data source for prediction
367
            When data type is string, it represents the path of txt file
wxchan's avatar
wxchan committed
368
        num_iteration : int
369
            Used iteration for prediction
wxchan's avatar
wxchan committed
370
371
372
373
        raw_score : bool
            True for predict raw score
        pred_leaf : bool
            True for predict leaf index
374
375
        pred_contrib : bool
            True for predict feature contributions
wxchan's avatar
wxchan committed
376
        data_has_header : bool
Guolin Ke's avatar
Guolin Ke committed
377
            Used for txt data, True if txt data has header
wxchan's avatar
wxchan committed
378
        is_reshape : bool
379
            Reshape to (nrow, ncol) if true
wxchan's avatar
wxchan committed
380
381
382
383
384

        Returns
        -------
        Prediction result
        """
wxchan's avatar
wxchan committed
385
        if isinstance(data, Dataset):
386
            raise TypeError("Cannot use Dataset instance for prediction, please use raw data instead")
387
        data = _data_from_pandas(data, None, None, self.pandas_categorical)[0]
wxchan's avatar
wxchan committed
388
389
390
391
392
        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
393
394
        if pred_contrib:
            predict_type = C_API_PREDICT_CONTRIB
wxchan's avatar
wxchan committed
395
        int_data_has_header = 1 if data_has_header else 0
396
397
        if num_iteration > self.num_total_iteration:
            num_iteration = self.num_total_iteration
cbecker's avatar
cbecker committed
398

wxchan's avatar
wxchan committed
399
        if isinstance(data, string_type):
400
            with _temp_file() as f:
wxchan's avatar
wxchan committed
401
402
403
                _safe_call(_LIB.LGBM_BoosterPredictForFile(
                    self.handle,
                    c_str(data),
Guolin Ke's avatar
Guolin Ke committed
404
405
406
                    ctypes.c_int(int_data_has_header),
                    ctypes.c_int(predict_type),
                    ctypes.c_int(num_iteration),
407
                    c_str(self.pred_parameter),
wxchan's avatar
wxchan committed
408
409
                    c_str(f.name)))
                lines = f.readlines()
410
411
                nrow = len(lines)
                preds = [float(token) for line in lines for token in line.split('\t')]
Guolin Ke's avatar
Guolin Ke committed
412
                preds = np.array(preds, dtype=np.float64, copy=False)
wxchan's avatar
wxchan committed
413
414
415
        elif isinstance(data, scipy.sparse.csr_matrix):
            preds, nrow = self.__pred_for_csr(data, num_iteration,
                                              predict_type)
Guolin Ke's avatar
Guolin Ke committed
416
417
418
        elif isinstance(data, scipy.sparse.csc_matrix):
            preds, nrow = self.__pred_for_csc(data, num_iteration,
                                              predict_type)
wxchan's avatar
wxchan committed
419
420
421
        elif isinstance(data, np.ndarray):
            preds, nrow = self.__pred_for_np2d(data, num_iteration,
                                               predict_type)
422
423
424
425
426
427
428
        elif isinstance(data, list):
            try:
                data = np.array(data)
            except:
                raise ValueError('Cannot convert data list to numpy array.')
            preds, nrow = self.__pred_for_np2d(data, num_iteration,
                                               predict_type)
wxchan's avatar
wxchan committed
429
430
        else:
            try:
431
                warnings.warn('Converting data to scipy sparse matrix.')
wxchan's avatar
wxchan committed
432
433
                csr = scipy.sparse.csr_matrix(data)
            except:
434
                raise TypeError('Cannot predict data for type {}'.format(type(data).__name__))
435
436
            preds, nrow = self.__pred_for_csr(csr, num_iteration,
                                              predict_type)
wxchan's avatar
wxchan committed
437
438
        if pred_leaf:
            preds = preds.astype(np.int32)
439
        if is_reshape and preds.size != nrow:
wxchan's avatar
wxchan committed
440
            if preds.size % nrow == 0:
441
                preds = preds.reshape(nrow, -1)
wxchan's avatar
wxchan committed
442
            else:
443
                raise ValueError('Length of predict result (%d) cannot be divide nrow (%d)'
wxchan's avatar
wxchan committed
444
445
446
447
                                 % (preds.size, nrow))
        return preds

    def __get_num_preds(self, num_iteration, nrow, predict_type):
Guolin Ke's avatar
Guolin Ke committed
448
449
450
        """
        Get size of prediction result
        """
Guolin Ke's avatar
Guolin Ke committed
451
452
453
        n_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterCalcNumPredict(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
454
455
456
            ctypes.c_int(nrow),
            ctypes.c_int(predict_type),
            ctypes.c_int(num_iteration),
Guolin Ke's avatar
Guolin Ke committed
457
458
            ctypes.byref(n_preds)))
        return n_preds.value
wxchan's avatar
wxchan committed
459
460
461
462
463
464

    def __pred_for_np2d(self, mat, num_iteration, predict_type):
        """
        Predict for a 2-D numpy matrix.
        """
        if len(mat.shape) != 2:
465
            raise ValueError('Input numpy.ndarray or list must be 2 dimensional')
wxchan's avatar
wxchan committed
466
467
468
469
470
471
472
473
474

        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)
Guolin Ke's avatar
Guolin Ke committed
475
        preds = np.zeros(n_preds, dtype=np.float64)
wxchan's avatar
wxchan committed
476
477
478
479
        out_num_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterPredictForMat(
            self.handle,
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
480
481
482
483
484
485
            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),
486
            c_str(self.pred_parameter),
wxchan's avatar
wxchan committed
487
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
488
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
489
        if n_preds != out_num_preds.value:
490
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
491
492
493
494
495
496
497
498
        return preds, mat.shape[0]

    def __pred_for_csr(self, csr, num_iteration, predict_type):
        """
        Predict for a csr data
        """
        nrow = len(csr.indptr) - 1
        n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
Guolin Ke's avatar
Guolin Ke committed
499
        preds = np.zeros(n_preds, dtype=np.float64)
wxchan's avatar
wxchan committed
500
501
502
503
504
505
506
507
        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,
Guolin Ke's avatar
Guolin Ke committed
508
            ctypes.c_int32(type_ptr_indptr),
wxchan's avatar
wxchan committed
509
510
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
511
512
513
514
515
516
            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),
517
            c_str(self.pred_parameter),
wxchan's avatar
wxchan committed
518
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
519
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
Guolin Ke's avatar
Guolin Ke committed
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
        if n_preds != out_num_preds.value:
            raise ValueError("Wrong length for predict results")
        return preds, nrow

    def __pred_for_csc(self, csc, num_iteration, predict_type):
        """
        Predict for a csc data
        """
        nrow = csc.shape[0]
        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)

        ptr_indptr, type_ptr_indptr = c_int_array(csc.indptr)
        ptr_data, type_ptr_data = c_float_array(csc.data)

        _safe_call(_LIB.LGBM_BoosterPredictForCSC(
            self.handle,
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
539
            ctypes.c_int32(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
540
541
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
542
543
544
545
546
547
            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),
548
            c_str(self.pred_parameter),
Guolin Ke's avatar
Guolin Ke committed
549
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
550
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
551
        if n_preds != out_num_preds.value:
552
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
553
554
        return preds, nrow

wxchan's avatar
wxchan committed
555

wxchan's avatar
wxchan committed
556
557
class Dataset(object):
    """Dataset in LightGBM."""
558
559
    def __init__(self, data, label=None, max_bin=None, reference=None,
                 weight=None, group=None, init_score=None, silent=False,
560
                 feature_name='auto', categorical_feature='auto', params=None,
wxchan's avatar
wxchan committed
561
                 free_raw_data=True):
562
563
        """Constract Dataset.

wxchan's avatar
wxchan committed
564
565
        Parameters
        ----------
566
        data : string, numpy array or scipy.sparse
wxchan's avatar
wxchan committed
567
            Data source of Dataset.
568
            If string, it represents the path to txt file.
569
        label : list, numpy 1-D array or None, optional (default=None)
570
            Label of the data.
571
        max_bin : int or None, optional (default=None)
572
            Max number of discrete bins for features.
573
            If None, default value from parameters of CLI-version will be used.
574
575
576
        reference : Dataset or None, optional (default=None)
            If this is Dataset for validation, training data should be used as reference.
        weight : list, numpy 1-D array or None, optional (default=None)
wxchan's avatar
wxchan committed
577
            Weight for each instance.
578
579
        group : list, numpy 1-D array or None, optional (default=None)
            Group/query size for Dataset.
580
581
        init_score : list, numpy 1-D array or None, optional (default=None)
            Init score for Dataset.
582
583
584
585
586
587
588
589
590
591
592
593
594
595
        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.
        params: dict or None, optional (default=None)
            Other parameters.
        free_raw_data: bool, optional (default=True)
            If True, raw data is freed after constructing inner Dataset.
wxchan's avatar
wxchan committed
596
        """
wxchan's avatar
wxchan committed
597
598
599
600
601
602
603
        self.handle = None
        self.data = data
        self.label = label
        self.max_bin = max_bin
        self.reference = reference
        self.weight = weight
        self.group = group
604
        self.init_score = init_score
wxchan's avatar
wxchan committed
605
606
        self.silent = silent
        self.feature_name = feature_name
607
        self.categorical_feature = categorical_feature
608
        self.params = copy.deepcopy(params)
wxchan's avatar
wxchan committed
609
610
611
        self.free_raw_data = free_raw_data
        self.used_indices = None
        self._predictor = None
612
        self.pandas_categorical = None
613
        self.params_back_up = None
wxchan's avatar
wxchan committed
614
615

    def __del__(self):
616
617
618
        self._free_handle()

    def _free_handle(self):
619
        if self.handle is not None:
620
            _safe_call(_LIB.LGBM_DatasetFree(self.handle))
621
            self.handle = None
wxchan's avatar
wxchan committed
622

623
624
    def _lazy_init(self, data, label=None, max_bin=None, reference=None,
                   weight=None, group=None, init_score=None, predictor=None,
wxchan's avatar
wxchan committed
625
                   silent=False, feature_name='auto',
626
                   categorical_feature='auto', params=None):
wxchan's avatar
wxchan committed
627
628
629
        if data is None:
            self.handle = None
            return
Guolin Ke's avatar
Guolin Ke committed
630
631
632
        if reference is not None:
            self.pandas_categorical = reference.pandas_categorical
            categorical_feature = reference.categorical_feature
633
        data, feature_name, categorical_feature, self.pandas_categorical = _data_from_pandas(data, feature_name, categorical_feature, self.pandas_categorical)
wxchan's avatar
wxchan committed
634
635
        label = _label_from_pandas(label)
        self.data_has_header = False
636
        # process for args
wxchan's avatar
wxchan committed
637
        params = {} if params is None else params
638
639
640
641
642
        args_names = getattr(self.__class__, '_lazy_init').__code__.co_varnames[:getattr(self.__class__, '_lazy_init').__code__.co_argcount]
        for key, _ in params.items():
            if key in args_names:
                warnings.warn('{0} keyword has been found in `params` and will be ignored. '
                              'Please use {0} argument of the Dataset constructor to pass this parameter.'.format(key))
wxchan's avatar
wxchan committed
643
644
        self.max_bin = max_bin
        self.predictor = predictor
645
646
        if self.max_bin is not None:
            params["max_bin"] = self.max_bin
647
648
            warnings.warn('The `max_bin` parameter is deprecated and will be removed in 2.0.12 version. '
                          'Please use `params` to pass this parameter.', LGBMDeprecationWarning)
649
650
        if "verbosity" in params:
            params.setdefault("verbose", params.pop("verbosity"))
wxchan's avatar
wxchan committed
651
652
653
654
        if silent:
            params["verbose"] = 0
        elif "verbose" not in params:
            params["verbose"] = 1
655
        # get categorical features
656
657
658
659
660
661
662
663
664
665
666
667
668
        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))
669
            if categorical_indices:
670
671
672
673
                if "categorical_feature" in params or "categorical_column" in params:
                    warnings.warn('categorical_feature in param dict is overrided.')
                    params.pop("categorical_feature", None)
                    params.pop("categorical_column", None)
674
                params['categorical_column'] = sorted(categorical_indices)
675

wxchan's avatar
wxchan committed
676
        params_str = param_dict_to_str(params)
677
        # process for reference dataset
wxchan's avatar
wxchan committed
678
        ref_dataset = None
wxchan's avatar
wxchan committed
679
        if isinstance(reference, Dataset):
680
            ref_dataset = reference.construct().handle
wxchan's avatar
wxchan committed
681
682
        elif reference is not None:
            raise TypeError('Reference dataset should be None or dataset instance')
683
        # start construct data
wxchan's avatar
wxchan committed
684
        if isinstance(data, string_type):
685
            # check data has header or not
Guolin Ke's avatar
Guolin Ke committed
686
            if str(params.get("has_header", "")).lower() == "true" \
wxchan's avatar
wxchan committed
687
                    or str(params.get("header", "")).lower() == "true":
688
                self.data_has_header = True
wxchan's avatar
wxchan committed
689
690
691
692
693
694
695
696
            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
697
698
        elif isinstance(data, scipy.sparse.csc_matrix):
            self.__init_from_csc(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
699
700
701
702
703
704
705
        elif isinstance(data, np.ndarray):
            self.__init_from_np2d(data, params_str, ref_dataset)
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                self.__init_from_csr(csr, params_str, ref_dataset)
            except:
wxchan's avatar
wxchan committed
706
                raise TypeError('Cannot initialize Dataset from {}'.format(type(data).__name__))
wxchan's avatar
wxchan committed
707
708
709
        if label is not None:
            self.set_label(label)
        if self.get_label() is None:
710
            raise ValueError("Label should not be None")
wxchan's avatar
wxchan committed
711
712
713
714
715
        if weight is not None:
            self.set_weight(weight)
        if group is not None:
            self.set_group(group)
        # load init score
716
717
718
719
720
        if init_score is not None:
            self.set_init_score(init_score)
            if self.predictor is not None:
                warnings.warn("The prediction of init_model will be overrided by init_score.")
        elif isinstance(self.predictor, _InnerPredictor):
wxchan's avatar
wxchan committed
721
722
723
724
725
726
            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
727
                new_init_score = np.zeros(init_score.size, dtype=np.float32)
wxchan's avatar
wxchan committed
728
                num_data = self.num_data()
wxchan's avatar
wxchan committed
729
730
                for i in range_(num_data):
                    for j in range_(self.predictor.num_class):
wxchan's avatar
wxchan committed
731
732
733
                        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
734
735
        elif self.predictor is not None:
            raise TypeError('wrong predictor type {}'.format(type(self.predictor).__name__))
Guolin Ke's avatar
Guolin Ke committed
736
737
        # set feature names
        self.set_feature_name(feature_name)
wxchan's avatar
wxchan committed
738
739
740
741
742
743
744
745
746
747
748
749

    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:
750
            # change non-float data to float data, need to copy
wxchan's avatar
wxchan committed
751
752
753
754
755
            data = np.array(mat.reshape(mat.size), dtype=np.float32)

        ptr_data, type_ptr_data = c_float_array(data)
        _safe_call(_LIB.LGBM_DatasetCreateFromMat(
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
756
757
758
759
            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
760
761
762
763
764
765
766
767
768
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

    def __init_from_csr(self, csr, params_str, ref_dataset):
        """
        Initialize data from a CSR matrix.
        """
        if len(csr.indices) != len(csr.data):
769
            raise ValueError('Length mismatch: {} vs {}'.format(len(csr.indices), len(csr.data)))
wxchan's avatar
wxchan committed
770
771
772
773
774
775
776
        self.handle = ctypes.c_void_p()

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

        _safe_call(_LIB.LGBM_DatasetCreateFromCSR(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
777
            ctypes.c_int(type_ptr_indptr),
wxchan's avatar
wxchan committed
778
779
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
780
781
782
783
            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
784
785
786
787
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

Guolin Ke's avatar
Guolin Ke committed
788
789
790
791
792
793
794
795
796
797
798
799
800
    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()

        ptr_indptr, type_ptr_indptr = c_int_array(csc.indptr)
        ptr_data, type_ptr_data = c_float_array(csc.data)

        _safe_call(_LIB.LGBM_DatasetCreateFromCSC(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
801
            ctypes.c_int(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
802
803
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
804
805
806
807
            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
808
809
810
811
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

wxchan's avatar
wxchan committed
812
    def construct(self):
813
814
815
816
817
818
819
        """Lazy init.

        Returns
        -------
        self : Dataset
            Returns self.
        """
820
        if self.handle is None:
wxchan's avatar
wxchan committed
821
822
            if self.reference is not None:
                if self.used_indices is None:
823
                    # create valid
wxchan's avatar
wxchan committed
824
                    self._lazy_init(self.data, label=self.label, max_bin=self.max_bin, reference=self.reference,
825
                                    weight=self.weight, group=self.group, init_score=self.init_score, predictor=self._predictor,
826
                                    silent=self.silent, feature_name=self.feature_name, params=self.params)
wxchan's avatar
wxchan committed
827
                else:
828
                    # construct subset
wxchan's avatar
wxchan committed
829
                    used_indices = list_to_1d_numpy(self.used_indices, np.int32, name='used_indices')
830
                    self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
831
832
                    params_str = param_dict_to_str(self.params)
                    _safe_call(_LIB.LGBM_DatasetGetSubset(
833
                        self.reference.construct().handle,
wxchan's avatar
wxchan committed
834
                        used_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
835
                        ctypes.c_int(used_indices.shape[0]),
wxchan's avatar
wxchan committed
836
837
838
839
840
                        c_str(params_str),
                        ctypes.byref(self.handle)))
                    if self.get_label() is None:
                        raise ValueError("Label should not be None.")
            else:
841
                # create train
wxchan's avatar
wxchan committed
842
                self._lazy_init(self.data, label=self.label, max_bin=self.max_bin,
843
844
                                weight=self.weight, group=self.group, init_score=self.init_score,
                                predictor=self._predictor, silent=self.silent, feature_name=self.feature_name,
845
                                categorical_feature=self.categorical_feature, params=self.params)
wxchan's avatar
wxchan committed
846
847
848
            if self.free_raw_data:
                self.data = None
        return self
wxchan's avatar
wxchan committed
849

wxchan's avatar
wxchan committed
850
    def create_valid(self, data, label=None, weight=None, group=None,
851
                     init_score=None, silent=False, params=None):
852
        """Create validation data align with current Dataset.
wxchan's avatar
wxchan committed
853
854
855

        Parameters
        ----------
856
        data : string, numpy array or scipy.sparse
wxchan's avatar
wxchan committed
857
            Data source of Dataset.
858
859
            If string, it represents the path to txt file.
        label : list or numpy 1-D array, optional (default=None)
wxchan's avatar
wxchan committed
860
            Label of the training data.
861
        weight : list, numpy 1-D array or None, optional (default=None)
wxchan's avatar
wxchan committed
862
            Weight for each instance.
863
864
        group : list, numpy 1-D array or None, optional (default=None)
            Group/query size for Dataset.
865
866
        init_score : list, numpy 1-D array or None, optional (default=None)
            Init score for Dataset.
867
868
869
870
871
872
873
874
875
        silent : bool, optional (default=False)
            Whether to print messages during construction.
        params: dict or None, optional (default=None)
            Other parameters.

        Returns
        -------
        self : Dataset
            Returns self.
wxchan's avatar
wxchan committed
876
877
        """
        ret = Dataset(data, label=label, max_bin=self.max_bin, reference=self,
878
879
                      weight=weight, group=group, init_score=init_score,
                      silent=silent, params=params, free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
880
        ret._predictor = self._predictor
881
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
882
        return ret
wxchan's avatar
wxchan committed
883

wxchan's avatar
wxchan committed
884
    def subset(self, used_indices, params=None):
885
        """Get subset of current Dataset.
wxchan's avatar
wxchan committed
886
887
888
889

        Parameters
        ----------
        used_indices : list of int
890
891
892
893
894
895
896
897
            Indices used to create the subset.
        params: dict or None, optional (default=None)
            Other parameters.

        Returns
        -------
        subset : Dataset
            Subset of the current Dataset.
wxchan's avatar
wxchan committed
898
        """
wxchan's avatar
wxchan committed
899
900
        if params is None:
            params = self.params
wxchan's avatar
wxchan committed
901
        ret = Dataset(None, reference=self, feature_name=self.feature_name,
902
                      categorical_feature=self.categorical_feature, params=params)
wxchan's avatar
wxchan committed
903
        ret._predictor = self._predictor
904
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
905
906
907
908
        ret.used_indices = used_indices
        return ret

    def save_binary(self, filename):
909
        """Save Dataset to binary file.
wxchan's avatar
wxchan committed
910
911
912
913
914
915
916
917
918
919
920
921
922

        Parameters
        ----------
        filename : string
            Name of the output file.
        """
        _safe_call(_LIB.LGBM_DatasetSaveBinary(
            self.construct().handle,
            c_str(filename)))

    def _update_params(self, params):
        if not self.params:
            self.params = params
wxchan's avatar
wxchan committed
923
        else:
924
            self.params_back_up = copy.deepcopy(self.params)
wxchan's avatar
wxchan committed
925
            self.params.update(params)
wxchan's avatar
wxchan committed
926

927
928
929
930
    def _reverse_update_params(self):
        self.params = copy.deepcopy(self.params_back_up)
        self.params_back_up = None

wxchan's avatar
wxchan committed
931
    def set_field(self, field_name, data):
wxchan's avatar
wxchan committed
932
        """Set property into the Dataset.
wxchan's avatar
wxchan committed
933
934
935

        Parameters
        ----------
936
937
938
939
        field_name: string
            The field name of the information.
        data: list, numpy array or None
            The array of data to be set.
wxchan's avatar
wxchan committed
940
        """
941
942
        if self.handle is None:
            raise Exception("Cannot set %s before construct dataset" % field_name)
wxchan's avatar
wxchan committed
943
        if data is None:
944
            # set to None
wxchan's avatar
wxchan committed
945
946
947
948
            _safe_call(_LIB.LGBM_DatasetSetField(
                self.handle,
                c_str(field_name),
                None,
Guolin Ke's avatar
Guolin Ke committed
949
950
                ctypes.c_int(0),
                ctypes.c_int(FIELD_TYPE_MAPPER[field_name])))
wxchan's avatar
wxchan committed
951
            return
Guolin Ke's avatar
Guolin Ke committed
952
953
954
955
956
        dtype = np.float32
        if field_name == 'group':
            dtype = np.int32
        elif field_name == 'init_score':
            dtype = np.float64
957
        data = list_to_1d_numpy(data, dtype, name=field_name)
wxchan's avatar
wxchan committed
958
959
960
        if data.dtype == np.float32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
            type_data = C_API_DTYPE_FLOAT32
Guolin Ke's avatar
Guolin Ke committed
961
962
963
        elif data.dtype == np.float64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
            type_data = C_API_DTYPE_FLOAT64
wxchan's avatar
wxchan committed
964
965
966
967
        elif data.dtype == np.int32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
            type_data = C_API_DTYPE_INT32
        else:
Guolin Ke's avatar
Guolin Ke committed
968
            raise TypeError("Excepted np.float32/64 or np.int32, meet type({})".format(data.dtype))
wxchan's avatar
wxchan committed
969
        if type_data != FIELD_TYPE_MAPPER[field_name]:
970
            raise TypeError("Input type error for set_field")
wxchan's avatar
wxchan committed
971
972
973
974
        _safe_call(_LIB.LGBM_DatasetSetField(
            self.handle,
            c_str(field_name),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
975
976
            ctypes.c_int(len(data)),
            ctypes.c_int(type_data)))
wxchan's avatar
wxchan committed
977

wxchan's avatar
wxchan committed
978
979
    def get_field(self, field_name):
        """Get property from the Dataset.
wxchan's avatar
wxchan committed
980
981
982

        Parameters
        ----------
983
984
        field_name: string
            The field name of the information.
wxchan's avatar
wxchan committed
985
986
987

        Returns
        -------
988
989
        info : numpy array
            A numpy array with information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
990
        """
991
992
        if self.handle is None:
            raise Exception("Cannot get %s before construct dataset" % field_name)
Guolin Ke's avatar
Guolin Ke committed
993
994
        tmp_out_len = ctypes.c_int()
        out_type = ctypes.c_int()
wxchan's avatar
wxchan committed
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
        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
1010
1011
        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)
1012
        else:
wxchan's avatar
wxchan committed
1013
            raise TypeError("Unknown type")
Guolin Ke's avatar
Guolin Ke committed
1014

1015
    def set_categorical_feature(self, categorical_feature):
1016
        """Set categorical features.
1017
1018
1019

        Parameters
        ----------
1020
1021
        categorical_feature : list of int or strings
            Names or indices of categorical features.
1022
1023
1024
1025
        """
        if self.categorical_feature == categorical_feature:
            return
        if self.data is not None:
1026
1027
1028
1029
1030
1031
1032
1033
1034
            if self.categorical_feature is None:
                self.categorical_feature = categorical_feature
                self._free_handle()
            elif categorical_feature == 'auto':
                warnings.warn('Using categorical_feature in Dataset.')
            else:
                warnings.warn('categorical_feature in Dataset is overrided. New categorical_feature is {}'.format(sorted(list(categorical_feature))))
                self.categorical_feature = categorical_feature
                self._free_handle()
1035
1036
1037
        else:
            raise LightGBMError("Cannot set categorical feature after freed raw data, set free_raw_data=False when construct Dataset to avoid this.")

Guolin Ke's avatar
Guolin Ke committed
1038
1039
1040
1041
1042
1043
1044
1045
1046
    def _set_predictor(self, predictor):
        """
        Set predictor for continued training, not recommand for user to call this function.
        Please set init_model in engine.train or engine.cv
        """
        if predictor is self._predictor:
            return
        if self.data is not None:
            self._predictor = predictor
1047
            self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1048
        else:
1049
            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
1050
1051

    def set_reference(self, reference):
1052
        """Set reference Dataset.
Guolin Ke's avatar
Guolin Ke committed
1053
1054
1055
1056

        Parameters
        ----------
        reference : Dataset
1057
            Reference that is used as a template to consturct the current Dataset.
Guolin Ke's avatar
Guolin Ke committed
1058
        """
1059
        self.set_categorical_feature(reference.categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
1060
1061
        self.set_feature_name(reference.feature_name)
        self._set_predictor(reference._predictor)
1062
1063
        # we're done if self and reference share a common upstrem reference
        if self.get_ref_chain().intersection(reference.get_ref_chain()):
Guolin Ke's avatar
Guolin Ke committed
1064
1065
1066
            return
        if self.data is not None:
            self.reference = reference
1067
            self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1068
        else:
1069
            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
1070
1071

    def set_feature_name(self, feature_name):
1072
        """Set feature name.
Guolin Ke's avatar
Guolin Ke committed
1073
1074
1075

        Parameters
        ----------
1076
1077
        feature_name : list of strings
            Feature names.
Guolin Ke's avatar
Guolin Ke committed
1078
        """
1079
1080
        if feature_name != 'auto':
            self.feature_name = feature_name
1081
        if self.handle is not None and feature_name is not None and feature_name != 'auto':
wxchan's avatar
wxchan committed
1082
1083
            if len(feature_name) != self.num_feature():
                raise ValueError("Length of feature_name({}) and num_feature({}) don't match".format(len(feature_name), self.num_feature()))
1084
            c_feature_name = [c_str(name) for name in feature_name]
wxchan's avatar
wxchan committed
1085
1086
1087
            _safe_call(_LIB.LGBM_DatasetSetFeatureNames(
                self.handle,
                c_array(ctypes.c_char_p, c_feature_name),
Guolin Ke's avatar
Guolin Ke committed
1088
                ctypes.c_int(len(feature_name))))
Guolin Ke's avatar
Guolin Ke committed
1089
1090

    def set_label(self, label):
1091
        """Set label of Dataset
Guolin Ke's avatar
Guolin Ke committed
1092
1093
1094

        Parameters
        ----------
1095
1096
        label: list, numpy array or None
            The label information to be set into Dataset.
Guolin Ke's avatar
Guolin Ke committed
1097
1098
        """
        self.label = label
1099
        if self.handle is not None:
wxchan's avatar
wxchan committed
1100
1101
            label = list_to_1d_numpy(label, name='label')
            self.set_field('label', label)
Guolin Ke's avatar
Guolin Ke committed
1102
1103

    def set_weight(self, weight):
1104
        """Set weight of each instance.
Guolin Ke's avatar
Guolin Ke committed
1105
1106
1107

        Parameters
        ----------
1108
1109
        weight : list, numpy array or None
            Weight to be set for each data point.
Guolin Ke's avatar
Guolin Ke committed
1110
1111
        """
        self.weight = weight
1112
        if self.handle is not None and weight is not None:
wxchan's avatar
wxchan committed
1113
1114
            weight = list_to_1d_numpy(weight, name='weight')
            self.set_field('weight', weight)
Guolin Ke's avatar
Guolin Ke committed
1115
1116

    def set_init_score(self, init_score):
1117
        """Set init score of Booster to start from.
Guolin Ke's avatar
Guolin Ke committed
1118
1119
1120

        Parameters
        ----------
1121
1122
        init_score : list, numpy array or None
            Init score for Booster.
Guolin Ke's avatar
Guolin Ke committed
1123
1124
        """
        self.init_score = init_score
1125
        if self.handle is not None and init_score is not None:
Guolin Ke's avatar
Guolin Ke committed
1126
            init_score = list_to_1d_numpy(init_score, np.float64, name='init_score')
wxchan's avatar
wxchan committed
1127
            self.set_field('init_score', init_score)
Guolin Ke's avatar
Guolin Ke committed
1128
1129

    def set_group(self, group):
1130
        """Set group size of Dataset (used for ranking).
Guolin Ke's avatar
Guolin Ke committed
1131
1132
1133

        Parameters
        ----------
1134
1135
        group : list, numpy array or None
            Group size of each group.
Guolin Ke's avatar
Guolin Ke committed
1136
1137
        """
        self.group = group
1138
        if self.handle is not None and group is not None:
wxchan's avatar
wxchan committed
1139
1140
            group = list_to_1d_numpy(group, np.int32, name='group')
            self.set_field('group', group)
Guolin Ke's avatar
Guolin Ke committed
1141
1142

    def get_label(self):
1143
        """Get the label of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1144
1145
1146

        Returns
        -------
1147
1148
        label : numpy array
            The label information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1149
        """
1150
        if self.label is None and self.handle is not None:
wxchan's avatar
wxchan committed
1151
            self.label = self.get_field('label')
Guolin Ke's avatar
Guolin Ke committed
1152
1153
1154
        return self.label

    def get_weight(self):
1155
        """Get the weight of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1156
1157
1158

        Returns
        -------
1159
1160
        weight : numpy array
            Weight for each data point from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1161
        """
1162
        if self.weight is None and self.handle is not None:
wxchan's avatar
wxchan committed
1163
            self.weight = self.get_field('weight')
Guolin Ke's avatar
Guolin Ke committed
1164
1165
1166
        return self.weight

    def get_init_score(self):
1167
        """Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1168
1169
1170

        Returns
        -------
1171
1172
        init_score : numpy array
            Init score of Booster.
Guolin Ke's avatar
Guolin Ke committed
1173
        """
1174
        if self.init_score is None and self.handle is not None:
wxchan's avatar
wxchan committed
1175
            self.init_score = self.get_field('init_score')
Guolin Ke's avatar
Guolin Ke committed
1176
1177
1178
        return self.init_score

    def get_group(self):
1179
        """Get the group of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1180
1181
1182

        Returns
        -------
1183
        group : numpy array
1184
            Group size of each group.
Guolin Ke's avatar
Guolin Ke committed
1185
        """
1186
        if self.group is None and self.handle is not None:
wxchan's avatar
wxchan committed
1187
            self.group = self.get_field('group')
Guolin Ke's avatar
Guolin Ke committed
1188
1189
1190
            if self.group is not None:
                # group data from LightGBM is boundaries data, need to convert to group size
                new_group = []
wxchan's avatar
wxchan committed
1191
                for i in range_(len(self.group) - 1):
Guolin Ke's avatar
Guolin Ke committed
1192
1193
                    new_group.append(self.group[i + 1] - self.group[i])
                self.group = new_group
Guolin Ke's avatar
Guolin Ke committed
1194
1195
1196
        return self.group

    def num_data(self):
1197
        """Get the number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1198
1199
1200

        Returns
        -------
1201
1202
        number_of_rows : int
            The number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1203
        """
1204
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1205
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1206
1207
1208
            _safe_call(_LIB.LGBM_DatasetGetNumData(self.handle,
                                                   ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1209
        else:
1210
            raise LightGBMError("Cannot get num_data before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1211
1212

    def num_feature(self):
1213
        """Get the number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1214
1215
1216

        Returns
        -------
1217
1218
        number_of_columns : int
            The number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1219
        """
1220
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1221
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1222
1223
1224
            _safe_call(_LIB.LGBM_DatasetGetNumFeature(self.handle,
                                                      ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1225
        else:
1226
            raise LightGBMError("Cannot get num_feature before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1227

1228
    def get_ref_chain(self, ref_limit=100):
1229
1230
1231
1232
1233
1234
1235
        """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.
1236
1237
1238

        Returns
        -------
1239
1240
1241
        ref_chain : set of Dataset
            Chain of references of the Datasets.
        """
1242
        head = self
1243
        ref_chain = set()
1244
1245
        while len(ref_chain) < ref_limit:
            if isinstance(head, Dataset):
1246
                ref_chain.add(head)
1247
1248
1249
1250
1251
1252
                if (head.reference is not None) and (head.reference not in ref_chain):
                    head = head.reference
                else:
                    break
            else:
                break
1253
        return(ref_chain)
1254

wxchan's avatar
wxchan committed
1255

wxchan's avatar
wxchan committed
1256
class Booster(object):
1257
    """Booster in LightGBM."""
wxchan's avatar
wxchan committed
1258
    def __init__(self, params=None, train_set=None, model_file=None, silent=False):
1259
        """Initialize the Booster.
wxchan's avatar
wxchan committed
1260
1261
1262

        Parameters
        ----------
1263
1264
1265
1266
1267
        params: dict or None, optional (default=None)
            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
1268
            Path to the model file.
1269
1270
        silent : bool, optional (default=False)
            Whether to print messages during construction.
wxchan's avatar
wxchan committed
1271
        """
1272
        self.handle = None
1273
        self.network = False
wxchan's avatar
wxchan committed
1274
1275
1276
1277
        self.__need_reload_eval_info = True
        self.__train_data_name = "training"
        self.__attr = {}
        self.best_iteration = -1
wxchan's avatar
wxchan committed
1278
        self.best_score = {}
wxchan's avatar
wxchan committed
1279
        params = {} if params is None else params
1280
1281
        if "verbosity" in params:
            params.setdefault("verbose", params.pop("verbosity"))
wxchan's avatar
wxchan committed
1282
1283
1284
1285
1286
        if silent:
            params["verbose"] = 0
        elif "verbose" not in params:
            params["verbose"] = 1
        if train_set is not None:
1287
            # Training task
wxchan's avatar
wxchan committed
1288
            if not isinstance(train_set, Dataset):
1289
                raise TypeError('Training data should be Dataset instance, met {}'.format(type(train_set).__name__))
wxchan's avatar
wxchan committed
1290
            params_str = param_dict_to_str(params)
1291
            # construct booster object
1292
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1293
            _safe_call(_LIB.LGBM_BoosterCreate(
wxchan's avatar
wxchan committed
1294
                train_set.construct().handle,
wxchan's avatar
wxchan committed
1295
1296
                c_str(params_str),
                ctypes.byref(self.handle)))
1297
            # save reference to data
wxchan's avatar
wxchan committed
1298
1299
1300
1301
            self.train_set = train_set
            self.valid_sets = []
            self.name_valid_sets = []
            self.__num_dataset = 1
Guolin Ke's avatar
Guolin Ke committed
1302
1303
            self.__init_predictor = train_set._predictor
            if self.__init_predictor is not None:
wxchan's avatar
wxchan committed
1304
1305
                _safe_call(_LIB.LGBM_BoosterMerge(
                    self.handle,
Guolin Ke's avatar
Guolin Ke committed
1306
                    self.__init_predictor.handle))
Guolin Ke's avatar
Guolin Ke committed
1307
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1308
1309
1310
1311
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1312
            # buffer for inner predict
wxchan's avatar
wxchan committed
1313
1314
1315
            self.__inner_predict_buffer = [None]
            self.__is_predicted_cur_iter = [False]
            self.__get_eval_info()
1316
            self.pandas_categorical = train_set.pandas_categorical
1317
            # set network if necessary
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
            if "machines" in params:
                machines = params["machines"]
                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))
wxchan's avatar
wxchan committed
1331
        elif model_file is not None:
1332
            # Prediction task
Guolin Ke's avatar
Guolin Ke committed
1333
            out_num_iterations = ctypes.c_int(0)
1334
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1335
1336
1337
1338
            _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
1339
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1340
1341
1342
1343
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1344
            self.pandas_categorical = _load_pandas_categorical(model_file)
1345
        elif 'model_str' in params:
1346
            self._load_model_from_string(params['model_str'])
wxchan's avatar
wxchan committed
1347
        else:
1348
            raise TypeError('Need at least one training dataset or model file to create booster instance')
wxchan's avatar
wxchan committed
1349
1350

    def __del__(self):
1351
1352
        if self.network:
            self.free_network()
Guolin Ke's avatar
Guolin Ke committed
1353
        if self.handle is not None:
wxchan's avatar
wxchan committed
1354
1355
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))

wxchan's avatar
wxchan committed
1356
1357
1358
1359
    def __copy__(self):
        return self.__deepcopy__(None)

    def __deepcopy__(self, _):
1360
        model_str = self._save_model_to_string()
1361
        booster = Booster({'model_str': model_str})
1362
        booster.pandas_categorical = self.pandas_categorical
1363
        return booster
wxchan's avatar
wxchan committed
1364
1365
1366
1367
1368
1369
1370

    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:
1371
            this["handle"] = self._save_model_to_string()
wxchan's avatar
wxchan committed
1372
1373
1374
        return this

    def __setstate__(self, state):
1375
1376
        model_str = state.get('handle', None)
        if model_str is not None:
wxchan's avatar
wxchan committed
1377
            handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
1378
            out_num_iterations = ctypes.c_int(0)
1379
1380
1381
1382
            _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
                c_str(model_str),
                ctypes.byref(out_num_iterations),
                ctypes.byref(handle)))
wxchan's avatar
wxchan committed
1383
1384
1385
            state['handle'] = handle
        self.__dict__.update(state)

wxchan's avatar
wxchan committed
1386
    def free_dataset(self):
1387
        """Free Booster's Datasets."""
wxchan's avatar
wxchan committed
1388
1389
        self.__dict__.pop('train_set', None)
        self.__dict__.pop('valid_sets', None)
1390
        self.__num_dataset = 0
wxchan's avatar
wxchan committed
1391

1392
1393
1394
1395
    def _free_buffer(self):
        self.__inner_predict_buffer = []
        self.__is_predicted_cur_iter = []

1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
    def set_network(self, machines, local_listen_port=12400,
                    listen_time_out=120, num_machines=1):
        """Set the network configuration.

        Parameters
        ----------
        machines: list, set or string
            Names of machines.
        local_listen_port: int, optional (default=12400)
            TCP listen port for local machines.
        listen_time_out: int, optional (default=120)
            Socket time-out in minutes.
        num_machines: int, optional (default=1)
            The number of machines for parallel learning application.
        """
        _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

    def free_network(self):
        """Free Network."""
        _safe_call(_LIB.LGBM_NetworkFree())
        self.network = False

wxchan's avatar
wxchan committed
1422
    def set_train_data_name(self, name):
1423
1424
1425
1426
1427
1428
1429
        """Set the name to the training Dataset.

        Parameters
        ----------
        name: string
            Name for training Dataset.
        """
wxchan's avatar
wxchan committed
1430
1431
1432
        self.__train_data_name = name

    def add_valid(self, data, name):
1433
        """Add validation data.
wxchan's avatar
wxchan committed
1434
1435
1436
1437

        Parameters
        ----------
        data : Dataset
1438
1439
1440
            Validation data.
        name : string
            Name of validation data.
wxchan's avatar
wxchan committed
1441
        """
Guolin Ke's avatar
Guolin Ke committed
1442
        if not isinstance(data, Dataset):
1443
            raise TypeError('Validation data should be Dataset instance, met {}'.format(type(data).__name__))
Guolin Ke's avatar
Guolin Ke committed
1444
1445
        if data._predictor is not self.__init_predictor:
            raise LightGBMError("Add validation data failed, you should use same predictor for these data")
wxchan's avatar
wxchan committed
1446
1447
        _safe_call(_LIB.LGBM_BoosterAddValidData(
            self.handle,
wxchan's avatar
wxchan committed
1448
            data.construct().handle))
wxchan's avatar
wxchan committed
1449
1450
1451
1452
1453
1454
1455
        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)

    def reset_parameter(self, params):
1456
        """Reset parameters of Booster.
wxchan's avatar
wxchan committed
1457
1458
1459
1460

        Parameters
        ----------
        params : dict
1461
            New parameters for Booster.
wxchan's avatar
wxchan committed
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
        """
        if 'metric' in params:
            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)))

    def update(self, train_set=None, fobj=None):
1472
        """Update for one iteration.
1473

wxchan's avatar
wxchan committed
1474
1475
        Parameters
        ----------
1476
1477
1478
1479
        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
1480
1481
            Customized objective function.

1482
1483
1484
1485
            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
1486
1487
        Returns
        -------
1488
1489
        is_finished : bool
            Whether the update was successfully finished.
wxchan's avatar
wxchan committed
1490
1491
        """

1492
        # need reset training data
wxchan's avatar
wxchan committed
1493
        if train_set is not None and train_set is not self.train_set:
Guolin Ke's avatar
Guolin Ke committed
1494
1495
            if not isinstance(train_set, Dataset):
                raise TypeError('Training data should be Dataset instance, met {}'.format(type(train_set).__name__))
Guolin Ke's avatar
Guolin Ke committed
1496
1497
            if train_set._predictor is not self.__init_predictor:
                raise LightGBMError("Replace training data failed, you should use same predictor for these data")
wxchan's avatar
wxchan committed
1498
1499
1500
            self.train_set = train_set
            _safe_call(_LIB.LGBM_BoosterResetTrainingData(
                self.handle,
wxchan's avatar
wxchan committed
1501
                self.train_set.construct().handle))
wxchan's avatar
wxchan committed
1502
1503
1504
1505
1506
1507
            self.__inner_predict_buffer[0] = None
        is_finished = ctypes.c_int(0)
        if fobj is None:
            _safe_call(_LIB.LGBM_BoosterUpdateOneIter(
                self.handle,
                ctypes.byref(is_finished)))
wxchan's avatar
wxchan committed
1508
            self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
            return is_finished.value == 1
        else:
            grad, hess = fobj(self.__inner_predict(0), self.train_set)
            return self.__boost(grad, hess)

    def __boost(self, grad, hess):
        """
        Boost the 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
1520

wxchan's avatar
wxchan committed
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
        Parameters
        ----------
        grad : 1d numpy or 1d list
            The first order of gradient.
        hess : 1d numpy or 1d list
            The second order of gradient.

        Returns
        -------
        is_finished, bool
        """
1532
1533
        grad = list_to_1d_numpy(grad, name='gradient')
        hess = list_to_1d_numpy(hess, name='hessian')
wxchan's avatar
wxchan committed
1534
        if len(grad) != len(hess):
1535
            raise ValueError("Lengths of gradient({}) and hessian({}) don't match".format(len(grad), len(hess)))
wxchan's avatar
wxchan committed
1536
1537
1538
1539
1540
1541
        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
1542
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1543
1544
1545
        return is_finished.value == 1

    def rollback_one_iter(self):
1546
        """Rollback one iteration."""
wxchan's avatar
wxchan committed
1547
1548
        _safe_call(_LIB.LGBM_BoosterRollbackOneIter(
            self.handle))
wxchan's avatar
wxchan committed
1549
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1550
1551

    def current_iteration(self):
1552
1553
1554
1555
1556
1557
1558
        """Get the index of the current iteration.

        Returns
        -------
        cur_iter : int
            The index of the current iteration.
        """
Guolin Ke's avatar
Guolin Ke committed
1559
        out_cur_iter = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1560
1561
1562
1563
1564
1565
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
            self.handle,
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

    def eval(self, data, name, feval=None):
1566
        """Evaluate for data.
wxchan's avatar
wxchan committed
1567
1568
1569

        Parameters
        ----------
1570
1571
1572
1573
1574
        data : Dataset
            Data for the evaluating.
        name : string
            Name of the data.
        feval : callable or None, optional (default=None)
wxchan's avatar
wxchan committed
1575
            Custom evaluation function.
1576

wxchan's avatar
wxchan committed
1577
1578
1579
        Returns
        -------
        result: list
1580
            List with evaluation results.
wxchan's avatar
wxchan committed
1581
        """
Guolin Ke's avatar
Guolin Ke committed
1582
1583
        if not isinstance(data, Dataset):
            raise TypeError("Can only eval for Dataset instance")
wxchan's avatar
wxchan committed
1584
1585
1586
1587
        data_idx = -1
        if data is self.train_set:
            data_idx = 0
        else:
wxchan's avatar
wxchan committed
1588
            for i in range_(len(self.valid_sets)):
wxchan's avatar
wxchan committed
1589
1590
1591
                if data is self.valid_sets[i]:
                    data_idx = i + 1
                    break
1592
        # need to push new valid data
wxchan's avatar
wxchan committed
1593
1594
1595
1596
1597
1598
1599
        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):
1600
        """Evaluate for training data.
wxchan's avatar
wxchan committed
1601
1602
1603

        Parameters
        ----------
1604
        feval : callable or None, optional (default=None)
wxchan's avatar
wxchan committed
1605
1606
1607
1608
            Custom evaluation function.

        Returns
        -------
1609
1610
        result: list
            List with evaluation results.
wxchan's avatar
wxchan committed
1611
1612
1613
1614
        """
        return self.__inner_eval(self.__train_data_name, 0, feval)

    def eval_valid(self, feval=None):
1615
        """Evaluate for validation data.
wxchan's avatar
wxchan committed
1616
1617
1618

        Parameters
        ----------
1619
        feval : callable or None, optional (default=None)
wxchan's avatar
wxchan committed
1620
1621
1622
1623
            Custom evaluation function.

        Returns
        -------
1624
1625
        result: list
            List with evaluation results.
wxchan's avatar
wxchan committed
1626
        """
wxchan's avatar
wxchan committed
1627
        return [item for i in range_(1, self.__num_dataset)
wxchan's avatar
wxchan committed
1628
                for item in self.__inner_eval(self.name_valid_sets[i - 1], i, feval)]
wxchan's avatar
wxchan committed
1629
1630

    def save_model(self, filename, num_iteration=-1):
1631
        """Save Booster to file.
wxchan's avatar
wxchan committed
1632
1633
1634

        Parameters
        ----------
1635
1636
1637
1638
1639
        filename : string
            Filename to save Booster.
        num_iteration: int, optional (default=-1)
            Index of the iteration that should to saved.
            If <0, the best iteration (if exists) is saved.
wxchan's avatar
wxchan committed
1640
        """
1641
1642
        if num_iteration <= 0:
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1643
1644
        _safe_call(_LIB.LGBM_BoosterSaveModel(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
1645
            ctypes.c_int(num_iteration),
wxchan's avatar
wxchan committed
1646
            c_str(filename)))
1647
        _save_pandas_categorical(filename, self.pandas_categorical)
wxchan's avatar
wxchan committed
1648

1649
    def _load_model_from_string(self, model_str):
1650
        """[Private] Load model from string"""
1651
1652
1653
1654
        if self.handle is not None:
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        self._free_buffer()
        self.handle = ctypes.c_void_p()
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
        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)))
        self.__num_class = out_num_class.value

1666
    def _save_model_to_string(self, num_iteration=-1):
1667
1668
1669
1670
        """[Private] Save model to string"""
        if num_iteration <= 0:
            num_iteration = self.best_iteration
        buffer_len = 1 << 20
1671
        tmp_out_len = ctypes.c_int64(0)
1672
1673
1674
1675
1676
        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,
            ctypes.c_int(num_iteration),
1677
            ctypes.c_int64(buffer_len),
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
            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,
                ctypes.c_int(num_iteration),
1688
                ctypes.c_int64(actual_len),
1689
1690
1691
1692
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
        return string_buffer.value.decode()

1693
    def dump_model(self, num_iteration=-1):
1694
        """Dump Booster to json format.
wxchan's avatar
wxchan committed
1695

1696
1697
        Parameters
        ----------
1698
1699
1700
        num_iteration: int, optional (default=-1)
            Index of the iteration that should to dumped.
            If <0, the best iteration (if exists) is dumped.
1701

wxchan's avatar
wxchan committed
1702
1703
        Returns
        -------
1704
1705
        json_repr : dict
            Json format of Booster.
wxchan's avatar
wxchan committed
1706
        """
1707
1708
        if num_iteration <= 0:
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1709
        buffer_len = 1 << 20
1710
        tmp_out_len = ctypes.c_int64(0)
wxchan's avatar
wxchan committed
1711
1712
1713
1714
        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,
Guolin Ke's avatar
Guolin Ke committed
1715
            ctypes.c_int(num_iteration),
1716
            ctypes.c_int64(buffer_len),
wxchan's avatar
wxchan committed
1717
            ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1718
            ptr_string_buffer))
wxchan's avatar
wxchan committed
1719
        actual_len = tmp_out_len.value
1720
        '''if buffer length is not long enough, reallocate a buffer'''
wxchan's avatar
wxchan committed
1721
1722
1723
1724
1725
        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,
Guolin Ke's avatar
Guolin Ke committed
1726
                ctypes.c_int(num_iteration),
1727
                ctypes.c_int64(actual_len),
wxchan's avatar
wxchan committed
1728
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1729
                ptr_string_buffer))
wxchan's avatar
wxchan committed
1730
1731
        return json.loads(string_buffer.value.decode())

1732
    def predict(self, data, num_iteration=-1, raw_score=False, pred_leaf=False, pred_contrib=False,
1733
1734
                data_has_header=False, is_reshape=True, pred_parameter=None):
        """Make a prediction.
wxchan's avatar
wxchan committed
1735
1736
1737

        Parameters
        ----------
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
        data : string, numpy array or scipy.sparse
            Data source for prediction.
            If string, it represents the path to txt file.
        num_iteration : int, optional (default=-1)
            Iteration used for prediction.
            If <0, the best iteration (if exists) is used for prediction.
        raw_score : bool, optional (default=False)
            Whether to predict raw scores.
        pred_leaf : bool, optional (default=False)
            Whether to predict leaf index.
1748
1749
        pred_contrib : bool, optional (default=False)
            Whether to predict feature contributions.
1750
1751
1752
1753
1754
1755
1756
        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].
        pred_parameter: dict or None, optional (default=None)
            Other parameters for the prediction.
wxchan's avatar
wxchan committed
1757
1758
1759

        Returns
        -------
1760
1761
        result : numpy array
            Prediction result.
wxchan's avatar
wxchan committed
1762
        """
1763
        predictor = self._to_predictor(pred_parameter)
1764
1765
        if num_iteration <= 0:
            num_iteration = self.best_iteration
1766
        return predictor.predict(data, num_iteration, raw_score, pred_leaf, pred_contrib, data_has_header, is_reshape)
wxchan's avatar
wxchan committed
1767

1768
    def get_leaf_output(self, tree_id, leaf_id):
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
        """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.
        """
1783
1784
1785
1786
1787
1788
1789
1790
        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

1791
    def _to_predictor(self, pred_parameter=None):
wxchan's avatar
wxchan committed
1792
        """Convert to predictor"""
1793
        predictor = _InnerPredictor(booster_handle=self.handle, pred_parameter=pred_parameter)
1794
        predictor.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1795
1796
        return predictor

1797
    def num_feature(self):
1798
1799
1800
1801
1802
1803
1804
        """Get number of features.

        Returns
        -------
        num_feature : int
            The number of features.
        """
1805
1806
1807
1808
1809
1810
        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
1811
    def feature_name(self):
1812
        """Get names of features.
wxchan's avatar
wxchan committed
1813
1814
1815

        Returns
        -------
1816
1817
        result : list
            List with names of features.
wxchan's avatar
wxchan committed
1818
        """
1819
        num_feature = self.num_feature()
1820
        # Get name of features
wxchan's avatar
wxchan committed
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
        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)]

1832
    def feature_importance(self, importance_type='split', iteration=-1):
1833
        """Get feature importances.
1834

1835
1836
        Parameters
        ----------
1837
1838
1839
1840
        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.
1841

1842
1843
        Returns
        -------
1844
1845
        result : numpy array
            Array with feature importances.
1846
        """
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
        if importance_type == "split":
            importance_type_int = 0
        elif importance_type == "gain":
            importance_type_int = 1
        else:
            importance_type_int = -1
        num_feature = self.num_feature()
        result = np.array([0 for _ in range_(num_feature)], dtype=np.float64)
        _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
1864

wxchan's avatar
wxchan committed
1865
1866
    def __inner_eval(self, data_name, data_idx, feval=None):
        """
1867
        Evaulate training or validation data
wxchan's avatar
wxchan committed
1868
1869
        """
        if data_idx >= self.__num_dataset:
1870
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
1871
1872
1873
        self.__get_eval_info()
        ret = []
        if self.__num_inner_eval > 0:
wxchan's avatar
wxchan committed
1874
            result = np.array([0.0 for _ in range_(self.__num_inner_eval)], dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
1875
            tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1876
1877
            _safe_call(_LIB.LGBM_BoosterGetEval(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
1878
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
1879
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1880
                result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
1881
            if tmp_out_len.value != self.__num_inner_eval:
1882
                raise ValueError("Wrong length of eval results")
wxchan's avatar
wxchan committed
1883
            for i in range_(self.__num_inner_eval):
wxchan's avatar
wxchan committed
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
                ret.append((data_name, self.__name_inner_eval[i], result[i], self.__higher_better_inner_eval[i]))
        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:
1904
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
1905
1906
1907
1908
1909
1910
        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
            self.__inner_predict_buffer[data_idx] = \
wxchan's avatar
wxchan committed
1911
                np.array([0.0 for _ in range_(n_preds)], dtype=np.float64, copy=False)
1912
        # avoid to predict many time in one iteration
wxchan's avatar
wxchan committed
1913
1914
        if not self.__is_predicted_cur_iter[data_idx]:
            tmp_out_len = ctypes.c_int64(0)
Guolin Ke's avatar
Guolin Ke committed
1915
            data_ptr = self.__inner_predict_buffer[data_idx].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
wxchan's avatar
wxchan committed
1916
1917
            _safe_call(_LIB.LGBM_BoosterGetPredict(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
1918
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
1919
1920
1921
                ctypes.byref(tmp_out_len),
                data_ptr))
            if tmp_out_len.value != len(self.__inner_predict_buffer[data_idx]):
1922
                raise ValueError("Wrong length of predict results for data %d" % (data_idx))
wxchan's avatar
wxchan committed
1923
1924
1925
1926
1927
1928
1929
1930
1931
            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
1932
            out_num_eval = ctypes.c_int(0)
1933
            # Get num of inner evals
wxchan's avatar
wxchan committed
1934
1935
1936
1937
1938
            _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:
1939
                # Get name of evals
Guolin Ke's avatar
Guolin Ke committed
1940
                tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1941
                string_buffers = [ctypes.create_string_buffer(255) for i in range_(self.__num_inner_eval)]
wxchan's avatar
wxchan committed
1942
                ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
wxchan's avatar
wxchan committed
1943
1944
1945
1946
1947
                _safe_call(_LIB.LGBM_BoosterGetEvalNames(
                    self.handle,
                    ctypes.byref(tmp_out_len),
                    ptr_string_buffers))
                if self.__num_inner_eval != tmp_out_len.value:
1948
                    raise ValueError("Length of eval names doesn't equal with num_evals")
1949
                self.__name_inner_eval = \
wxchan's avatar
wxchan committed
1950
                    [string_buffers[i].value.decode() for i in range_(self.__num_inner_eval)]
1951
                self.__higher_better_inner_eval = \
Guolin Ke's avatar
Guolin Ke committed
1952
                    [name.startswith(('auc', 'ndcg', 'map')) for name in self.__name_inner_eval]
1953

wxchan's avatar
wxchan committed
1954
    def attr(self, key):
1955
        """Get attribute string from the Booster.
wxchan's avatar
wxchan committed
1956
1957
1958

        Parameters
        ----------
1959
1960
        key : string
            The name of the attribute.
wxchan's avatar
wxchan committed
1961
1962
1963

        Returns
        -------
1964
1965
1966
        value : string or None
            The attribute value.
            Returns None if attribute do not exist.
wxchan's avatar
wxchan committed
1967
        """
1968
        return self.__attr.get(key, None)
wxchan's avatar
wxchan committed
1969
1970

    def set_attr(self, **kwargs):
1971
        """Set the attribute of the Booster.
wxchan's avatar
wxchan committed
1972
1973
1974
1975

        Parameters
        ----------
        **kwargs
1976
1977
            The attributes to set.
            Setting a value to None deletes an attribute.
wxchan's avatar
wxchan committed
1978
1979
1980
        """
        for key, value in kwargs.items():
            if value is not None:
wxchan's avatar
wxchan committed
1981
                if not isinstance(value, string_type):
1982
                    raise ValueError("Set attr only accepts strings")
wxchan's avatar
wxchan committed
1983
1984
1985
                self.__attr[key] = value
            else:
                self.__attr.pop(key, None)