basic.py 95.6 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
4
"""Wrapper for C API of LightGBM."""
wxchan's avatar
wxchan committed
5
6
from __future__ import absolute_import

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

import numpy as np
import scipy.sparse

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

wxchan's avatar
wxchan committed
23

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

wxchan's avatar
wxchan committed
33

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

wxchan's avatar
wxchan committed
36

wxchan's avatar
wxchan committed
37
def _safe_call(ret):
38
39
    """Check the return value from C API call.

wxchan's avatar
wxchan committed
40
41
42
    Parameters
    ----------
    ret : int
43
        The return value from C API calls.
wxchan's avatar
wxchan committed
44
45
    """
    if ret != 0:
46
        raise LightGBMError(decode_string(_LIB.LGBM_GetLastError()))
wxchan's avatar
wxchan committed
47

wxchan's avatar
wxchan committed
48

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

wxchan's avatar
wxchan committed
59

wxchan's avatar
wxchan committed
60
def is_numpy_1d_array(data):
61
    """Check whether data is a 1-D numpy array."""
62
    return isinstance(data, np.ndarray) and len(data.shape) == 1
wxchan's avatar
wxchan committed
63

wxchan's avatar
wxchan committed
64

wxchan's avatar
wxchan committed
65
def is_1d_list(data):
66
67
    """Check whether data is a 1-D list."""
    return isinstance(data, list) and (not data or is_numeric(data[0]))
wxchan's avatar
wxchan committed
68

wxchan's avatar
wxchan committed
69

70
def list_to_1d_numpy(data, dtype=np.float32, name='list'):
71
    """Convert data to 1-D numpy array."""
wxchan's avatar
wxchan committed
72
73
74
75
76
77
78
    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)
79
80
    elif isinstance(data, Series):
        return data.values.astype(dtype)
wxchan's avatar
wxchan committed
81
    else:
82
83
        raise TypeError("Wrong type({0}) for {1}.\n"
                        "It should be list, numpy 1-D array or pandas Series".format(type(data).__name__, name))
wxchan's avatar
wxchan committed
84

wxchan's avatar
wxchan committed
85

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

Guolin Ke's avatar
Guolin Ke committed
93

Guolin Ke's avatar
Guolin Ke committed
94
def cfloat64_array_to_numpy(cptr, length):
95
    """Convert a ctypes double pointer array to a numpy array."""
Guolin Ke's avatar
Guolin Ke committed
96
97
98
99
100
    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
101

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

wxchan's avatar
wxchan committed
109

wxchan's avatar
wxchan committed
110
def c_str(string):
111
    """Convert a Python string to C string."""
wxchan's avatar
wxchan committed
112
113
    return ctypes.c_char_p(string.encode('utf-8'))

wxchan's avatar
wxchan committed
114

wxchan's avatar
wxchan committed
115
def c_array(ctype, values):
116
    """Convert a Python array to C array."""
wxchan's avatar
wxchan committed
117
118
    return (ctype * len(values))(*values)

wxchan's avatar
wxchan committed
119

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

wxchan's avatar
wxchan committed
135

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

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

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

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

wxchan's avatar
wxchan committed
155

156
class LightGBMError(Exception):
157
158
    """Error thrown by LightGBM."""

159
160
161
162
163
    pass


MAX_INT32 = (1 << 31) - 1

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

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

173
"""Macro definition of prediction type in C API of LightGBM"""
wxchan's avatar
wxchan committed
174
175
176
C_API_PREDICT_NORMAL = 0
C_API_PREDICT_RAW_SCORE = 1
C_API_PREDICT_LEAF_INDEX = 2
177
C_API_PREDICT_CONTRIB = 3
wxchan's avatar
wxchan committed
178

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

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

wxchan's avatar
wxchan committed
190

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


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

wxchan's avatar
wxchan committed
221

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

wxchan's avatar
wxchan committed
242

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

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


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


298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
def _dump_pandas_categorical(pandas_categorical, file_name=None):
    pandas_str = ('\npandas_categorical:'
                  + json.dumps(pandas_categorical, default=json_default_with_numpy)
                  + '\n')
    if file_name is not None:
        with open(file_name, 'a') as f:
            f.write(pandas_str)
    return pandas_str


def _load_pandas_categorical(file_name=None, model_str=None):
    if file_name is not None:
        with open(file_name, 'r') as f:
            lines = f.readlines()
            last_line = lines[-1]
            if last_line.strip() == "":
                last_line = lines[-2]
            if last_line.startswith('pandas_categorical:'):
                return json.loads(last_line[len('pandas_categorical:'):])
    elif model_str is not None:
        lines = model_str.split('\n')
Guolin Ke's avatar
Guolin Ke committed
319
320
321
        last_line = lines[-1]
        if last_line.strip() == "":
            last_line = lines[-2]
322
323
324
325
326
        if last_line.startswith('pandas_categorical:'):
            return json.loads(last_line[len('pandas_categorical:'):])
    return None


Guolin Ke's avatar
Guolin Ke committed
327
class _InnerPredictor(object):
328
329
330
331
332
333
334
335
    """_InnerPredictor of LightGBM.

    Not exposed to user.
    Used only for prediction, usually used for continued training.

    Note
    ----
    Can be converted from Booster, but cannot be converted to Booster.
Guolin Ke's avatar
Guolin Ke committed
336
    """
337

338
    def __init__(self, model_file=None, booster_handle=None, pred_parameter=None):
339
        """Initialize the _InnerPredictor.
wxchan's avatar
wxchan committed
340
341
342

        Parameters
        ----------
343
        model_file : string or None, optional (default=None)
wxchan's avatar
wxchan committed
344
            Path to the model file.
345
346
347
348
        booster_handle : object or None, optional (default=None)
            Handle of Booster.
        pred_parameter: dict or None, optional (default=None)
            Other parameters for the prediciton.
wxchan's avatar
wxchan committed
349
350
351
352
353
        """
        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
354
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
355
356
357
358
            _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
359
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
360
361
362
363
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
364
            self.num_total_iteration = out_num_iterations.value
365
            self.pandas_categorical = _load_pandas_categorical(file_name=model_file)
wxchan's avatar
wxchan committed
366
        elif booster_handle is not None:
Guolin Ke's avatar
Guolin Ke committed
367
            self.__is_manage_handle = False
wxchan's avatar
wxchan committed
368
            self.handle = booster_handle
Guolin Ke's avatar
Guolin Ke committed
369
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
370
371
372
373
            _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
374
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
375
376
377
            _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
                self.handle,
                ctypes.byref(out_num_iterations)))
378
            self.num_total_iteration = out_num_iterations.value
379
            self.pandas_categorical = None
wxchan's avatar
wxchan committed
380
        else:
Guolin Ke's avatar
Guolin Ke committed
381
            raise TypeError('Need Model file or Booster handle to create a predictor')
wxchan's avatar
wxchan committed
382

383
384
        pred_parameter = {} if pred_parameter is None else pred_parameter
        self.pred_parameter = param_dict_to_str(pred_parameter)
cbecker's avatar
cbecker committed
385

wxchan's avatar
wxchan committed
386
    def __del__(self):
387
388
389
390
391
        try:
            if self.__is_manage_handle:
                _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
392

393
394
395
396
397
    def __getstate__(self):
        this = self.__dict__.copy()
        this.pop('handle', None)
        return this

wxchan's avatar
wxchan committed
398
    def predict(self, data, num_iteration=-1,
399
                raw_score=False, pred_leaf=False, pred_contrib=False, data_has_header=False,
wxchan's avatar
wxchan committed
400
                is_reshape=True):
401
        """Predict logic.
wxchan's avatar
wxchan committed
402
403
404

        Parameters
        ----------
405
        data : string, numpy array, pandas DataFrame or scipy.sparse
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
            Data source for prediction.
            When data type is string, it represents the path of txt file.
        num_iteration : int, optional (default=-1)
            Iteration 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.
        pred_contrib : bool, optional (default=False)
            Whether to predict feature contributions.
        data_has_header : bool, optional (default=False)
            Whether data has header.
            Used only for txt data.
        is_reshape : bool, optional (default=True)
            Whether to reshape to (nrow, ncol).
wxchan's avatar
wxchan committed
421
422
423

        Returns
        -------
424
425
        result : numpy array
            Prediction result.
wxchan's avatar
wxchan committed
426
        """
wxchan's avatar
wxchan committed
427
        if isinstance(data, Dataset):
428
            raise TypeError("Cannot use Dataset instance for prediction, please use raw data instead")
429
        data = _data_from_pandas(data, None, None, self.pandas_categorical)[0]
wxchan's avatar
wxchan committed
430
431
432
433
434
        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
435
436
        if pred_contrib:
            predict_type = C_API_PREDICT_CONTRIB
wxchan's avatar
wxchan committed
437
        int_data_has_header = 1 if data_has_header else 0
438
439
        if num_iteration > self.num_total_iteration:
            num_iteration = self.num_total_iteration
cbecker's avatar
cbecker committed
440

wxchan's avatar
wxchan committed
441
        if isinstance(data, string_type):
442
            with _TempFile() as f:
wxchan's avatar
wxchan committed
443
444
445
                _safe_call(_LIB.LGBM_BoosterPredictForFile(
                    self.handle,
                    c_str(data),
Guolin Ke's avatar
Guolin Ke committed
446
447
448
                    ctypes.c_int(int_data_has_header),
                    ctypes.c_int(predict_type),
                    ctypes.c_int(num_iteration),
449
                    c_str(self.pred_parameter),
wxchan's avatar
wxchan committed
450
451
                    c_str(f.name)))
                lines = f.readlines()
452
453
                nrow = len(lines)
                preds = [float(token) for line in lines for token in line.split('\t')]
Guolin Ke's avatar
Guolin Ke committed
454
                preds = np.array(preds, dtype=np.float64, copy=False)
wxchan's avatar
wxchan committed
455
        elif isinstance(data, scipy.sparse.csr_matrix):
