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

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

import numpy as np
import scipy.sparse

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

wxchan's avatar
wxchan committed
20

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

wxchan's avatar
wxchan committed
30

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

wxchan's avatar
wxchan committed
33

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

wxchan's avatar
wxchan committed
38

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

wxchan's avatar
wxchan committed
49

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

wxchan's avatar
wxchan committed
60

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

wxchan's avatar
wxchan committed
65

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

wxchan's avatar
wxchan committed
71

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

wxchan's avatar
wxchan committed
86

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

Guolin Ke's avatar
Guolin Ke committed
95

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

wxchan's avatar
wxchan committed
104

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

wxchan's avatar
wxchan committed
113

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

wxchan's avatar
wxchan committed
118

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

wxchan's avatar
wxchan committed
123

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

wxchan's avatar
wxchan committed
138

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

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

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

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

wxchan's avatar
wxchan committed
158

wxchan's avatar
wxchan committed
159
160
161
162
163
"""marco definition of data type in c_api of LightGBM"""
C_API_DTYPE_FLOAT32 = 0
C_API_DTYPE_FLOAT64 = 1
C_API_DTYPE_INT32 = 2
C_API_DTYPE_INT64 = 3
Guolin Ke's avatar
Guolin Ke committed
164

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

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

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

wxchan's avatar
wxchan committed
179

wxchan's avatar
wxchan committed
180
def c_float_array(data):
Guolin Ke's avatar
Guolin Ke committed
181
    """get pointer of float numpy array / list"""
wxchan's avatar
wxchan committed
182
183
184
185
186
187
188
189
190
191
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
        if data.dtype == np.float32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
            type_data = C_API_DTYPE_FLOAT32
        elif data.dtype == np.float64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
            type_data = C_API_DTYPE_FLOAT64
        else:
192
            raise TypeError("Expected np.float32 or np.float64, met type({})"
wxchan's avatar
wxchan committed
193
194
                            .format(data.dtype))
    else:
195
        raise TypeError("Unknown type({})".format(type(data).__name__))
wxchan's avatar
wxchan committed
196
197
    return (ptr_data, type_data)

wxchan's avatar
wxchan committed
198

wxchan's avatar
wxchan committed
199
def c_int_array(data):
Guolin Ke's avatar
Guolin Ke committed
200
    """get pointer of int numpy array / list"""
wxchan's avatar
wxchan committed
201
202
203
204
205
206
207
208
209
210
    if is_1d_list(data):
        data = np.array(data, copy=False)
    if is_numpy_1d_array(data):
        if data.dtype == np.int32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
            type_data = C_API_DTYPE_INT32
        elif data.dtype == np.int64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int64))
            type_data = C_API_DTYPE_INT64
        else:
211
            raise TypeError("Expected np.int32 or np.int64, met type({})"
wxchan's avatar
wxchan committed
212
213
                            .format(data.dtype))
    else:
214
        raise TypeError("Unknown type({})".format(type(data).__name__))
wxchan's avatar
wxchan committed
215
216
    return (ptr_data, type_data)

wxchan's avatar
wxchan committed
217

218
219
220
221
222
223
PANDAS_DTYPE_MAPPER = {'int8': 'int', 'int16': 'int', 'int32': 'int',
                       'int64': 'int', 'uint8': 'int', 'uint16': 'int',
                       'uint32': 'int', 'uint64': 'int', 'float16': 'float',
                       'float32': 'float', 'float64': 'float', 'bool': 'int'}


224
def _data_from_pandas(data, feature_name, categorical_feature, pandas_categorical):
225
    if isinstance(data, DataFrame):
226
        if feature_name == 'auto' or feature_name is None:
wxchan's avatar
wxchan committed
227
228
229
230
            if all([isinstance(name, integer_types + (np.integer, )) for name in data.columns]):
                msg = """Using Pandas (default) integer column names, not column indexes. You can use indexes with DataFrame.values."""
                warnings.filterwarnings('once')
                warnings.warn(msg, stacklevel=5)
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
            data = data.rename(columns=str)
        cat_cols = data.select_dtypes(include=['category']).columns
        if pandas_categorical is None:  # train dataset
            pandas_categorical = [list(data[col].cat.categories) for col in cat_cols]
        else:
            if len(cat_cols) != len(pandas_categorical):
                raise ValueError('train and valid dataset categorical_feature do not match.')
            for col, category in zip(cat_cols, pandas_categorical):
                if list(data[col].cat.categories) != list(category):
                    data[col] = data[col].cat.set_categories(category)
        if len(cat_cols):  # cat_cols is pandas Index object
            data = data.copy()  # not alter origin DataFrame
            data[cat_cols] = data[cat_cols].apply(lambda x: x.cat.codes)
        if categorical_feature is not None:
            if feature_name is None:
                feature_name = list(data.columns)
            if categorical_feature == 'auto':
                categorical_feature = list(cat_cols)
            else:
                categorical_feature = list(categorical_feature) + list(cat_cols)
        if feature_name == 'auto':
            feature_name = list(data.columns)
        data_dtypes = data.dtypes
        if not all(dtype.name in PANDAS_DTYPE_MAPPER for dtype in data_dtypes):
            bad_fields = [data.columns[i] for i, dtype in
                          enumerate(data_dtypes) if dtype.name not in PANDAS_DTYPE_MAPPER]

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


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


280
281
282
283
284
285
286
287
288
289
290
291
292
def _save_pandas_categorical(file_name, pandas_categorical):
    with open(file_name, 'a') as f:
        f.write('\npandas_categorical:' + json.dumps(pandas_categorical, default=json_default_with_numpy))


def _load_pandas_categorical(file_name):
    with open(file_name, 'r') as f:
        last_line = f.readlines()[-1]
        if last_line.startswith('pandas_categorical:'):
            return json.loads(last_line[len('pandas_categorical:'):])
    return None


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

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

    def __del__(self):
        if self.__is_manage_handle:
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))

    def predict(self, data, num_iteration=-1,
                raw_score=False, pred_leaf=False, data_has_header=False,
                is_reshape=True):
        """
        Predict logic

        Parameters
        ----------
        data : string/numpy array/scipy.sparse
            Data source for prediction
356
            When data type is string, it represents the path of txt file
wxchan's avatar
wxchan committed
357
        num_iteration : int
358
            Used iteration for prediction
wxchan's avatar
wxchan committed
359
360
361
362
363
        raw_score : bool
            True for predict raw score
        pred_leaf : bool
            True for predict leaf index
        data_has_header : bool
Guolin Ke's avatar
Guolin Ke committed
364
            Used for txt data, True if txt data has header
wxchan's avatar
wxchan committed
365
        is_reshape : bool
366
            Reshape to (nrow, ncol) if true
wxchan's avatar
wxchan committed
367
368
369
370
371

        Returns
        -------
        Prediction result
        """
wxchan's avatar
wxchan committed
372
        if isinstance(data, Dataset):
373
            raise TypeError("Cannot use Dataset instance for prediction, please use raw data instead")
374
        data = _data_from_pandas(data, None, None, self.pandas_categorical)[0]
wxchan's avatar
wxchan committed
375
376
377
378
379
380
        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
        int_data_has_header = 1 if data_has_header else 0
381
382
        if num_iteration > self.num_total_iteration:
            num_iteration = self.num_total_iteration
wxchan's avatar
wxchan committed
383
        if isinstance(data, string_type):
384
            with _temp_file() as f:
wxchan's avatar
wxchan committed
385
386
387
                _safe_call(_LIB.LGBM_BoosterPredictForFile(
                    self.handle,
                    c_str(data),
Guolin Ke's avatar
Guolin Ke committed
388
389
390
                    ctypes.c_int(int_data_has_header),
                    ctypes.c_int(predict_type),
                    ctypes.c_int(num_iteration),
wxchan's avatar
wxchan committed
391
392
                    c_str(f.name)))
                lines = f.readlines()
393
394
                nrow = len(lines)
                preds = [float(token) for line in lines for token in line.split('\t')]
Guolin Ke's avatar
Guolin Ke committed
395
                preds = np.array(preds, dtype=np.float64, copy=False)
wxchan's avatar
wxchan committed
396
397
398
        elif isinstance(data, scipy.sparse.csr_matrix):
            preds, nrow = self.__pred_for_csr(data, num_iteration,
                                              predict_type)
Guolin Ke's avatar
Guolin Ke committed
399
400
401
        elif isinstance(data, scipy.sparse.csc_matrix):
            preds, nrow = self.__pred_for_csc(data, num_iteration,
                                              predict_type)
wxchan's avatar
wxchan committed
402
403
404
        elif isinstance(data, np.ndarray):
            preds, nrow = self.__pred_for_np2d(data, num_iteration,
                                               predict_type)
405
406
407
        elif isinstance(data, DataFrame):
            preds, nrow = self.__pred_for_np2d(data.values, num_iteration,
                                               predict_type)
wxchan's avatar
wxchan committed
408
409
410
411
412
413
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                preds, nrow = self.__pred_for_csr(csr, num_iteration,
                                                  predict_type)
            except:
414
                raise TypeError('Cannot predict data for type {}'.format(type(data).__name__))
wxchan's avatar
wxchan committed
415
416
        if pred_leaf:
            preds = preds.astype(np.int32)
417
        if is_reshape and preds.size != nrow:
wxchan's avatar
wxchan committed
418
            if preds.size % nrow == 0:
419
                preds = preds.reshape(nrow, -1)
wxchan's avatar
wxchan committed
420
            else:
421
                raise ValueError('Length of predict result (%d) cannot be divide nrow (%d)'
wxchan's avatar
wxchan committed
422
423
424
425
                                 % (preds.size, nrow))
        return preds

    def __get_num_preds(self, num_iteration, nrow, predict_type):
Guolin Ke's avatar
Guolin Ke committed
426
427
428
        """
        Get size of prediction result
        """
Guolin Ke's avatar
Guolin Ke committed
429
430
431
        n_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterCalcNumPredict(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
432
433
434
            ctypes.c_int(nrow),
            ctypes.c_int(predict_type),
            ctypes.c_int(num_iteration),
Guolin Ke's avatar
Guolin Ke committed
435
436
            ctypes.byref(n_preds)))
        return n_preds.value
wxchan's avatar
wxchan committed
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452

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

        if mat.dtype == np.float32 or mat.dtype == np.float64:
            data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
        else:
            """change non-float data to float data, need to copy"""
            data = np.array(mat.reshape(mat.size), dtype=np.float32)
        ptr_data, type_ptr_data = c_float_array(data)
        n_preds = self.__get_num_preds(num_iteration, mat.shape[0],
                                       predict_type)
Guolin Ke's avatar
Guolin Ke committed
453
        preds = np.zeros(n_preds, dtype=np.float64)
wxchan's avatar
wxchan committed
454
455
456
457
        out_num_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterPredictForMat(
            self.handle,
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
458
459
460
461
462
463
            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),
wxchan's avatar
wxchan committed
464
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
465
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
466
        if n_preds != out_num_preds.value:
467
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
468
469
470
471
472
473
474
475
        return preds, mat.shape[0]

    def __pred_for_csr(self, csr, num_iteration, predict_type):
        """
        Predict for a csr data
        """
        nrow = len(csr.indptr) - 1
        n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
Guolin Ke's avatar
Guolin Ke committed
476
        preds = np.zeros(n_preds, dtype=np.float64)