456
            preds, nrow = self.__pred_for_csr(data, num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
457
        elif isinstance(data, scipy.sparse.csc_matrix):
458
            preds, nrow = self.__pred_for_csc(data, num_iteration, predict_type)
wxchan's avatar
wxchan committed
459
        elif isinstance(data, np.ndarray):
460
            preds, nrow = self.__pred_for_np2d(data, num_iteration, predict_type)
461
462
463
        elif isinstance(data, list):
            try:
                data = np.array(data)
464
            except BaseException:
465
                raise ValueError('Cannot convert data list to numpy array.')
466
            preds, nrow = self.__pred_for_np2d(data, num_iteration, predict_type)
wxchan's avatar
wxchan committed
467
468
        else:
            try:
469
                warnings.warn('Converting data to scipy sparse matrix.')
wxchan's avatar
wxchan committed
470
                csr = scipy.sparse.csr_matrix(data)
471
            except BaseException:
472
                raise TypeError('Cannot predict data for type {}'.format(type(data).__name__))
473
            preds, nrow = self.__pred_for_csr(csr, num_iteration, predict_type)
wxchan's avatar
wxchan committed
474
475
        if pred_leaf:
            preds = preds.astype(np.int32)
476
        if is_reshape and preds.size != nrow:
wxchan's avatar
wxchan committed
477
            if preds.size % nrow == 0:
478
                preds = preds.reshape(nrow, -1)
wxchan's avatar
wxchan committed
479
            else:
480
                raise ValueError('Length of predict result (%d) cannot be divide nrow (%d)'
wxchan's avatar
wxchan committed
481
482
483
484
                                 % (preds.size, nrow))
        return preds

    def __get_num_preds(self, num_iteration, nrow, predict_type):
485
        """Get size of prediction result."""
486
487
488
489
490
        if nrow > MAX_INT32:
            raise LightGBMError('LightGBM cannot perform prediction for data'
                                'with number of rows greater than MAX_INT32 (%d).\n'
                                'You can split your data into chunks'
                                'and then concatenate predictions for them' % MAX_INT32)
Guolin Ke's avatar
Guolin Ke committed
491
492
493
        n_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterCalcNumPredict(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
494
495
496
            ctypes.c_int(nrow),
            ctypes.c_int(predict_type),
            ctypes.c_int(num_iteration),
Guolin Ke's avatar
Guolin Ke committed
497
498
            ctypes.byref(n_preds)))
        return n_preds.value
wxchan's avatar
wxchan committed
499
500

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

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
        def inner_predict(mat, num_iteration, predict_type, preds=None):
            if mat.dtype == np.float32 or mat.dtype == np.float64:
                data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
            else:
                """change non-float data to float data, need to copy"""
                data = np.array(mat.reshape(mat.size), dtype=np.float32)
            ptr_data, type_ptr_data, _ = c_float_array(data)
            n_preds = self.__get_num_preds(num_iteration, mat.shape[0], predict_type)
            if preds is None:
                preds = np.zeros(n_preds, dtype=np.float64)
            elif len(preds.shape) != 1 or len(preds) != n_preds:
                raise ValueError("Wrong length of pre-allocated predict array")
            out_num_preds = ctypes.c_int64(0)
            _safe_call(_LIB.LGBM_BoosterPredictForMat(
                self.handle,
                ptr_data,
                ctypes.c_int(type_ptr_data),
                ctypes.c_int(mat.shape[0]),
                ctypes.c_int(mat.shape[1]),
                ctypes.c_int(C_API_IS_ROW_MAJOR),
                ctypes.c_int(predict_type),
                ctypes.c_int(num_iteration),
                c_str(self.pred_parameter),
                ctypes.byref(out_num_preds),
                preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
            if n_preds != out_num_preds.value:
                raise ValueError("Wrong length for predict results")
            return preds, mat.shape[0]

        nrow = mat.shape[0]
        if nrow > MAX_INT32:
            sections = np.arange(start=MAX_INT32, stop=nrow, step=MAX_INT32)
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
            n_preds = [self.__get_num_preds(num_iteration, i, predict_type) for i in np.diff([0] + list(sections) + [nrow])]
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
            preds = np.zeros(sum(n_preds), dtype=np.float64)
541
542
            for chunk, (start_idx_pred, end_idx_pred) in zip_(np.array_split(mat, sections),
                                                              zip_(n_preds_sections, n_preds_sections[1:])):
543
544
545
                # avoid memory consumption by arrays concatenation operations
                inner_predict(chunk, num_iteration, predict_type, preds[start_idx_pred:end_idx_pred])
            return preds, nrow
wxchan's avatar
wxchan committed
546
        else:
547
            return inner_predict(mat, num_iteration, predict_type)
wxchan's avatar
wxchan committed
548
549

    def __pred_for_csr(self, csr, num_iteration, predict_type):
550
        """Predict for a CSR data."""
551
552
553
554
555
556
557
558
559
560
561
562
        def inner_predict(csr, num_iteration, predict_type, preds=None):
            nrow = len(csr.indptr) - 1
            n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
            if preds is None:
                preds = np.zeros(n_preds, dtype=np.float64)
            elif len(preds.shape) != 1 or len(preds) != n_preds:
                raise ValueError("Wrong length of pre-allocated predict array")
            out_num_preds = ctypes.c_int64(0)

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

563
564
565
            assert csr.shape[1] <= MAX_INT32
            csr.indices = csr.indices.astype(np.int32, copy=False)

566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
            _safe_call(_LIB.LGBM_BoosterPredictForCSR(
                self.handle,
                ptr_indptr,
                ctypes.c_int32(type_ptr_indptr),
                csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
                ptr_data,
                ctypes.c_int(type_ptr_data),
                ctypes.c_int64(len(csr.indptr)),
                ctypes.c_int64(len(csr.data)),
                ctypes.c_int64(csr.shape[1]),
                ctypes.c_int(predict_type),
                ctypes.c_int(num_iteration),
                c_str(self.pred_parameter),
                ctypes.byref(out_num_preds),
                preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
            if n_preds != out_num_preds.value:
                raise ValueError("Wrong length for predict results")
            return preds, nrow
wxchan's avatar
wxchan committed
584

585
586
587
588
589
590
591
592
593
594
595
596
597
598
        nrow = len(csr.indptr) - 1
        if nrow > MAX_INT32:
            sections = [0] + list(np.arange(start=MAX_INT32, stop=nrow, step=MAX_INT32)) + [nrow]
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
            n_preds = [self.__get_num_preds(num_iteration, i, predict_type) for i in np.diff(sections)]
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
            preds = np.zeros(sum(n_preds), dtype=np.float64)
            for (start_idx, end_idx), (start_idx_pred, end_idx_pred) in zip_(zip_(sections, sections[1:]),
                                                                             zip_(n_preds_sections, n_preds_sections[1:])):
                # avoid memory consumption by arrays concatenation operations
                inner_predict(csr[start_idx:end_idx], num_iteration, predict_type, preds[start_idx_pred:end_idx_pred])
            return preds, nrow
        else:
            return inner_predict(csr, num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
599
600

    def __pred_for_csc(self, csc, num_iteration, predict_type):
601
        """Predict for a CSC data."""
Guolin Ke's avatar
Guolin Ke committed
602
        nrow = csc.shape[0]
603
604
        if nrow > MAX_INT32:
            return self.__pred_for_csr(csc.tocsr(), num_iteration, predict_type)
Guolin Ke's avatar
Guolin Ke committed
605
606
607
608
        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)

609
610
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csc.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csc.data)
Guolin Ke's avatar
Guolin Ke committed
611

612
613
614
        assert csc.shape[0] <= MAX_INT32
        csc.indices = csc.indices.astype(np.int32, copy=False)

Guolin Ke's avatar
Guolin Ke committed
615
616
617
        _safe_call(_LIB.LGBM_BoosterPredictForCSC(
            self.handle,
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
618
            ctypes.c_int32(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
619
620
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
621
622
623
624
625
626
            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),
627
            c_str(self.pred_parameter),
Guolin Ke's avatar
Guolin Ke committed
628
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
629
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
630
        if n_preds != out_num_preds.value:
631
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
632
633
        return preds, nrow

wxchan's avatar
wxchan committed
634

wxchan's avatar
wxchan committed
635
636
class Dataset(object):
    """Dataset in LightGBM."""
637

638
    def __init__(self, data, label=None, reference=None,
639
                 weight=None, group=None, init_score=None, silent=False,
640
                 feature_name='auto', categorical_feature='auto', params=None,
wxchan's avatar
wxchan committed
641
                 free_raw_data=True):
642
        """Initialize Dataset.
643

wxchan's avatar
wxchan committed
644
645
        Parameters
        ----------
646
        data : string, numpy array, pandas DataFrame, scipy.sparse or list of numpy arrays
wxchan's avatar
wxchan committed
647
            Data source of Dataset.
648
            If string, it represents the path to txt file.
649
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
650
651
652
            Label of the data.
        reference : Dataset or None, optional (default=None)
            If this is Dataset for validation, training data should be used as reference.
653
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
wxchan's avatar
wxchan committed
654
            Weight for each instance.
655
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
656
            Group/query size for Dataset.
657
        init_score : list, numpy 1-D array, pandas Series or None, optional (default=None)
658
            Init score for Dataset.
659
660
661
662
663
664
665
666
667
668
        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.
669
            All values in categorical features should be less than int32 max value (2147483647).
670
            Large values could be memory consuming. Consider using consecutive integers starting from zero.
671
            All negative values in categorical features will be treated as missing values.
Nikita Titov's avatar
Nikita Titov committed
672
        params : dict or None, optional (default=None)
673
            Other parameters for Dataset.
Nikita Titov's avatar
Nikita Titov committed
674
        free_raw_data : bool, optional (default=True)
675
            If True, raw data is freed after constructing inner Dataset.
wxchan's avatar
wxchan committed
676
        """
wxchan's avatar
wxchan committed
677
678
679
680
681
682
        self.handle = None
        self.data = data
        self.label = label
        self.reference = reference
        self.weight = weight
        self.group = group
683
        self.init_score = init_score
wxchan's avatar
wxchan committed
684
685
        self.silent = silent
        self.feature_name = feature_name
686
        self.categorical_feature = categorical_feature
687
        self.params = copy.deepcopy(params)
wxchan's avatar
wxchan committed
688
689
        self.free_raw_data = free_raw_data
        self.used_indices = None
690
        self.need_slice = True
wxchan's avatar
wxchan committed
691
        self._predictor = None
692
        self.pandas_categorical = None
693
        self.params_back_up = None
wxchan's avatar
wxchan committed
694
695

    def __del__(self):
696
697
698
699
        try:
            self._free_handle()
        except AttributeError:
            pass
700
701

    def _free_handle(self):
702
        if self.handle is not None:
703
            _safe_call(_LIB.LGBM_DatasetFree(self.handle))
704
            self.handle = None
Nikita Titov's avatar
Nikita Titov committed
705
        return self
wxchan's avatar
wxchan committed
706

707
    def _lazy_init(self, data, label=None, reference=None,
708
                   weight=None, group=None, init_score=None, predictor=None,
wxchan's avatar
wxchan committed
709
                   silent=False, feature_name='auto',
710
                   categorical_feature='auto', params=None):
wxchan's avatar
wxchan committed
711
712
        if data is None:
            self.handle = None
Nikita Titov's avatar
Nikita Titov committed
713
            return self
Guolin Ke's avatar
Guolin Ke committed
714
715
716
        if reference is not None:
            self.pandas_categorical = reference.pandas_categorical
            categorical_feature = reference.categorical_feature
717
718
719
720
        data, feature_name, categorical_feature, self.pandas_categorical = _data_from_pandas(data,
                                                                                             feature_name,
                                                                                             categorical_feature,
                                                                                             self.pandas_categorical)
wxchan's avatar
wxchan committed
721
722
        label = _label_from_pandas(label)
        self.data_has_header = False
723
        # process for args
wxchan's avatar
wxchan committed
724
        params = {} if params is None else params
725
726
727
        args_names = (getattr(self.__class__, '_lazy_init')
                      .__code__
                      .co_varnames[:getattr(self.__class__, '_lazy_init').__code__.co_argcount])
728
729
        for key, _ in params.items():
            if key in args_names:
730
731
732
                warnings.warn('{0} keyword has been found in `params` and will be ignored.\n'
                              'Please use {0} argument of the Dataset constructor to pass this parameter.'
                              .format(key))
wxchan's avatar
wxchan committed
733
        self.predictor = predictor
734
735
736
        # user can set verbose with params, it has higher priority
        if not any(verbose_alias in params for verbose_alias in ('verbose', 'verbosity')) and silent:
            params["verbose"] = -1
737
        # get categorical features
738
739
740
741
742
743
744
745
746
747
748
749
750
        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))
751
            if categorical_indices:
752
                if "categorical_feature" in params or "categorical_column" in params:
753
                    warnings.warn('categorical_feature in param dict is overridden.')
754
755
                    params.pop("categorical_feature", None)
                    params.pop("categorical_column", None)
756
                params['categorical_column'] = sorted(categorical_indices)
757

wxchan's avatar
wxchan committed
758
        params_str = param_dict_to_str(params)
759
        # process for reference dataset
wxchan's avatar
wxchan committed
760
        ref_dataset = None
wxchan's avatar
wxchan committed
761
        if isinstance(reference, Dataset):
762
            ref_dataset = reference.construct().handle
wxchan's avatar
wxchan committed
763
764
        elif reference is not None:
            raise TypeError('Reference dataset should be None or dataset instance')
765
        # start construct data
wxchan's avatar
wxchan committed
766
        if isinstance(data, string_type):
767
            # check data has header or not
Guolin Ke's avatar
Guolin Ke committed
768
            if str(params.get("has_header", "")).lower() == "true" \
wxchan's avatar
wxchan committed
769
                    or str(params.get("header", "")).lower() == "true":
770
                self.data_has_header = True
wxchan's avatar
wxchan committed
771
772
773
774
775
776
777
778
            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
779
780
        elif isinstance(data, scipy.sparse.csc_matrix):
            self.__init_from_csc(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
781
782
        elif isinstance(data, np.ndarray):
            self.__init_from_np2d(data, params_str, ref_dataset)
783
784
        elif isinstance(data, list) and len(data) > 0 and all(isinstance(x, np.ndarray) for x in data):
            self.__init_from_list_np2d(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
785
786
787
788
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                self.__init_from_csr(csr, params_str, ref_dataset)
789
            except BaseException:
wxchan's avatar
wxchan committed
790
                raise TypeError('Cannot initialize Dataset from {}'.format(type(data).__name__))
wxchan's avatar
wxchan committed
791
792
793
        if label is not None:
            self.set_label(label)
        if self.get_label() is None:
794
            raise ValueError("Label should not be None")
wxchan's avatar
wxchan committed
795
796
797
798
799
        if weight is not None:
            self.set_weight(weight)
        if group is not None:
            self.set_group(group)
        # load init score
800
801
802
        if init_score is not None:
            self.set_init_score(init_score)
            if self.predictor is not None:
803
                warnings.warn("The prediction of init_model will be overridden by init_score.")
804
        elif isinstance(self.predictor, _InnerPredictor):
wxchan's avatar
wxchan committed
805
806
807
808
809
810
            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
811
                new_init_score = np.zeros(init_score.size, dtype=np.float32)
wxchan's avatar
wxchan committed
812
                num_data = self.num_data()
wxchan's avatar
wxchan committed
813
814
                for i in range_(num_data):
                    for j in range_(self.predictor.num_class):
wxchan's avatar
wxchan committed
815
816
817
                        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
818
819
        elif self.predictor is not None:
            raise TypeError('wrong predictor type {}'.format(type(self.predictor).__name__))
Guolin Ke's avatar
Guolin Ke committed
820
        # set feature names
Nikita Titov's avatar
Nikita Titov committed
821
        return self.set_feature_name(feature_name)
wxchan's avatar
wxchan committed
822
823

    def __init_from_np2d(self, mat, params_str, ref_dataset):
824
        """Initialize data from a 2-D numpy matrix."""
wxchan's avatar
wxchan committed
825
826
827
828
829
830
831
        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:
832
            # change non-float data to float data, need to copy
wxchan's avatar
wxchan committed
833
834
            data = np.array(mat.reshape(mat.size), dtype=np.float32)

835
        ptr_data, type_ptr_data, _ = c_float_array(data)
wxchan's avatar
wxchan committed
836
837
        _safe_call(_LIB.LGBM_DatasetCreateFromMat(
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
838
839
840
841
            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
842
843
844
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
845
        return self
wxchan's avatar
wxchan committed
846

847
    def __init_from_list_np2d(self, mats, params_str, ref_dataset):
848
        """Initialize data from a list of 2-D numpy matrices."""
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
        ncol = mats[0].shape[1]
        nrow = np.zeros((len(mats),), np.int32)
        if mats[0].dtype == np.float64:
            ptr_data = (ctypes.POINTER(ctypes.c_double) * len(mats))()
        else:
            ptr_data = (ctypes.POINTER(ctypes.c_float) * len(mats))()

        holders = []
        type_ptr_data = None

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

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

            nrow[i] = mat.shape[0]

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

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

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

wxchan's avatar
wxchan committed
894
    def __init_from_csr(self, csr, params_str, ref_dataset):
895
        """Initialize data from a CSR matrix."""
wxchan's avatar
wxchan committed
896
        if len(csr.indices) != len(csr.data):
897
            raise ValueError('Length mismatch: {} vs {}'.format(len(csr.indices), len(csr.data)))
wxchan's avatar
wxchan committed
898
899
        self.handle = ctypes.c_void_p()

900
901
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csr.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csr.data)
wxchan's avatar
wxchan committed
902

903
904
905
        assert csr.shape[1] <= MAX_INT32
        csr.indices = csr.indices.astype(np.int32, copy=False)

wxchan's avatar
wxchan committed
906
907
        _safe_call(_LIB.LGBM_DatasetCreateFromCSR(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
908
            ctypes.c_int(type_ptr_indptr),
wxchan's avatar
wxchan committed
909
910
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
911
912
913
914
            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
915
916
917
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
918
        return self
wxchan's avatar
wxchan committed
919

Guolin Ke's avatar
Guolin Ke committed
920
    def __init_from_csc(self, csc, params_str, ref_dataset):
921
        """Initialize data from a CSC matrix."""
Guolin Ke's avatar
Guolin Ke committed
922
923
924
925
        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()

926
927
        ptr_indptr, type_ptr_indptr, __ = c_int_array(csc.indptr)
        ptr_data, type_ptr_data, _ = c_float_array(csc.data)
Guolin Ke's avatar
Guolin Ke committed
928

929
930
931
        assert csc.shape[0] <= MAX_INT32
        csc.indices = csc.indices.astype(np.int32, copy=False)

Guolin Ke's avatar
Guolin Ke committed
932
933
        _safe_call(_LIB.LGBM_DatasetCreateFromCSC(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
934
            ctypes.c_int(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
935
936
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
937
938
939
940
            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
941
942
943
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))
Nikita Titov's avatar
Nikita Titov committed
944
        return self
Guolin Ke's avatar
Guolin Ke committed
945

wxchan's avatar
wxchan committed
946
    def construct(self):
947
948
949
950
951
        """Lazy init.

        Returns
        -------
        self : Dataset
Nikita Titov's avatar
Nikita Titov committed
952
            Constructed Dataset object.
953
        """
954
        if self.handle is None:
wxchan's avatar
wxchan committed
955
956
            if self.reference is not None:
                if self.used_indices is None:
957
                    # create valid
958
                    self._lazy_init(self.data, label=self.label, reference=self.reference,
959
960
                                    weight=self.weight, group=self.group,
                                    init_score=self.init_score, predictor=self._predictor,
961
                                    silent=self.silent, feature_name=self.feature_name, params=self.params)
wxchan's avatar
wxchan committed
962
                else:
963
                    # construct subset
wxchan's avatar
wxchan committed
964
                    used_indices = list_to_1d_numpy(self.used_indices, np.int32, name='used_indices')
965
                    assert used_indices.flags.c_contiguous
Guolin Ke's avatar
Guolin Ke committed
966
967
                    if self.reference.group is not None:
                        group_info = np.array(self.reference.group).astype(int)
968
969
                        _, self.group = np.unique(np.repeat(range_(len(group_info)), repeats=group_info)[self.used_indices],
                                                  return_counts=True)
970
                    self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
971
972
                    params_str = param_dict_to_str(self.params)
                    _safe_call(_LIB.LGBM_DatasetGetSubset(
973
                        self.reference.construct().handle,
wxchan's avatar
wxchan committed
974
                        used_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
975
                        ctypes.c_int(used_indices.shape[0]),
wxchan's avatar
wxchan committed
976
977
                        c_str(params_str),
                        ctypes.byref(self.handle)))
978
979
                    self.data = self.reference.data
                    self.get_data()
Guolin Ke's avatar
Guolin Ke committed
980
981
                    if self.group is not None:
                        self.set_group(self.group)
wxchan's avatar
wxchan committed
982
983
984
                    if self.get_label() is None:
                        raise ValueError("Label should not be None.")
            else:
985
                # create train
986
                self._lazy_init(self.data, label=self.label,
987
988
989
                                weight=self.weight, group=self.group,
                                init_score=self.init_score, predictor=self._predictor,
                                silent=self.silent, feature_name=self.feature_name,
990
                                categorical_feature=self.categorical_feature, params=self.params)
wxchan's avatar
wxchan committed
991
992
993
            if self.free_raw_data:
                self.data = None
        return self
wxchan's avatar
wxchan committed
994

wxchan's avatar
wxchan committed
995
    def create_valid(self, data, label=None, weight=None, group=None,
996
                     init_score=None, silent=False, params=None):
997
        """Create validation data align with current Dataset.
wxchan's avatar
wxchan committed
998
999
1000

        Parameters
        ----------
1001
        data : string, numpy array, pandas DataFrame, scipy.sparse or list of numpy arrays
wxchan's avatar
wxchan committed
1002
            Data source of Dataset.
1003
            If string, it represents the path to txt file.
1004
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
1005
1006
            Label of the data.
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
wxchan's avatar
wxchan committed
1007
            Weight for each instance.
1008
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
1009
            Group/query size for Dataset.
1010
        init_score : list, numpy 1-D array, pandas Series or None, optional (default=None)
1011
            Init score for Dataset.
1012
1013
        silent : bool, optional (default=False)
            Whether to print messages during construction.
Nikita Titov's avatar
Nikita Titov committed
1014
        params : dict or None, optional (default=None)
1015
            Other parameters for validation Dataset.
1016
1017
1018

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1019
1020
        valid : Dataset
            Validation Dataset with reference to self.
wxchan's avatar
wxchan committed
1021
        """
1022
        ret = Dataset(data, label=label, reference=self,
1023
1024
                      weight=weight, group=group, init_score=init_score,
                      silent=silent, params=params, free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
1025
        ret._predictor = self._predictor
1026
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1027
        return ret
wxchan's avatar
wxchan committed
1028

wxchan's avatar
wxchan committed
1029
    def subset(self, used_indices, params=None):
1030
        """Get subset of current Dataset.
wxchan's avatar
wxchan committed
1031
1032
1033
1034

        Parameters
        ----------
        used_indices : list of int
1035
            Indices used to create the subset.
Nikita Titov's avatar
Nikita Titov committed
1036
        params : dict or None, optional (default=None)
1037
            These parameters will be passed to Dataset constructor.
1038
1039
1040
1041
1042

        Returns
        -------
        subset : Dataset
            Subset of the current Dataset.
wxchan's avatar
wxchan committed
1043
        """
wxchan's avatar
wxchan committed
1044
1045
        if params is None:
            params = self.params
wxchan's avatar
wxchan committed
1046
        ret = Dataset(None, reference=self, feature_name=self.feature_name,
1047
1048
                      categorical_feature=self.categorical_feature, params=params,
                      free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
1049
        ret._predictor = self._predictor
1050
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1051
1052
1053
1054
        ret.used_indices = used_indices
        return ret

    def save_binary(self, filename):
1055
        """Save Dataset to a binary file.
wxchan's avatar
wxchan committed
1056
1057
1058
1059
1060

        Parameters
        ----------
        filename : string
            Name of the output file.
Nikita Titov's avatar
Nikita Titov committed
1061
1062
1063
1064
1065

        Returns
        -------
        self : Dataset
            Returns self.
wxchan's avatar
wxchan committed
1066
1067
1068
1069
        """
        _safe_call(_LIB.LGBM_DatasetSaveBinary(
            self.construct().handle,
            c_str(filename)))
Nikita Titov's avatar
Nikita Titov committed
1070
        return self
wxchan's avatar
wxchan committed
1071
1072

    def _update_params(self, params):
1073
1074
        if self.handle is not None and params is not None:
            _safe_call(_LIB.LGBM_DatasetUpdateParam(self.handle, c_str(param_dict_to_str(params))))
wxchan's avatar
wxchan committed
1075
1076
        if not self.params:
            self.params = params
wxchan's avatar
wxchan committed
1077
        else:
1078
            self.params_back_up = copy.deepcopy(self.params)
wxchan's avatar
wxchan committed
1079
            self.params.update(params)
Nikita Titov's avatar
Nikita Titov committed
1080
        return self
wxchan's avatar
wxchan committed
1081

1082
1083
1084
    def _reverse_update_params(self):
        self.params = copy.deepcopy(self.params_back_up)
        self.params_back_up = None
1085
1086
        if self.handle is not None and self.params is not None:
            _safe_call(_LIB.LGBM_DatasetUpdateParam(self.handle, c_str(param_dict_to_str(self.params))))
Nikita Titov's avatar
Nikita Titov committed
1087
        return self
1088

wxchan's avatar
wxchan committed
1089
    def set_field(self, field_name, data):
wxchan's avatar
wxchan committed
1090
        """Set property into the Dataset.
wxchan's avatar
wxchan committed
1091
1092
1093

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1094
        field_name : string
1095
            The field name of the information.
1096
        data : list, numpy 1-D array, pandas Series or None
1097
            The array of data to be set.
Nikita Titov's avatar
Nikita Titov committed
1098
1099
1100
1101
1102

        Returns
        -------
        self : Dataset
            Dataset with set property.
wxchan's avatar
wxchan committed
1103
        """
1104
1105
        if self.handle is None:
            raise Exception("Cannot set %s before construct dataset" % field_name)
wxchan's avatar
wxchan committed
1106
        if data is None:
1107
            # set to None
wxchan's avatar
wxchan committed
1108
1109
1110
1111
            _safe_call(_LIB.LGBM_DatasetSetField(
                self.handle,
                c_str(field_name),
                None,
Guolin Ke's avatar
Guolin Ke committed
1112
1113
                ctypes.c_int(0),
                ctypes.c_int(FIELD_TYPE_MAPPER[field_name])))
Nikita Titov's avatar
Nikita Titov committed
1114
            return self
Guolin Ke's avatar
Guolin Ke committed
1115
1116
1117
1118
1119
        dtype = np.float32
        if field_name == 'group':
            dtype = np.int32
        elif field_name == 'init_score':
            dtype = np.float64
1120
        data = list_to_1d_numpy(data, dtype, name=field_name)
1121
1122
        if data.dtype == np.float32 or data.dtype == np.float64:
            ptr_data, type_data, _ = c_float_array(data)
wxchan's avatar
wxchan committed
1123
        elif data.dtype == np.int32:
1124
            ptr_data, type_data, _ = c_int_array(data)
wxchan's avatar
wxchan committed
1125
        else:
Guolin Ke's avatar
Guolin Ke committed
1126
            raise TypeError("Excepted np.float32/64 or np.int32, meet type({})".format(data.dtype))
wxchan's avatar
wxchan committed
1127
        if type_data != FIELD_TYPE_MAPPER[field_name]:
1128
            raise TypeError("Input type error for set_field")
wxchan's avatar
wxchan committed
1129
1130
1131
1132
        _safe_call(_LIB.LGBM_DatasetSetField(
            self.handle,
            c_str(field_name),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
1133
1134
            ctypes.c_int(len(data)),
            ctypes.c_int(type_data)))
Nikita Titov's avatar
Nikita Titov committed
1135
        return self
wxchan's avatar
wxchan committed
1136

wxchan's avatar
wxchan committed
1137
1138
    def get_field(self, field_name):
        """Get property from the Dataset.
wxchan's avatar
wxchan committed
1139
1140
1141

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1142
        field_name : string
1143
            The field name of the information.
wxchan's avatar
wxchan committed
1144
1145
1146

        Returns
        -------
1147
1148
        info : numpy array
            A numpy array with information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1149
        """
1150
        if self.handle is None:
1151
            raise Exception("Cannot get %s before construct Dataset" % field_name)
Guolin Ke's avatar
Guolin Ke committed
1152
1153
        tmp_out_len = ctypes.c_int()
        out_type = ctypes.c_int()
wxchan's avatar
wxchan committed
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
        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
1169
1170
        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)
1171
        else:
wxchan's avatar
wxchan committed
1172
            raise TypeError("Unknown type")
Guolin Ke's avatar
Guolin Ke committed
1173

1174
    def set_categorical_feature(self, categorical_feature):
1175
        """Set categorical features.
1176
1177
1178

        Parameters
        ----------
1179
1180
        categorical_feature : list of int or strings
            Names or indices of categorical features.
Nikita Titov's avatar
Nikita Titov committed
1181
1182
1183
1184
1185

        Returns
        -------
        self : Dataset
            Dataset with set categorical features.
1186
1187
        """
        if self.categorical_feature == categorical_feature:
Nikita Titov's avatar
Nikita Titov committed
1188
            return self
1189
        if self.data is not None:
1190
1191
            if self.categorical_feature is None:
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
1192
                return self._free_handle()
1193
1194
            elif categorical_feature == 'auto':
                warnings.warn('Using categorical_feature in Dataset.')
Nikita Titov's avatar
Nikita Titov committed
1195
                return self
1196
            else:
1197
1198
                warnings.warn('categorical_feature in Dataset is overridden.\n'
                              'New categorical_feature is {}'.format(sorted(list(categorical_feature))))
1199
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
1200
                return self._free_handle()
1201
        else:
1202
1203
            raise LightGBMError("Cannot set categorical feature after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
1204

Guolin Ke's avatar
Guolin Ke committed
1205
    def _set_predictor(self, predictor):
1206
1207
1208
1209
        """Set predictor for continued training.

        It is not recommended for user to call this function.
        Please use init_model argument in engine.train() or engine.cv() instead.
Guolin Ke's avatar
Guolin Ke committed
1210
1211
        """
        if predictor is self._predictor:
Nikita Titov's avatar
Nikita Titov committed
1212
            return self
Guolin Ke's avatar
Guolin Ke committed
1213
1214
        if self.data is not None:
            self._predictor = predictor
Nikita Titov's avatar
Nikita Titov committed
1215
            return self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1216
        else:
1217
1218
            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
1219
1220

    def set_reference(self, reference):
1221
        """Set reference Dataset.
Guolin Ke's avatar
Guolin Ke committed
1222
1223
1224
1225

        Parameters
        ----------
        reference : Dataset
1226
            Reference that is used as a template to construct the current Dataset.
Nikita Titov's avatar
Nikita Titov committed
1227
1228
1229
1230
1231

        Returns
        -------
        self : Dataset
            Dataset with set reference.
Guolin Ke's avatar
Guolin Ke committed
1232
        """
1233
1234
1235
        self.set_categorical_feature(reference.categorical_feature) \
            .set_feature_name(reference.feature_name) \
            ._set_predictor(reference._predictor)
1236
1237
        # we're done if self and reference share a common upstrem reference
        if self.get_ref_chain().intersection(reference.get_ref_chain()):
Nikita Titov's avatar
Nikita Titov committed
1238
            return self
Guolin Ke's avatar
Guolin Ke committed
1239
1240
        if self.data is not None:
            self.reference = reference
Nikita Titov's avatar
Nikita Titov committed
1241
            return self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
1242
        else:
1243
1244
            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
1245
1246

    def set_feature_name(self, feature_name):
1247
        """Set feature name.
Guolin Ke's avatar
Guolin Ke committed
1248
1249
1250

        Parameters
        ----------
1251
1252
        feature_name : list of strings
            Feature names.
Nikita Titov's avatar
Nikita Titov committed
1253
1254
1255
1256
1257

        Returns
        -------
        self : Dataset
            Dataset with set feature name.
Guolin Ke's avatar
Guolin Ke committed
1258
        """
1259
1260
        if feature_name != 'auto':
            self.feature_name = feature_name
1261
        if self.handle is not None and feature_name is not None and feature_name != 'auto':
wxchan's avatar
wxchan committed
1262
            if len(feature_name) != self.num_feature():
1263
1264
                raise ValueError("Length of feature_name({}) and num_feature({}) don't match"
                                 .format(len(feature_name), self.num_feature()))
1265
            c_feature_name = [c_str(name) for name in feature_name]
wxchan's avatar
wxchan committed
1266
1267
1268
            _safe_call(_LIB.LGBM_DatasetSetFeatureNames(
                self.handle,
                c_array(ctypes.c_char_p, c_feature_name),
Guolin Ke's avatar
Guolin Ke committed
1269
                ctypes.c_int(len(feature_name))))
Nikita Titov's avatar
Nikita Titov committed
1270
        return self
Guolin Ke's avatar
Guolin Ke committed
1271
1272

    def set_label(self, label):
1273
        """Set label of Dataset.
Guolin Ke's avatar
Guolin Ke committed
1274
1275
1276

        Parameters
        ----------
1277
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None
1278
            The label information to be set into Dataset.
Nikita Titov's avatar
Nikita Titov committed
1279
1280
1281
1282
1283

        Returns
        -------
        self : Dataset
            Dataset with set label.
Guolin Ke's avatar
Guolin Ke committed
1284
1285
        """
        self.label = label
1286
        if self.handle is not None:
1287
            label = list_to_1d_numpy(_label_from_pandas(label), name='label')
wxchan's avatar
wxchan committed
1288
            self.set_field('label', label)
Nikita Titov's avatar
Nikita Titov committed
1289
        return self
Guolin Ke's avatar
Guolin Ke committed
1290
1291

    def set_weight(self, weight):
1292
        """Set weight of each instance.
Guolin Ke's avatar
Guolin Ke committed
1293
1294
1295

        Parameters
        ----------
1296
        weight : list, numpy 1-D array, pandas Series or None
1297
            Weight to be set for each data point.
Nikita Titov's avatar
Nikita Titov committed
1298
1299
1300
1301
1302

        Returns
        -------
        self : Dataset
            Dataset with set weight.
Guolin Ke's avatar
Guolin Ke committed
1303
        """
1304
1305
        if weight is not None and np.all(weight == 1):
            weight = None
Guolin Ke's avatar
Guolin Ke committed
1306
        self.weight = weight
1307
        if self.handle is not None and weight is not None:
wxchan's avatar
wxchan committed
1308
1309
            weight = list_to_1d_numpy(weight, name='weight')
            self.set_field('weight', weight)
Nikita Titov's avatar
Nikita Titov committed
1310
        return self
Guolin Ke's avatar
Guolin Ke committed
1311
1312

    def set_init_score(self, init_score):
1313
        """Set init score of Booster to start from.
Guolin Ke's avatar
Guolin Ke committed
1314
1315
1316

        Parameters
        ----------
1317
        init_score : list, numpy 1-D array, pandas Series or None
1318
            Init score for Booster.
Nikita Titov's avatar
Nikita Titov committed
1319
1320
1321
1322
1323

        Returns
        -------
        self : Dataset
            Dataset with set init score.
Guolin Ke's avatar
Guolin Ke committed
1324
1325
        """
        self.init_score = init_score
1326
        if self.handle is not None and init_score is not None:
Guolin Ke's avatar
Guolin Ke committed
1327
            init_score = list_to_1d_numpy(init_score, np.float64, name='init_score')
wxchan's avatar
wxchan committed
1328
            self.set_field('init_score', init_score)
Nikita Titov's avatar
Nikita Titov committed
1329
        return self
Guolin Ke's avatar
Guolin Ke committed
1330
1331

    def set_group(self, group):
1332
        """Set group size of Dataset (used for ranking).
Guolin Ke's avatar
Guolin Ke committed
1333
1334
1335

        Parameters
        ----------
1336
        group : list, numpy 1-D array, pandas Series or None
1337
            Group size of each group.
Nikita Titov's avatar
Nikita Titov committed
1338
1339
1340
1341
1342

        Returns
        -------
        self : Dataset
            Dataset with set group.
Guolin Ke's avatar
Guolin Ke committed
1343
1344
        """
        self.group = group
1345
        if self.handle is not None and group is not None:
wxchan's avatar
wxchan committed
1346
1347
            group = list_to_1d_numpy(group, np.int32, name='group')
            self.set_field('group', group)
Nikita Titov's avatar
Nikita Titov committed
1348
        return self
Guolin Ke's avatar
Guolin Ke committed
1349
1350

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

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1355
        label : numpy array or None
1356
            The label information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1357
        """
1358
        if self.label is None:
wxchan's avatar
wxchan committed
1359
            self.label = self.get_field('label')
Guolin Ke's avatar
Guolin Ke committed
1360
1361
1362
        return self.label

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

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1367
        weight : numpy array or None
1368
            Weight for each data point from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1369
        """
1370
        if self.weight is None:
wxchan's avatar
wxchan committed
1371
            self.weight = self.get_field('weight')
Guolin Ke's avatar
Guolin Ke committed
1372
1373
1374
        return self.weight

    def get_init_score(self):
1375
        """Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1376
1377
1378

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1379
        init_score : numpy array or None
1380
            Init score of Booster.
Guolin Ke's avatar
Guolin Ke committed
1381
        """
1382
        if self.init_score is None:
wxchan's avatar
wxchan committed
1383
            self.init_score = self.get_field('init_score')
Guolin Ke's avatar
Guolin Ke committed
1384
1385
        return self.init_score

1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
    def get_data(self):
        """Get the raw data of the Dataset.

        Returns
        -------
        data : string, numpy array, pandas DataFrame, scipy.sparse, list of numpy arrays or None
            Raw data used in the Dataset construction.
        """
        if self.handle is None:
            raise Exception("Cannot get data before construct Dataset")
        if self.data is not None and self.used_indices is not None and self.need_slice:
            if isinstance(self.data, np.ndarray) or scipy.sparse.issparse(self.data):
                self.data = self.data[self.used_indices, :]
            elif isinstance(self.data, DataFrame):
                self.data = self.data.iloc[self.used_indices].copy()
            else:
                warnings.warn("Cannot subset {} type of raw data.\n"
                              "Returning original raw data".format(type(self.data).__name__))
            self.need_slice = False
        return self.data

Guolin Ke's avatar
Guolin Ke committed
1407
    def get_group(self):
1408
        """Get the group of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1409
1410
1411

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1412
        group : numpy array or None
1413
            Group size of each group.
Guolin Ke's avatar
Guolin Ke committed
1414
        """
1415
        if self.group is None:
wxchan's avatar
wxchan committed
1416
            self.group = self.get_field('group')
Guolin Ke's avatar
Guolin Ke committed
1417
1418
            if self.group is not None:
                # group data from LightGBM is boundaries data, need to convert to group size
Nikita Titov's avatar
Nikita Titov committed
1419
                self.group = np.diff(self.group)
Guolin Ke's avatar
Guolin Ke committed
1420
1421
1422
        return self.group

    def num_data(self):
1423
        """Get the number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1424
1425
1426

        Returns
        -------
1427
1428
        number_of_rows : int
            The number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1429
        """
1430
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1431
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1432
1433
1434
            _safe_call(_LIB.LGBM_DatasetGetNumData(self.handle,
                                                   ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1435
        else:
1436
            raise LightGBMError("Cannot get num_data before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1437
1438

    def num_feature(self):
1439
        """Get the number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1440
1441
1442

        Returns
        -------
1443
1444
        number_of_columns : int
            The number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1445
        """
1446
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1447
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1448
1449
1450
            _safe_call(_LIB.LGBM_DatasetGetNumFeature(self.handle,
                                                      ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1451
        else:
1452
            raise LightGBMError("Cannot get num_feature before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1453

1454
    def get_ref_chain(self, ref_limit=100):
1455
1456
1457
1458
1459
        """Get a chain of Dataset objects.

        Starts with r, then goes to r.reference (if exists),
        then to r.reference.reference, etc.
        until we hit ``ref_limit`` or a reference loop.
1460
1461
1462
1463
1464

        Parameters
        ----------
        ref_limit : int, optional (default=100)
            The limit number of references.
1465
1466
1467

        Returns
        -------
1468
1469
1470
        ref_chain : set of Dataset
            Chain of references of the Datasets.
        """
1471
        head = self
1472
        ref_chain = set()
1473
1474
        while len(ref_chain) < ref_limit:
            if isinstance(head, Dataset):
1475
                ref_chain.add(head)
1476
1477
1478
1479
1480
1481
                if (head.reference is not None) and (head.reference not in ref_chain):
                    head = head.reference
                else:
                    break
            else:
                break
Nikita Titov's avatar
Nikita Titov committed
1482
        return ref_chain
1483

wxchan's avatar
wxchan committed
1484

wxchan's avatar
wxchan committed
1485
class Booster(object):
1486
    """Booster in LightGBM."""
1487

wxchan's avatar
wxchan committed
1488
    def __init__(self, params=None, train_set=None, model_file=None, silent=False):
1489
        """Initialize the Booster.
wxchan's avatar
wxchan committed
1490
1491
1492

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1493
        params : dict or None, optional (default=None)
1494
1495
1496
1497
            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
1498
            Path to the model file.
1499
1500
        silent : bool, optional (default=False)
            Whether to print messages during construction.
wxchan's avatar
wxchan committed
1501
        """
1502
        self.handle = None
1503
        self.network = False
wxchan's avatar
wxchan committed
1504
1505
1506
        self.__need_reload_eval_info = True
        self.__train_data_name = "training"
        self.__attr = {}
1507
        self.__set_objective_to_none = False
wxchan's avatar
wxchan committed
1508
        self.best_iteration = -1
wxchan's avatar
wxchan committed
1509
        self.best_score = {}
1510
        params = {} if params is None else copy.deepcopy(params)
1511
1512
1513
        # user can set verbose with params, it has higher priority
        if not any(verbose_alias in params for verbose_alias in ('verbose', 'verbosity')) and silent:
            params["verbose"] = -1
wxchan's avatar
wxchan committed
1514
        if train_set is not None:
1515
            # Training task
wxchan's avatar
wxchan committed
1516
            if not isinstance(train_set, Dataset):
1517
1518
                raise TypeError('Training data should be Dataset instance, met {}'
                                .format(type(train_set).__name__))
wxchan's avatar
wxchan committed
1519
            params_str = param_dict_to_str(params)
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
            # set network if necessary
            for alias in ["machines", "workers", "nodes"]:
                if alias in params:
                    machines = params[alias]
                    if isinstance(machines, string_type):
                        num_machines = len(machines.split(','))
                    elif isinstance(machines, (list, set)):
                        num_machines = len(machines)
                        machines = ','.join(machines)
                    else:
                        raise ValueError("Invalid machines in params.")
                    self.set_network(machines,
                                     local_listen_port=params.get("local_listen_port", 12400),
                                     listen_time_out=params.get("listen_time_out", 120),
                                     num_machines=params.get("num_machines", num_machines))
                    break
1536
            # construct booster object
1537
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1538
            _safe_call(_LIB.LGBM_BoosterCreate(
wxchan's avatar
wxchan committed
1539
                train_set.construct().handle,
wxchan's avatar
wxchan committed
1540
1541
                c_str(params_str),
                ctypes.byref(self.handle)))
1542
            # save reference to data
wxchan's avatar
wxchan committed
1543
1544
1545
1546
            self.train_set = train_set
            self.valid_sets = []
            self.name_valid_sets = []
            self.__num_dataset = 1
Guolin Ke's avatar
Guolin Ke committed
1547
1548
            self.__init_predictor = train_set._predictor
            if self.__init_predictor is not None:
wxchan's avatar
wxchan committed
1549
1550
                _safe_call(_LIB.LGBM_BoosterMerge(
                    self.handle,
Guolin Ke's avatar
Guolin Ke committed
1551
                    self.__init_predictor.handle))
Guolin Ke's avatar
Guolin Ke committed
1552
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1553
1554
1555
1556
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1557
            # buffer for inner predict
wxchan's avatar
wxchan committed
1558
1559
1560
            self.__inner_predict_buffer = [None]
            self.__is_predicted_cur_iter = [False]
            self.__get_eval_info()
1561
            self.pandas_categorical = train_set.pandas_categorical
wxchan's avatar
wxchan committed
1562
        elif model_file is not None:
1563
            # Prediction task
Guolin Ke's avatar
Guolin Ke committed
1564
            out_num_iterations = ctypes.c_int(0)
1565
            self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1566
1567
1568
1569
            _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
1570
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1571
1572
1573
1574
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1575
            self.pandas_categorical = _load_pandas_categorical(file_name=model_file)
1576
        elif 'model_str' in params:
1577
            self.model_from_string(params['model_str'], False)
wxchan's avatar
wxchan committed
1578
        else:
1579
            raise TypeError('Need at least one training dataset or model file to create booster instance')
1580
        self.params = params
wxchan's avatar
wxchan committed
1581
1582

    def __del__(self):
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
        try:
            if self.network:
                self.free_network()
        except AttributeError:
            pass
        try:
            if self.handle is not None:
                _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
1593

wxchan's avatar
wxchan committed
1594
1595
1596
1597
    def __copy__(self):
        return self.__deepcopy__(None)

    def __deepcopy__(self, _):
1598
        model_str = self.model_to_string(num_iteration=-1)
1599
1600
        booster = Booster({'model_str': model_str})
        return booster
wxchan's avatar
wxchan committed
1601
1602
1603
1604
1605
1606
1607

    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:
1608
            this["handle"] = self.model_to_string(num_iteration=-1)
wxchan's avatar
wxchan committed
1609
1610
1611
        return this

    def __setstate__(self, state):
1612
1613
        model_str = state.get('handle', None)
        if model_str is not None:
wxchan's avatar
wxchan committed
1614
            handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
1615
            out_num_iterations = ctypes.c_int(0)
1616
1617
1618
1619
            _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
                c_str(model_str),
                ctypes.byref(out_num_iterations),
                ctypes.byref(handle)))
wxchan's avatar
wxchan committed
1620
1621
1622
            state['handle'] = handle
        self.__dict__.update(state)

wxchan's avatar
wxchan committed
1623
    def free_dataset(self):
Nikita Titov's avatar
Nikita Titov committed
1624
1625
1626
1627
1628
1629
1630
        """Free Booster's Datasets.

        Returns
        -------
        self : Booster
            Booster without Datasets.
        """
wxchan's avatar
wxchan committed
1631
1632
        self.__dict__.pop('train_set', None)
        self.__dict__.pop('valid_sets', None)
1633
        self.__num_dataset = 0
Nikita Titov's avatar
Nikita Titov committed
1634
        return self
wxchan's avatar
wxchan committed
1635

1636
1637
1638
    def _free_buffer(self):
        self.__inner_predict_buffer = []
        self.__is_predicted_cur_iter = []
Nikita Titov's avatar
Nikita Titov committed
1639
        return self
1640

1641
1642
1643
1644
1645
1646
    def set_network(self, machines, local_listen_port=12400,
                    listen_time_out=120, num_machines=1):
        """Set the network configuration.

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1647
        machines : list, set or string
1648
            Names of machines.
Nikita Titov's avatar
Nikita Titov committed
1649
        local_listen_port : int, optional (default=12400)
1650
            TCP listen port for local machines.
Nikita Titov's avatar
Nikita Titov committed
1651
        listen_time_out : int, optional (default=120)
1652
            Socket time-out in minutes.
Nikita Titov's avatar
Nikita Titov committed
1653
        num_machines : int, optional (default=1)
1654
            The number of machines for parallel learning application.
Nikita Titov's avatar
Nikita Titov committed
1655
1656
1657
1658
1659

        Returns
        -------
        self : Booster
            Booster with set network.
1660
1661
1662
1663
1664
1665
        """
        _safe_call(_LIB.LGBM_NetworkInit(c_str(machines),
                                         ctypes.c_int(local_listen_port),
                                         ctypes.c_int(listen_time_out),
                                         ctypes.c_int(num_machines)))
        self.network = True
Nikita Titov's avatar
Nikita Titov committed
1666
        return self
1667
1668

    def free_network(self):
Nikita Titov's avatar
Nikita Titov committed
1669
1670
1671
1672
1673
1674
1675
        """Free Booster's network.

        Returns
        -------
        self : Booster
            Booster with freed network.
        """
1676
1677
        _safe_call(_LIB.LGBM_NetworkFree())
        self.network = False
Nikita Titov's avatar
Nikita Titov committed
1678
        return self
1679

wxchan's avatar
wxchan committed
1680
    def set_train_data_name(self, name):
1681
1682
1683
1684
        """Set the name to the training Dataset.

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
1685
1686
1687
1688
1689
1690
1691
        name : string
            Name for the training Dataset.

        Returns
        -------
        self : Booster
            Booster with set training Dataset name.
1692
        """
wxchan's avatar
wxchan committed
1693
        self.__train_data_name = name
Nikita Titov's avatar
Nikita Titov committed
1694
        return self
wxchan's avatar
wxchan committed
1695
1696

    def add_valid(self, data, name):
1697
        """Add validation data.
wxchan's avatar
wxchan committed
1698
1699
1700
1701

        Parameters
        ----------
        data : Dataset
1702
1703
1704
            Validation data.
        name : string
            Name of validation data.
Nikita Titov's avatar
Nikita Titov committed
1705
1706
1707
1708
1709

        Returns
        -------
        self : Booster
            Booster with set validation data.
wxchan's avatar
wxchan committed
1710
        """
Guolin Ke's avatar
Guolin Ke committed
1711
        if not isinstance(data, Dataset):
1712
1713
            raise TypeError('Validation data should be Dataset instance, met {}'
                            .format(type(data).__name__))
Guolin Ke's avatar
Guolin Ke committed
1714
        if data._predictor is not self.__init_predictor:
1715
1716
            raise LightGBMError("Add validation data failed, "
                                "you should use same predictor for these data")
wxchan's avatar
wxchan committed
1717
1718
        _safe_call(_LIB.LGBM_BoosterAddValidData(
            self.handle,
wxchan's avatar
wxchan committed
1719
            data.construct().handle))
wxchan's avatar
wxchan committed
1720
1721
1722
1723
1724
        self.valid_sets.append(data)
        self.name_valid_sets.append(name)
        self.__num_dataset += 1
        self.__inner_predict_buffer.append(None)
        self.__is_predicted_cur_iter.append(False)
Nikita Titov's avatar
Nikita Titov committed
1725
        return self
wxchan's avatar
wxchan committed
1726
1727

    def reset_parameter(self, params):
1728
        """Reset parameters of Booster.
wxchan's avatar
wxchan committed
1729
1730
1731
1732

        Parameters
        ----------
        params : dict
1733
            New parameters for Booster.
Nikita Titov's avatar
Nikita Titov committed
1734
1735
1736
1737
1738

        Returns
        -------
        self : Booster
            Booster with new parameters.
wxchan's avatar
wxchan committed
1739
        """
1740
        if any(metric_alias in params for metric_alias in ('metric', 'metrics', 'metric_types')):
wxchan's avatar
wxchan committed
1741
1742
1743
1744
1745
1746
            self.__need_reload_eval_info = True
        params_str = param_dict_to_str(params)
        if params_str:
            _safe_call(_LIB.LGBM_BoosterResetParameter(
                self.handle,
                c_str(params_str)))
Guolin Ke's avatar
Guolin Ke committed
1747
        self.params.update(params)
Nikita Titov's avatar
Nikita Titov committed
1748
        return self
wxchan's avatar
wxchan committed
1749
1750

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

wxchan's avatar
wxchan committed
1753
1754
        Parameters
        ----------
1755
1756
1757
1758
        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
1759
1760
            Customized objective function.

1761
1762
1763
1764
            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
1765
1766
        Returns
        -------
1767
1768
        is_finished : bool
            Whether the update was successfully finished.
wxchan's avatar
wxchan committed
1769
        """
1770
        # need reset training data
wxchan's avatar
wxchan committed
1771
        if train_set is not None and train_set is not self.train_set:
Guolin Ke's avatar
Guolin Ke committed
1772
            if not isinstance(train_set, Dataset):
1773
1774
                raise TypeError('Training data should be Dataset instance, met {}'
                                .format(type(train_set).__name__))
Guolin Ke's avatar
Guolin Ke committed
1775
            if train_set._predictor is not self.__init_predictor:
1776
1777
                raise LightGBMError("Replace training data failed, "
                                    "you should use same predictor for these data")
wxchan's avatar
wxchan committed
1778
1779
1780
            self.train_set = train_set
            _safe_call(_LIB.LGBM_BoosterResetTrainingData(
                self.handle,
wxchan's avatar
wxchan committed
1781
                self.train_set.construct().handle))
wxchan's avatar
wxchan committed
1782
1783
1784
            self.__inner_predict_buffer[0] = None
        is_finished = ctypes.c_int(0)
        if fobj is None:
1785
            if self.__set_objective_to_none:
1786
                raise LightGBMError('Cannot update due to null objective function.')
wxchan's avatar
wxchan committed
1787
1788
1789
            _safe_call(_LIB.LGBM_BoosterUpdateOneIter(
                self.handle,
                ctypes.byref(is_finished)))
wxchan's avatar
wxchan committed
1790
            self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1791
1792
            return is_finished.value == 1
        else:
1793
            if not self.__set_objective_to_none:
Nikita Titov's avatar
Nikita Titov committed
1794
                self.reset_parameter({"objective": "none"}).__set_objective_to_none = True
wxchan's avatar
wxchan committed
1795
1796
1797
1798
            grad, hess = fobj(self.__inner_predict(0), self.train_set)
            return self.__boost(grad, hess)

    def __boost(self, grad, hess):
1799
        """Boost Booster for one iteration with customized gradient statistics.
Nikita Titov's avatar
Nikita Titov committed
1800

1801
1802
1803
1804
1805
        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.
1806

wxchan's avatar
wxchan committed
1807
1808
        Parameters
        ----------
1809
        grad : 1-D numpy array or 1-D list
Nikita Titov's avatar
Nikita Titov committed
1810
            The first order derivative (gradient).
1811
        hess : 1-D numpy array or 1-D list
Nikita Titov's avatar
Nikita Titov committed
1812
            The second order derivative (Hessian).
wxchan's avatar
wxchan committed
1813
1814
1815

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1816
1817
        is_finished : bool
            Whether the boost was successfully finished.
wxchan's avatar
wxchan committed
1818
        """
1819
1820
        grad = list_to_1d_numpy(grad, name='gradient')
        hess = list_to_1d_numpy(hess, name='hessian')
1821
1822
        assert grad.flags.c_contiguous
        assert hess.flags.c_contiguous
wxchan's avatar
wxchan committed
1823
        if len(grad) != len(hess):
1824
1825
            raise ValueError("Lengths of gradient({}) and hessian({}) don't match"
                             .format(len(grad), len(hess)))
wxchan's avatar
wxchan committed
1826
1827
1828
1829
1830
1831
        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
1832
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1833
1834
1835
        return is_finished.value == 1

    def rollback_one_iter(self):
Nikita Titov's avatar
Nikita Titov committed
1836
1837
1838
1839
1840
1841
1842
        """Rollback one iteration.

        Returns
        -------
        self : Booster
            Booster with rolled back one iteration.
        """
wxchan's avatar
wxchan committed
1843
1844
        _safe_call(_LIB.LGBM_BoosterRollbackOneIter(
            self.handle))
wxchan's avatar
wxchan committed
1845
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
Nikita Titov's avatar
Nikita Titov committed
1846
        return self
wxchan's avatar
wxchan committed
1847
1848

    def current_iteration(self):
1849
1850
1851
1852
1853
1854
1855
        """Get the index of the current iteration.

        Returns
        -------
        cur_iter : int
            The index of the current iteration.
        """
Guolin Ke's avatar
Guolin Ke committed
1856
        out_cur_iter = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1857
1858
1859
1860
1861
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
            self.handle,
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
    def num_model_per_iteration(self):
        """Get number of models per iteration.

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

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

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

wxchan's avatar
wxchan committed
1890
    def eval(self, data, name, feval=None):
1891
        """Evaluate for data.
wxchan's avatar
wxchan committed
1892
1893
1894

        Parameters
        ----------
1895
1896
1897
1898
1899
        data : Dataset
            Data for the evaluating.
        name : string
            Name of the data.
        feval : callable or None, optional (default=None)
1900
            Customized evaluation function.
1901
1902
            Should accept two parameters: preds, train_data,
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
1903
1904
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
1905

wxchan's avatar
wxchan committed
1906
1907
        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1908
        result : list
1909
            List with evaluation results.
wxchan's avatar
wxchan committed
1910
        """
Guolin Ke's avatar
Guolin Ke committed
1911
1912
        if not isinstance(data, Dataset):
            raise TypeError("Can only eval for Dataset instance")
wxchan's avatar
wxchan committed
1913
1914
1915
1916
        data_idx = -1
        if data is self.train_set:
            data_idx = 0
        else:
wxchan's avatar
wxchan committed
1917
            for i in range_(len(self.valid_sets)):
wxchan's avatar
wxchan committed
1918
1919
1920
                if data is self.valid_sets[i]:
                    data_idx = i + 1
                    break
1921
        # need to push new valid data
wxchan's avatar
wxchan committed
1922
1923
1924
1925
1926
1927
1928
        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):
1929
        """Evaluate for training data.
wxchan's avatar
wxchan committed
1930
1931
1932

        Parameters
        ----------
1933
        feval : callable or None, optional (default=None)
1934
            Customized evaluation function.
1935
1936
            Should accept two parameters: preds, train_data,
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
1937
1938
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
wxchan's avatar
wxchan committed
1939
1940
1941

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1942
        result : list
1943
            List with evaluation results.
wxchan's avatar
wxchan committed
1944
1945
1946
1947
        """
        return self.__inner_eval(self.__train_data_name, 0, feval)

    def eval_valid(self, feval=None):
1948
        """Evaluate for validation data.
wxchan's avatar
wxchan committed
1949
1950
1951

        Parameters
        ----------
1952
        feval : callable or None, optional (default=None)
1953
            Customized evaluation function.
1954
1955
            Should accept two parameters: preds, train_data,
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
1956
1957
            For multi-class task, the preds is group by class_id first, then group by row_id.
            If you want to get i-th row preds in j-th class, the access way is preds[j * num_data + i].
wxchan's avatar
wxchan committed
1958
1959
1960

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
1961
        result : list
1962
            List with evaluation results.
wxchan's avatar
wxchan committed
1963
        """
wxchan's avatar
wxchan committed
1964
        return [item for i in range_(1, self.__num_dataset)
wxchan's avatar
wxchan committed
1965
                for item in self.__inner_eval(self.name_valid_sets[i - 1], i, feval)]
wxchan's avatar
wxchan committed
1966

1967
    def save_model(self, filename, num_iteration=None, start_iteration=0):
1968
        """Save Booster to file.
wxchan's avatar
wxchan committed
1969
1970
1971

        Parameters
        ----------
1972
1973
        filename : string
            Filename to save Booster.
1974
1975
1976
1977
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
Nikita Titov's avatar
Nikita Titov committed
1978
        start_iteration : int, optional (default=0)
1979
            Start index of the iteration that should be saved.
Nikita Titov's avatar
Nikita Titov committed
1980
1981
1982
1983
1984

        Returns
        -------
        self : Booster
            Returns self.
wxchan's avatar
wxchan committed
1985
        """
1986
        if num_iteration is None:
1987
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1988
1989
        _safe_call(_LIB.LGBM_BoosterSaveModel(
            self.handle,
1990
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
1991
            ctypes.c_int(num_iteration),
wxchan's avatar
wxchan committed
1992
            c_str(filename)))
1993
        _dump_pandas_categorical(self.pandas_categorical, filename)
Nikita Titov's avatar
Nikita Titov committed
1994
        return self
wxchan's avatar
wxchan committed
1995

1996
    def shuffle_models(self, start_iteration=0, end_iteration=-1):
1997
        """Shuffle models.
Nikita Titov's avatar
Nikita Titov committed
1998

1999
2000
2001
        Parameters
        ----------
        start_iteration : int, optional (default=0)
2002
            The first iteration that will be shuffled.
2003
2004
        end_iteration : int, optional (default=-1)
            The last iteration that will be shuffled.
2005
            If <= 0, means the last available iteration.
2006

Nikita Titov's avatar
Nikita Titov committed
2007
2008
2009
2010
        Returns
        -------
        self : Booster
            Booster with shuffled models.
2011
        """
2012
2013
        _safe_call(_LIB.LGBM_BoosterShuffleModels(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
2014
2015
            ctypes.c_int(start_iteration),
            ctypes.c_int(end_iteration)))
Nikita Titov's avatar
Nikita Titov committed
2016
        return self
2017
2018
2019
2020
2021
2022

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

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
2023
        model_str : string
2024
            Model will be loaded from this string.
Nikita Titov's avatar
Nikita Titov committed
2025
2026
        verbose : bool, optional (default=True)
            Whether to print messages while loading model.
2027
2028
2029

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2030
        self : Booster
2031
2032
            Loaded Booster object.
        """
2033
2034
2035
2036
        if self.handle is not None:
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))
        self._free_buffer()
        self.handle = ctypes.c_void_p()
2037
2038
2039
2040
2041
2042
2043
2044
2045
        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)))
2046
        if verbose:
Nikita Titov's avatar
Nikita Titov committed
2047
            print('Finished loading model, total used %d iterations' % int(out_num_iterations.value))
2048
        self.__num_class = out_num_class.value
2049
        self.pandas_categorical = _load_pandas_categorical(model_str=model_str)
2050
2051
2052
2053
        return self

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

2055
2056
2057
2058
2059
2060
        Parameters
        ----------
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be saved.
            If None, if the best iteration exists, it is saved; otherwise, all iterations are saved.
            If <= 0, all iterations are saved.
Nikita Titov's avatar
Nikita Titov committed
2061
        start_iteration : int, optional (default=0)
2062
2063
2064
2065
            Start index of the iteration that should be saved.

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2066
        str_repr : string
2067
2068
            String representation of Booster.
        """
2069
        if num_iteration is None:
2070
2071
            num_iteration = self.best_iteration
        buffer_len = 1 << 20
2072
        tmp_out_len = ctypes.c_int64(0)
2073
2074
2075
2076
        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,
2077
            ctypes.c_int(start_iteration),
2078
            ctypes.c_int(num_iteration),
2079
            ctypes.c_int64(buffer_len),
2080
2081
2082
            ctypes.byref(tmp_out_len),
            ptr_string_buffer))
        actual_len = tmp_out_len.value
2083
        # if buffer length is not long enough, re-allocate a buffer
2084
2085
2086
2087
2088
        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,
2089
                ctypes.c_int(start_iteration),
2090
                ctypes.c_int(num_iteration),
2091
                ctypes.c_int64(actual_len),
2092
2093
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
2094
2095
2096
        ret = string_buffer.value.decode()
        ret += _dump_pandas_categorical(self.pandas_categorical)
        return ret
2097

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

2101
2102
        Parameters
        ----------
2103
2104
2105
2106
        num_iteration : int or None, optional (default=None)
            Index of the iteration that should be dumped.
            If None, if the best iteration exists, it is dumped; otherwise, all iterations are dumped.
            If <= 0, all iterations are dumped.
Nikita Titov's avatar
Nikita Titov committed
2107
        start_iteration : int, optional (default=0)
2108
            Start index of the iteration that should be dumped.
2109

wxchan's avatar
wxchan committed
2110
2111
        Returns
        -------
2112
        json_repr : dict
Nikita Titov's avatar
Nikita Titov committed
2113
            JSON format of Booster.
wxchan's avatar
wxchan committed
2114
        """
2115
        if num_iteration is None:
2116
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
2117
        buffer_len = 1 << 20
2118
        tmp_out_len = ctypes.c_int64(0)
wxchan's avatar
wxchan committed
2119
2120
2121
2122
        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,
2123
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
2124
            ctypes.c_int(num_iteration),
2125
            ctypes.c_int64(buffer_len),
wxchan's avatar
wxchan committed
2126
            ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2127
            ptr_string_buffer))
wxchan's avatar
wxchan committed
2128
        actual_len = tmp_out_len.value
2129
        # if buffer length is not long enough, reallocate a buffer
wxchan's avatar
wxchan committed
2130
2131
2132
2133
2134
        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,
2135
                ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
2136
                ctypes.c_int(num_iteration),
2137
                ctypes.c_int64(actual_len),
wxchan's avatar
wxchan committed
2138
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2139
                ptr_string_buffer))
2140
2141
2142
2143
        ret = json.loads(string_buffer.value.decode())
        ret['pandas_categorical'] = json.loads(json.dumps(self.pandas_categorical,
                                                          default=json_default_with_numpy))
        return ret
wxchan's avatar
wxchan committed
2144

2145
2146
    def predict(self, data, num_iteration=None,
                raw_score=False, pred_leaf=False, pred_contrib=False,
2147
                data_has_header=False, is_reshape=True, **kwargs):
2148
        """Make a prediction.
wxchan's avatar
wxchan committed
2149
2150
2151

        Parameters
        ----------
2152
        data : string, numpy array, pandas DataFrame or scipy.sparse
2153
2154
            Data source for prediction.
            If string, it represents the path to txt file.
2155
2156
2157
2158
        num_iteration : int or None, optional (default=None)
            Limit number of iterations in the prediction.
            If None, if the best iteration exists, it is used; otherwise, all iterations are used.
            If <= 0, all iterations are used (no limits).
2159
2160
2161
2162
        raw_score : bool, optional (default=False)
            Whether to predict raw scores.
        pred_leaf : bool, optional (default=False)
            Whether to predict leaf index.
2163
2164
        pred_contrib : bool, optional (default=False)
            Whether to predict feature contributions.
2165
2166
2167
2168
2169
2170
2171

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

2172
2173
2174
2175
2176
        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].
2177
2178
        **kwargs
            Other parameters for the prediction.
wxchan's avatar
wxchan committed
2179
2180
2181

        Returns
        -------
2182
2183
        result : numpy array
            Prediction result.
wxchan's avatar
wxchan committed
2184
        """
2185
        predictor = self._to_predictor(copy.deepcopy(kwargs))
2186
        if num_iteration is None:
2187
            num_iteration = self.best_iteration
2188
2189
2190
        return predictor.predict(data, num_iteration,
                                 raw_score, pred_leaf, pred_contrib,
                                 data_has_header, is_reshape)
wxchan's avatar
wxchan committed
2191

2192
    def refit(self, data, label, decay_rate=0.9, **kwargs):
Guolin Ke's avatar
Guolin Ke committed
2193
2194
2195
2196
        """Refit the existing Booster by new data.

        Parameters
        ----------
2197
        data : string, numpy array, pandas DataFrame or scipy.sparse
Guolin Ke's avatar
Guolin Ke committed
2198
2199
            Data source for refit.
            If string, it represents the path to txt file.
2200
        label : list, numpy 1-D array or pandas Series / one-column DataFrame
Guolin Ke's avatar
Guolin Ke committed
2201
2202
            Label for refit.
        decay_rate : float, optional (default=0.9)
2203
2204
            Decay rate of refit,
            will use ``leaf_output = decay_rate * old_leaf_output + (1.0 - decay_rate) * new_leaf_output`` to refit trees.
2205
2206
        **kwargs
            Other parameters for refit.
2207
            These parameters will be passed to ``predict`` method.
Guolin Ke's avatar
Guolin Ke committed
2208
2209
2210
2211
2212
2213

        Returns
        -------
        result : Booster
            Refitted Booster.
        """
2214
2215
        if self.__set_objective_to_none:
            raise LightGBMError('Cannot refit due to null objective function.')
2216
        predictor = self._to_predictor(copy.deepcopy(kwargs))
2217
        leaf_preds = predictor.predict(data, -1, pred_leaf=True)
2218
        nrow, ncol = leaf_preds.shape
2219
        train_set = Dataset(data, label, silent=True)
Guolin Ke's avatar
Guolin Ke committed
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
        new_booster = Booster(self.params, train_set, silent=True)
        # Copy models
        _safe_call(_LIB.LGBM_BoosterMerge(
            new_booster.handle,
            predictor.handle))
        leaf_preds = leaf_preds.reshape(-1)
        ptr_data, type_ptr_data, _ = c_int_array(leaf_preds)
        _safe_call(_LIB.LGBM_BoosterRefit(
            new_booster.handle,
            ptr_data,
            ctypes.c_int(nrow),
            ctypes.c_int(ncol)))
2232
2233
        new_booster.network = self.network
        new_booster.__attr = self.__attr.copy()
Guolin Ke's avatar
Guolin Ke committed
2234
2235
        return new_booster

2236
    def get_leaf_output(self, tree_id, leaf_id):
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
        """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.
        """
2251
2252
2253
2254
2255
2256
2257
2258
        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

2259
    def _to_predictor(self, pred_parameter=None):
2260
        """Convert to predictor."""
2261
        predictor = _InnerPredictor(booster_handle=self.handle, pred_parameter=pred_parameter)
2262
        predictor.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
2263
2264
        return predictor

2265
    def num_feature(self):
2266
2267
2268
2269
2270
2271
2272
        """Get number of features.

        Returns
        -------
        num_feature : int
            The number of features.
        """
2273
2274
2275
2276
2277
2278
        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
2279
    def feature_name(self):
2280
        """Get names of features.
wxchan's avatar
wxchan committed
2281
2282
2283

        Returns
        -------
2284
2285
        result : list
            List with names of features.
wxchan's avatar
wxchan committed
2286
        """
2287
        num_feature = self.num_feature()
2288
        # Get name of features
wxchan's avatar
wxchan committed
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
        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)]

2300
    def feature_importance(self, importance_type='split', iteration=None):
2301
        """Get feature importances.
2302

2303
2304
        Parameters
        ----------
2305
2306
2307
2308
        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.
2309
2310
2311
2312
        iteration : int or None, optional (default=None)
            Limit number of iterations in the feature importance calculation.
            If None, if the best iteration exists, it is used; otherwise, all trees are used.
            If <= 0, all trees are used (no limits).
2313

2314
2315
        Returns
        -------
2316
2317
        result : numpy array
            Array with feature importances.
2318
        """
2319
2320
        if iteration is None:
            iteration = self.best_iteration
2321
2322
2323
2324
2325
2326
        if importance_type == "split":
            importance_type_int = 0
        elif importance_type == "gain":
            importance_type_int = 1
        else:
            importance_type_int = -1
Nikita Titov's avatar
Nikita Titov committed
2327
        result = np.zeros(self.num_feature(), dtype=np.float64)
2328
2329
2330
2331
2332
2333
2334
2335
2336
        _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
2337

wxchan's avatar
wxchan committed
2338
    def __inner_eval(self, data_name, data_idx, feval=None):
2339
        """Evaluate training or validation data."""
wxchan's avatar
wxchan committed
2340
        if data_idx >= self.__num_dataset:
2341
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
2342
2343
2344
        self.__get_eval_info()
        ret = []
        if self.__num_inner_eval > 0:
2345
            result = np.zeros(self.__num_inner_eval, dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
2346
            tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
2347
2348
            _safe_call(_LIB.LGBM_BoosterGetEval(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
2349
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
2350
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
2351
                result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
2352
            if tmp_out_len.value != self.__num_inner_eval:
2353
                raise ValueError("Wrong length of eval results")
wxchan's avatar
wxchan committed
2354
            for i in range_(self.__num_inner_eval):
2355
2356
                ret.append((data_name, self.__name_inner_eval[i],
                            result[i], self.__higher_better_inner_eval[i]))
wxchan's avatar
wxchan committed
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
        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):
2372
        """Predict for training and validation dataset."""
wxchan's avatar
wxchan committed
2373
        if data_idx >= self.__num_dataset:
2374
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
2375
2376
2377
2378
2379
        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
2380
            self.__inner_predict_buffer[data_idx] = np.zeros(n_preds, dtype=np.float64)
2381
        # avoid to predict many time in one iteration
wxchan's avatar
wxchan committed
2382
2383
        if not self.__is_predicted_cur_iter[data_idx]:
            tmp_out_len = ctypes.c_int64(0)
Guolin Ke's avatar
Guolin Ke committed
2384
            data_ptr = self.__inner_predict_buffer[data_idx].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
wxchan's avatar
wxchan committed
2385
2386
            _safe_call(_LIB.LGBM_BoosterGetPredict(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
2387
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
2388
2389
2390
                ctypes.byref(tmp_out_len),
                data_ptr))
            if tmp_out_len.value != len(self.__inner_predict_buffer[data_idx]):
2391
                raise ValueError("Wrong length of predict results for data %d" % (data_idx))
wxchan's avatar
wxchan committed
2392
2393
2394
2395
            self.__is_predicted_cur_iter[data_idx] = True
        return self.__inner_predict_buffer[data_idx]

    def __get_eval_info(self):
2396
        """Get inner evaluation count and names."""
wxchan's avatar
wxchan committed
2397
2398
        if self.__need_reload_eval_info:
            self.__need_reload_eval_info = False
Guolin Ke's avatar
Guolin Ke committed
2399
            out_num_eval = ctypes.c_int(0)
2400
            # Get num of inner evals
wxchan's avatar
wxchan committed
2401
2402
2403
2404
2405
            _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:
2406
                # Get name of evals
Guolin Ke's avatar
Guolin Ke committed
2407
                tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
2408
                string_buffers = [ctypes.create_string_buffer(255) for i in range_(self.__num_inner_eval)]
wxchan's avatar
wxchan committed
2409
                ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
wxchan's avatar
wxchan committed
2410
2411
2412
2413
2414
                _safe_call(_LIB.LGBM_BoosterGetEvalNames(
                    self.handle,
                    ctypes.byref(tmp_out_len),
                    ptr_string_buffers))
                if self.__num_inner_eval != tmp_out_len.value:
2415
                    raise ValueError("Length of eval names doesn't equal with num_evals")
2416
                self.__name_inner_eval = \
wxchan's avatar
wxchan committed
2417
                    [string_buffers[i].value.decode() for i in range_(self.__num_inner_eval)]
2418
                self.__higher_better_inner_eval = \
2419
                    [name.startswith(('auc', 'ndcg@', 'map@')) for name in self.__name_inner_eval]
2420

wxchan's avatar
wxchan committed
2421
    def attr(self, key):
2422
        """Get attribute string from the Booster.
wxchan's avatar
wxchan committed
2423
2424
2425

        Parameters
        ----------
2426
2427
        key : string
            The name of the attribute.
wxchan's avatar
wxchan committed
2428
2429
2430

        Returns
        -------
2431
2432
        value : string or None
            The attribute value.
Nikita Titov's avatar
Nikita Titov committed
2433
            Returns None if attribute does not exist.
wxchan's avatar
wxchan committed
2434
        """
2435
        return self.__attr.get(key, None)
wxchan's avatar
wxchan committed
2436
2437

    def set_attr(self, **kwargs):
2438
        """Set attributes to the Booster.
wxchan's avatar
wxchan committed
2439
2440
2441
2442

        Parameters
        ----------
        **kwargs
2443
2444
            The attributes to set.
            Setting a value to None deletes an attribute.
Nikita Titov's avatar
Nikita Titov committed
2445
2446
2447
2448

        Returns
        -------
        self : Booster
2449
            Booster with set attributes.
wxchan's avatar
wxchan committed
2450
2451
2452
        """
        for key, value in kwargs.items():
            if value is not None:
wxchan's avatar
wxchan committed
2453
                if not isinstance(value, string_type):
Nikita Titov's avatar
Nikita Titov committed
2454
                    raise ValueError("Only string values are accepted")
wxchan's avatar
wxchan committed
2455
2456
2457
                self.__attr[key] = value
            else:
                self.__attr.pop(key, None)
Nikita Titov's avatar
Nikita Titov committed
2458
        return self