wxchan's avatar
wxchan committed
477
478
479
480
481
482
483
484
        out_num_preds = ctypes.c_int64(0)

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

        _safe_call(_LIB.LGBM_BoosterPredictForCSR(
            self.handle,
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
485
            ctypes.c_int32(type_ptr_indptr),
wxchan's avatar
wxchan committed
486
487
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
488
489
490
491
492
493
            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),
wxchan's avatar
wxchan committed
494
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
495
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
Guolin Ke's avatar
Guolin Ke committed
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
        if n_preds != out_num_preds.value:
            raise ValueError("Wrong length for predict results")
        return preds, nrow

    def __pred_for_csc(self, csc, num_iteration, predict_type):
        """
        Predict for a csc data
        """
        nrow = csc.shape[0]
        n_preds = self.__get_num_preds(num_iteration, nrow, predict_type)
        preds = np.zeros(n_preds, dtype=np.float64)
        out_num_preds = ctypes.c_int64(0)

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

        _safe_call(_LIB.LGBM_BoosterPredictForCSC(
            self.handle,
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
515
            ctypes.c_int32(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
516
517
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
518
519
520
521
522
523
            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),
Guolin Ke's avatar
Guolin Ke committed
524
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
525
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
526
        if n_preds != out_num_preds.value:
527
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
528
529
        return preds, nrow

wxchan's avatar
wxchan committed
530

wxchan's avatar
wxchan committed
531
532
class Dataset(object):
    """Dataset in LightGBM."""
wxchan's avatar
wxchan committed
533
    def __init__(self, data, label=None, max_bin=255, reference=None,
wxchan's avatar
wxchan committed
534
                 weight=None, group=None, silent=False,
535
                 feature_name='auto', categorical_feature='auto', params=None,
wxchan's avatar
wxchan committed
536
                 free_raw_data=True):
wxchan's avatar
wxchan committed
537
538
539
540
        """
        Parameters
        ----------
        data : string/numpy array/scipy.sparse
wxchan's avatar
wxchan committed
541
            Data source of Dataset.
542
            When data type is string, it represents the path of txt file
wxchan's avatar
wxchan committed
543
544
545
        label : list or numpy 1-D array, optional
            Label of the data
        max_bin : int, required
546
            Max number of discrete bin for features
wxchan's avatar
wxchan committed
547
        reference : Other Dataset, optional
wxchan's avatar
wxchan committed
548
549
550
551
            If this dataset validation, need to use training data as reference
        weight : list or numpy 1-D array , optional
            Weight for each instance.
        group : list or numpy 1-D array , optional
552
            Group/query size for dataset
wxchan's avatar
wxchan committed
553
554
        silent : boolean, optional
            Whether print messages during construction
wxchan's avatar
wxchan committed
555
        feature_name : list of str, or 'auto'
556
            Feature names
wxchan's avatar
wxchan committed
557
            If 'auto' and data is pandas DataFrame, use data columns name
558
559
560
561
562
        categorical_feature : list of str or int, or 'auto'
            Categorical features,
            type int represents index,
            type str represents feature names (need to specify feature_name as well)
            If 'auto' and data is pandas DataFrame, use pandas categorical columns
wxchan's avatar
wxchan committed
563
        params: dict, optional
564
            Other parameters
wxchan's avatar
wxchan committed
565
566
        free_raw_data: Bool
            True if need to free raw data after construct inner dataset
wxchan's avatar
wxchan committed
567
        """
wxchan's avatar
wxchan committed
568
569
570
571
572
573
574
575
576
        self.handle = None
        self.data = data
        self.label = label
        self.max_bin = max_bin
        self.reference = reference
        self.weight = weight
        self.group = group
        self.silent = silent
        self.feature_name = feature_name
577
        self.categorical_feature = categorical_feature
wxchan's avatar
wxchan committed
578
579
580
581
        self.params = params
        self.free_raw_data = free_raw_data
        self.used_indices = None
        self._predictor = None
582
        self.pandas_categorical = None
wxchan's avatar
wxchan committed
583
584

    def __del__(self):
585
586
587
        self._free_handle()

    def _free_handle(self):
588
        if self.handle is not None:
589
            _safe_call(_LIB.LGBM_DatasetFree(self.handle))
590
            self.handle = None
wxchan's avatar
wxchan committed
591
592
593

    def _lazy_init(self, data, label=None, max_bin=255, reference=None,
                   weight=None, group=None, predictor=None,
wxchan's avatar
wxchan committed
594
                   silent=False, feature_name='auto',
595
                   categorical_feature='auto', params=None):
wxchan's avatar
wxchan committed
596
597
598
        if data is None:
            self.handle = None
            return
599
        data, feature_name, categorical_feature, self.pandas_categorical = _data_from_pandas(data, feature_name, categorical_feature, self.pandas_categorical)
wxchan's avatar
wxchan committed
600
601
602
603
604
605
606
607
608
609
610
        label = _label_from_pandas(label)
        self.data_has_header = False
        """process for args"""
        params = {} if params is None else params
        self.max_bin = max_bin
        self.predictor = predictor
        params["max_bin"] = max_bin
        if silent:
            params["verbose"] = 0
        elif "verbose" not in params:
            params["verbose"] = 1
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
        """get categorical features"""
        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))

            params['categorical_column'] = sorted(categorical_indices)

wxchan's avatar
wxchan committed
628
629
630
        params_str = param_dict_to_str(params)
        """process for reference dataset"""
        ref_dataset = None
wxchan's avatar
wxchan committed
631
        if isinstance(reference, Dataset):
632
            ref_dataset = reference.construct().handle
wxchan's avatar
wxchan committed
633
634
635
        elif reference is not None:
            raise TypeError('Reference dataset should be None or dataset instance')
        """start construct data"""
wxchan's avatar
wxchan committed
636
        if isinstance(data, string_type):
wxchan's avatar
wxchan committed
637
            """check data has header or not"""
Guolin Ke's avatar
Guolin Ke committed
638
            if str(params.get("has_header", "")).lower() == "true" \
wxchan's avatar
wxchan committed
639
                    or str(params.get("header", "")).lower() == "true":
640
                self.data_has_header = True
wxchan's avatar
wxchan committed
641
642
643
644
645
646
647
648
            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
649
650
        elif isinstance(data, scipy.sparse.csc_matrix):
            self.__init_from_csc(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
651
652
653
654
655
656
657
        elif isinstance(data, np.ndarray):
            self.__init_from_np2d(data, params_str, ref_dataset)
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                self.__init_from_csr(csr, params_str, ref_dataset)
            except:
wxchan's avatar
wxchan committed
658
                raise TypeError('Cannot initialize Dataset from {}'.format(type(data).__name__))
wxchan's avatar
wxchan committed
659
660
661
        if label is not None:
            self.set_label(label)
        if self.get_label() is None:
662
            raise ValueError("Label should not be None")
wxchan's avatar
wxchan committed
663
664
665
666
667
        if weight is not None:
            self.set_weight(weight)
        if group is not None:
            self.set_group(group)
        # load init score
Guolin Ke's avatar
Guolin Ke committed
668
        if isinstance(self.predictor, _InnerPredictor):
wxchan's avatar
wxchan committed
669
670
671
672
673
674
            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
675
                new_init_score = np.zeros(init_score.size, dtype=np.float32)
wxchan's avatar
wxchan committed
676
                num_data = self.num_data()
wxchan's avatar
wxchan committed
677
678
                for i in range_(num_data):
                    for j in range_(self.predictor.num_class):
wxchan's avatar
wxchan committed
679
680
681
                        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
682
683
        elif self.predictor is not None:
            raise TypeError('wrong predictor type {}'.format(type(self.predictor).__name__))
Guolin Ke's avatar
Guolin Ke committed
684
685
        # set feature names
        self.set_feature_name(feature_name)
wxchan's avatar
wxchan committed
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703

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

        self.handle = ctypes.c_void_p()
        if mat.dtype == np.float32 or mat.dtype == np.float64:
            data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
        else:
            """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)
        _safe_call(_LIB.LGBM_DatasetCreateFromMat(
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
704
705
706
707
            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
708
709
710
711
712
713
714
715
716
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

    def __init_from_csr(self, csr, params_str, ref_dataset):
        """
        Initialize data from a CSR matrix.
        """
        if len(csr.indices) != len(csr.data):
717
            raise ValueError('Length mismatch: {} vs {}'.format(len(csr.indices), len(csr.data)))
wxchan's avatar
wxchan committed
718
719
720
721
722
723
724
        self.handle = ctypes.c_void_p()

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

        _safe_call(_LIB.LGBM_DatasetCreateFromCSR(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
725
            ctypes.c_int(type_ptr_indptr),
wxchan's avatar
wxchan committed
726
727
            csr.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
728
729
730
731
            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
732
733
734
735
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

Guolin Ke's avatar
Guolin Ke committed
736
737
738
739
740
741
742
743
744
745
746
747
748
    def __init_from_csc(self, csc, params_str, ref_dataset):
        """
        Initialize data from a csc matrix.
        """
        if len(csc.indices) != len(csc.data):
            raise ValueError('Length mismatch: {} vs {}'.format(len(csc.indices), len(csc.data)))
        self.handle = ctypes.c_void_p()

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

        _safe_call(_LIB.LGBM_DatasetCreateFromCSC(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
749
            ctypes.c_int(type_ptr_indptr),
Guolin Ke's avatar
Guolin Ke committed
750
751
            csc.indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
752
753
754
755
            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
756
757
758
759
            c_str(params_str),
            ref_dataset,
            ctypes.byref(self.handle)))

wxchan's avatar
wxchan committed
760
761
    def construct(self):
        """Lazy init"""
762
        if self.handle is None:
wxchan's avatar
wxchan committed
763
764
765
766
767
768
769
770
771
            if self.reference is not None:
                if self.used_indices is None:
                    """create valid"""
                    self._lazy_init(self.data, label=self.label, max_bin=self.max_bin, reference=self.reference,
                                    weight=self.weight, group=self.group, predictor=self._predictor,
                                    silent=self.silent, params=self.params)
                else:
                    """construct subset"""
                    used_indices = list_to_1d_numpy(self.used_indices, np.int32, name='used_indices')
772
                    self.handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
773
774
                    params_str = param_dict_to_str(self.params)
                    _safe_call(_LIB.LGBM_DatasetGetSubset(
775
                        self.reference.construct().handle,
wxchan's avatar
wxchan committed
776
                        used_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
777
                        ctypes.c_int(used_indices.shape[0]),
wxchan's avatar
wxchan committed
778
779
780
781
782
783
784
785
786
                        c_str(params_str),
                        ctypes.byref(self.handle)))
                    if self.get_label() is None:
                        raise ValueError("Label should not be None.")
            else:
                """create train"""
                self._lazy_init(self.data, label=self.label, max_bin=self.max_bin,
                                weight=self.weight, group=self.group, predictor=self._predictor,
                                silent=self.silent, feature_name=self.feature_name,
787
                                categorical_feature=self.categorical_feature, params=self.params)
wxchan's avatar
wxchan committed
788
789
790
            if self.free_raw_data:
                self.data = None
        return self
wxchan's avatar
wxchan committed
791

wxchan's avatar
wxchan committed
792
793
794
795
    def create_valid(self, data, label=None, weight=None, group=None,
                     silent=False, params=None):
        """
        Create validation data align with current dataset
wxchan's avatar
wxchan committed
796
797
798

        Parameters
        ----------
wxchan's avatar
wxchan committed
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
        data : string/numpy array/scipy.sparse
            Data source of Dataset.
            When data type is string, it represents the path of txt file
        label : list or numpy 1-D array, optional
            Label of the training data.
        weight : list or numpy 1-D array , optional
            Weight for each instance.
        group : list or numpy 1-D array , optional
            Group/query size for dataset
        silent : boolean, optional
            Whether print messages during construction
        params: dict, optional
            Other parameters
        """
        ret = Dataset(data, label=label, max_bin=self.max_bin, reference=self,
                      weight=weight, group=group, silent=silent, params=params,
                      free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
816
        ret._predictor = self._predictor
817
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
818
        return ret
wxchan's avatar
wxchan committed
819

wxchan's avatar
wxchan committed
820
    def subset(self, used_indices, params=None):
wxchan's avatar
wxchan committed
821
        """
wxchan's avatar
wxchan committed
822
823
824
825
826
827
828
829
830
831
        Get subset of current dataset

        Parameters
        ----------
        used_indices : list of int
            Used indices of this subset
        params : dict
            Other parameters
        """
        ret = Dataset(None, reference=self, feature_name=self.feature_name,
832
                      categorical_feature=self.categorical_feature, params=params)
wxchan's avatar
wxchan committed
833
        ret._predictor = self._predictor
834
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
        ret.used_indices = used_indices
        return ret

    def save_binary(self, filename):
        """
        Save Dataset to binary file

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

    def _update_params(self, params):
        if not self.params:
            self.params = params
wxchan's avatar
wxchan committed
854
        else:
wxchan's avatar
wxchan committed
855
            self.params.update(params)
wxchan's avatar
wxchan committed
856
857

    def set_field(self, field_name, data):
wxchan's avatar
wxchan committed
858
        """Set property into the Dataset.
wxchan's avatar
wxchan committed
859
860
861
862
863
864
865
866
867

        Parameters
        ----------
        field_name: str
            The field name of the information

        data: numpy array or list or None
            The array ofdata to be set
        """
868
869
        if self.handle is None:
            raise Exception("Cannot set %s before construct dataset" % field_name)
wxchan's avatar
wxchan committed
870
871
872
873
874
875
        if data is None:
            """set to None"""
            _safe_call(_LIB.LGBM_DatasetSetField(
                self.handle,
                c_str(field_name),
                None,
Guolin Ke's avatar
Guolin Ke committed
876
877
                ctypes.c_int(0),
                ctypes.c_int(FIELD_TYPE_MAPPER[field_name])))
wxchan's avatar
wxchan committed
878
            return
Guolin Ke's avatar
Guolin Ke committed
879
880
881
882
883
        dtype = np.float32
        if field_name == 'group':
            dtype = np.int32
        elif field_name == 'init_score':
            dtype = np.float64
884
        data = list_to_1d_numpy(data, dtype, name=field_name)
wxchan's avatar
wxchan committed
885
886
887
        if data.dtype == np.float32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
            type_data = C_API_DTYPE_FLOAT32
Guolin Ke's avatar
Guolin Ke committed
888
889
890
        elif data.dtype == np.float64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
            type_data = C_API_DTYPE_FLOAT64
wxchan's avatar
wxchan committed
891
892
893
894
        elif data.dtype == np.int32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
            type_data = C_API_DTYPE_INT32
        else:
Guolin Ke's avatar
Guolin Ke committed
895
            raise TypeError("Excepted np.float32/64 or np.int32, meet type({})".format(data.dtype))
wxchan's avatar
wxchan committed
896
        if type_data != FIELD_TYPE_MAPPER[field_name]:
897
            raise TypeError("Input type error for set_field")
wxchan's avatar
wxchan committed
898
899
900
901
        _safe_call(_LIB.LGBM_DatasetSetField(
            self.handle,
            c_str(field_name),
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
902
903
            ctypes.c_int(len(data)),
            ctypes.c_int(type_data)))
wxchan's avatar
wxchan committed
904

wxchan's avatar
wxchan committed
905
906
    def get_field(self, field_name):
        """Get property from the Dataset.
wxchan's avatar
wxchan committed
907
908
909

        Parameters
        ----------
wxchan's avatar
wxchan committed
910
911
        field_name: str
            The field name of the information
wxchan's avatar
wxchan committed
912
913
914

        Returns
        -------
wxchan's avatar
wxchan committed
915
916
        info : array
            A numpy array of information of the data
Guolin Ke's avatar
Guolin Ke committed
917
        """
918
919
        if self.handle is None:
            raise Exception("Cannot get %s before construct dataset" % field_name)
Guolin Ke's avatar
Guolin Ke committed
920
921
        tmp_out_len = ctypes.c_int()
        out_type = ctypes.c_int()
wxchan's avatar
wxchan committed
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
        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
937
938
        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)
939
        else:
wxchan's avatar
wxchan committed
940
            raise TypeError("Unknown type")
Guolin Ke's avatar
Guolin Ke committed
941

942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
    def set_categorical_feature(self, categorical_feature):
        """
        Set categorical features

        Parameters
        ----------
        categorical_feature : list of int or str
            Name/index of categorical features

        """
        if self.categorical_feature == categorical_feature:
            return
        if self.data is not None:
            self.categorical_feature = categorical_feature
            self._free_handle()
        else:
            raise LightGBMError("Cannot set categorical feature after freed raw data, set free_raw_data=False when construct Dataset to avoid this.")

Guolin Ke's avatar
Guolin Ke committed
960
961
962
963
964
965
966
967
968
    def _set_predictor(self, predictor):
        """
        Set predictor for continued training, not recommand for user to call this function.
        Please set init_model in engine.train or engine.cv
        """
        if predictor is self._predictor:
            return
        if self.data is not None:
            self._predictor = predictor
969
            self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
970
        else:
971
            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
972
973
974
975
976
977
978
979

    def set_reference(self, reference):
        """
        Set reference dataset

        Parameters
        ----------
        reference : Dataset
980
            Will use reference as template to consturct current dataset
Guolin Ke's avatar
Guolin Ke committed
981
        """
982
        self.set_categorical_feature(reference.categorical_feature)
Guolin Ke's avatar
Guolin Ke committed
983
984
985
986
987
988
        self.set_feature_name(reference.feature_name)
        self._set_predictor(reference._predictor)
        if self.reference is reference:
            return
        if self.data is not None:
            self.reference = reference
989
            self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
990
        else:
991
            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
992
993
994
995
996
997
998
999

    def set_feature_name(self, feature_name):
        """
        Set feature name

        Parameters
        ----------
        feature_name : list of str
1000
            Feature names
Guolin Ke's avatar
Guolin Ke committed
1001
1002
        """
        self.feature_name = feature_name
1003
        if self.handle is not None and feature_name is not None and feature_name != 'auto':
wxchan's avatar
wxchan committed
1004
1005
            if len(feature_name) != self.num_feature():
                raise ValueError("Length of feature_name({}) and num_feature({}) don't match".format(len(feature_name), self.num_feature()))
1006
            c_feature_name = [c_str(name) for name in feature_name]
wxchan's avatar
wxchan committed
1007
1008
1009
            _safe_call(_LIB.LGBM_DatasetSetFeatureNames(
                self.handle,
                c_array(ctypes.c_char_p, c_feature_name),
Guolin Ke's avatar
Guolin Ke committed
1010
                ctypes.c_int(len(feature_name))))
Guolin Ke's avatar
Guolin Ke committed
1011
1012

    def set_label(self, label):
wxchan's avatar
wxchan committed
1013
1014
        """
        Set label of Dataset
Guolin Ke's avatar
Guolin Ke committed
1015
1016
1017
1018
1019
1020
1021

        Parameters
        ----------
        label: numpy array or list or None
            The label information to be set into Dataset
        """
        self.label = label
1022
        if self.handle is not None:
wxchan's avatar
wxchan committed
1023
1024
            label = list_to_1d_numpy(label, name='label')
            self.set_field('label', label)
Guolin Ke's avatar
Guolin Ke committed
1025
1026

    def set_weight(self, weight):
wxchan's avatar
wxchan committed
1027
1028
        """
        Set weight of each instance.
Guolin Ke's avatar
Guolin Ke committed
1029
1030
1031
1032
1033
1034
1035

        Parameters
        ----------
        weight : numpy array or list or None
            Weight for each data point
        """
        self.weight = weight
1036
        if self.handle is not None and weight is not None:
wxchan's avatar
wxchan committed
1037
1038
            weight = list_to_1d_numpy(weight, name='weight')
            self.set_field('weight', weight)
Guolin Ke's avatar
Guolin Ke committed
1039
1040

    def set_init_score(self, init_score):
wxchan's avatar
wxchan committed
1041
1042
        """
        Set init score of booster to start from.
Guolin Ke's avatar
Guolin Ke committed
1043
1044
1045
1046
1047
1048
1049

        Parameters
        ----------
        init_score: numpy array or list or None
            Init score for booster
        """
        self.init_score = init_score
1050
        if self.handle is not None and init_score is not None:
Guolin Ke's avatar
Guolin Ke committed
1051
            init_score = list_to_1d_numpy(init_score, np.float64, name='init_score')
wxchan's avatar
wxchan committed
1052
            self.set_field('init_score', init_score)
Guolin Ke's avatar
Guolin Ke committed
1053
1054

    def set_group(self, group):
wxchan's avatar
wxchan committed
1055
1056
        """
        Set group size of Dataset (used for ranking).
Guolin Ke's avatar
Guolin Ke committed
1057
1058
1059
1060
1061
1062
1063

        Parameters
        ----------
        group : numpy array or list or None
            Group size of each group
        """
        self.group = group
1064
        if self.handle is not None and group is not None:
wxchan's avatar
wxchan committed
1065
1066
            group = list_to_1d_numpy(group, np.int32, name='group')
            self.set_field('group', group)
Guolin Ke's avatar
Guolin Ke committed
1067
1068

    def get_label(self):
wxchan's avatar
wxchan committed
1069
1070
        """
        Get the label of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1071
1072
1073
1074
1075

        Returns
        -------
        label : array
        """
1076
        if self.label is None and self.handle is not None:
wxchan's avatar
wxchan committed
1077
            self.label = self.get_field('label')
Guolin Ke's avatar
Guolin Ke committed
1078
1079
1080
        return self.label

    def get_weight(self):
wxchan's avatar
wxchan committed
1081
1082
        """
        Get the weight of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1083
1084
1085
1086
1087

        Returns
        -------
        weight : array
        """
1088
        if self.weight is None and self.handle is not None:
wxchan's avatar
wxchan committed
1089
            self.weight = self.get_field('weight')
Guolin Ke's avatar
Guolin Ke committed
1090
1091
1092
        return self.weight

    def get_init_score(self):
wxchan's avatar
wxchan committed
1093
1094
        """
        Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1095
1096
1097
1098
1099

        Returns
        -------
        init_score : array
        """
1100
        if self.init_score is None and self.handle is not None:
wxchan's avatar
wxchan committed
1101
            self.init_score = self.get_field('init_score')
Guolin Ke's avatar
Guolin Ke committed
1102
1103
1104
        return self.init_score

    def get_group(self):
wxchan's avatar
wxchan committed
1105
1106
        """
        Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1107
1108
1109
1110
1111

        Returns
        -------
        init_score : array
        """
1112
        if self.group is None and self.handle is not None:
wxchan's avatar
wxchan committed
1113
            self.group = self.get_field('group')
Guolin Ke's avatar
Guolin Ke committed
1114
1115
1116
            if self.group is not None:
                # group data from LightGBM is boundaries data, need to convert to group size
                new_group = []
wxchan's avatar
wxchan committed
1117
                for i in range_(len(self.group) - 1):
Guolin Ke's avatar
Guolin Ke committed
1118
1119
                    new_group.append(self.group[i + 1] - self.group[i])
                self.group = new_group
Guolin Ke's avatar
Guolin Ke committed
1120
1121
1122
        return self.group

    def num_data(self):
wxchan's avatar
wxchan committed
1123
1124
        """
        Get the number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1125
1126
1127
1128
1129

        Returns
        -------
        number of rows : int
        """
1130
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1131
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1132
1133
1134
            _safe_call(_LIB.LGBM_DatasetGetNumData(self.handle,
                                                   ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1135
        else:
1136
            raise LightGBMError("Cannot get num_data before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1137
1138

    def num_feature(self):
wxchan's avatar
wxchan committed
1139
1140
        """
        Get the number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
1141
1142
1143
1144
1145

        Returns
        -------
        number of columns : int
        """
1146
        if self.handle is not None:
Guolin Ke's avatar
Guolin Ke committed
1147
            ret = ctypes.c_int()
wxchan's avatar
wxchan committed
1148
1149
1150
            _safe_call(_LIB.LGBM_DatasetGetNumFeature(self.handle,
                                                      ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
1151
        else:
1152
            raise LightGBMError("Cannot get num_feature before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
1153

wxchan's avatar
wxchan committed
1154

wxchan's avatar
wxchan committed
1155
class Booster(object):
wxchan's avatar
wxchan committed
1156
    """"Booster in LightGBM."""
wxchan's avatar
wxchan committed
1157
    def __init__(self, params=None, train_set=None, model_file=None, silent=False):
wxchan's avatar
wxchan committed
1158
1159
        """
        Initialize the Booster.
wxchan's avatar
wxchan committed
1160
1161
1162
1163
1164
1165

        Parameters
        ----------
        params : dict
            Parameters for boosters.
        train_set : Dataset
1166
            Training dataset
wxchan's avatar
wxchan committed
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
        model_file : string
            Path to the model file.
        silent : boolean, optional
            Whether print messages during construction
        """
        self.handle = ctypes.c_void_p()
        self.__need_reload_eval_info = True
        self.__train_data_name = "training"
        self.__attr = {}
        self.best_iteration = -1
        params = {} if params is None else params
        if silent:
            params["verbose"] = 0
        elif "verbose" not in params:
            params["verbose"] = 1
        if train_set is not None:
            """Training task"""
            if not isinstance(train_set, Dataset):
1185
                raise TypeError('Training data should be Dataset instance, met {}'.format(type(train_set).__name__))
wxchan's avatar
wxchan committed
1186
1187
1188
            params_str = param_dict_to_str(params)
            """construct booster object"""
            _safe_call(_LIB.LGBM_BoosterCreate(
wxchan's avatar
wxchan committed
1189
                train_set.construct().handle,
wxchan's avatar
wxchan committed
1190
1191
1192
1193
1194
1195
1196
                c_str(params_str),
                ctypes.byref(self.handle)))
            """save reference to data"""
            self.train_set = train_set
            self.valid_sets = []
            self.name_valid_sets = []
            self.__num_dataset = 1
Guolin Ke's avatar
Guolin Ke committed
1197
1198
            self.__init_predictor = train_set._predictor
            if self.__init_predictor is not None:
wxchan's avatar
wxchan committed
1199
1200
                _safe_call(_LIB.LGBM_BoosterMerge(
                    self.handle,
Guolin Ke's avatar
Guolin Ke committed
1201
                    self.__init_predictor.handle))
Guolin Ke's avatar
Guolin Ke committed
1202
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1203
1204
1205
1206
1207
1208
1209
1210
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
            """buffer for inner predict"""
            self.__inner_predict_buffer = [None]
            self.__is_predicted_cur_iter = [False]
            self.__get_eval_info()
1211
            self.pandas_categorical = train_set.pandas_categorical
wxchan's avatar
wxchan committed
1212
1213
        elif model_file is not None:
            """Prediction task"""
Guolin Ke's avatar
Guolin Ke committed
1214
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1215
1216
1217
1218
            _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
1219
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1220
1221
1222
1223
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
                self.handle,
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
1224
            self.pandas_categorical = _load_pandas_categorical(model_file)
1225
1226
        elif 'model_str' in params:
            self.__load_model_from_string(params['model_str'])
wxchan's avatar
wxchan committed
1227
        else:
1228
            raise TypeError('Need at least one training dataset or model file to create booster instance')
wxchan's avatar
wxchan committed
1229
1230

    def __del__(self):
Guolin Ke's avatar
Guolin Ke committed
1231
        if self.handle is not None:
wxchan's avatar
wxchan committed
1232
1233
            _safe_call(_LIB.LGBM_BoosterFree(self.handle))

wxchan's avatar
wxchan committed
1234
1235
1236
1237
    def __copy__(self):
        return self.__deepcopy__(None)

    def __deepcopy__(self, _):
1238
1239
        model_str = self.__save_model_to_string()
        booster = Booster({'model_str': model_str})
1240
        booster.pandas_categorical = self.pandas_categorical
1241
        return booster
wxchan's avatar
wxchan committed
1242
1243
1244
1245
1246
1247
1248

    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:
1249
            this["handle"] = self.__save_model_to_string()
wxchan's avatar
wxchan committed
1250
1251
1252
        return this

    def __setstate__(self, state):
1253
1254
        model_str = state.get('handle', None)
        if model_str is not None:
wxchan's avatar
wxchan committed
1255
            handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
1256
            out_num_iterations = ctypes.c_int(0)
1257
1258
1259
1260
            _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
                c_str(model_str),
                ctypes.byref(out_num_iterations),
                ctypes.byref(handle)))
wxchan's avatar
wxchan committed
1261
1262
1263
            state['handle'] = handle
        self.__dict__.update(state)

wxchan's avatar
wxchan committed
1264
1265
1266
1267
    def set_train_data_name(self, name):
        self.__train_data_name = name

    def add_valid(self, data, name):
wxchan's avatar
wxchan committed
1268
1269
        """
        Add an validation data
wxchan's avatar
wxchan committed
1270
1271
1272
1273

        Parameters
        ----------
        data : Dataset
1274
            Validation data
wxchan's avatar
wxchan committed
1275
        name : String
1276
            Name of validation data
wxchan's avatar
wxchan committed
1277
        """
Guolin Ke's avatar
Guolin Ke committed
1278
        if not isinstance(data, Dataset):
wxchan's avatar
wxchan committed
1279
            raise TypeError('valid data should be Dataset instance, met {}'.format(type(data).__name__))
Guolin Ke's avatar
Guolin Ke committed
1280
1281
        if data._predictor is not self.__init_predictor:
            raise LightGBMError("Add validation data failed, you should use same predictor for these data")
wxchan's avatar
wxchan committed
1282
1283
        _safe_call(_LIB.LGBM_BoosterAddValidData(
            self.handle,
wxchan's avatar
wxchan committed
1284
            data.construct().handle))
wxchan's avatar
wxchan committed
1285
1286
1287
1288
1289
1290
1291
        self.valid_sets.append(data)
        self.name_valid_sets.append(name)
        self.__num_dataset += 1
        self.__inner_predict_buffer.append(None)
        self.__is_predicted_cur_iter.append(False)

    def reset_parameter(self, params):
wxchan's avatar
wxchan committed
1292
1293
        """
        Reset parameters for booster
wxchan's avatar
wxchan committed
1294
1295
1296
1297

        Parameters
        ----------
        params : dict
1298
            New parameters for boosters
wxchan's avatar
wxchan committed
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
        silent : boolean, optional
            Whether print messages during construction
        """
        if 'metric' in params:
            self.__need_reload_eval_info = True
        params_str = param_dict_to_str(params)
        if params_str:
            _safe_call(_LIB.LGBM_BoosterResetParameter(
                self.handle,
                c_str(params_str)))

    def update(self, train_set=None, fobj=None):
        """
        Update for one iteration
        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
1316

wxchan's avatar
wxchan committed
1317
1318
        Parameters
        ----------
1319
1320
        train_set :
            Training data, None means use last training data
wxchan's avatar
wxchan committed
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
        fobj : function
            Customized objective function.

        Returns
        -------
        is_finished, bool
        """

        """need reset training data"""
        if train_set is not None and train_set is not self.train_set:
Guolin Ke's avatar
Guolin Ke committed
1331
1332
            if not isinstance(train_set, Dataset):
                raise TypeError('Training data should be Dataset instance, met {}'.format(type(train_set).__name__))
Guolin Ke's avatar
Guolin Ke committed
1333
1334
            if train_set._predictor is not self.__init_predictor:
                raise LightGBMError("Replace training data failed, you should use same predictor for these data")
wxchan's avatar
wxchan committed
1335
1336
1337
            self.train_set = train_set
            _safe_call(_LIB.LGBM_BoosterResetTrainingData(
                self.handle,
wxchan's avatar
wxchan committed
1338
                self.train_set.construct().handle))
wxchan's avatar
wxchan committed
1339
1340
1341
1342
1343
1344
            self.__inner_predict_buffer[0] = None
        is_finished = ctypes.c_int(0)
        if fobj is None:
            _safe_call(_LIB.LGBM_BoosterUpdateOneIter(
                self.handle,
                ctypes.byref(is_finished)))
wxchan's avatar
wxchan committed
1345
            self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
            return is_finished.value == 1
        else:
            grad, hess = fobj(self.__inner_predict(0), self.train_set)
            return self.__boost(grad, hess)

    def __boost(self, grad, hess):
        """
        Boost the booster for one iteration, with customized gradient statistics.
        Note: for multi-class task, the score is group by class_id first, then group by row_id
              if you want to get i-th row score in j-th class, the access way is score[j*num_data+i]
              and you should group grad and hess in this way as well
1357

wxchan's avatar
wxchan committed
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
        Parameters
        ----------
        grad : 1d numpy or 1d list
            The first order of gradient.
        hess : 1d numpy or 1d list
            The second order of gradient.

        Returns
        -------
        is_finished, bool
        """
1369
1370
        grad = list_to_1d_numpy(grad, name='gradient')
        hess = list_to_1d_numpy(hess, name='hessian')
wxchan's avatar
wxchan committed
1371
        if len(grad) != len(hess):
1372
            raise ValueError("Lengths of gradient({}) and hessian({}) don't match".format(len(grad), len(hess)))
wxchan's avatar
wxchan committed
1373
1374
1375
1376
1377
1378
        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
1379
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1380
1381
1382
1383
1384
1385
1386
1387
        return is_finished.value == 1

    def rollback_one_iter(self):
        """
        Rollback one iteration
        """
        _safe_call(_LIB.LGBM_BoosterRollbackOneIter(
            self.handle))
wxchan's avatar
wxchan committed
1388
        self.__is_predicted_cur_iter = [False for _ in range_(self.__num_dataset)]
wxchan's avatar
wxchan committed
1389
1390

    def current_iteration(self):
Guolin Ke's avatar
Guolin Ke committed
1391
        out_cur_iter = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1392
1393
1394
1395
1396
1397
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
            self.handle,
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

    def eval(self, data, name, feval=None):
wxchan's avatar
wxchan committed
1398
1399
        """
        Evaluate for data
wxchan's avatar
wxchan committed
1400
1401
1402

        Parameters
        ----------
Guolin Ke's avatar
Guolin Ke committed
1403
        data : Dataset object
1404
1405
        name :
            Name of data
wxchan's avatar
wxchan committed
1406
1407
1408
1409
1410
1411
1412
        feval : function
            Custom evaluation function.
        Returns
        -------
        result: list
            Evaluation result list.
        """
Guolin Ke's avatar
Guolin Ke committed
1413
1414
        if not isinstance(data, Dataset):
            raise TypeError("Can only eval for Dataset instance")
wxchan's avatar
wxchan committed
1415
1416
1417
1418
        data_idx = -1
        if data is self.train_set:
            data_idx = 0
        else:
wxchan's avatar
wxchan committed
1419
            for i in range_(len(self.valid_sets)):
wxchan's avatar
wxchan committed
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
                if data is self.valid_sets[i]:
                    data_idx = i + 1
                    break
        """need to push new valid data"""
        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):
wxchan's avatar
wxchan committed
1431
1432
        """
        Evaluate for training data
wxchan's avatar
wxchan committed
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446

        Parameters
        ----------
        feval : function
            Custom evaluation function.

        Returns
        -------
        result: str
            Evaluation result list.
        """
        return self.__inner_eval(self.__train_data_name, 0, feval)

    def eval_valid(self, feval=None):
wxchan's avatar
wxchan committed
1447
1448
        """
        Evaluate for validation data
wxchan's avatar
wxchan committed
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459

        Parameters
        ----------
        feval : function
            Custom evaluation function.

        Returns
        -------
        result: str
            Evaluation result list.
        """
wxchan's avatar
wxchan committed
1460
        return [item for i in range_(1, self.__num_dataset)
wxchan's avatar
wxchan committed
1461
                for item in self.__inner_eval(self.name_valid_sets[i - 1], i, feval)]
wxchan's avatar
wxchan committed
1462
1463

    def save_model(self, filename, num_iteration=-1):
wxchan's avatar
wxchan committed
1464
1465
        """
        Save model of booster to file
wxchan's avatar
wxchan committed
1466
1467
1468
1469

        Parameters
        ----------
        filename : str
1470
            Filename to save
wxchan's avatar
wxchan committed
1471
        num_iteration: int
1472
            Number of iteration that want to save. < 0 means save the best iteration(if have)
wxchan's avatar
wxchan committed
1473
        """
1474
1475
        if num_iteration <= 0:
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1476
1477
        _safe_call(_LIB.LGBM_BoosterSaveModel(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
1478
            ctypes.c_int(num_iteration),
wxchan's avatar
wxchan committed
1479
            c_str(filename)))
1480
        _save_pandas_categorical(filename, self.pandas_categorical)
wxchan's avatar
wxchan committed
1481

1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
    def __load_model_from_string(self, model_str):
        """[Private] Load model from string"""
        out_num_iterations = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
            c_str(model_str),
            ctypes.byref(out_num_iterations),
            ctypes.byref(self.handle)))
        out_num_class = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetNumClasses(
            self.handle,
            ctypes.byref(out_num_class)))
        self.__num_class = out_num_class.value

    def __save_model_to_string(self, num_iteration=-1):
        """[Private] Save model to string"""
        if num_iteration <= 0:
            num_iteration = self.best_iteration
        buffer_len = 1 << 20
        tmp_out_len = ctypes.c_int(0)
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterSaveModelToString(
            self.handle,
            ctypes.c_int(num_iteration),
            ctypes.c_int(buffer_len),
            ctypes.byref(tmp_out_len),
            ptr_string_buffer))
        actual_len = tmp_out_len.value
        '''if buffer length is not long enough, re-allocate a buffer'''
        if actual_len > buffer_len:
            string_buffer = ctypes.create_string_buffer(actual_len)
            ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
            _safe_call(_LIB.LGBM_BoosterSaveModelToString(
                self.handle,
                ctypes.c_int(num_iteration),
                ctypes.c_int(actual_len),
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
        return string_buffer.value.decode()

1522
    def dump_model(self, num_iteration=-1):
wxchan's avatar
wxchan committed
1523
1524
        """
        Dump model to json format
wxchan's avatar
wxchan committed
1525

1526
1527
1528
1529
1530
        Parameters
        ----------
        num_iteration: int
            Number of iteration that want to dump. < 0 means dump to best iteration(if have)

wxchan's avatar
wxchan committed
1531
1532
1533
1534
        Returns
        -------
        Json format of model
        """
1535
1536
        if num_iteration <= 0:
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1537
        buffer_len = 1 << 20
Guolin Ke's avatar
Guolin Ke committed
1538
        tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1539
1540
1541
1542
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterDumpModel(
            self.handle,
Guolin Ke's avatar
Guolin Ke committed
1543
1544
            ctypes.c_int(num_iteration),
            ctypes.c_int(buffer_len),
wxchan's avatar
wxchan committed
1545
            ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1546
            ptr_string_buffer))
wxchan's avatar
wxchan committed
1547
        actual_len = tmp_out_len.value
1548
        '''if buffer length is not long enough, reallocate a buffer'''
wxchan's avatar
wxchan committed
1549
1550
1551
1552
1553
        if actual_len > buffer_len:
            string_buffer = ctypes.create_string_buffer(actual_len)
            ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
            _safe_call(_LIB.LGBM_BoosterDumpModel(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
1554
1555
                ctypes.c_int(num_iteration),
                ctypes.c_int(actual_len),
wxchan's avatar
wxchan committed
1556
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1557
                ptr_string_buffer))
wxchan's avatar
wxchan committed
1558
1559
1560
        return json.loads(string_buffer.value.decode())

    def predict(self, data, num_iteration=-1, raw_score=False, pred_leaf=False, data_has_header=False, is_reshape=True):
wxchan's avatar
wxchan committed
1561
1562
        """
        Predict logic
wxchan's avatar
wxchan committed
1563
1564
1565
1566
1567

        Parameters
        ----------
        data : string/numpy array/scipy.sparse
            Data source for prediction
1568
            When data type is string, it represents the path of txt file
wxchan's avatar
wxchan committed
1569
        num_iteration : int
1570
            Used iteration for prediction, < 0 means predict for best iteration(if have)
wxchan's avatar
wxchan committed
1571
1572
1573
1574
1575
1576
1577
        raw_score : bool
            True for predict raw score
        pred_leaf : bool
            True for predict leaf index
        data_has_header : bool
            Used for txt data
        is_reshape : bool
1578
            Reshape to (nrow, ncol) if true
wxchan's avatar
wxchan committed
1579
1580
1581
1582
1583

        Returns
        -------
        Prediction result
        """
1584
        predictor = self._to_predictor()
1585
1586
        if num_iteration <= 0:
            num_iteration = self.best_iteration
wxchan's avatar
wxchan committed
1587
1588
        return predictor.predict(data, num_iteration, raw_score, pred_leaf, data_has_header, is_reshape)

Guolin Ke's avatar
Guolin Ke committed
1589
    def _to_predictor(self):
wxchan's avatar
wxchan committed
1590
        """Convert to predictor"""
Guolin Ke's avatar
Guolin Ke committed
1591
        predictor = _InnerPredictor(booster_handle=self.handle)
1592
        predictor.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
1593
1594
        return predictor

wxchan's avatar
wxchan committed
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
    def feature_name(self):
        """
        Get feature names.

        Returns
        -------
        result : array
            Array of feature names.
        """
        out_num_feature = ctypes.c_int(0)
        """Get num of features"""
        _safe_call(_LIB.LGBM_BoosterGetNumFeature(
            self.handle,
            ctypes.byref(out_num_feature)))
        num_feature = out_num_feature.value
        """Get name of features"""
        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)]

1622
    def feature_importance(self, importance_type='split'):
wxchan's avatar
wxchan committed
1623
        """
wxchan's avatar
wxchan committed
1624
        Get feature importances
1625

1626
1627
1628
1629
1630
1631
1632
        Parameters
        ----------
        importance_type : str, default "split"
        How the importance is calculated: "split" or "gain"
        "split" is the number of times a feature is used in a model
        "gain" is the total gain of splits which use the feature

1633
1634
        Returns
        -------
wxchan's avatar
wxchan committed
1635
1636
        result : array
            Array of feature importances.
1637
1638
1639
1640
1641
        """
        if importance_type not in ["split", "gain"]:
            raise KeyError("importance_type must be split or gain")
        dump_model = self.dump_model()
        ret = [0] * (dump_model["max_feature_idx"] + 1)
wxchan's avatar
wxchan committed
1642

1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
        def dfs(root):
            if "split_feature" in root:
                if importance_type == 'split':
                    ret[root["split_feature"]] += 1
                elif importance_type == 'gain':
                    ret[root["split_feature"]] += root["split_gain"]
                dfs(root["left_child"])
                dfs(root["right_child"])
        for tree in dump_model["tree_info"]:
            dfs(tree["tree_structure"])
        return np.array(ret)

wxchan's avatar
wxchan committed
1655
1656
    def __inner_eval(self, data_name, data_idx, feval=None):
        """
1657
        Evaulate training or validation data
wxchan's avatar
wxchan committed
1658
1659
        """
        if data_idx >= self.__num_dataset:
1660
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
1661
1662
1663
        self.__get_eval_info()
        ret = []
        if self.__num_inner_eval > 0:
wxchan's avatar
wxchan committed
1664
            result = np.array([0.0 for _ in range_(self.__num_inner_eval)], dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
1665
            tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1666
1667
            _safe_call(_LIB.LGBM_BoosterGetEval(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
1668
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
1669
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
1670
                result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
1671
            if tmp_out_len.value != self.__num_inner_eval:
1672
                raise ValueError("Wrong length of eval results")
wxchan's avatar
wxchan committed
1673
            for i in range_(self.__num_inner_eval):
wxchan's avatar
wxchan committed
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
                ret.append((data_name, self.__name_inner_eval[i], result[i], self.__higher_better_inner_eval[i]))
        if feval is not None:
            if data_idx == 0:
                cur_data = self.train_set
            else:
                cur_data = self.valid_sets[data_idx - 1]
            feval_ret = feval(self.__inner_predict(data_idx), cur_data)
            if isinstance(feval_ret, list):
                for eval_name, val, is_higher_better in feval_ret:
                    ret.append((data_name, eval_name, val, is_higher_better))
            else:
                eval_name, val, is_higher_better = feval_ret
                ret.append((data_name, eval_name, val, is_higher_better))
        return ret

    def __inner_predict(self, data_idx):
        """
        Predict for training and validation dataset
        """
        if data_idx >= self.__num_dataset:
1694
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
1695
1696
1697
1698
1699
1700
        if self.__inner_predict_buffer[data_idx] is None:
            if data_idx == 0:
                n_preds = self.train_set.num_data() * self.__num_class
            else:
                n_preds = self.valid_sets[data_idx - 1].num_data() * self.__num_class
            self.__inner_predict_buffer[data_idx] = \
wxchan's avatar
wxchan committed
1701
                np.array([0.0 for _ in range_(n_preds)], dtype=np.float64, copy=False)
wxchan's avatar
wxchan committed
1702
1703
1704
        """avoid to predict many time in one iteration"""
        if not self.__is_predicted_cur_iter[data_idx]:
            tmp_out_len = ctypes.c_int64(0)
Guolin Ke's avatar
Guolin Ke committed
1705
            data_ptr = self.__inner_predict_buffer[data_idx].ctypes.data_as(ctypes.POINTER(ctypes.c_double))
wxchan's avatar
wxchan committed
1706
1707
            _safe_call(_LIB.LGBM_BoosterGetPredict(
                self.handle,
Guolin Ke's avatar
Guolin Ke committed
1708
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
1709
1710
1711
                ctypes.byref(tmp_out_len),
                data_ptr))
            if tmp_out_len.value != len(self.__inner_predict_buffer[data_idx]):
1712
                raise ValueError("Wrong length of predict results for data %d" % (data_idx))
wxchan's avatar
wxchan committed
1713
1714
1715
1716
1717
1718
1719
1720
1721
            self.__is_predicted_cur_iter[data_idx] = True
        return self.__inner_predict_buffer[data_idx]

    def __get_eval_info(self):
        """
        Get inner evaluation count and names
        """
        if self.__need_reload_eval_info:
            self.__need_reload_eval_info = False
Guolin Ke's avatar
Guolin Ke committed
1722
            out_num_eval = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1723
1724
1725
1726
1727
1728
1729
            """Get num of inner evals"""
            _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:
                """Get name of evals"""
Guolin Ke's avatar
Guolin Ke committed
1730
                tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
1731
                string_buffers = [ctypes.create_string_buffer(255) for i in range_(self.__num_inner_eval)]
wxchan's avatar
wxchan committed
1732
                ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
wxchan's avatar
wxchan committed
1733
1734
1735
1736
1737
                _safe_call(_LIB.LGBM_BoosterGetEvalNames(
                    self.handle,
                    ctypes.byref(tmp_out_len),
                    ptr_string_buffers))
                if self.__num_inner_eval != tmp_out_len.value:
1738
                    raise ValueError("Length of eval names doesn't equal with num_evals")
1739
                self.__name_inner_eval = \
wxchan's avatar
wxchan committed
1740
                    [string_buffers[i].value.decode() for i in range_(self.__num_inner_eval)]
1741
1742
1743
                self.__higher_better_inner_eval = \
                    [name.startswith(('auc', 'ndcg')) for name in self.__name_inner_eval]

wxchan's avatar
wxchan committed
1744
    def attr(self, key):
wxchan's avatar
wxchan committed
1745
1746
        """
        Get attribute string from the Booster.
wxchan's avatar
wxchan committed
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757

        Parameters
        ----------
        key : str
            The key to get attribute from.

        Returns
        -------
        value : str
            The attribute value of the key, returns None if attribute do not exist.
        """
1758
        return self.__attr.get(key, None)
wxchan's avatar
wxchan committed
1759
1760

    def set_attr(self, **kwargs):
wxchan's avatar
wxchan committed
1761
1762
        """
        Set the attribute of the Booster.
wxchan's avatar
wxchan committed
1763
1764
1765
1766
1767
1768
1769
1770

        Parameters
        ----------
        **kwargs
            The attributes to set. Setting a value to None deletes an attribute.
        """
        for key, value in kwargs.items():
            if value is not None:
wxchan's avatar
wxchan committed
1771
                if not isinstance(value, string_type):
1772
                    raise ValueError("Set attr only accepts strings")
wxchan's avatar
wxchan committed
1773
1774
1775
                self.__attr[key] = value
            else:
                self.__attr.pop(key, None)