basic.py 183 KB
Newer Older
wxchan's avatar
wxchan committed
1
# coding: utf-8
2
"""Wrapper for C API of LightGBM."""
3
import abc
wxchan's avatar
wxchan committed
4
import ctypes
5
import inspect
6
import json
wxchan's avatar
wxchan committed
7
import warnings
8
from collections import OrderedDict
9
from copy import deepcopy
10
from enum import Enum
11
from functools import wraps
12
from os import SEEK_END, environ
13
14
from os.path import getsize
from pathlib import Path
15
from tempfile import NamedTemporaryFile
16
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
wxchan's avatar
wxchan committed
17
18
19
20

import numpy as np
import scipy.sparse

21
from .compat import PANDAS_INSTALLED, concat, dt_DataTable, pd_CategoricalDtype, pd_DataFrame, pd_Series
wxchan's avatar
wxchan committed
22
23
from .libpath import find_lib_path

24
25
26
if TYPE_CHECKING:
    from typing import Literal

27
28
29
30
31
32
33
34
35
__all__ = [
    'Booster',
    'Dataset',
    'LGBMDeprecationWarning',
    'LightGBMError',
    'register_logger',
    'Sequence',
]

36
_DatasetHandle = ctypes.c_void_p
37
38
39
40
_ctypes_int_ptr = Union[
    "ctypes._Pointer[ctypes.c_int32]",
    "ctypes._Pointer[ctypes.c_int64]"
]
41
42
43
44
_ctypes_int_array = Union[
    "ctypes.Array[ctypes._Pointer[ctypes.c_int32]]",
    "ctypes.Array[ctypes._Pointer[ctypes.c_int64]]"
]
45
46
47
48
49
50
51
52
_ctypes_float_ptr = Union[
    "ctypes._Pointer[ctypes.c_float]",
    "ctypes._Pointer[ctypes.c_double]"
]
_ctypes_float_array = Union[
    "ctypes.Array[ctypes._Pointer[ctypes.c_float]]",
    "ctypes.Array[ctypes._Pointer[ctypes.c_double]]"
]
53
_LGBM_EvalFunctionResultType = Tuple[str, float, bool]
54
_LGBM_BoosterBestScoreType = Dict[str, Dict[str, float]]
55
_LGBM_BoosterEvalMethodResultType = Tuple[str, str, float, bool]
56
57
_LGBM_CategoricalFeatureConfiguration = Union[List[str], List[int], "Literal['auto']"]
_LGBM_FeatureNameConfiguration = Union[List[str], "Literal['auto']"]
58
59
60
61
62
63
64
65
66
67
68
69
70
_LGBM_GroupType = Union[
    List[float],
    List[int],
    np.ndarray,
    pd_Series
]
_LGBM_InitScoreType = Union[
    List[float],
    List[List[float]],
    np.ndarray,
    pd_Series,
    pd_DataFrame,
]
71
72
73
74
75
76
77
78
79
80
81
_LGBM_TrainDataType = Union[
    str,
    Path,
    np.ndarray,
    pd_DataFrame,
    dt_DataTable,
    scipy.sparse.spmatrix,
    "Sequence",
    List["Sequence"],
    List[np.ndarray]
]
82
_LGBM_LabelType = Union[
83
84
    List[float],
    List[int],
85
86
87
88
    np.ndarray,
    pd_Series,
    pd_DataFrame
]
89
90
91
92
93
94
95
96
_LGBM_PredictDataType = Union[
    str,
    Path,
    np.ndarray,
    pd_DataFrame,
    dt_DataTable,
    scipy.sparse.spmatrix
]
97
98
99
100
101
102
_LGBM_WeightType = Union[
    List[float],
    List[int],
    np.ndarray,
    pd_Series
]
103
104
105
ZERO_THRESHOLD = 1e-35


106
107
108
109
def _is_zero(x: float) -> bool:
    return -ZERO_THRESHOLD <= x <= ZERO_THRESHOLD


110
def _get_sample_count(total_nrow: int, params: str) -> int:
111
112
113
    sample_cnt = ctypes.c_int(0)
    _safe_call(_LIB.LGBM_GetSampleCount(
        ctypes.c_int32(total_nrow),
114
        _c_str(params),
115
116
117
118
        ctypes.byref(sample_cnt),
    ))
    return sample_cnt.value

wxchan's avatar
wxchan committed
119

120
121
122
123
124
125
class _MissingType(Enum):
    NONE = 'None'
    NAN = 'NaN'
    ZERO = 'Zero'


126
class _DummyLogger:
127
    def info(self, msg: str) -> None:
128
129
        print(msg)

130
    def warning(self, msg: str) -> None:
131
132
133
        warnings.warn(msg, stacklevel=3)


134
135
136
_LOGGER: Any = _DummyLogger()
_INFO_METHOD_NAME = "info"
_WARNING_METHOD_NAME = "warning"
137
138


139
140
141
142
def _has_method(logger: Any, method_name: str) -> bool:
    return callable(getattr(logger, method_name, None))


143
144
145
def register_logger(
    logger: Any, info_method_name: str = "info", warning_method_name: str = "warning"
) -> None:
146
147
148
149
    """Register custom logger.

    Parameters
    ----------
150
    logger : Any
151
        Custom logger.
152
153
154
155
    info_method_name : str, optional (default="info")
        Method used to log info messages.
    warning_method_name : str, optional (default="warning")
        Method used to log warning messages.
156
    """
157
158
159
160
161
162
    if not _has_method(logger, info_method_name) or not _has_method(logger, warning_method_name):
        raise TypeError(
            f"Logger must provide '{info_method_name}' and '{warning_method_name}' method"
        )

    global _LOGGER, _INFO_METHOD_NAME, _WARNING_METHOD_NAME
163
    _LOGGER = logger
164
165
    _INFO_METHOD_NAME = info_method_name
    _WARNING_METHOD_NAME = warning_method_name
166
167


168
def _normalize_native_string(func: Callable[[str], None]) -> Callable[[str], None]:
169
    """Join log messages from native library which come by chunks."""
170
    msg_normalized: List[str] = []
171
172

    @wraps(func)
173
    def wrapper(msg: str) -> None:
174
175
176
177
178
179
180
181
182
183
184
        nonlocal msg_normalized
        if msg.strip() == '':
            msg = ''.join(msg_normalized)
            msg_normalized = []
            return func(msg)
        else:
            msg_normalized.append(msg)

    return wrapper


185
def _log_info(msg: str) -> None:
186
    getattr(_LOGGER, _INFO_METHOD_NAME)(msg)
187
188


189
def _log_warning(msg: str) -> None:
190
    getattr(_LOGGER, _WARNING_METHOD_NAME)(msg)
191
192
193


@_normalize_native_string
194
def _log_native(msg: str) -> None:
195
    getattr(_LOGGER, _INFO_METHOD_NAME)(msg)
196
197


198
def _log_callback(msg: bytes) -> None:
199
    """Redirect logs from native library into Python."""
200
    _log_native(str(msg.decode('utf-8')))
201
202


203
def _load_lib() -> ctypes.CDLL:
204
    """Load LightGBM library."""
wxchan's avatar
wxchan committed
205
206
207
    lib_path = find_lib_path()
    lib = ctypes.cdll.LoadLibrary(lib_path[0])
    lib.LGBM_GetLastError.restype = ctypes.c_char_p
208
    callback = ctypes.CFUNCTYPE(None, ctypes.c_char_p)
209
    lib.callback = callback(_log_callback)  # type: ignore[attr-defined]
210
    if lib.LGBM_RegisterLogCallback(lib.callback) != 0:
211
        raise LightGBMError(lib.LGBM_GetLastError().decode('utf-8'))
wxchan's avatar
wxchan committed
212
213
    return lib

wxchan's avatar
wxchan committed
214

215
216
217
218
219
220
221
# we don't need lib_lightgbm while building docs
_LIB: ctypes.CDLL
if environ.get('LIGHTGBM_BUILD_DOC', False):
    from unittest.mock import Mock  # isort: skip
    _LIB = Mock(ctypes.CDLL)  # type: ignore
else:
    _LIB = _load_lib()
wxchan's avatar
wxchan committed
222

wxchan's avatar
wxchan committed
223

224
_NUMERIC_TYPES = (int, float, bool)
225
_ArrayLike = Union[List, np.ndarray, pd_Series]
226
227


228
def _safe_call(ret: int) -> None:
229
230
    """Check the return value from C API call.

wxchan's avatar
wxchan committed
231
232
233
    Parameters
    ----------
    ret : int
234
        The return value from C API calls.
wxchan's avatar
wxchan committed
235
236
    """
    if ret != 0:
237
        raise LightGBMError(_LIB.LGBM_GetLastError().decode('utf-8'))
wxchan's avatar
wxchan committed
238

wxchan's avatar
wxchan committed
239

240
def _is_numeric(obj: Any) -> bool:
241
    """Check whether object is a number or not, include numpy number, etc."""
wxchan's avatar
wxchan committed
242
243
244
    try:
        float(obj)
        return True
wxchan's avatar
wxchan committed
245
246
247
    except (TypeError, ValueError):
        # TypeError: obj is not a string or a number
        # ValueError: invalid literal
wxchan's avatar
wxchan committed
248
249
        return False

wxchan's avatar
wxchan committed
250

251
def _is_numpy_1d_array(data: Any) -> bool:
252
    """Check whether data is a numpy 1-D array."""
253
    return isinstance(data, np.ndarray) and len(data.shape) == 1
wxchan's avatar
wxchan committed
254

wxchan's avatar
wxchan committed
255

256
def _is_numpy_column_array(data: Any) -> bool:
257
258
259
260
261
262
263
    """Check whether data is a column numpy array."""
    if not isinstance(data, np.ndarray):
        return False
    shape = data.shape
    return len(shape) == 2 and shape[1] == 1


264
def _cast_numpy_array_to_dtype(array: np.ndarray, dtype: "np.typing.DTypeLike") -> np.ndarray:
265
    """Cast numpy array to given dtype."""
266
267
268
269
270
    if array.dtype == dtype:
        return array
    return array.astype(dtype=dtype, copy=False)


271
def _is_1d_list(data: Any) -> bool:
272
    """Check whether data is a 1-D list."""
273
    return isinstance(data, list) and (not data or _is_numeric(data[0]))
wxchan's avatar
wxchan committed
274

wxchan's avatar
wxchan committed
275

276
277
278
def _is_1d_collection(data: Any) -> bool:
    """Check whether data is a 1-D collection."""
    return (
279
        _is_numpy_1d_array(data)
280
        or _is_numpy_column_array(data)
281
        or _is_1d_list(data)
282
283
284
285
        or isinstance(data, pd_Series)
    )


286
287
def _list_to_1d_numpy(
    data: Any,
288
289
    dtype: "np.typing.DTypeLike",
    name: str
290
) -> np.ndarray:
291
    """Convert data to numpy 1-D array."""
292
    if _is_numpy_1d_array(data):
293
        return _cast_numpy_array_to_dtype(data, dtype)
294
    elif _is_numpy_column_array(data):
295
296
        _log_warning('Converting column-vector to 1d array')
        array = data.ravel()
297
        return _cast_numpy_array_to_dtype(array, dtype)
298
    elif _is_1d_list(data):
wxchan's avatar
wxchan committed
299
        return np.array(data, dtype=dtype, copy=False)
300
    elif isinstance(data, pd_Series):
301
        _check_for_bad_pandas_dtypes(data.to_frame().dtypes)
302
        return np.array(data, dtype=dtype, copy=False)  # SparseArray should be supported as well
wxchan's avatar
wxchan committed
303
    else:
304
305
        raise TypeError(f"Wrong type({type(data).__name__}) for {name}.\n"
                        "It should be list, numpy 1-D array or pandas Series")
wxchan's avatar
wxchan committed
306

wxchan's avatar
wxchan committed
307

308
309
310
311
312
313
314
def _is_numpy_2d_array(data: Any) -> bool:
    """Check whether data is a numpy 2-D array."""
    return isinstance(data, np.ndarray) and len(data.shape) == 2 and data.shape[1] > 1


def _is_2d_list(data: Any) -> bool:
    """Check whether data is a 2-D list."""
315
    return isinstance(data, list) and len(data) > 0 and _is_1d_list(data[0])
316
317
318
319
320
321
322
323
324
325
326


def _is_2d_collection(data: Any) -> bool:
    """Check whether data is a 2-D collection."""
    return (
        _is_numpy_2d_array(data)
        or _is_2d_list(data)
        or isinstance(data, pd_DataFrame)
    )


327
328
def _data_to_2d_numpy(
    data: Any,
329
330
    dtype: "np.typing.DTypeLike",
    name: str
331
) -> np.ndarray:
332
333
    """Convert data to numpy 2-D array."""
    if _is_numpy_2d_array(data):
334
        return _cast_numpy_array_to_dtype(data, dtype)
335
336
337
    if _is_2d_list(data):
        return np.array(data, dtype=dtype)
    if isinstance(data, pd_DataFrame):
338
        _check_for_bad_pandas_dtypes(data.dtypes)
339
        return _cast_numpy_array_to_dtype(data.values, dtype)
340
341
342
343
    raise TypeError(f"Wrong type({type(data).__name__}) for {name}.\n"
                    "It should be list of lists, numpy 2-D array or pandas DataFrame")


344
def _cfloat32_array_to_numpy(cptr: "ctypes._Pointer", length: int) -> np.ndarray:
345
    """Convert a ctypes float pointer array to a numpy array."""
wxchan's avatar
wxchan committed
346
    if isinstance(cptr, ctypes.POINTER(ctypes.c_float)):
347
        return np.ctypeslib.as_array(cptr, shape=(length,)).copy()
wxchan's avatar
wxchan committed
348
    else:
349
        raise RuntimeError('Expected float pointer')
wxchan's avatar
wxchan committed
350

Guolin Ke's avatar
Guolin Ke committed
351

352
def _cfloat64_array_to_numpy(cptr: "ctypes._Pointer", length: int) -> np.ndarray:
353
    """Convert a ctypes double pointer array to a numpy array."""
Guolin Ke's avatar
Guolin Ke committed
354
    if isinstance(cptr, ctypes.POINTER(ctypes.c_double)):
355
        return np.ctypeslib.as_array(cptr, shape=(length,)).copy()
Guolin Ke's avatar
Guolin Ke committed
356
357
358
    else:
        raise RuntimeError('Expected double pointer')

wxchan's avatar
wxchan committed
359

360
def _cint32_array_to_numpy(cptr: "ctypes._Pointer", length: int) -> np.ndarray:
361
    """Convert a ctypes int pointer array to a numpy array."""
wxchan's avatar
wxchan committed
362
    if isinstance(cptr, ctypes.POINTER(ctypes.c_int32)):
363
        return np.ctypeslib.as_array(cptr, shape=(length,)).copy()
wxchan's avatar
wxchan committed
364
    else:
365
366
367
        raise RuntimeError('Expected int32 pointer')


368
def _cint64_array_to_numpy(cptr: "ctypes._Pointer", length: int) -> np.ndarray:
369
370
    """Convert a ctypes int pointer array to a numpy array."""
    if isinstance(cptr, ctypes.POINTER(ctypes.c_int64)):
371
        return np.ctypeslib.as_array(cptr, shape=(length,)).copy()
372
373
    else:
        raise RuntimeError('Expected int64 pointer')
wxchan's avatar
wxchan committed
374

wxchan's avatar
wxchan committed
375

376
def _c_str(string: str) -> ctypes.c_char_p:
377
    """Convert a Python string to C string."""
wxchan's avatar
wxchan committed
378
379
    return ctypes.c_char_p(string.encode('utf-8'))

wxchan's avatar
wxchan committed
380

381
def _c_array(ctype: type, values: List[Any]) -> ctypes.Array:
382
    """Convert a Python array to C array."""
383
    return (ctype * len(values))(*values)  # type: ignore[operator]
wxchan's avatar
wxchan committed
384

wxchan's avatar
wxchan committed
385

386
def _json_default_with_numpy(obj: Any) -> Any:
387
388
389
390
391
392
393
394
395
    """Convert numpy classes to JSON serializable objects."""
    if isinstance(obj, (np.integer, np.floating, np.bool_)):
        return obj.item()
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    else:
        return obj


396
397
398
399
400
401
402
403
def _to_string(x: Union[int, float, str, List]) -> str:
    if isinstance(x, list):
        val_list = ",".join(str(val) for val in x)
        return f"[{val_list}]"
    else:
        return str(x)


404
def _param_dict_to_str(data: Optional[Dict[str, Any]]) -> str:
405
    """Convert Python dictionary to string, which is passed to C API."""
406
    if data is None or not data:
wxchan's avatar
wxchan committed
407
408
409
        return ""
    pairs = []
    for key, val in data.items():
410
        if isinstance(val, (list, tuple, set)) or _is_numpy_1d_array(val):
411
            pairs.append(f"{key}={','.join(map(_to_string, val))}")
412
        elif isinstance(val, (str, Path, _NUMERIC_TYPES)) or _is_numeric(val):
413
            pairs.append(f"{key}={val}")
414
        elif val is not None:
415
            raise TypeError(f'Unknown type of parameter:{key}, got:{type(val).__name__}')
wxchan's avatar
wxchan committed
416
    return ' '.join(pairs)
417

wxchan's avatar
wxchan committed
418

419
class _TempFile:
420
421
    """Proxy class to workaround errors on Windows."""

422
423
424
    def __enter__(self):
        with NamedTemporaryFile(prefix="lightgbm_tmp_", delete=True) as f:
            self.name = f.name
425
            self.path = Path(self.name)
426
        return self
wxchan's avatar
wxchan committed
427

428
    def __exit__(self, exc_type, exc_val, exc_tb):
429
430
        if self.path.is_file():
            self.path.unlink()
431

wxchan's avatar
wxchan committed
432

433
class LightGBMError(Exception):
434
435
    """Error thrown by LightGBM."""

436
437
438
    pass


439
440
441
442
443
444
445
446
# DeprecationWarning is not shown by default, so let's create our own with higher level
class LGBMDeprecationWarning(UserWarning):
    """Custom deprecation warning."""

    pass


class _ConfigAliases:
447
448
449
450
    # lazy evaluation to allow import without dynamic library, e.g., for docs generation
    aliases = None

    @staticmethod
451
    def _get_all_param_aliases() -> Dict[str, List[str]]:
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
        buffer_len = 1 << 20
        tmp_out_len = ctypes.c_int64(0)
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_DumpParamAliases(
            ctypes.c_int64(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_DumpParamAliases(
                ctypes.c_int64(actual_len),
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
        aliases = json.loads(
            string_buffer.value.decode('utf-8'),
471
            object_hook=lambda obj: {k: [k] + v for k, v in obj.items()}
472
473
        )
        return aliases
474
475

    @classmethod
476
477
478
    def get(cls, *args) -> Set[str]:
        if cls.aliases is None:
            cls.aliases = cls._get_all_param_aliases()
479
480
        ret = set()
        for i in args:
481
            ret.update(cls.get_sorted(i))
482
483
        return ret

484
485
486
487
488
489
    @classmethod
    def get_sorted(cls, name: str) -> List[str]:
        if cls.aliases is None:
            cls.aliases = cls._get_all_param_aliases()
        return cls.aliases.get(name, [name])

490
    @classmethod
491
492
493
    def get_by_alias(cls, *args) -> Set[str]:
        if cls.aliases is None:
            cls.aliases = cls._get_all_param_aliases()
494
495
496
497
        ret = set(args)
        for arg in args:
            for aliases in cls.aliases.values():
                if arg in aliases:
498
                    ret.update(aliases)
499
500
501
                    break
        return ret

502

503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
def _choose_param_value(main_param_name: str, params: Dict[str, Any], default_value: Any) -> Dict[str, Any]:
    """Get a single parameter value, accounting for aliases.

    Parameters
    ----------
    main_param_name : str
        Name of the main parameter to get a value for. One of the keys of ``_ConfigAliases``.
    params : dict
        Dictionary of LightGBM parameters.
    default_value : Any
        Default value to use for the parameter, if none is found in ``params``.

    Returns
    -------
    params : dict
        A ``params`` dict with exactly one value for ``main_param_name``, and all aliases ``main_param_name`` removed.
        If both ``main_param_name`` and one or more aliases for it are found, the value of ``main_param_name`` will be preferred.
    """
    # avoid side effects on passed-in parameters
    params = deepcopy(params)

524
525
    aliases = _ConfigAliases.get_sorted(main_param_name)
    aliases = [a for a in aliases if a != main_param_name]
526
527

    # if main_param_name was provided, keep that value and remove all aliases
528
    if main_param_name in params.keys():
529
530
531
        for param in aliases:
            params.pop(param, None)
        return params
532

533
534
535
536
537
    # if main param name was not found, search for an alias
    for param in aliases:
        if param in params.keys():
            params[main_param_name] = params[param]
            break
538

539
540
541
542
543
544
545
    if main_param_name in params.keys():
        for param in aliases:
            params.pop(param, None)
        return params

    # neither of main_param_name, aliases were found
    params[main_param_name] = default_value
546
547
548
549

    return params


550
_MAX_INT32 = (1 << 31) - 1
551

552
"""Macro definition of data type in C API of LightGBM"""
553
554
555
556
_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
557

558
"""Matrix is row major in Python"""
559
_C_API_IS_ROW_MAJOR = 1
wxchan's avatar
wxchan committed
560

561
"""Macro definition of prediction type in C API of LightGBM"""
562
563
564
565
_C_API_PREDICT_NORMAL = 0
_C_API_PREDICT_RAW_SCORE = 1
_C_API_PREDICT_LEAF_INDEX = 2
_C_API_PREDICT_CONTRIB = 3
wxchan's avatar
wxchan committed
566

567
"""Macro definition of sparse matrix type"""
568
569
_C_API_MATRIX_TYPE_CSR = 0
_C_API_MATRIX_TYPE_CSC = 1
570

571
"""Macro definition of feature importance type"""
572
573
_C_API_FEATURE_IMPORTANCE_SPLIT = 0
_C_API_FEATURE_IMPORTANCE_GAIN = 1
574

575
"""Data type of data field"""
576
577
578
579
580
581
_FIELD_TYPE_MAPPER = {
    "label": _C_API_DTYPE_FLOAT32,
    "weight": _C_API_DTYPE_FLOAT32,
    "init_score": _C_API_DTYPE_FLOAT64,
    "group": _C_API_DTYPE_INT32
}
wxchan's avatar
wxchan committed
582

583
"""String name to int feature importance type mapper"""
584
585
586
587
_FEATURE_IMPORTANCE_TYPE_MAPPER = {
    "split": _C_API_FEATURE_IMPORTANCE_SPLIT,
    "gain": _C_API_FEATURE_IMPORTANCE_GAIN
}
588

wxchan's avatar
wxchan committed
589

590
def _convert_from_sliced_object(data: np.ndarray) -> np.ndarray:
591
    """Fix the memory of multi-dimensional sliced object."""
592
    if isinstance(data, np.ndarray) and isinstance(data.base, np.ndarray):
593
        if not data.flags.c_contiguous:
594
595
            _log_warning("Usage of np.ndarray subset (sliced data) is not recommended "
                         "due to it will double the peak memory cost in LightGBM.")
596
597
598
599
            return np.copy(data)
    return data


600
601
602
def _c_float_array(
    data: np.ndarray
) -> Tuple[_ctypes_float_ptr, int, np.ndarray]:
603
    """Get pointer of float numpy array / list."""
604
    if _is_1d_list(data):
wxchan's avatar
wxchan committed
605
        data = np.array(data, copy=False)
606
    if _is_numpy_1d_array(data):
607
        data = _convert_from_sliced_object(data)
608
        assert data.flags.c_contiguous
609
        ptr_data: _ctypes_float_ptr
wxchan's avatar
wxchan committed
610
611
        if data.dtype == np.float32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
612
            type_data = _C_API_DTYPE_FLOAT32
wxchan's avatar
wxchan committed
613
614
        elif data.dtype == np.float64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
615
            type_data = _C_API_DTYPE_FLOAT64
wxchan's avatar
wxchan committed
616
        else:
617
            raise TypeError(f"Expected np.float32 or np.float64, met type({data.dtype})")
wxchan's avatar
wxchan committed
618
    else:
619
        raise TypeError(f"Unknown type({type(data).__name__})")
620
    return (ptr_data, type_data, data)  # return `data` to avoid the temporary copy is freed
wxchan's avatar
wxchan committed
621

wxchan's avatar
wxchan committed
622

623
624
625
def _c_int_array(
    data: np.ndarray
) -> Tuple[_ctypes_int_ptr, int, np.ndarray]:
626
    """Get pointer of int numpy array / list."""
627
    if _is_1d_list(data):
wxchan's avatar
wxchan committed
628
        data = np.array(data, copy=False)
629
    if _is_numpy_1d_array(data):
630
        data = _convert_from_sliced_object(data)
631
        assert data.flags.c_contiguous
632
        ptr_data: _ctypes_int_ptr
wxchan's avatar
wxchan committed
633
634
        if data.dtype == np.int32:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
635
            type_data = _C_API_DTYPE_INT32
wxchan's avatar
wxchan committed
636
637
        elif data.dtype == np.int64:
            ptr_data = data.ctypes.data_as(ctypes.POINTER(ctypes.c_int64))
638
            type_data = _C_API_DTYPE_INT64
wxchan's avatar
wxchan committed
639
        else:
640
            raise TypeError(f"Expected np.int32 or np.int64, met type({data.dtype})")
wxchan's avatar
wxchan committed
641
    else:
642
        raise TypeError(f"Unknown type({type(data).__name__})")
643
    return (ptr_data, type_data, data)  # return `data` to avoid the temporary copy is freed
wxchan's avatar
wxchan committed
644

wxchan's avatar
wxchan committed
645

646
def _is_allowed_numpy_dtype(dtype: type) -> bool:
647
    float128 = getattr(np, 'float128', type(None))
648
649
650
651
    return (
        issubclass(dtype, (np.integer, np.floating, np.bool_))
        and not issubclass(dtype, (np.timedelta64, float128))
    )
652
653


654
def _check_for_bad_pandas_dtypes(pandas_dtypes_series: pd_Series) -> None:
655
656
    bad_pandas_dtypes = [
        f'{column_name}: {pandas_dtype}'
657
        for column_name, pandas_dtype in pandas_dtypes_series.items()
658
        if not _is_allowed_numpy_dtype(pandas_dtype.type)
659
660
661
662
    ]
    if bad_pandas_dtypes:
        raise ValueError('pandas dtypes must be int, float or bool.\n'
                         f'Fields with bad pandas dtypes: {", ".join(bad_pandas_dtypes)}')
663
664


665
666
667
668
669
670
def _data_from_pandas(
    data,
    feature_name: Optional[_LGBM_FeatureNameConfiguration],
    categorical_feature: Optional[_LGBM_CategoricalFeatureConfiguration],
    pandas_categorical: Optional[List[List]]
):
671
    if isinstance(data, pd_DataFrame):
672
673
        if len(data.shape) != 2 or data.shape[0] < 1:
            raise ValueError('Input data must be 2 dimensional and non empty.')
674
        if feature_name == 'auto' or feature_name is None:
675
            data = data.rename(columns=str, copy=False)
676
        cat_cols = [col for col, dtype in zip(data.columns, data.dtypes) if isinstance(dtype, pd_CategoricalDtype)]
677
        cat_cols_not_ordered = [col for col in cat_cols if not data[col].cat.ordered]
678
679
680
681
682
        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.')
683
            for col, category in zip(cat_cols, pandas_categorical):
684
685
                if list(data[col].cat.categories) != list(category):
                    data[col] = data[col].cat.set_categories(category)
686
        if len(cat_cols):  # cat_cols is list
687
            data = data.copy(deep=False)  # not alter origin DataFrame
688
            data[cat_cols] = data[cat_cols].apply(lambda x: x.cat.codes).replace({-1: np.nan})
689
690
691
        if categorical_feature is not None:
            if feature_name is None:
                feature_name = list(data.columns)
692
            if categorical_feature == 'auto':  # use cat cols from DataFrame
693
                categorical_feature = cat_cols_not_ordered
694
            else:  # use cat cols specified by user
695
                categorical_feature = list(categorical_feature)  # type: ignore[assignment]
696
697
        if feature_name == 'auto':
            feature_name = list(data.columns)
698
        _check_for_bad_pandas_dtypes(data.dtypes)
699
700
701
        df_dtypes = [dtype.type for dtype in data.dtypes]
        df_dtypes.append(np.float32)  # so that the target dtype considers floats
        target_dtype = np.find_common_type(df_dtypes, [])
702
703
704
705
706
707
708
709
710
711
        try:
            # most common case (no nullable dtypes)
            data = data.to_numpy(dtype=target_dtype, copy=False)
        except TypeError:
            # 1.0 <= pd version < 1.1 and nullable dtypes, least common case
            # raises error because array is casted to type(pd.NA) and there's no na_value argument
            data = data.astype(target_dtype, copy=False).values
        except ValueError:
            # data has nullable dtypes, but we can specify na_value argument and copy will be made
            data = data.to_numpy(dtype=target_dtype, na_value=np.nan)
712
713
714
715
716
717
    else:
        if feature_name == 'auto':
            feature_name = None
        if categorical_feature == 'auto':
            categorical_feature = None
    return data, feature_name, categorical_feature, pandas_categorical
718
719


720
721
722
723
def _dump_pandas_categorical(
    pandas_categorical: Optional[List[List]],
    file_name: Optional[Union[str, Path]] = None
) -> str:
724
    categorical_json = json.dumps(pandas_categorical, default=_json_default_with_numpy)
725
    pandas_str = f'\npandas_categorical:{categorical_json}\n'
726
727
728
729
730
731
    if file_name is not None:
        with open(file_name, 'a') as f:
            f.write(pandas_str)
    return pandas_str


732
733
734
def _load_pandas_categorical(
    file_name: Optional[Union[str, Path]] = None,
    model_str: Optional[str] = None
735
) -> Optional[List[List]]:
736
737
    pandas_key = 'pandas_categorical:'
    offset = -len(pandas_key)
738
    if file_name is not None:
739
        max_offset = -getsize(file_name)
740
741
742
743
        with open(file_name, 'rb') as f:
            while True:
                if offset < max_offset:
                    offset = max_offset
744
                f.seek(offset, SEEK_END)
745
746
747
748
                lines = f.readlines()
                if len(lines) >= 2:
                    break
                offset *= 2
749
        last_line = lines[-1].decode('utf-8').strip()
750
        if not last_line.startswith(pandas_key):
751
            last_line = lines[-2].decode('utf-8').strip()
752
    elif model_str is not None:
753
754
755
756
757
758
        idx = model_str.rfind('\n', 0, offset)
        last_line = model_str[idx:].strip()
    if last_line.startswith(pandas_key):
        return json.loads(last_line[len(pandas_key):])
    else:
        return None
759
760


761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
class Sequence(abc.ABC):
    """
    Generic data access interface.

    Object should support the following operations:

    .. code-block::

        # Get total row number.
        >>> len(seq)
        # Random access by row index. Used for data sampling.
        >>> seq[10]
        # Range data access. Used to read data in batch when constructing Dataset.
        >>> seq[0:100]
        # Optionally specify batch_size to control range data read size.
        >>> seq.batch_size

    - With random access, **data sampling does not need to go through all data**.
    - With range data access, there's **no need to read all data into memory thus reduce memory usage**.

781
782
    .. versionadded:: 3.3.0

783
784
785
786
787
788
789
790
791
    Attributes
    ----------
    batch_size : int
        Default size of a batch.
    """

    batch_size = 4096  # Defaults to read 4K rows in each batch.

    @abc.abstractmethod
792
    def __getitem__(self, idx: Union[int, slice, List[int]]) -> np.ndarray:
793
794
795
796
797
798
799
        """Return data for given row index.

        A basic implementation should look like this:

        .. code-block:: python

            if isinstance(idx, numbers.Integral):
800
                return self._get_one_line(idx)
801
            elif isinstance(idx, slice):
802
803
                return np.stack([self._get_one_line(i) for i in range(idx.start, idx.stop)])
            elif isinstance(idx, list):
804
                # Only required if using ``Dataset.subset()``.
805
                return np.array([self._get_one_line(i) for i in idx])
806
            else:
807
                raise TypeError(f"Sequence index must be integer, slice or list, got {type(idx).__name__}")
808
809
810

        Parameters
        ----------
811
        idx : int, slice[int], list[int]
812
813
814
815
            Item index.

        Returns
        -------
816
        result : numpy 1-D array or numpy 2-D array
817
            1-D array if idx is int, 2-D array if idx is slice or list.
818
819
820
821
822
823
824
825
826
        """
        raise NotImplementedError("Sub-classes of lightgbm.Sequence must implement __getitem__()")

    @abc.abstractmethod
    def __len__(self) -> int:
        """Return row count of this sequence."""
        raise NotImplementedError("Sub-classes of lightgbm.Sequence must implement __len__()")


827
class _InnerPredictor:
828
829
830
831
832
    """_InnerPredictor of LightGBM.

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

Nikita Titov's avatar
Nikita Titov committed
833
834
835
    .. note::

        Can be converted from Booster, but cannot be converted to Booster.
Guolin Ke's avatar
Guolin Ke committed
836
    """
837

838
839
840
841
842
843
    def __init__(
        self,
        model_file: Optional[Union[str, Path]] = None,
        booster_handle: Optional[ctypes.c_void_p] = None,
        pred_parameter: Optional[Dict[str, Any]] = None
    ):
844
        """Initialize the _InnerPredictor.
wxchan's avatar
wxchan committed
845
846
847

        Parameters
        ----------
848
        model_file : str, pathlib.Path or None, optional (default=None)
wxchan's avatar
wxchan committed
849
            Path to the model file.
850
851
852
        booster_handle : object or None, optional (default=None)
            Handle of Booster.
        pred_parameter: dict or None, optional (default=None)
853
            Other parameters for the prediction.
wxchan's avatar
wxchan committed
854
        """
855
        self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
856
857
858
        self.__is_manage_handle = True
        if model_file is not None:
            """Prediction task"""
Guolin Ke's avatar
Guolin Ke committed
859
            out_num_iterations = ctypes.c_int(0)
wxchan's avatar
wxchan committed
860
            _safe_call(_LIB.LGBM_BoosterCreateFromModelfile(
861
                _c_str(str(model_file)),
wxchan's avatar
wxchan committed
862
                ctypes.byref(out_num_iterations),
863
                ctypes.byref(self._handle)))
Guolin Ke's avatar
Guolin Ke committed
864
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
865
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
866
                self._handle,
wxchan's avatar
wxchan committed
867
868
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
869
            self.num_total_iteration = out_num_iterations.value
870
            self.pandas_categorical = _load_pandas_categorical(file_name=model_file)
wxchan's avatar
wxchan committed
871
        elif booster_handle is not None:
Guolin Ke's avatar
Guolin Ke committed
872
            self.__is_manage_handle = False
873
            self._handle = booster_handle
Guolin Ke's avatar
Guolin Ke committed
874
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
875
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
876
                self._handle,
wxchan's avatar
wxchan committed
877
878
                ctypes.byref(out_num_class)))
            self.num_class = out_num_class.value
879
            self.num_total_iteration = self.current_iteration()
880
            self.pandas_categorical = None
wxchan's avatar
wxchan committed
881
        else:
882
            raise TypeError('Need model_file or booster_handle to create a predictor')
wxchan's avatar
wxchan committed
883

884
        pred_parameter = {} if pred_parameter is None else pred_parameter
885
        self.pred_parameter = _param_dict_to_str(pred_parameter)
cbecker's avatar
cbecker committed
886

887
    def __del__(self) -> None:
888
889
        try:
            if self.__is_manage_handle:
890
                _safe_call(_LIB.LGBM_BoosterFree(self._handle))
891
892
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
893

894
    def __getstate__(self) -> Dict[str, Any]:
895
896
        this = self.__dict__.copy()
        this.pop('handle', None)
897
        this.pop('_handle', None)
898
899
        return this

900
901
    def predict(
        self,
902
        data: _LGBM_PredictDataType,
903
904
905
906
907
908
909
        start_iteration: int = 0,
        num_iteration: int = -1,
        raw_score: bool = False,
        pred_leaf: bool = False,
        pred_contrib: bool = False,
        data_has_header: bool = False,
        validate_features: bool = False
910
    ) -> Union[np.ndarray, scipy.sparse.spmatrix, List[scipy.sparse.spmatrix]]:
911
        """Predict logic.
wxchan's avatar
wxchan committed
912
913
914

        Parameters
        ----------
915
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame or scipy.sparse
916
            Data source for prediction.
917
            If str or pathlib.Path, it represents the path to a text file (CSV, TSV, or LibSVM).
918
919
        start_iteration : int, optional (default=0)
            Start index of the iteration to predict.
920
921
922
923
924
925
926
927
928
929
930
        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.
931
932
933
        validate_features : bool, optional (default=False)
            If True, ensure that the features used to predict match the ones used to train.
            Used only if data is pandas DataFrame.
wxchan's avatar
wxchan committed
934

935
936
            .. versionadded:: 4.0.0

wxchan's avatar
wxchan committed
937
938
        Returns
        -------
939
        result : numpy array, scipy.sparse or list of scipy.sparse
940
            Prediction result.
941
            Can be sparse or a list of sparse objects (each element represents predictions for one class) for feature contributions (when ``pred_contrib=True``).
wxchan's avatar
wxchan committed
942
        """
wxchan's avatar
wxchan committed
943
        if isinstance(data, Dataset):
944
            raise TypeError("Cannot use Dataset instance for prediction, please use raw data instead")
945
946
947
948
949
950
        elif isinstance(data, pd_DataFrame) and validate_features:
            data_names = [str(x) for x in data.columns]
            ptr_names = (ctypes.c_char_p * len(data_names))()
            ptr_names[:] = [x.encode('utf-8') for x in data_names]
            _safe_call(
                _LIB.LGBM_BoosterValidateFeatureNames(
951
                    self._handle,
952
953
954
955
                    ptr_names,
                    ctypes.c_int(len(data_names)),
                )
            )
956
        data = _data_from_pandas(data, None, None, self.pandas_categorical)[0]
957
        predict_type = _C_API_PREDICT_NORMAL
wxchan's avatar
wxchan committed
958
        if raw_score:
959
            predict_type = _C_API_PREDICT_RAW_SCORE
wxchan's avatar
wxchan committed
960
        if pred_leaf:
961
            predict_type = _C_API_PREDICT_LEAF_INDEX
962
        if pred_contrib:
963
            predict_type = _C_API_PREDICT_CONTRIB
wxchan's avatar
wxchan committed
964
        int_data_has_header = 1 if data_has_header else 0
cbecker's avatar
cbecker committed
965

966
        if isinstance(data, (str, Path)):
967
            with _TempFile() as f:
wxchan's avatar
wxchan committed
968
                _safe_call(_LIB.LGBM_BoosterPredictForFile(
969
                    self._handle,
970
                    _c_str(str(data)),
Guolin Ke's avatar
Guolin Ke committed
971
972
                    ctypes.c_int(int_data_has_header),
                    ctypes.c_int(predict_type),
973
                    ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
974
                    ctypes.c_int(num_iteration),
975
976
                    _c_str(self.pred_parameter),
                    _c_str(f.name)))
977
978
                preds = np.loadtxt(f.name, dtype=np.float64)
                nrow = preds.shape[0]
wxchan's avatar
wxchan committed
979
        elif isinstance(data, scipy.sparse.csr_matrix):
980
981
982
983
984
985
            preds, nrow = self.__pred_for_csr(
                csr=data,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
Guolin Ke's avatar
Guolin Ke committed
986
        elif isinstance(data, scipy.sparse.csc_matrix):
987
988
989
990
991
992
            preds, nrow = self.__pred_for_csc(
                csc=data,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
wxchan's avatar
wxchan committed
993
        elif isinstance(data, np.ndarray):
994
995
996
997
998
999
            preds, nrow = self.__pred_for_np2d(
                mat=data,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
1000
1001
1002
        elif isinstance(data, list):
            try:
                data = np.array(data)
1003
1004
            except BaseException as err:
                raise ValueError('Cannot convert data list to numpy array.') from err
1005
1006
1007
1008
1009
1010
            preds, nrow = self.__pred_for_np2d(
                mat=data,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
1011
        elif isinstance(data, dt_DataTable):
1012
1013
1014
1015
1016
1017
            preds, nrow = self.__pred_for_np2d(
                mat=data.to_numpy(),
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
wxchan's avatar
wxchan committed
1018
1019
        else:
            try:
1020
                _log_warning('Converting data to scipy sparse matrix.')
wxchan's avatar
wxchan committed
1021
                csr = scipy.sparse.csr_matrix(data)
1022
1023
            except BaseException as err:
                raise TypeError(f'Cannot predict data for type {type(data).__name__}') from err
1024
1025
1026
1027
1028
1029
            preds, nrow = self.__pred_for_csr(
                csr=csr,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
wxchan's avatar
wxchan committed
1030
1031
        if pred_leaf:
            preds = preds.astype(np.int32)
1032
        is_sparse = isinstance(preds, scipy.sparse.spmatrix) or isinstance(preds, list)
1033
        if not is_sparse and preds.size != nrow:
wxchan's avatar
wxchan committed
1034
            if preds.size % nrow == 0:
1035
                preds = preds.reshape(nrow, -1)
wxchan's avatar
wxchan committed
1036
            else:
1037
                raise ValueError(f'Length of predict result ({preds.size}) cannot be divide nrow ({nrow})')
wxchan's avatar
wxchan committed
1038
1039
        return preds

1040
1041
1042
1043
1044
1045
1046
    def __get_num_preds(
        self,
        start_iteration: int,
        num_iteration: int,
        nrow: int,
        predict_type: int
    ) -> int:
1047
        """Get size of prediction result."""
1048
        if nrow > _MAX_INT32:
1049
            raise LightGBMError('LightGBM cannot perform prediction for data '
1050
                                f'with number of rows greater than MAX_INT32 ({_MAX_INT32}).\n'
1051
                                'You can split your data into chunks '
1052
                                'and then concatenate predictions for them')
Guolin Ke's avatar
Guolin Ke committed
1053
1054
        n_preds = ctypes.c_int64(0)
        _safe_call(_LIB.LGBM_BoosterCalcNumPredict(
1055
            self._handle,
Guolin Ke's avatar
Guolin Ke committed
1056
1057
            ctypes.c_int(nrow),
            ctypes.c_int(predict_type),
1058
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
1059
            ctypes.c_int(num_iteration),
Guolin Ke's avatar
Guolin Ke committed
1060
1061
            ctypes.byref(n_preds)))
        return n_preds.value
wxchan's avatar
wxchan committed
1062

1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
    def __inner_predict_np2d(
        self,
        mat: np.ndarray,
        start_iteration: int,
        num_iteration: int,
        predict_type: int,
        preds: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, int]:
        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)
1076
1077
1078
1079
1080
1081
        n_preds = self.__get_num_preds(
            start_iteration=start_iteration,
            num_iteration=num_iteration,
            nrow=mat.shape[0],
            predict_type=predict_type
        )
1082
1083
1084
1085
1086
1087
        if preds is None:
            preds = np.empty(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(
1088
            self._handle,
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
            ptr_data,
            ctypes.c_int(type_ptr_data),
            ctypes.c_int32(mat.shape[0]),
            ctypes.c_int32(mat.shape[1]),
            ctypes.c_int(_C_API_IS_ROW_MAJOR),
            ctypes.c_int(predict_type),
            ctypes.c_int(start_iteration),
            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]

    def __pred_for_np2d(
        self,
        mat: np.ndarray,
        start_iteration: int,
        num_iteration: int,
        predict_type: int
    ) -> Tuple[np.ndarray, int]:
1111
        """Predict for a 2-D numpy matrix."""
wxchan's avatar
wxchan committed
1112
        if len(mat.shape) != 2:
1113
            raise ValueError('Input numpy.ndarray or list must be 2 dimensional')
wxchan's avatar
wxchan committed
1114

1115
        nrow = mat.shape[0]
1116
1117
        if nrow > _MAX_INT32:
            sections = np.arange(start=_MAX_INT32, stop=nrow, step=_MAX_INT32)
1118
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
1119
            n_preds = [self.__get_num_preds(start_iteration, num_iteration, i, predict_type) for i in np.diff([0] + list(sections) + [nrow])]
1120
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
1121
            preds = np.empty(sum(n_preds), dtype=np.float64)
1122
1123
            for chunk, (start_idx_pred, end_idx_pred) in zip(np.array_split(mat, sections),
                                                             zip(n_preds_sections, n_preds_sections[1:])):
1124
                # avoid memory consumption by arrays concatenation operations
1125
1126
1127
1128
1129
1130
1131
                self.__inner_predict_np2d(
                    mat=chunk,
                    start_iteration=start_iteration,
                    num_iteration=num_iteration,
                    predict_type=predict_type,
                    preds=preds[start_idx_pred:end_idx_pred]
                )
1132
            return preds, nrow
wxchan's avatar
wxchan committed
1133
        else:
1134
1135
1136
1137
1138
1139
1140
            return self.__inner_predict_np2d(
                mat=mat,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type,
                preds=None
            )
wxchan's avatar
wxchan committed
1141

1142
1143
1144
    def __create_sparse_native(
        self,
        cs: Union[scipy.sparse.csc_matrix, scipy.sparse.csr_matrix],
1145
1146
1147
1148
1149
1150
        out_shape: np.ndarray,
        out_ptr_indptr: "ctypes._Pointer",
        out_ptr_indices: "ctypes._Pointer",
        out_ptr_data: "ctypes._Pointer",
        indptr_type: int,
        data_type: int,
1151
        is_csr: bool
1152
    ) -> Union[List[scipy.sparse.csc_matrix], List[scipy.sparse.csr_matrix]]:
1153
1154
1155
        # create numpy array from output arrays
        data_indices_len = out_shape[0]
        indptr_len = out_shape[1]
1156
        if indptr_type == _C_API_DTYPE_INT32:
1157
            out_indptr = _cint32_array_to_numpy(out_ptr_indptr, indptr_len)
1158
        elif indptr_type == _C_API_DTYPE_INT64:
1159
            out_indptr = _cint64_array_to_numpy(out_ptr_indptr, indptr_len)
1160
1161
        else:
            raise TypeError("Expected int32 or int64 type for indptr")
1162
        if data_type == _C_API_DTYPE_FLOAT32:
1163
            out_data = _cfloat32_array_to_numpy(out_ptr_data, data_indices_len)
1164
        elif data_type == _C_API_DTYPE_FLOAT64:
1165
            out_data = _cfloat64_array_to_numpy(out_ptr_data, data_indices_len)
1166
1167
        else:
            raise TypeError("Expected float32 or float64 type for data")
1168
        out_indices = _cint32_array_to_numpy(out_ptr_indices, data_indices_len)
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
        # break up indptr based on number of rows (note more than one matrix in multiclass case)
        per_class_indptr_shape = cs.indptr.shape[0]
        # for CSC there is extra column added
        if not is_csr:
            per_class_indptr_shape += 1
        out_indptr_arrays = np.split(out_indptr, out_indptr.shape[0] / per_class_indptr_shape)
        # reformat output into a csr or csc matrix or list of csr or csc matrices
        cs_output_matrices = []
        offset = 0
        for cs_indptr in out_indptr_arrays:
            matrix_indptr_len = cs_indptr[cs_indptr.shape[0] - 1]
            cs_indices = out_indices[offset + cs_indptr[0]:offset + matrix_indptr_len]
            cs_data = out_data[offset + cs_indptr[0]:offset + matrix_indptr_len]
            offset += matrix_indptr_len
            # same shape as input csr or csc matrix except extra column for expected value
            cs_shape = [cs.shape[0], cs.shape[1] + 1]
            # note: make sure we copy data as it will be deallocated next
            if is_csr:
                cs_output_matrices.append(scipy.sparse.csr_matrix((cs_data, cs_indices, cs_indptr), cs_shape))
            else:
                cs_output_matrices.append(scipy.sparse.csc_matrix((cs_data, cs_indices, cs_indptr), cs_shape))
        # free the temporary native indptr, indices, and data
        _safe_call(_LIB.LGBM_BoosterFreePredictSparse(out_ptr_indptr, out_ptr_indices, out_ptr_data,
                                                      ctypes.c_int(indptr_type), ctypes.c_int(data_type)))
        if len(cs_output_matrices) == 1:
            return cs_output_matrices[0]
        return cs_output_matrices

1197
1198
1199
1200
1201
1202
1203
1204
1205
    def __inner_predict_csr(
        self,
        csr: scipy.sparse.csr_matrix,
        start_iteration: int,
        num_iteration: int,
        predict_type: int,
        preds: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, int]:
        nrow = len(csr.indptr) - 1
1206
1207
1208
1209
1210
1211
        n_preds = self.__get_num_preds(
            start_iteration=start_iteration,
            num_iteration=num_iteration,
            nrow=nrow,
            predict_type=predict_type
        )
1212
1213
1214
1215
1216
        if preds is None:
            preds = np.empty(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)
wxchan's avatar
wxchan committed
1217

1218
1219
        ptr_indptr, type_ptr_indptr, _ = _c_int_array(csr.indptr)
        ptr_data, type_ptr_data, _ = _c_float_array(csr.data)
1220

1221
1222
1223
1224
        assert csr.shape[1] <= _MAX_INT32
        csr_indices = csr.indices.astype(np.int32, copy=False)

        _safe_call(_LIB.LGBM_BoosterPredictForCSR(
1225
            self._handle,
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
            ptr_indptr,
            ctypes.c_int(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(start_iteration),
            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

    def __inner_predict_csr_sparse(
        self,
        csr: scipy.sparse.csr_matrix,
        start_iteration: int,
        num_iteration: int,
        predict_type: int
1250
    ) -> Tuple[Union[List[scipy.sparse.csc_matrix], List[scipy.sparse.csr_matrix]], int]:
1251
1252
1253
1254
        ptr_indptr, type_ptr_indptr, __ = _c_int_array(csr.indptr)
        ptr_data, type_ptr_data, _ = _c_float_array(csr.data)
        csr_indices = csr.indices.astype(np.int32, copy=False)
        matrix_type = _C_API_MATRIX_TYPE_CSR
1255
        out_ptr_indptr: _ctypes_int_ptr
1256
1257
1258
1259
1260
        if type_ptr_indptr == _C_API_DTYPE_INT32:
            out_ptr_indptr = ctypes.POINTER(ctypes.c_int32)()
        else:
            out_ptr_indptr = ctypes.POINTER(ctypes.c_int64)()
        out_ptr_indices = ctypes.POINTER(ctypes.c_int32)()
1261
        out_ptr_data: _ctypes_float_ptr
1262
1263
1264
1265
1266
1267
        if type_ptr_data == _C_API_DTYPE_FLOAT32:
            out_ptr_data = ctypes.POINTER(ctypes.c_float)()
        else:
            out_ptr_data = ctypes.POINTER(ctypes.c_double)()
        out_shape = np.empty(2, dtype=np.int64)
        _safe_call(_LIB.LGBM_BoosterPredictSparseOutput(
1268
            self._handle,
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
            ptr_indptr,
            ctypes.c_int(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(start_iteration),
            ctypes.c_int(num_iteration),
            _c_str(self.pred_parameter),
            ctypes.c_int(matrix_type),
            out_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
            ctypes.byref(out_ptr_indptr),
            ctypes.byref(out_ptr_indices),
            ctypes.byref(out_ptr_data)))
        matrices = self.__create_sparse_native(
            cs=csr,
            out_shape=out_shape,
            out_ptr_indptr=out_ptr_indptr,
            out_ptr_indices=out_ptr_indices,
            out_ptr_data=out_ptr_data,
            indptr_type=type_ptr_indptr,
            data_type=type_ptr_data,
            is_csr=True
        )
        nrow = len(csr.indptr) - 1
        return matrices, nrow

1299
1300
1301
1302
1303
1304
1305
    def __pred_for_csr(
        self,
        csr: scipy.sparse.csr_matrix,
        start_iteration: int,
        num_iteration: int,
        predict_type: int
    ) -> Tuple[np.ndarray, int]:
1306
        """Predict for a CSR data."""
1307
        if predict_type == _C_API_PREDICT_CONTRIB:
1308
1309
1310
1311
1312
1313
            return self.__inner_predict_csr_sparse(
                csr=csr,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
1314
        nrow = len(csr.indptr) - 1
1315
1316
        if nrow > _MAX_INT32:
            sections = [0] + list(np.arange(start=_MAX_INT32, stop=nrow, step=_MAX_INT32)) + [nrow]
1317
            # __get_num_preds() cannot work with nrow > MAX_INT32, so calculate overall number of predictions piecemeal
1318
            n_preds = [self.__get_num_preds(start_iteration, num_iteration, i, predict_type) for i in np.diff(sections)]
1319
            n_preds_sections = np.array([0] + n_preds, dtype=np.intp).cumsum()
1320
            preds = np.empty(sum(n_preds), dtype=np.float64)
1321
1322
            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:])):
1323
                # avoid memory consumption by arrays concatenation operations
1324
1325
1326
1327
1328
1329
1330
                self.__inner_predict_csr(
                    csr=csr[start_idx:end_idx],
                    start_iteration=start_iteration,
                    num_iteration=num_iteration,
                    predict_type=predict_type,
                    preds=preds[start_idx_pred:end_idx_pred]
                )
1331
1332
            return preds, nrow
        else:
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
            return self.__inner_predict_csr(
                csr=csr,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type,
                preds=None
            )

    def __inner_predict_sparse_csc(
        self,
1343
1344
1345
1346
        csc: scipy.sparse.csc_matrix,
        start_iteration: int,
        num_iteration: int,
        predict_type: int
1347
1348
1349
1350
1351
    ):
        ptr_indptr, type_ptr_indptr, __ = _c_int_array(csc.indptr)
        ptr_data, type_ptr_data, _ = _c_float_array(csc.data)
        csc_indices = csc.indices.astype(np.int32, copy=False)
        matrix_type = _C_API_MATRIX_TYPE_CSC
1352
        out_ptr_indptr: _ctypes_int_ptr
1353
1354
1355
1356
1357
        if type_ptr_indptr == _C_API_DTYPE_INT32:
            out_ptr_indptr = ctypes.POINTER(ctypes.c_int32)()
        else:
            out_ptr_indptr = ctypes.POINTER(ctypes.c_int64)()
        out_ptr_indices = ctypes.POINTER(ctypes.c_int32)()
1358
        out_ptr_data: _ctypes_float_ptr
1359
1360
1361
1362
1363
1364
        if type_ptr_data == _C_API_DTYPE_FLOAT32:
            out_ptr_data = ctypes.POINTER(ctypes.c_float)()
        else:
            out_ptr_data = ctypes.POINTER(ctypes.c_double)()
        out_shape = np.empty(2, dtype=np.int64)
        _safe_call(_LIB.LGBM_BoosterPredictSparseOutput(
1365
            self._handle,
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
            ptr_indptr,
            ctypes.c_int(type_ptr_indptr),
            csc_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ptr_data,
            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(start_iteration),
            ctypes.c_int(num_iteration),
            _c_str(self.pred_parameter),
            ctypes.c_int(matrix_type),
            out_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_int64)),
            ctypes.byref(out_ptr_indptr),
            ctypes.byref(out_ptr_indices),
            ctypes.byref(out_ptr_data)))
        matrices = self.__create_sparse_native(
            cs=csc,
            out_shape=out_shape,
            out_ptr_indptr=out_ptr_indptr,
            out_ptr_indices=out_ptr_indices,
            out_ptr_data=out_ptr_data,
            indptr_type=type_ptr_indptr,
            data_type=type_ptr_data,
            is_csr=False
        )
        nrow = csc.shape[0]
        return matrices, nrow
Guolin Ke's avatar
Guolin Ke committed
1395

1396
1397
1398
1399
1400
1401
1402
    def __pred_for_csc(
        self,
        csc: scipy.sparse.csc_matrix,
        start_iteration: int,
        num_iteration: int,
        predict_type: int
    ) -> Tuple[np.ndarray, int]:
1403
        """Predict for a CSC data."""
Guolin Ke's avatar
Guolin Ke committed
1404
        nrow = csc.shape[0]
1405
        if nrow > _MAX_INT32:
1406
1407
1408
1409
1410
1411
            return self.__pred_for_csr(
                csr=csc.tocsr(),
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
1412
        if predict_type == _C_API_PREDICT_CONTRIB:
1413
1414
1415
1416
1417
1418
            return self.__inner_predict_sparse_csc(
                csc=csc,
                start_iteration=start_iteration,
                num_iteration=num_iteration,
                predict_type=predict_type
            )
1419
1420
1421
1422
1423
1424
        n_preds = self.__get_num_preds(
            start_iteration=start_iteration,
            num_iteration=num_iteration,
            nrow=nrow,
            predict_type=predict_type
        )
1425
        preds = np.empty(n_preds, dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
1426
1427
        out_num_preds = ctypes.c_int64(0)

1428
1429
        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
1430

1431
        assert csc.shape[0] <= _MAX_INT32
1432
        csc_indices = csc.indices.astype(np.int32, copy=False)
1433

Guolin Ke's avatar
Guolin Ke committed
1434
        _safe_call(_LIB.LGBM_BoosterPredictForCSC(
1435
            self._handle,
Guolin Ke's avatar
Guolin Ke committed
1436
            ptr_indptr,
1437
            ctypes.c_int(type_ptr_indptr),
1438
            csc_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
1439
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
1440
1441
1442
1443
1444
            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),
1445
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
1446
            ctypes.c_int(num_iteration),
1447
            _c_str(self.pred_parameter),
Guolin Ke's avatar
Guolin Ke committed
1448
            ctypes.byref(out_num_preds),
wxchan's avatar
wxchan committed
1449
            preds.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
1450
        if n_preds != out_num_preds.value:
1451
            raise ValueError("Wrong length for predict results")
wxchan's avatar
wxchan committed
1452
1453
        return preds, nrow

1454
    def current_iteration(self) -> int:
1455
1456
1457
1458
1459
1460
1461
1462
1463
        """Get the index of the current iteration.

        Returns
        -------
        cur_iter : int
            The index of the current iteration.
        """
        out_cur_iter = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
1464
            self._handle,
1465
1466
1467
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

wxchan's avatar
wxchan committed
1468

1469
class Dataset:
wxchan's avatar
wxchan committed
1470
    """Dataset in LightGBM."""
1471

1472
1473
    def __init__(
        self,
1474
        data: _LGBM_TrainDataType,
1475
        label: Optional[_LGBM_LabelType] = None,
1476
        reference: Optional["Dataset"] = None,
1477
1478
1479
        weight: Optional[_LGBM_WeightType] = None,
        group: Optional[_LGBM_GroupType] = None,
        init_score: Optional[_LGBM_InitScoreType] = None,
1480
1481
        feature_name: _LGBM_FeatureNameConfiguration = 'auto',
        categorical_feature: _LGBM_CategoricalFeatureConfiguration = 'auto',
1482
1483
1484
        params: Optional[Dict[str, Any]] = None,
        free_raw_data: bool = True
    ):
1485
        """Initialize Dataset.
1486

wxchan's avatar
wxchan committed
1487
1488
        Parameters
        ----------
1489
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse, Sequence, list of Sequence or list of numpy array
wxchan's avatar
wxchan committed
1490
            Data source of Dataset.
1491
            If str or pathlib.Path, it represents the path to a text file (CSV, TSV, or LibSVM) or a LightGBM Dataset binary file.
1492
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
1493
1494
1495
            Label of the data.
        reference : Dataset or None, optional (default=None)
            If this is Dataset for validation, training data should be used as reference.
1496
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
1497
            Weight for each instance. Weights should be non-negative.
1498
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
1499
1500
1501
            Group/query data.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
1502
1503
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
1504
        init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), or None, optional (default=None)
1505
            Init score for Dataset.
1506
        feature_name : list of str, or 'auto', optional (default="auto")
1507
1508
            Feature names.
            If 'auto' and data is pandas DataFrame, data columns names are used.
1509
        categorical_feature : list of str or int, or 'auto', optional (default="auto")
1510
1511
            Categorical features.
            If list of int, interpreted as indices.
1512
            If list of str, interpreted as feature names (need to specify ``feature_name`` as well).
1513
            If 'auto' and data is pandas DataFrame, pandas unordered categorical columns are used.
1514
            All values in categorical features will be cast to int32 and thus should be less than int32 max value (2147483647).
1515
            Large values could be memory consuming. Consider using consecutive integers starting from zero.
1516
            All negative values in categorical features will be treated as missing values.
1517
            The output cannot be monotonically constrained with respect to a categorical feature.
1518
            Floating point numbers in categorical features will be rounded towards 0.
Nikita Titov's avatar
Nikita Titov committed
1519
        params : dict or None, optional (default=None)
1520
            Other parameters for Dataset.
Nikita Titov's avatar
Nikita Titov committed
1521
        free_raw_data : bool, optional (default=True)
1522
            If True, raw data is freed after constructing inner Dataset.
wxchan's avatar
wxchan committed
1523
        """
1524
        self._handle: Optional[_DatasetHandle] = None
wxchan's avatar
wxchan committed
1525
1526
1527
1528
1529
        self.data = data
        self.label = label
        self.reference = reference
        self.weight = weight
        self.group = group
1530
        self.init_score = init_score
1531
1532
        self.feature_name: _LGBM_FeatureNameConfiguration = feature_name
        self.categorical_feature: _LGBM_CategoricalFeatureConfiguration = categorical_feature
1533
        self.params = deepcopy(params)
wxchan's avatar
wxchan committed
1534
        self.free_raw_data = free_raw_data
1535
        self.used_indices: Optional[List[int]] = None
1536
        self._need_slice = True
1537
        self._predictor: Optional[_InnerPredictor] = None
1538
        self.pandas_categorical = None
1539
        self._params_back_up = None
1540
        self.version = 0
1541
        self._start_row = 0  # Used when pushing rows one by one.
wxchan's avatar
wxchan committed
1542

1543
    def __del__(self) -> None:
1544
1545
1546
1547
        try:
            self._free_handle()
        except AttributeError:
            pass
1548

1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
    def _create_sample_indices(self, total_nrow: int) -> np.ndarray:
        """Get an array of randomly chosen indices from this ``Dataset``.

        Indices are sampled without replacement.

        Parameters
        ----------
        total_nrow : int
            Total number of rows to sample from.
            If this value is greater than the value of parameter ``bin_construct_sample_cnt``, only ``bin_construct_sample_cnt`` indices will be used.
            If Dataset has multiple input data, this should be the sum of rows of every file.

        Returns
        -------
        indices : numpy array
            Indices for sampled data.
        """
1566
        param_str = _param_dict_to_str(self.get_params())
1567
1568
        sample_cnt = _get_sample_count(total_nrow, param_str)
        indices = np.empty(sample_cnt, dtype=np.int32)
1569
        ptr_data, _, _ = _c_int_array(indices)
1570
1571
1572
1573
        actual_sample_cnt = ctypes.c_int32(0)

        _safe_call(_LIB.LGBM_SampleIndices(
            ctypes.c_int32(total_nrow),
1574
            _c_str(param_str),
1575
1576
1577
            ptr_data,
            ctypes.byref(actual_sample_cnt),
        ))
1578
1579
        assert sample_cnt == actual_sample_cnt.value
        return indices
1580

1581
1582
1583
1584
1585
    def _init_from_ref_dataset(
        self,
        total_nrow: int,
        ref_dataset: _DatasetHandle
    ) -> 'Dataset':
1586
1587
1588
1589
1590
1591
        """Create dataset from a reference dataset.

        Parameters
        ----------
        total_nrow : int
            Number of rows expected to add to dataset.
1592
1593
        ref_dataset : object
            Handle of reference dataset to extract metadata from.
1594
1595
1596
1597
1598
1599

        Returns
        -------
        self : Dataset
            Constructed Dataset object.
        """
1600
        self._handle = ctypes.c_void_p()
1601
1602
1603
        _safe_call(_LIB.LGBM_DatasetCreateByReference(
            ref_dataset,
            ctypes.c_int64(total_nrow),
1604
            ctypes.byref(self._handle),
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
        ))
        return self

    def _init_from_sample(
        self,
        sample_data: List[np.ndarray],
        sample_indices: List[np.ndarray],
        sample_cnt: int,
        total_nrow: int,
    ) -> "Dataset":
        """Create Dataset from sampled data structures.

        Parameters
        ----------
1619
        sample_data : list of numpy array
1620
            Sample data for each column.
1621
        sample_indices : list of numpy array
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
            Sample data row index for each column.
        sample_cnt : int
            Number of samples.
        total_nrow : int
            Total number of rows for all input files.

        Returns
        -------
        self : Dataset
            Constructed Dataset object.
        """
        ncol = len(sample_indices)
        assert len(sample_data) == ncol, "#sample data column != #column indices"

        for i in range(ncol):
            if sample_data[i].dtype != np.double:
                raise ValueError(f"sample_data[{i}] type {sample_data[i].dtype} is not double")
            if sample_indices[i].dtype != np.int32:
                raise ValueError(f"sample_indices[{i}] type {sample_indices[i].dtype} is not int32")

        # c type: double**
        # each double* element points to start of each column of sample data.
1644
        sample_col_ptr: _ctypes_float_array = (ctypes.POINTER(ctypes.c_double) * ncol)()
1645
1646
        # c type int**
        # each int* points to start of indices for each column
1647
        indices_col_ptr: _ctypes_int_array = (ctypes.POINTER(ctypes.c_int32) * ncol)()
1648
        for i in range(ncol):
1649
1650
            sample_col_ptr[i] = _c_float_array(sample_data[i])[0]
            indices_col_ptr[i] = _c_int_array(sample_indices[i])[0]
1651
1652

        num_per_col = np.array([len(d) for d in sample_indices], dtype=np.int32)
1653
        num_per_col_ptr, _, _ = _c_int_array(num_per_col)
1654

1655
        self._handle = ctypes.c_void_p()
1656
        params_str = _param_dict_to_str(self.get_params())
1657
1658
1659
1660
1661
1662
1663
        _safe_call(_LIB.LGBM_DatasetCreateFromSampledColumn(
            ctypes.cast(sample_col_ptr, ctypes.POINTER(ctypes.POINTER(ctypes.c_double))),
            ctypes.cast(indices_col_ptr, ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))),
            ctypes.c_int32(ncol),
            num_per_col_ptr,
            ctypes.c_int32(sample_cnt),
            ctypes.c_int32(total_nrow),
1664
            ctypes.c_int64(total_nrow),
1665
            _c_str(params_str),
1666
            ctypes.byref(self._handle),
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
        ))
        return self

    def _push_rows(self, data: np.ndarray) -> 'Dataset':
        """Add rows to Dataset.

        Parameters
        ----------
        data : numpy 1-D array
            New data to add to the Dataset.

        Returns
        -------
        self : Dataset
            Dataset object.
        """
        nrow, ncol = data.shape
        data = data.reshape(data.size)
1685
        data_ptr, data_type, _ = _c_float_array(data)
1686
1687

        _safe_call(_LIB.LGBM_DatasetPushRows(
1688
            self._handle,
1689
1690
1691
1692
1693
1694
1695
1696
1697
            data_ptr,
            data_type,
            ctypes.c_int32(nrow),
            ctypes.c_int32(ncol),
            ctypes.c_int32(self._start_row),
        ))
        self._start_row += nrow
        return self

1698
    def get_params(self) -> Dict[str, Any]:
1699
1700
1701
1702
        """Get the used parameters in the Dataset.

        Returns
        -------
1703
        params : dict
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
            The used parameters in this Dataset object.
        """
        if self.params is not None:
            # no min_data, nthreads and verbose in this function
            dataset_params = _ConfigAliases.get("bin_construct_sample_cnt",
                                                "categorical_feature",
                                                "data_random_seed",
                                                "enable_bundle",
                                                "feature_pre_filter",
                                                "forcedbins_filename",
                                                "group_column",
                                                "header",
                                                "ignore_column",
                                                "is_enable_sparse",
                                                "label_column",
1719
                                                "linear_tree",
1720
1721
1722
1723
                                                "max_bin",
                                                "max_bin_by_feature",
                                                "min_data_in_bin",
                                                "pre_partition",
Nikita Titov's avatar
Nikita Titov committed
1724
                                                "precise_float_parser",
1725
1726
1727
1728
1729
                                                "two_round",
                                                "use_missing",
                                                "weight_column",
                                                "zero_as_missing")
            return {k: v for k, v in self.params.items() if k in dataset_params}
1730
1731
        else:
            return {}
1732

1733
    def _free_handle(self) -> "Dataset":
1734
1735
1736
        if self._handle is not None:
            _safe_call(_LIB.LGBM_DatasetFree(self._handle))
            self._handle = None
1737
        self._need_slice = True
Guolin Ke's avatar
Guolin Ke committed
1738
1739
        if self.used_indices is not None:
            self.data = None
Nikita Titov's avatar
Nikita Titov committed
1740
        return self
wxchan's avatar
wxchan committed
1741

1742
1743
1744
    def _set_init_score_by_predictor(
        self,
        predictor: Optional[_InnerPredictor],
1745
        data: _LGBM_TrainDataType,
1746
        used_indices: Optional[Union[List[int], np.ndarray]]
1747
    ) -> "Dataset":
Guolin Ke's avatar
Guolin Ke committed
1748
        data_has_header = False
1749
        if isinstance(data, (str, Path)) and self.params is not None:
Guolin Ke's avatar
Guolin Ke committed
1750
            # check data has header or not
1751
            data_has_header = any(self.params.get(alias, False) for alias in _ConfigAliases.get("header"))
Guolin Ke's avatar
Guolin Ke committed
1752
        num_data = self.num_data()
1753
        if predictor is not None:
1754
1755
1756
1757
1758
            init_score: Union[np.ndarray, scipy.sparse.spmatrix] = predictor.predict(
                data=data,
                raw_score=True,
                data_has_header=data_has_header
            )
1759
            init_score = init_score.ravel()
1760
            if used_indices is not None:
1761
                assert not self._need_slice
1762
                if isinstance(data, (str, Path)):
1763
                    sub_init_score = np.empty(num_data * predictor.num_class, dtype=np.float64)
1764
                    assert num_data == len(used_indices)
1765
1766
                    for i in range(len(used_indices)):
                        for j in range(predictor.num_class):
1767
1768
1769
1770
                            sub_init_score[i * predictor.num_class + j] = init_score[used_indices[i] * predictor.num_class + j]
                    init_score = sub_init_score
            if predictor.num_class > 1:
                # need to regroup init_score
1771
                new_init_score = np.empty(init_score.size, dtype=np.float64)
1772
1773
                for i in range(num_data):
                    for j in range(predictor.num_class):
1774
1775
1776
                        new_init_score[j * num_data + i] = init_score[i * predictor.num_class + j]
                init_score = new_init_score
        elif self.init_score is not None:
1777
            init_score = np.full_like(self.init_score, fill_value=0.0, dtype=np.float64)
1778
1779
        else:
            return self
Guolin Ke's avatar
Guolin Ke committed
1780
        self.set_init_score(init_score)
1781
        return self
Guolin Ke's avatar
Guolin Ke committed
1782

1783
1784
    def _lazy_init(
        self,
1785
        data: Optional[_LGBM_TrainDataType],
1786
1787
1788
1789
1790
1791
1792
1793
1794
        label: Optional[_LGBM_LabelType],
        reference: Optional["Dataset"],
        weight: Optional[_LGBM_WeightType],
        group: Optional[_LGBM_GroupType],
        init_score: Optional[_LGBM_InitScoreType],
        predictor: Optional[_InnerPredictor],
        feature_name: _LGBM_FeatureNameConfiguration,
        categorical_feature: _LGBM_CategoricalFeatureConfiguration,
        params: Optional[Dict[str, Any]]
1795
    ) -> "Dataset":
wxchan's avatar
wxchan committed
1796
        if data is None:
1797
            self._handle = None
Nikita Titov's avatar
Nikita Titov committed
1798
            return self
Guolin Ke's avatar
Guolin Ke committed
1799
1800
1801
        if reference is not None:
            self.pandas_categorical = reference.pandas_categorical
            categorical_feature = reference.categorical_feature
1802
1803
1804
1805
        data, feature_name, categorical_feature, self.pandas_categorical = _data_from_pandas(data=data,
                                                                                             feature_name=feature_name,
                                                                                             categorical_feature=categorical_feature,
                                                                                             pandas_categorical=self.pandas_categorical)
Guolin Ke's avatar
Guolin Ke committed
1806

1807
        # process for args
wxchan's avatar
wxchan committed
1808
        params = {} if params is None else params
1809
        args_names = inspect.signature(self.__class__._lazy_init).parameters.keys()
1810
        for key in params.keys():
1811
            if key in args_names:
1812
1813
                _log_warning(f'{key} keyword has been found in `params` and will be ignored.\n'
                             f'Please use {key} argument of the Dataset constructor to pass this parameter.')
1814
        # get categorical features
1815
1816
1817
1818
1819
1820
        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:
1821
                if isinstance(name, str) and name in feature_dict:
1822
                    categorical_indices.add(feature_dict[name])
1823
                elif isinstance(name, int):
1824
1825
                    categorical_indices.add(name)
                else:
1826
                    raise TypeError(f"Wrong type({type(name).__name__}) or unknown name({name}) in categorical_feature")
1827
            if categorical_indices:
1828
1829
                for cat_alias in _ConfigAliases.get("categorical_feature"):
                    if cat_alias in params:
1830
                        # If the params[cat_alias] is equal to categorical_indices, do not report the warning.
1831
                        if not (isinstance(params[cat_alias], list) and set(params[cat_alias]) == categorical_indices):
1832
                            _log_warning(f'{cat_alias} in param dict is overridden.')
1833
                        params.pop(cat_alias, None)
1834
                params['categorical_column'] = sorted(categorical_indices)
1835

1836
        params_str = _param_dict_to_str(params)
1837
        self.params = params
1838
        # process for reference dataset
wxchan's avatar
wxchan committed
1839
        ref_dataset = None
wxchan's avatar
wxchan committed
1840
        if isinstance(reference, Dataset):
1841
            ref_dataset = reference.construct()._handle
wxchan's avatar
wxchan committed
1842
1843
        elif reference is not None:
            raise TypeError('Reference dataset should be None or dataset instance')
1844
        # start construct data
1845
        if isinstance(data, (str, Path)):
1846
            self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1847
            _safe_call(_LIB.LGBM_DatasetCreateFromFile(
1848
1849
                _c_str(str(data)),
                _c_str(params_str),
wxchan's avatar
wxchan committed
1850
                ref_dataset,
1851
                ctypes.byref(self._handle)))
wxchan's avatar
wxchan committed
1852
1853
        elif isinstance(data, scipy.sparse.csr_matrix):
            self.__init_from_csr(data, params_str, ref_dataset)
Guolin Ke's avatar
Guolin Ke committed
1854
1855
        elif isinstance(data, scipy.sparse.csc_matrix):
            self.__init_from_csc(data, params_str, ref_dataset)
wxchan's avatar
wxchan committed
1856
1857
        elif isinstance(data, np.ndarray):
            self.__init_from_np2d(data, params_str, ref_dataset)
1858
1859
1860
1861
1862
1863
1864
1865
1866
        elif isinstance(data, list) and len(data) > 0:
            if all(isinstance(x, np.ndarray) for x in data):
                self.__init_from_list_np2d(data, params_str, ref_dataset)
            elif all(isinstance(x, Sequence) for x in data):
                self.__init_from_seqs(data, ref_dataset)
            else:
                raise TypeError('Data list can only be of ndarray or Sequence')
        elif isinstance(data, Sequence):
            self.__init_from_seqs([data], ref_dataset)
1867
        elif isinstance(data, dt_DataTable):
1868
            self.__init_from_np2d(data.to_numpy(), params_str, ref_dataset)
wxchan's avatar
wxchan committed
1869
1870
1871
1872
        else:
            try:
                csr = scipy.sparse.csr_matrix(data)
                self.__init_from_csr(csr, params_str, ref_dataset)
1873
1874
            except BaseException as err:
                raise TypeError(f'Cannot initialize Dataset from {type(data).__name__}') from err
wxchan's avatar
wxchan committed
1875
1876
1877
        if label is not None:
            self.set_label(label)
        if self.get_label() is None:
1878
            raise ValueError("Label should not be None")
wxchan's avatar
wxchan committed
1879
1880
1881
1882
        if weight is not None:
            self.set_weight(weight)
        if group is not None:
            self.set_group(group)
1883
1884
        if isinstance(predictor, _InnerPredictor):
            if self._predictor is None and init_score is not None:
1885
                _log_warning("The init_score will be overridden by the prediction of init_model.")
1886
1887
1888
1889
1890
            self._set_init_score_by_predictor(
                predictor=predictor,
                data=data,
                used_indices=None
            )
1891
1892
        elif init_score is not None:
            self.set_init_score(init_score)
Guolin Ke's avatar
Guolin Ke committed
1893
        elif predictor is not None:
1894
            raise TypeError(f'Wrong predictor type {type(predictor).__name__}')
Guolin Ke's avatar
Guolin Ke committed
1895
        # set feature names
Nikita Titov's avatar
Nikita Titov committed
1896
        return self.set_feature_name(feature_name)
wxchan's avatar
wxchan committed
1897

1898
1899
    @staticmethod
    def _yield_row_from_seqlist(seqs: List[Sequence], indices: Iterable[int]):
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
        offset = 0
        seq_id = 0
        seq = seqs[seq_id]
        for row_id in indices:
            assert row_id >= offset, "sample indices are expected to be monotonic"
            while row_id >= offset + len(seq):
                offset += len(seq)
                seq_id += 1
                seq = seqs[seq_id]
            id_in_seq = row_id - offset
            row = seq[id_in_seq]
            yield row if row.flags['OWNDATA'] else row.copy()

    def __sample(self, seqs: List[Sequence], total_nrow: int) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        """Sample data from seqs.

        Mimics behavior in c_api.cpp:LGBM_DatasetCreateFromMats()

        Returns
        -------
            sampled_rows, sampled_row_indices
        """
        indices = self._create_sample_indices(total_nrow)

        # Select sampled rows, transpose to column order.
1925
        sampled = np.array(list(self._yield_row_from_seqlist(seqs, indices)))
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
        sampled = sampled.T

        filtered = []
        filtered_idx = []
        sampled_row_range = np.arange(len(indices), dtype=np.int32)
        for col in sampled:
            col_predicate = (np.abs(col) > ZERO_THRESHOLD) | np.isnan(col)
            filtered_col = col[col_predicate]
            filtered_row_idx = sampled_row_range[col_predicate]

            filtered.append(filtered_col)
            filtered_idx.append(filtered_row_idx)

        return filtered, filtered_idx

1941
1942
1943
    def __init_from_seqs(
        self,
        seqs: List[Sequence],
1944
        ref_dataset: Optional[_DatasetHandle]
1945
    ) -> "Dataset":
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
        """
        Initialize data from list of Sequence objects.

        Sequence: Generic Data Access Object
            Supports random access and access by batch if properly defined by user

        Data scheme uniformity are trusted, not checked
        """
        total_nrow = sum(len(seq) for seq in seqs)

        # create validation dataset from ref_dataset
        if ref_dataset is not None:
            self._init_from_ref_dataset(total_nrow, ref_dataset)
        else:
1960
            param_str = _param_dict_to_str(self.get_params())
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
            sample_cnt = _get_sample_count(total_nrow, param_str)

            sample_data, col_indices = self.__sample(seqs, total_nrow)
            self._init_from_sample(sample_data, col_indices, sample_cnt, total_nrow)

        for seq in seqs:
            nrow = len(seq)
            batch_size = getattr(seq, 'batch_size', None) or Sequence.batch_size
            for start in range(0, nrow, batch_size):
                end = min(start + batch_size, nrow)
                self._push_rows(seq[start:end])
        return self

1974
1975
1976
1977
1978
1979
    def __init_from_np2d(
        self,
        mat: np.ndarray,
        params_str: str,
        ref_dataset: Optional[_DatasetHandle]
    ) -> "Dataset":
1980
        """Initialize data from a 2-D numpy matrix."""
wxchan's avatar
wxchan committed
1981
1982
1983
        if len(mat.shape) != 2:
            raise ValueError('Input numpy.ndarray must be 2 dimensional')

1984
        self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
1985
1986
        if mat.dtype == np.float32 or mat.dtype == np.float64:
            data = np.array(mat.reshape(mat.size), dtype=mat.dtype, copy=False)
1987
        else:  # change non-float data to float data, need to copy
wxchan's avatar
wxchan committed
1988
1989
            data = np.array(mat.reshape(mat.size), dtype=np.float32)

1990
        ptr_data, type_ptr_data, _ = _c_float_array(data)
wxchan's avatar
wxchan committed
1991
1992
        _safe_call(_LIB.LGBM_DatasetCreateFromMat(
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
1993
            ctypes.c_int(type_ptr_data),
1994
1995
            ctypes.c_int32(mat.shape[0]),
            ctypes.c_int32(mat.shape[1]),
1996
            ctypes.c_int(_C_API_IS_ROW_MAJOR),
1997
            _c_str(params_str),
wxchan's avatar
wxchan committed
1998
            ref_dataset,
1999
            ctypes.byref(self._handle)))
Nikita Titov's avatar
Nikita Titov committed
2000
        return self
wxchan's avatar
wxchan committed
2001

2002
2003
2004
2005
2006
2007
    def __init_from_list_np2d(
        self,
        mats: List[np.ndarray],
        params_str: str,
        ref_dataset: Optional[_DatasetHandle]
    ) -> "Dataset":
2008
        """Initialize data from a list of 2-D numpy matrices."""
2009
        ncol = mats[0].shape[1]
2010
        nrow = np.empty((len(mats),), np.int32)
2011
        ptr_data: _ctypes_float_array
2012
2013
2014
2015
2016
2017
        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 = []
2018
        type_ptr_data = -1
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030

        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)
2031
            else:  # change non-float data to float data, need to copy
2032
2033
                mats[i] = np.array(mat.reshape(mat.size), dtype=np.float32)

2034
            chunk_ptr_data, chunk_type_ptr_data, holder = _c_float_array(mats[i])
2035
            if type_ptr_data != -1 and chunk_type_ptr_data != type_ptr_data:
2036
2037
2038
2039
2040
                raise ValueError('Input chunks must have same type')
            ptr_data[i] = chunk_ptr_data
            type_ptr_data = chunk_type_ptr_data
            holders.append(holder)

2041
        self._handle = ctypes.c_void_p()
2042
        _safe_call(_LIB.LGBM_DatasetCreateFromMats(
2043
            ctypes.c_int32(len(mats)),
2044
2045
2046
            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)),
2047
            ctypes.c_int32(ncol),
2048
            ctypes.c_int(_C_API_IS_ROW_MAJOR),
2049
            _c_str(params_str),
2050
            ref_dataset,
2051
            ctypes.byref(self._handle)))
Nikita Titov's avatar
Nikita Titov committed
2052
        return self
2053

2054
2055
2056
2057
2058
2059
    def __init_from_csr(
        self,
        csr: scipy.sparse.csr_matrix,
        params_str: str,
        ref_dataset: Optional[_DatasetHandle]
    ) -> "Dataset":
2060
        """Initialize data from a CSR matrix."""
wxchan's avatar
wxchan committed
2061
        if len(csr.indices) != len(csr.data):
2062
            raise ValueError(f'Length mismatch: {len(csr.indices)} vs {len(csr.data)}')
2063
        self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
2064

2065
2066
        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
2067

2068
        assert csr.shape[1] <= _MAX_INT32
2069
        csr_indices = csr.indices.astype(np.int32, copy=False)
2070

wxchan's avatar
wxchan committed
2071
2072
        _safe_call(_LIB.LGBM_DatasetCreateFromCSR(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
2073
            ctypes.c_int(type_ptr_indptr),
2074
            csr_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
wxchan's avatar
wxchan committed
2075
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
2076
2077
2078
2079
            ctypes.c_int(type_ptr_data),
            ctypes.c_int64(len(csr.indptr)),
            ctypes.c_int64(len(csr.data)),
            ctypes.c_int64(csr.shape[1]),
2080
            _c_str(params_str),
wxchan's avatar
wxchan committed
2081
            ref_dataset,
2082
            ctypes.byref(self._handle)))
Nikita Titov's avatar
Nikita Titov committed
2083
        return self
wxchan's avatar
wxchan committed
2084

2085
2086
2087
2088
2089
2090
    def __init_from_csc(
        self,
        csc: scipy.sparse.csc_matrix,
        params_str: str,
        ref_dataset: Optional[_DatasetHandle]
    ) -> "Dataset":
2091
        """Initialize data from a CSC matrix."""
Guolin Ke's avatar
Guolin Ke committed
2092
        if len(csc.indices) != len(csc.data):
2093
            raise ValueError(f'Length mismatch: {len(csc.indices)} vs {len(csc.data)}')
2094
        self._handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
2095

2096
2097
        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
2098

2099
        assert csc.shape[0] <= _MAX_INT32
2100
        csc_indices = csc.indices.astype(np.int32, copy=False)
2101

Guolin Ke's avatar
Guolin Ke committed
2102
2103
        _safe_call(_LIB.LGBM_DatasetCreateFromCSC(
            ptr_indptr,
Guolin Ke's avatar
Guolin Ke committed
2104
            ctypes.c_int(type_ptr_indptr),
2105
            csc_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
Guolin Ke's avatar
Guolin Ke committed
2106
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
2107
2108
2109
2110
            ctypes.c_int(type_ptr_data),
            ctypes.c_int64(len(csc.indptr)),
            ctypes.c_int64(len(csc.data)),
            ctypes.c_int64(csc.shape[0]),
2111
            _c_str(params_str),
Guolin Ke's avatar
Guolin Ke committed
2112
            ref_dataset,
2113
            ctypes.byref(self._handle)))
Nikita Titov's avatar
Nikita Titov committed
2114
        return self
Guolin Ke's avatar
Guolin Ke committed
2115

2116
    @staticmethod
2117
2118
2119
2120
2121
2122
    def _compare_params_for_warning(
        params: Optional[Dict[str, Any]],
        other_params: Optional[Dict[str, Any]],
        ignore_keys: Set[str]
    ) -> bool:
        """Compare two dictionaries with params ignoring some keys.
2123

2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
        It is only for the warning purpose.

        Parameters
        ----------
        params : dict or None
            One dictionary with parameters to compare.
        other_params : dict or None
            Another dictionary with parameters to compare.
        ignore_keys : set
            Keys that should be ignored during comparing two dictionaries.
2134
2135
2136

        Returns
        -------
2137
2138
        compare_result : bool
          Returns whether two dictionaries with params are equal.
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
        """
        if params is None:
            params = {}
        if other_params is None:
            other_params = {}
        for k in other_params:
            if k not in ignore_keys:
                if k not in params or params[k] != other_params[k]:
                    return False
        for k in params:
            if k not in ignore_keys:
                if k not in other_params or params[k] != other_params[k]:
                    return False
        return True

2154
    def construct(self) -> "Dataset":
2155
2156
2157
2158
2159
        """Lazy init.

        Returns
        -------
        self : Dataset
Nikita Titov's avatar
Nikita Titov committed
2160
            Constructed Dataset object.
2161
        """
2162
        if self._handle is None:
wxchan's avatar
wxchan committed
2163
            if self.reference is not None:
2164
                reference_params = self.reference.get_params()
2165
2166
                params = self.get_params()
                if params != reference_params:
2167
2168
2169
2170
2171
                    if not self._compare_params_for_warning(
                        params=params,
                        other_params=reference_params,
                        ignore_keys=_ConfigAliases.get("categorical_feature")
                    ):
2172
                        _log_warning('Overriding the parameters from Reference Dataset.')
2173
                    self._update_params(reference_params)
wxchan's avatar
wxchan committed
2174
                if self.used_indices is None:
2175
                    # create valid
2176
                    self._lazy_init(data=self.data, label=self.label, reference=self.reference,
2177
2178
                                    weight=self.weight, group=self.group,
                                    init_score=self.init_score, predictor=self._predictor,
2179
                                    feature_name=self.feature_name, categorical_feature='auto', params=self.params)
wxchan's avatar
wxchan committed
2180
                else:
2181
                    # construct subset
2182
                    used_indices = _list_to_1d_numpy(self.used_indices, dtype=np.int32, name='used_indices')
2183
                    assert used_indices.flags.c_contiguous
Guolin Ke's avatar
Guolin Ke committed
2184
                    if self.reference.group is not None:
2185
                        group_info = np.array(self.reference.group).astype(np.int32, copy=False)
2186
                        _, self.group = np.unique(np.repeat(range(len(group_info)), repeats=group_info)[self.used_indices],
2187
                                                  return_counts=True)
2188
                    self._handle = ctypes.c_void_p()
2189
                    params_str = _param_dict_to_str(self.params)
wxchan's avatar
wxchan committed
2190
                    _safe_call(_LIB.LGBM_DatasetGetSubset(
2191
                        self.reference.construct()._handle,
wxchan's avatar
wxchan committed
2192
                        used_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
2193
                        ctypes.c_int32(used_indices.shape[0]),
2194
                        _c_str(params_str),
2195
                        ctypes.byref(self._handle)))
Guolin Ke's avatar
Guolin Ke committed
2196
2197
                    if not self.free_raw_data:
                        self.get_data()
Guolin Ke's avatar
Guolin Ke committed
2198
2199
                    if self.group is not None:
                        self.set_group(self.group)
wxchan's avatar
wxchan committed
2200
2201
                    if self.get_label() is None:
                        raise ValueError("Label should not be None.")
Guolin Ke's avatar
Guolin Ke committed
2202
2203
                    if isinstance(self._predictor, _InnerPredictor) and self._predictor is not self.reference._predictor:
                        self.get_data()
2204
2205
2206
2207
2208
                        self._set_init_score_by_predictor(
                            predictor=self._predictor,
                            data=self.data,
                            used_indices=used_indices
                        )
wxchan's avatar
wxchan committed
2209
            else:
2210
                # create train
2211
                self._lazy_init(data=self.data, label=self.label, reference=None,
2212
2213
                                weight=self.weight, group=self.group,
                                init_score=self.init_score, predictor=self._predictor,
2214
                                feature_name=self.feature_name, categorical_feature=self.categorical_feature, params=self.params)
wxchan's avatar
wxchan committed
2215
2216
            if self.free_raw_data:
                self.data = None
2217
            self.feature_name = self.get_feature_name()
wxchan's avatar
wxchan committed
2218
        return self
wxchan's avatar
wxchan committed
2219

2220
2221
    def create_valid(
        self,
2222
        data: _LGBM_TrainDataType,
2223
        label: Optional[_LGBM_LabelType] = None,
2224
2225
2226
        weight: Optional[_LGBM_WeightType] = None,
        group: Optional[_LGBM_GroupType] = None,
        init_score: Optional[_LGBM_InitScoreType] = None,
2227
2228
        params: Optional[Dict[str, Any]] = None
    ) -> "Dataset":
2229
        """Create validation data align with current Dataset.
wxchan's avatar
wxchan committed
2230
2231
2232

        Parameters
        ----------
2233
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse, Sequence, list of Sequence or list of numpy array
wxchan's avatar
wxchan committed
2234
            Data source of Dataset.
2235
            If str or pathlib.Path, it represents the path to a text file (CSV, TSV, or LibSVM) or a LightGBM Dataset binary file.
2236
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None, optional (default=None)
2237
2238
            Label of the data.
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
2239
            Weight for each instance. Weights should be non-negative.
2240
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
2241
2242
2243
            Group/query data.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
2244
2245
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
2246
        init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), or None, optional (default=None)
2247
            Init score for Dataset.
Nikita Titov's avatar
Nikita Titov committed
2248
        params : dict or None, optional (default=None)
2249
            Other parameters for validation Dataset.
2250
2251
2252

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2253
2254
        valid : Dataset
            Validation Dataset with reference to self.
wxchan's avatar
wxchan committed
2255
        """
2256
        ret = Dataset(data, label=label, reference=self,
2257
                      weight=weight, group=group, init_score=init_score,
2258
                      params=params, free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
2259
        ret._predictor = self._predictor
2260
        ret.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
2261
        return ret
wxchan's avatar
wxchan committed
2262

2263
2264
2265
2266
2267
    def subset(
        self,
        used_indices: List[int],
        params: Optional[Dict[str, Any]] = None
    ) -> "Dataset":
2268
        """Get subset of current Dataset.
wxchan's avatar
wxchan committed
2269
2270
2271
2272

        Parameters
        ----------
        used_indices : list of int
2273
            Indices used to create the subset.
Nikita Titov's avatar
Nikita Titov committed
2274
        params : dict or None, optional (default=None)
2275
            These parameters will be passed to Dataset constructor.
2276
2277
2278
2279
2280

        Returns
        -------
        subset : Dataset
            Subset of the current Dataset.
wxchan's avatar
wxchan committed
2281
        """
wxchan's avatar
wxchan committed
2282
2283
        if params is None:
            params = self.params
wxchan's avatar
wxchan committed
2284
        ret = Dataset(None, reference=self, feature_name=self.feature_name,
2285
2286
                      categorical_feature=self.categorical_feature, params=params,
                      free_raw_data=self.free_raw_data)
wxchan's avatar
wxchan committed
2287
        ret._predictor = self._predictor
2288
        ret.pandas_categorical = self.pandas_categorical
2289
        ret.used_indices = sorted(used_indices)
wxchan's avatar
wxchan committed
2290
2291
        return ret

2292
    def save_binary(self, filename: Union[str, Path]) -> "Dataset":
2293
        """Save Dataset to a binary file.
wxchan's avatar
wxchan committed
2294

2295
2296
2297
2298
2299
        .. note::

            Please note that `init_score` is not saved in binary file.
            If you need it, please set it again after loading Dataset.

wxchan's avatar
wxchan committed
2300
2301
        Parameters
        ----------
2302
        filename : str or pathlib.Path
wxchan's avatar
wxchan committed
2303
            Name of the output file.
Nikita Titov's avatar
Nikita Titov committed
2304
2305
2306
2307
2308

        Returns
        -------
        self : Dataset
            Returns self.
wxchan's avatar
wxchan committed
2309
2310
        """
        _safe_call(_LIB.LGBM_DatasetSaveBinary(
2311
            self.construct()._handle,
2312
            _c_str(str(filename))))
Nikita Titov's avatar
Nikita Titov committed
2313
        return self
wxchan's avatar
wxchan committed
2314

2315
    def _update_params(self, params: Optional[Dict[str, Any]]) -> "Dataset":
2316
2317
        if not params:
            return self
2318
        params = deepcopy(params)
2319
2320
2321
2322
2323

        def update():
            if not self.params:
                self.params = params
            else:
2324
                self._params_back_up = deepcopy(self.params)
2325
2326
                self.params.update(params)

2327
        if self._handle is None:
2328
2329
2330
            update()
        elif params is not None:
            ret = _LIB.LGBM_DatasetUpdateParamChecking(
2331
2332
                _c_str(_param_dict_to_str(self.params)),
                _c_str(_param_dict_to_str(params)))
2333
2334
2335
2336
2337
2338
            if ret != 0:
                # could be updated if data is not freed
                if self.data is not None:
                    update()
                    self._free_handle()
                else:
2339
                    raise LightGBMError(_LIB.LGBM_GetLastError().decode('utf-8'))
Nikita Titov's avatar
Nikita Titov committed
2340
        return self
wxchan's avatar
wxchan committed
2341

2342
    def _reverse_update_params(self) -> "Dataset":
2343
        if self._handle is None:
2344
2345
            self.params = deepcopy(self._params_back_up)
            self._params_back_up = None
Nikita Titov's avatar
Nikita Titov committed
2346
        return self
2347

2348
2349
2350
    def set_field(
        self,
        field_name: str,
2351
        data: Optional[Union[List[List[float]], List[List[int]], List[float], List[int], np.ndarray, pd_Series, pd_DataFrame]]
2352
    ) -> "Dataset":
wxchan's avatar
wxchan committed
2353
        """Set property into the Dataset.
wxchan's avatar
wxchan committed
2354
2355
2356

        Parameters
        ----------
2357
        field_name : str
2358
            The field name of the information.
2359
2360
        data : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), or None
            The data to be set.
Nikita Titov's avatar
Nikita Titov committed
2361
2362
2363
2364
2365

        Returns
        -------
        self : Dataset
            Dataset with set property.
wxchan's avatar
wxchan committed
2366
        """
2367
        if self._handle is None:
2368
            raise Exception(f"Cannot set {field_name} before construct dataset")
wxchan's avatar
wxchan committed
2369
        if data is None:
2370
            # set to None
wxchan's avatar
wxchan committed
2371
            _safe_call(_LIB.LGBM_DatasetSetField(
2372
                self._handle,
2373
                _c_str(field_name),
wxchan's avatar
wxchan committed
2374
                None,
Guolin Ke's avatar
Guolin Ke committed
2375
                ctypes.c_int(0),
2376
                ctypes.c_int(_FIELD_TYPE_MAPPER[field_name])))
Nikita Titov's avatar
Nikita Titov committed
2377
            return self
2378
        dtype: "np.typing.DTypeLike"
2379
        if field_name == 'init_score':
Guolin Ke's avatar
Guolin Ke committed
2380
            dtype = np.float64
2381
            if _is_1d_collection(data):
2382
                data = _list_to_1d_numpy(data, dtype=dtype, name=field_name)
2383
            elif _is_2d_collection(data):
2384
                data = _data_to_2d_numpy(data, dtype=dtype, name=field_name)
2385
2386
2387
2388
2389
2390
2391
2392
                data = data.ravel(order='F')
            else:
                raise TypeError(
                    'init_score must be list, numpy 1-D array or pandas Series.\n'
                    'In multiclass classification init_score can also be a list of lists, numpy 2-D array or pandas DataFrame.'
                )
        else:
            dtype = np.int32 if field_name == 'group' else np.float32
2393
            data = _list_to_1d_numpy(data, dtype=dtype, name=field_name)
2394

2395
        ptr_data: Union[_ctypes_float_ptr, _ctypes_int_ptr]
2396
        if data.dtype == np.float32 or data.dtype == np.float64:
2397
            ptr_data, type_data, _ = _c_float_array(data)
wxchan's avatar
wxchan committed
2398
        elif data.dtype == np.int32:
2399
            ptr_data, type_data, _ = _c_int_array(data)
wxchan's avatar
wxchan committed
2400
        else:
2401
            raise TypeError(f"Expected np.float32/64 or np.int32, met type({data.dtype})")
2402
        if type_data != _FIELD_TYPE_MAPPER[field_name]:
2403
            raise TypeError("Input type error for set_field")
wxchan's avatar
wxchan committed
2404
        _safe_call(_LIB.LGBM_DatasetSetField(
2405
            self._handle,
2406
            _c_str(field_name),
wxchan's avatar
wxchan committed
2407
            ptr_data,
Guolin Ke's avatar
Guolin Ke committed
2408
2409
            ctypes.c_int(len(data)),
            ctypes.c_int(type_data)))
2410
        self.version += 1
Nikita Titov's avatar
Nikita Titov committed
2411
        return self
wxchan's avatar
wxchan committed
2412

2413
    def get_field(self, field_name: str) -> Optional[np.ndarray]:
wxchan's avatar
wxchan committed
2414
        """Get property from the Dataset.
wxchan's avatar
wxchan committed
2415
2416
2417

        Parameters
        ----------
2418
        field_name : str
2419
            The field name of the information.
wxchan's avatar
wxchan committed
2420
2421
2422

        Returns
        -------
2423
        info : numpy array or None
2424
            A numpy array with information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2425
        """
2426
        if self._handle is None:
2427
            raise Exception(f"Cannot get {field_name} before construct Dataset")
2428
2429
        tmp_out_len = ctypes.c_int(0)
        out_type = ctypes.c_int(0)
wxchan's avatar
wxchan committed
2430
2431
        ret = ctypes.POINTER(ctypes.c_void_p)()
        _safe_call(_LIB.LGBM_DatasetGetField(
2432
            self._handle,
2433
            _c_str(field_name),
wxchan's avatar
wxchan committed
2434
2435
2436
            ctypes.byref(tmp_out_len),
            ctypes.byref(ret),
            ctypes.byref(out_type)))
2437
        if out_type.value != _FIELD_TYPE_MAPPER[field_name]:
wxchan's avatar
wxchan committed
2438
2439
2440
            raise TypeError("Return type error for get_field")
        if tmp_out_len.value == 0:
            return None
2441
        if out_type.value == _C_API_DTYPE_INT32:
2442
            arr = _cint32_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_int32)), tmp_out_len.value)
2443
        elif out_type.value == _C_API_DTYPE_FLOAT32:
2444
            arr = _cfloat32_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_float)), tmp_out_len.value)
2445
        elif out_type.value == _C_API_DTYPE_FLOAT64:
2446
            arr = _cfloat64_array_to_numpy(ctypes.cast(ret, ctypes.POINTER(ctypes.c_double)), tmp_out_len.value)
2447
        else:
wxchan's avatar
wxchan committed
2448
            raise TypeError("Unknown type")
2449
2450
2451
2452
2453
2454
        if field_name == 'init_score':
            num_data = self.num_data()
            num_classes = arr.size // num_data
            if num_classes > 1:
                arr = arr.reshape((num_data, num_classes), order='F')
        return arr
Guolin Ke's avatar
Guolin Ke committed
2455

2456
2457
    def set_categorical_feature(
        self,
2458
        categorical_feature: _LGBM_CategoricalFeatureConfiguration
2459
    ) -> "Dataset":
2460
        """Set categorical features.
2461
2462
2463

        Parameters
        ----------
2464
        categorical_feature : list of str or int, or 'auto'
2465
            Names or indices of categorical features.
Nikita Titov's avatar
Nikita Titov committed
2466
2467
2468
2469
2470

        Returns
        -------
        self : Dataset
            Dataset with set categorical features.
2471
2472
        """
        if self.categorical_feature == categorical_feature:
Nikita Titov's avatar
Nikita Titov committed
2473
            return self
2474
        if self.data is not None:
2475
2476
            if self.categorical_feature is None:
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
2477
                return self._free_handle()
2478
            elif categorical_feature == 'auto':
Nikita Titov's avatar
Nikita Titov committed
2479
                return self
2480
            else:
2481
2482
                if self.categorical_feature != 'auto':
                    _log_warning('categorical_feature in Dataset is overridden.\n'
2483
                                 f'New categorical_feature is {list(categorical_feature)}')
2484
                self.categorical_feature = categorical_feature
Nikita Titov's avatar
Nikita Titov committed
2485
                return self._free_handle()
2486
        else:
2487
2488
            raise LightGBMError("Cannot set categorical feature after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
2489

2490
2491
2492
2493
    def _set_predictor(
        self,
        predictor: Optional[_InnerPredictor]
    ) -> "Dataset":
2494
2495
2496
2497
        """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
2498
        """
2499
        if predictor is None and self._predictor is None:
Nikita Titov's avatar
Nikita Titov committed
2500
            return self
2501
2502
2503
        elif isinstance(predictor, _InnerPredictor) and isinstance(self._predictor, _InnerPredictor):
            if (predictor == self._predictor) and (predictor.current_iteration() == self._predictor.current_iteration()):
                return self
2504
        if self._handle is None:
Guolin Ke's avatar
Guolin Ke committed
2505
            self._predictor = predictor
2506
2507
        elif self.data is not None:
            self._predictor = predictor
2508
2509
2510
2511
2512
            self._set_init_score_by_predictor(
                predictor=self._predictor,
                data=self.data,
                used_indices=None
            )
2513
2514
        elif self.used_indices is not None and self.reference is not None and self.reference.data is not None:
            self._predictor = predictor
2515
2516
2517
2518
2519
            self._set_init_score_by_predictor(
                predictor=self._predictor,
                data=self.reference.data,
                used_indices=self.used_indices
            )
Guolin Ke's avatar
Guolin Ke committed
2520
        else:
2521
2522
            raise LightGBMError("Cannot set predictor after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
2523
        return self
Guolin Ke's avatar
Guolin Ke committed
2524

2525
    def set_reference(self, reference: "Dataset") -> "Dataset":
2526
        """Set reference Dataset.
Guolin Ke's avatar
Guolin Ke committed
2527
2528
2529
2530

        Parameters
        ----------
        reference : Dataset
2531
            Reference that is used as a template to construct the current Dataset.
Nikita Titov's avatar
Nikita Titov committed
2532
2533
2534
2535
2536

        Returns
        -------
        self : Dataset
            Dataset with set reference.
Guolin Ke's avatar
Guolin Ke committed
2537
        """
2538
2539
2540
        self.set_categorical_feature(reference.categorical_feature) \
            .set_feature_name(reference.feature_name) \
            ._set_predictor(reference._predictor)
2541
        # we're done if self and reference share a common upstream reference
2542
        if self.get_ref_chain().intersection(reference.get_ref_chain()):
Nikita Titov's avatar
Nikita Titov committed
2543
            return self
Guolin Ke's avatar
Guolin Ke committed
2544
2545
        if self.data is not None:
            self.reference = reference
Nikita Titov's avatar
Nikita Titov committed
2546
            return self._free_handle()
Guolin Ke's avatar
Guolin Ke committed
2547
        else:
2548
2549
            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
2550

2551
    def set_feature_name(self, feature_name: _LGBM_FeatureNameConfiguration) -> "Dataset":
2552
        """Set feature name.
Guolin Ke's avatar
Guolin Ke committed
2553
2554
2555

        Parameters
        ----------
2556
        feature_name : list of str
2557
            Feature names.
Nikita Titov's avatar
Nikita Titov committed
2558
2559
2560
2561
2562

        Returns
        -------
        self : Dataset
            Dataset with set feature name.
Guolin Ke's avatar
Guolin Ke committed
2563
        """
2564
2565
        if feature_name != 'auto':
            self.feature_name = feature_name
2566
        if self._handle is not None and feature_name is not None and feature_name != 'auto':
wxchan's avatar
wxchan committed
2567
            if len(feature_name) != self.num_feature():
2568
                raise ValueError(f"Length of feature_name({len(feature_name)}) and num_feature({self.num_feature()}) don't match")
2569
            c_feature_name = [_c_str(name) for name in feature_name]
wxchan's avatar
wxchan committed
2570
            _safe_call(_LIB.LGBM_DatasetSetFeatureNames(
2571
                self._handle,
2572
                _c_array(ctypes.c_char_p, c_feature_name),
Guolin Ke's avatar
Guolin Ke committed
2573
                ctypes.c_int(len(feature_name))))
Nikita Titov's avatar
Nikita Titov committed
2574
        return self
Guolin Ke's avatar
Guolin Ke committed
2575

2576
    def set_label(self, label: Optional[_LGBM_LabelType]) -> "Dataset":
2577
        """Set label of Dataset.
Guolin Ke's avatar
Guolin Ke committed
2578
2579
2580

        Parameters
        ----------
2581
        label : list, numpy 1-D array, pandas Series / one-column DataFrame or None
2582
            The label information to be set into Dataset.
Nikita Titov's avatar
Nikita Titov committed
2583
2584
2585
2586
2587

        Returns
        -------
        self : Dataset
            Dataset with set label.
Guolin Ke's avatar
Guolin Ke committed
2588
2589
        """
        self.label = label
2590
        if self._handle is not None:
2591
2592
2593
2594
            if isinstance(label, pd_DataFrame):
                if len(label.columns) > 1:
                    raise ValueError('DataFrame for label cannot have multiple columns')
                _check_for_bad_pandas_dtypes(label.dtypes)
2595
2596
2597
2598
2599
2600
2601
2602
2603
2604
2605
                try:
                    # most common case (no nullable dtypes)
                    label = label.to_numpy(dtype=np.float32, copy=False)
                except TypeError:
                    # 1.0 <= pd version < 1.1 and nullable dtypes, least common case
                    # raises error because array is casted to type(pd.NA) and there's no na_value argument
                    label = label.astype(np.float32, copy=False).values
                except ValueError:
                    # data has nullable dtypes, but we can specify na_value argument and copy will be made
                    label = label.to_numpy(dtype=np.float32, na_value=np.nan)
                label_array = np.ravel(label)
2606
            else:
2607
                label_array = _list_to_1d_numpy(label, dtype=np.float32, name='label')
2608
            self.set_field('label', label_array)
2609
            self.label = self.get_field('label')  # original values can be modified at cpp side
Nikita Titov's avatar
Nikita Titov committed
2610
        return self
Guolin Ke's avatar
Guolin Ke committed
2611

2612
2613
2614
2615
    def set_weight(
        self,
        weight: Optional[_LGBM_WeightType]
    ) -> "Dataset":
2616
        """Set weight of each instance.
Guolin Ke's avatar
Guolin Ke committed
2617
2618
2619

        Parameters
        ----------
2620
        weight : list, numpy 1-D array, pandas Series or None
2621
            Weight to be set for each data point. Weights should be non-negative.
Nikita Titov's avatar
Nikita Titov committed
2622
2623
2624
2625
2626

        Returns
        -------
        self : Dataset
            Dataset with set weight.
Guolin Ke's avatar
Guolin Ke committed
2627
        """
2628
2629
        if weight is not None and np.all(weight == 1):
            weight = None
Guolin Ke's avatar
Guolin Ke committed
2630
        self.weight = weight
2631
        if self._handle is not None and weight is not None:
2632
            weight = _list_to_1d_numpy(weight, dtype=np.float32, name='weight')
wxchan's avatar
wxchan committed
2633
            self.set_field('weight', weight)
2634
            self.weight = self.get_field('weight')  # original values can be modified at cpp side
Nikita Titov's avatar
Nikita Titov committed
2635
        return self
Guolin Ke's avatar
Guolin Ke committed
2636

2637
2638
2639
2640
    def set_init_score(
        self,
        init_score: Optional[_LGBM_InitScoreType]
    ) -> "Dataset":
2641
        """Set init score of Booster to start from.
Guolin Ke's avatar
Guolin Ke committed
2642
2643
2644

        Parameters
        ----------
2645
        init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), or None
2646
            Init score for Booster.
Nikita Titov's avatar
Nikita Titov committed
2647
2648
2649
2650
2651

        Returns
        -------
        self : Dataset
            Dataset with set init score.
Guolin Ke's avatar
Guolin Ke committed
2652
2653
        """
        self.init_score = init_score
2654
        if self._handle is not None and init_score is not None:
wxchan's avatar
wxchan committed
2655
            self.set_field('init_score', init_score)
2656
            self.init_score = self.get_field('init_score')  # original values can be modified at cpp side
Nikita Titov's avatar
Nikita Titov committed
2657
        return self
Guolin Ke's avatar
Guolin Ke committed
2658

2659
2660
2661
2662
    def set_group(
        self,
        group: Optional[_LGBM_GroupType]
    ) -> "Dataset":
2663
        """Set group size of Dataset (used for ranking).
Guolin Ke's avatar
Guolin Ke committed
2664
2665
2666

        Parameters
        ----------
2667
        group : list, numpy 1-D array, pandas Series or None
2668
2669
2670
            Group/query data.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
2671
2672
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
Nikita Titov's avatar
Nikita Titov committed
2673
2674
2675
2676
2677

        Returns
        -------
        self : Dataset
            Dataset with set group.
Guolin Ke's avatar
Guolin Ke committed
2678
2679
        """
        self.group = group
2680
        if self._handle is not None and group is not None:
2681
            group = _list_to_1d_numpy(group, dtype=np.int32, name='group')
wxchan's avatar
wxchan committed
2682
            self.set_field('group', group)
Nikita Titov's avatar
Nikita Titov committed
2683
        return self
Guolin Ke's avatar
Guolin Ke committed
2684

2685
    def get_feature_name(self) -> List[str]:
2686
2687
2688
2689
        """Get the names of columns (features) in the Dataset.

        Returns
        -------
2690
        feature_names : list of str
2691
2692
            The names of columns (features) in the Dataset.
        """
2693
        if self._handle is None:
2694
2695
2696
2697
2698
            raise LightGBMError("Cannot get feature_name before construct dataset")
        num_feature = self.num_feature()
        tmp_out_len = ctypes.c_int(0)
        reserved_string_buffer_size = 255
        required_string_buffer_size = ctypes.c_size_t(0)
2699
        string_buffers = [ctypes.create_string_buffer(reserved_string_buffer_size) for _ in range(num_feature)]
2700
2701
        ptr_string_buffers = (ctypes.c_char_p * num_feature)(*map(ctypes.addressof, string_buffers))
        _safe_call(_LIB.LGBM_DatasetGetFeatureNames(
2702
            self._handle,
2703
            ctypes.c_int(num_feature),
2704
            ctypes.byref(tmp_out_len),
2705
            ctypes.c_size_t(reserved_string_buffer_size),
2706
2707
2708
2709
            ctypes.byref(required_string_buffer_size),
            ptr_string_buffers))
        if num_feature != tmp_out_len.value:
            raise ValueError("Length of feature names doesn't equal with num_feature")
2710
2711
2712
2713
2714
2715
        actual_string_buffer_size = required_string_buffer_size.value
        # if buffer length is not long enough, reallocate buffers
        if reserved_string_buffer_size < actual_string_buffer_size:
            string_buffers = [ctypes.create_string_buffer(actual_string_buffer_size) for _ in range(num_feature)]
            ptr_string_buffers = (ctypes.c_char_p * num_feature)(*map(ctypes.addressof, string_buffers))
            _safe_call(_LIB.LGBM_DatasetGetFeatureNames(
2716
                self._handle,
2717
2718
2719
2720
2721
                ctypes.c_int(num_feature),
                ctypes.byref(tmp_out_len),
                ctypes.c_size_t(actual_string_buffer_size),
                ctypes.byref(required_string_buffer_size),
                ptr_string_buffers))
2722
        return [string_buffers[i].value.decode('utf-8') for i in range(num_feature)]
2723

2724
    def get_label(self) -> Optional[np.ndarray]:
2725
        """Get the label of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2726
2727
2728

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2729
        label : numpy array or None
2730
            The label information from the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2731
        """
2732
        if self.label is None:
wxchan's avatar
wxchan committed
2733
            self.label = self.get_field('label')
Guolin Ke's avatar
Guolin Ke committed
2734
2735
        return self.label

2736
    def get_weight(self) -> Optional[np.ndarray]:
2737
        """Get the weight of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2738
2739
2740

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2741
        weight : numpy array or None
2742
            Weight for each data point from the Dataset. Weights should be non-negative.
Guolin Ke's avatar
Guolin Ke committed
2743
        """
2744
        if self.weight is None:
wxchan's avatar
wxchan committed
2745
            self.weight = self.get_field('weight')
Guolin Ke's avatar
Guolin Ke committed
2746
2747
        return self.weight

2748
    def get_init_score(self) -> Optional[np.ndarray]:
2749
        """Get the initial score of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2750
2751
2752

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2753
        init_score : numpy array or None
2754
            Init score of Booster.
Guolin Ke's avatar
Guolin Ke committed
2755
        """
2756
        if self.init_score is None:
wxchan's avatar
wxchan committed
2757
            self.init_score = self.get_field('init_score')
Guolin Ke's avatar
Guolin Ke committed
2758
2759
        return self.init_score

2760
    def get_data(self) -> Optional[_LGBM_TrainDataType]:
2761
2762
2763
2764
        """Get the raw data of the Dataset.

        Returns
        -------
2765
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse, Sequence, list of Sequence or list of numpy array or None
2766
2767
            Raw data used in the Dataset construction.
        """
2768
        if self._handle is None:
2769
            raise Exception("Cannot get data before construct Dataset")
2770
        if self._need_slice and self.used_indices is not None and self.reference is not None:
Guolin Ke's avatar
Guolin Ke committed
2771
2772
            self.data = self.reference.data
            if self.data is not None:
2773
                if isinstance(self.data, np.ndarray) or isinstance(self.data, scipy.sparse.spmatrix):
Guolin Ke's avatar
Guolin Ke committed
2774
                    self.data = self.data[self.used_indices, :]
2775
                elif isinstance(self.data, pd_DataFrame):
Guolin Ke's avatar
Guolin Ke committed
2776
                    self.data = self.data.iloc[self.used_indices].copy()
2777
                elif isinstance(self.data, dt_DataTable):
Guolin Ke's avatar
Guolin Ke committed
2778
                    self.data = self.data[self.used_indices, :]
2779
2780
2781
                elif isinstance(self.data, Sequence):
                    self.data = self.data[self.used_indices]
                elif isinstance(self.data, list) and len(self.data) > 0 and all(isinstance(x, Sequence) for x in self.data):
2782
                    self.data = np.array(list(self._yield_row_from_seqlist(self.data, self.used_indices)))
Guolin Ke's avatar
Guolin Ke committed
2783
                else:
2784
2785
                    _log_warning(f"Cannot subset {type(self.data).__name__} type of raw data.\n"
                                 "Returning original raw data")
2786
            self._need_slice = False
Guolin Ke's avatar
Guolin Ke committed
2787
2788
2789
        if self.data is None:
            raise LightGBMError("Cannot call `get_data` after freed raw data, "
                                "set free_raw_data=False when construct Dataset to avoid this.")
2790
2791
        return self.data

2792
    def get_group(self) -> Optional[np.ndarray]:
2793
        """Get the group of the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2794
2795
2796

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
2797
        group : numpy array or None
2798
2799
2800
            Group/query data.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
2801
2802
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
Guolin Ke's avatar
Guolin Ke committed
2803
        """
2804
        if self.group is None:
wxchan's avatar
wxchan committed
2805
            self.group = self.get_field('group')
Guolin Ke's avatar
Guolin Ke committed
2806
2807
            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
2808
                self.group = np.diff(self.group)
Guolin Ke's avatar
Guolin Ke committed
2809
2810
        return self.group

2811
    def num_data(self) -> int:
2812
        """Get the number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2813
2814
2815

        Returns
        -------
2816
2817
        number_of_rows : int
            The number of rows in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2818
        """
2819
        if self._handle is not None:
2820
            ret = ctypes.c_int(0)
2821
            _safe_call(_LIB.LGBM_DatasetGetNumData(self._handle,
wxchan's avatar
wxchan committed
2822
2823
                                                   ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
2824
        else:
2825
            raise LightGBMError("Cannot get num_data before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
2826

2827
    def num_feature(self) -> int:
2828
        """Get the number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2829
2830
2831

        Returns
        -------
2832
2833
        number_of_columns : int
            The number of columns (features) in the Dataset.
Guolin Ke's avatar
Guolin Ke committed
2834
        """
2835
        if self._handle is not None:
2836
            ret = ctypes.c_int(0)
2837
            _safe_call(_LIB.LGBM_DatasetGetNumFeature(self._handle,
wxchan's avatar
wxchan committed
2838
2839
                                                      ctypes.byref(ret)))
            return ret.value
Guolin Ke's avatar
Guolin Ke committed
2840
        else:
2841
            raise LightGBMError("Cannot get num_feature before construct dataset")
Guolin Ke's avatar
Guolin Ke committed
2842

2843
    def feature_num_bin(self, feature: Union[int, str]) -> int:
2844
2845
        """Get the number of bins for a feature.

2846
2847
        .. versionadded:: 4.0.0

2848
2849
        Parameters
        ----------
2850
2851
        feature : int or str
            Index or name of the feature.
2852
2853
2854
2855
2856
2857

        Returns
        -------
        number_of_bins : int
            The number of constructed bins for the feature in the Dataset.
        """
2858
        if self._handle is not None:
2859
            if isinstance(feature, str):
2860
2861
2862
                feature_index = self.feature_name.index(feature)
            else:
                feature_index = feature
2863
            ret = ctypes.c_int(0)
2864
            _safe_call(_LIB.LGBM_DatasetGetFeatureNumBin(self._handle,
2865
                                                         ctypes.c_int(feature_index),
2866
2867
2868
2869
2870
                                                         ctypes.byref(ret)))
            return ret.value
        else:
            raise LightGBMError("Cannot get feature_num_bin before construct dataset")

2871
    def get_ref_chain(self, ref_limit: int = 100) -> Set["Dataset"]:
2872
2873
2874
2875
2876
        """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.
2877
2878
2879
2880
2881

        Parameters
        ----------
        ref_limit : int, optional (default=100)
            The limit number of references.
2882
2883
2884

        Returns
        -------
2885
2886
2887
        ref_chain : set of Dataset
            Chain of references of the Datasets.
        """
2888
        head = self
2889
        ref_chain: Set[Dataset] = set()
2890
2891
        while len(ref_chain) < ref_limit:
            if isinstance(head, Dataset):
2892
                ref_chain.add(head)
2893
2894
2895
2896
2897
2898
                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
2899
        return ref_chain
2900

2901
    def add_features_from(self, other: "Dataset") -> "Dataset":
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
        """Add features from other Dataset to the current Dataset.

        Both Datasets must be constructed before calling this method.

        Parameters
        ----------
        other : Dataset
            The Dataset to take features from.

        Returns
        -------
        self : Dataset
            Dataset with the new features added.
        """
2916
        if self._handle is None or other._handle is None:
2917
            raise ValueError('Both source and target Datasets must be constructed before adding features')
2918
        _safe_call(_LIB.LGBM_DatasetAddFeaturesFrom(self._handle, other._handle))
Guolin Ke's avatar
Guolin Ke committed
2919
2920
2921
2922
2923
2924
2925
2926
        was_none = self.data is None
        old_self_data_type = type(self.data).__name__
        if other.data is None:
            self.data = None
        elif self.data is not None:
            if isinstance(self.data, np.ndarray):
                if isinstance(other.data, np.ndarray):
                    self.data = np.hstack((self.data, other.data))
2927
                elif isinstance(other.data, scipy.sparse.spmatrix):
Guolin Ke's avatar
Guolin Ke committed
2928
                    self.data = np.hstack((self.data, other.data.toarray()))
2929
                elif isinstance(other.data, pd_DataFrame):
Guolin Ke's avatar
Guolin Ke committed
2930
                    self.data = np.hstack((self.data, other.data.values))
2931
                elif isinstance(other.data, dt_DataTable):
Guolin Ke's avatar
Guolin Ke committed
2932
2933
2934
                    self.data = np.hstack((self.data, other.data.to_numpy()))
                else:
                    self.data = None
2935
            elif isinstance(self.data, scipy.sparse.spmatrix):
Guolin Ke's avatar
Guolin Ke committed
2936
                sparse_format = self.data.getformat()
2937
                if isinstance(other.data, np.ndarray) or isinstance(other.data, scipy.sparse.spmatrix):
Guolin Ke's avatar
Guolin Ke committed
2938
                    self.data = scipy.sparse.hstack((self.data, other.data), format=sparse_format)
2939
                elif isinstance(other.data, pd_DataFrame):
Guolin Ke's avatar
Guolin Ke committed
2940
                    self.data = scipy.sparse.hstack((self.data, other.data.values), format=sparse_format)
2941
                elif isinstance(other.data, dt_DataTable):
Guolin Ke's avatar
Guolin Ke committed
2942
2943
2944
                    self.data = scipy.sparse.hstack((self.data, other.data.to_numpy()), format=sparse_format)
                else:
                    self.data = None
2945
            elif isinstance(self.data, pd_DataFrame):
Guolin Ke's avatar
Guolin Ke committed
2946
2947
                if not PANDAS_INSTALLED:
                    raise LightGBMError("Cannot add features to DataFrame type of raw data "
2948
2949
                                        "without pandas installed. "
                                        "Install pandas and restart your session.")
Guolin Ke's avatar
Guolin Ke committed
2950
                if isinstance(other.data, np.ndarray):
2951
                    self.data = concat((self.data, pd_DataFrame(other.data)),
Guolin Ke's avatar
Guolin Ke committed
2952
                                       axis=1, ignore_index=True)
2953
                elif isinstance(other.data, scipy.sparse.spmatrix):
2954
                    self.data = concat((self.data, pd_DataFrame(other.data.toarray())),
Guolin Ke's avatar
Guolin Ke committed
2955
                                       axis=1, ignore_index=True)
2956
                elif isinstance(other.data, pd_DataFrame):
Guolin Ke's avatar
Guolin Ke committed
2957
2958
                    self.data = concat((self.data, other.data),
                                       axis=1, ignore_index=True)
2959
2960
                elif isinstance(other.data, dt_DataTable):
                    self.data = concat((self.data, pd_DataFrame(other.data.to_numpy())),
Guolin Ke's avatar
Guolin Ke committed
2961
2962
2963
                                       axis=1, ignore_index=True)
                else:
                    self.data = None
2964
            elif isinstance(self.data, dt_DataTable):
Guolin Ke's avatar
Guolin Ke committed
2965
                if isinstance(other.data, np.ndarray):
2966
                    self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data)))
2967
                elif isinstance(other.data, scipy.sparse.spmatrix):
2968
2969
2970
2971
2972
                    self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data.toarray())))
                elif isinstance(other.data, pd_DataFrame):
                    self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data.values)))
                elif isinstance(other.data, dt_DataTable):
                    self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data.to_numpy())))
Guolin Ke's avatar
Guolin Ke committed
2973
2974
2975
2976
2977
                else:
                    self.data = None
            else:
                self.data = None
        if self.data is None:
2978
2979
            err_msg = (f"Cannot add features from {type(other.data).__name__} type of raw data to "
                       f"{old_self_data_type} type of raw data.\n")
Guolin Ke's avatar
Guolin Ke committed
2980
2981
            err_msg += ("Set free_raw_data=False when construct Dataset to avoid this"
                        if was_none else "Freeing raw data")
2982
            _log_warning(err_msg)
Guolin Ke's avatar
Guolin Ke committed
2983
        self.feature_name = self.get_feature_name()
2984
2985
        _log_warning("Reseting categorical features.\n"
                     "You can set new categorical features via ``set_categorical_feature`` method")
Guolin Ke's avatar
Guolin Ke committed
2986
2987
        self.categorical_feature = "auto"
        self.pandas_categorical = None
2988
2989
        return self

2990
    def _dump_text(self, filename: Union[str, Path]) -> "Dataset":
2991
2992
2993
2994
2995
2996
        """Save Dataset to a text file.

        This format cannot be loaded back in by LightGBM, but is useful for debugging purposes.

        Parameters
        ----------
2997
        filename : str or pathlib.Path
2998
2999
3000
3001
3002
3003
3004
3005
            Name of the output file.

        Returns
        -------
        self : Dataset
            Returns self.
        """
        _safe_call(_LIB.LGBM_DatasetDumpText(
3006
            self.construct()._handle,
3007
            _c_str(str(filename))))
3008
3009
        return self

wxchan's avatar
wxchan committed
3010

3011
3012
3013
3014
_LGBM_CustomObjectiveFunction = Callable[
    [np.ndarray, Dataset],
    Tuple[np.ndarray, np.ndarray]
]
3015
3016
3017
3018
3019
3020
3021
3022
3023
3024
_LGBM_CustomEvalFunction = Union[
    Callable[
        [np.ndarray, Dataset],
        _LGBM_EvalFunctionResultType
    ],
    Callable[
        [np.ndarray, Dataset],
        List[_LGBM_EvalFunctionResultType]
    ]
]
3025
3026


3027
class Booster:
3028
    """Booster in LightGBM."""
3029

3030
3031
3032
3033
3034
3035
3036
    def __init__(
        self,
        params: Optional[Dict[str, Any]] = None,
        train_set: Optional[Dataset] = None,
        model_file: Optional[Union[str, Path]] = None,
        model_str: Optional[str] = None
    ):
3037
        """Initialize the Booster.
wxchan's avatar
wxchan committed
3038
3039
3040

        Parameters
        ----------
Nikita Titov's avatar
Nikita Titov committed
3041
        params : dict or None, optional (default=None)
3042
3043
3044
            Parameters for Booster.
        train_set : Dataset or None, optional (default=None)
            Training dataset.
3045
        model_file : str, pathlib.Path or None, optional (default=None)
wxchan's avatar
wxchan committed
3046
            Path to the model file.
3047
        model_str : str or None, optional (default=None)
3048
            Model will be loaded from this string.
wxchan's avatar
wxchan committed
3049
        """
3050
        self._handle = None
3051
        self._network = False
wxchan's avatar
wxchan committed
3052
        self.__need_reload_eval_info = True
3053
        self._train_data_name = "training"
3054
        self.__set_objective_to_none = False
wxchan's avatar
wxchan committed
3055
        self.best_iteration = -1
3056
        self.best_score: _LGBM_BoosterBestScoreType = {}
3057
        params = {} if params is None else deepcopy(params)
wxchan's avatar
wxchan committed
3058
        if train_set is not None:
3059
            # Training task
wxchan's avatar
wxchan committed
3060
            if not isinstance(train_set, Dataset):
3061
                raise TypeError(f'Training data should be Dataset instance, met {type(train_set).__name__}')
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
3082
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
            params = _choose_param_value(
                main_param_name="machines",
                params=params,
                default_value=None
            )
            # if "machines" is given, assume user wants to do distributed learning, and set up network
            if params["machines"] is None:
                params.pop("machines", None)
            else:
                machines = params["machines"]
                if isinstance(machines, str):
                    num_machines_from_machine_list = len(machines.split(','))
                elif isinstance(machines, (list, set)):
                    num_machines_from_machine_list = len(machines)
                    machines = ','.join(machines)
                else:
                    raise ValueError("Invalid machines in params.")

                params = _choose_param_value(
                    main_param_name="num_machines",
                    params=params,
                    default_value=num_machines_from_machine_list
                )
                params = _choose_param_value(
                    main_param_name="local_listen_port",
                    params=params,
                    default_value=12400
                )
                self.set_network(
                    machines=machines,
                    local_listen_port=params["local_listen_port"],
                    listen_time_out=params.get("time_out", 120),
                    num_machines=params["num_machines"]
                )
3096
            # construct booster object
3097
3098
3099
            train_set.construct()
            # copy the parameters from train_set
            params.update(train_set.get_params())
3100
            params_str = _param_dict_to_str(params)
3101
            self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
3102
            _safe_call(_LIB.LGBM_BoosterCreate(
3103
                train_set._handle,
3104
                _c_str(params_str),
3105
                ctypes.byref(self._handle)))
3106
            # save reference to data
wxchan's avatar
wxchan committed
3107
            self.train_set = train_set
3108
3109
            self.valid_sets: List[Dataset] = []
            self.name_valid_sets: List[str] = []
wxchan's avatar
wxchan committed
3110
            self.__num_dataset = 1
Guolin Ke's avatar
Guolin Ke committed
3111
3112
            self.__init_predictor = train_set._predictor
            if self.__init_predictor is not None:
wxchan's avatar
wxchan committed
3113
                _safe_call(_LIB.LGBM_BoosterMerge(
3114
3115
                    self._handle,
                    self.__init_predictor._handle))
Guolin Ke's avatar
Guolin Ke committed
3116
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
3117
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
3118
                self._handle,
wxchan's avatar
wxchan committed
3119
3120
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
3121
            # buffer for inner predict
3122
            self.__inner_predict_buffer: List[Optional[np.ndarray]] = [None]
wxchan's avatar
wxchan committed
3123
3124
            self.__is_predicted_cur_iter = [False]
            self.__get_eval_info()
3125
            self.pandas_categorical = train_set.pandas_categorical
3126
            self.train_set_version = train_set.version
wxchan's avatar
wxchan committed
3127
        elif model_file is not None:
3128
            # Prediction task
Guolin Ke's avatar
Guolin Ke committed
3129
            out_num_iterations = ctypes.c_int(0)
3130
            self._handle = ctypes.c_void_p()
wxchan's avatar
wxchan committed
3131
            _safe_call(_LIB.LGBM_BoosterCreateFromModelfile(
3132
                _c_str(str(model_file)),
wxchan's avatar
wxchan committed
3133
                ctypes.byref(out_num_iterations),
3134
                ctypes.byref(self._handle)))
Guolin Ke's avatar
Guolin Ke committed
3135
            out_num_class = ctypes.c_int(0)
wxchan's avatar
wxchan committed
3136
            _safe_call(_LIB.LGBM_BoosterGetNumClasses(
3137
                self._handle,
wxchan's avatar
wxchan committed
3138
3139
                ctypes.byref(out_num_class)))
            self.__num_class = out_num_class.value
3140
            self.pandas_categorical = _load_pandas_categorical(file_name=model_file)
3141
3142
3143
            if params:
                _log_warning('Ignoring params argument, using parameters from model file.')
            params = self._get_loaded_param()
3144
        elif model_str is not None:
3145
            self.model_from_string(model_str)
wxchan's avatar
wxchan committed
3146
        else:
3147
3148
            raise TypeError('Need at least one training dataset or model file or model string '
                            'to create Booster instance')
3149
        self.params = params
wxchan's avatar
wxchan committed
3150

3151
    def __del__(self) -> None:
3152
        try:
3153
            if self._network:
3154
3155
3156
3157
                self.free_network()
        except AttributeError:
            pass
        try:
3158
3159
            if self._handle is not None:
                _safe_call(_LIB.LGBM_BoosterFree(self._handle))
3160
3161
        except AttributeError:
            pass
wxchan's avatar
wxchan committed
3162

3163
    def __copy__(self) -> "Booster":
wxchan's avatar
wxchan committed
3164
3165
        return self.__deepcopy__(None)

3166
    def __deepcopy__(self, _) -> "Booster":
3167
        model_str = self.model_to_string(num_iteration=-1)
3168
        booster = Booster(model_str=model_str)
3169
        return booster
wxchan's avatar
wxchan committed
3170

3171
    def __getstate__(self) -> Dict[str, Any]:
wxchan's avatar
wxchan committed
3172
        this = self.__dict__.copy()
3173
        handle = this['_handle']
wxchan's avatar
wxchan committed
3174
3175
3176
        this.pop('train_set', None)
        this.pop('valid_sets', None)
        if handle is not None:
3177
            this["_handle"] = self.model_to_string(num_iteration=-1)
wxchan's avatar
wxchan committed
3178
3179
        return this

3180
    def __setstate__(self, state: Dict[str, Any]) -> None:
3181
        model_str = state.get('_handle', state.get('handle', None))
3182
        if model_str is not None:
wxchan's avatar
wxchan committed
3183
            handle = ctypes.c_void_p()
Guolin Ke's avatar
Guolin Ke committed
3184
            out_num_iterations = ctypes.c_int(0)
3185
            _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
3186
                _c_str(model_str),
3187
3188
                ctypes.byref(out_num_iterations),
                ctypes.byref(handle)))
3189
            state['_handle'] = handle
wxchan's avatar
wxchan committed
3190
3191
        self.__dict__.update(state)

3192
3193
3194
3195
3196
3197
    def _get_loaded_param(self) -> Dict[str, Any]:
        buffer_len = 1 << 20
        tmp_out_len = ctypes.c_int64(0)
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterGetLoadedParam(
3198
            self._handle,
3199
3200
3201
3202
3203
3204
3205
3206
3207
            ctypes.c_int64(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_BoosterGetLoadedParam(
3208
                self._handle,
3209
3210
3211
3212
3213
                ctypes.c_int64(actual_len),
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
        return json.loads(string_buffer.value.decode('utf-8'))

3214
    def free_dataset(self) -> "Booster":
Nikita Titov's avatar
Nikita Titov committed
3215
3216
3217
3218
3219
3220
3221
        """Free Booster's Datasets.

        Returns
        -------
        self : Booster
            Booster without Datasets.
        """
wxchan's avatar
wxchan committed
3222
3223
        self.__dict__.pop('train_set', None)
        self.__dict__.pop('valid_sets', None)
3224
        self.__num_dataset = 0
Nikita Titov's avatar
Nikita Titov committed
3225
        return self
wxchan's avatar
wxchan committed
3226

3227
    def _free_buffer(self) -> "Booster":
3228
3229
        self.__inner_predict_buffer = []
        self.__is_predicted_cur_iter = []
Nikita Titov's avatar
Nikita Titov committed
3230
        return self
3231

3232
3233
3234
3235
3236
3237
3238
    def set_network(
        self,
        machines: Union[List[str], Set[str], str],
        local_listen_port: int = 12400,
        listen_time_out: int = 120,
        num_machines: int = 1
    ) -> "Booster":
3239
3240
3241
3242
        """Set the network configuration.

        Parameters
        ----------
3243
        machines : list, set or str
3244
            Names of machines.
Nikita Titov's avatar
Nikita Titov committed
3245
        local_listen_port : int, optional (default=12400)
3246
            TCP listen port for local machines.
Nikita Titov's avatar
Nikita Titov committed
3247
        listen_time_out : int, optional (default=120)
3248
            Socket time-out in minutes.
Nikita Titov's avatar
Nikita Titov committed
3249
        num_machines : int, optional (default=1)
3250
            The number of machines for distributed learning application.
Nikita Titov's avatar
Nikita Titov committed
3251
3252
3253
3254
3255

        Returns
        -------
        self : Booster
            Booster with set network.
3256
        """
3257
3258
        if isinstance(machines, (list, set)):
            machines = ','.join(machines)
3259
        _safe_call(_LIB.LGBM_NetworkInit(_c_str(machines),
3260
3261
3262
                                         ctypes.c_int(local_listen_port),
                                         ctypes.c_int(listen_time_out),
                                         ctypes.c_int(num_machines)))
3263
        self._network = True
Nikita Titov's avatar
Nikita Titov committed
3264
        return self
3265

3266
    def free_network(self) -> "Booster":
Nikita Titov's avatar
Nikita Titov committed
3267
3268
3269
3270
3271
3272
3273
        """Free Booster's network.

        Returns
        -------
        self : Booster
            Booster with freed network.
        """
3274
        _safe_call(_LIB.LGBM_NetworkFree())
3275
        self._network = False
Nikita Titov's avatar
Nikita Titov committed
3276
        return self
3277

3278
    def trees_to_dataframe(self) -> pd_DataFrame:
3279
3280
        """Parse the fitted model and return in an easy-to-read pandas DataFrame.

3281
3282
3283
3284
        The returned DataFrame has the following columns.

            - ``tree_index`` : int64, which tree a node belongs to. 0-based, so a value of ``6``, for example, means "this node is in the 7th tree".
            - ``node_depth`` : int64, how far a node is from the root of the tree. The root node has a value of ``1``, its direct children are ``2``, etc.
3285
3286
3287
3288
3289
            - ``node_index`` : str, unique identifier for a node.
            - ``left_child`` : str, ``node_index`` of the child node to the left of a split. ``None`` for leaf nodes.
            - ``right_child`` : str, ``node_index`` of the child node to the right of a split. ``None`` for leaf nodes.
            - ``parent_index`` : str, ``node_index`` of this node's parent. ``None`` for the root node.
            - ``split_feature`` : str, name of the feature used for splitting. ``None`` for leaf nodes.
3290
3291
            - ``split_gain`` : float64, gain from adding this split to the tree. ``NaN`` for leaf nodes.
            - ``threshold`` : float64, value of the feature used to decide which side of the split a record will go down. ``NaN`` for leaf nodes.
3292
            - ``decision_type`` : str, logical operator describing how to compare a value to ``threshold``.
3293
3294
              For example, ``split_feature = "Column_10", threshold = 15, decision_type = "<="`` means that
              records where ``Column_10 <= 15`` follow the left side of the split, otherwise follows the right side of the split. ``None`` for leaf nodes.
3295
3296
            - ``missing_direction`` : str, split direction that missing values should go to. ``None`` for leaf nodes.
            - ``missing_type`` : str, describes what types of values are treated as missing.
3297
            - ``value`` : float64, predicted value for this leaf node, multiplied by the learning rate.
3298
            - ``weight`` : float64 or int64, sum of Hessian (second-order derivative of objective), summed over observations that fall in this node.
3299
3300
            - ``count`` : int64, number of records in the training data that fall into this node.

3301
3302
3303
3304
3305
3306
        Returns
        -------
        result : pandas DataFrame
            Returns a pandas DataFrame of the parsed model.
        """
        if not PANDAS_INSTALLED:
3307
3308
            raise LightGBMError('This method cannot be run without pandas installed. '
                                'You must install pandas and restart your session to use this method.')
3309
3310
3311
3312

        if self.num_trees() == 0:
            raise LightGBMError('There are no trees in this Booster and thus nothing to parse')

3313
        def _is_split_node(tree: Dict[str, Any]) -> bool:
3314
3315
            return 'split_index' in tree.keys()

3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
        def create_node_record(
            tree: Dict[str, Any],
            node_depth: int = 1,
            tree_index: Optional[int] = None,
            feature_names: Optional[List[str]] = None,
            parent_node: Optional[str] = None
        ) -> Dict[str, Any]:

            def _get_node_index(
                tree: Dict[str, Any],
                tree_index: Optional[int]
            ) -> str:
3328
                tree_num = f'{tree_index}-' if tree_index is not None else ''
3329
3330
3331
                is_split = _is_split_node(tree)
                node_type = 'S' if is_split else 'L'
                # if a single node tree it won't have `leaf_index` so return 0
3332
3333
                node_num = tree.get('split_index' if is_split else 'leaf_index', 0)
                return f"{tree_num}{node_type}{node_num}"
3334

3335
3336
3337
3338
            def _get_split_feature(
                tree: Dict[str, Any],
                feature_names: Optional[List[str]]
            ) -> Optional[str]:
3339
3340
3341
3342
3343
3344
3345
3346
3347
                if _is_split_node(tree):
                    if feature_names is not None:
                        feature_name = feature_names[tree['split_feature']]
                    else:
                        feature_name = tree['split_feature']
                else:
                    feature_name = None
                return feature_name

3348
            def _is_single_node_tree(tree: Dict[str, Any]) -> bool:
3349
                return set(tree.keys()) == {'leaf_value'}
3350
3351

            # Create the node record, and populate universal data members
3352
            node: Dict[str, Union[int, str, None]] = OrderedDict()
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
            node['tree_index'] = tree_index
            node['node_depth'] = node_depth
            node['node_index'] = _get_node_index(tree, tree_index)
            node['left_child'] = None
            node['right_child'] = None
            node['parent_index'] = parent_node
            node['split_feature'] = _get_split_feature(tree, feature_names)
            node['split_gain'] = None
            node['threshold'] = None
            node['decision_type'] = None
            node['missing_direction'] = None
            node['missing_type'] = None
            node['value'] = None
            node['weight'] = None
            node['count'] = None

            # Update values to reflect node type (leaf or split)
            if _is_split_node(tree):
                node['left_child'] = _get_node_index(tree['left_child'], tree_index)
                node['right_child'] = _get_node_index(tree['right_child'], tree_index)
                node['split_gain'] = tree['split_gain']
                node['threshold'] = tree['threshold']
                node['decision_type'] = tree['decision_type']
                node['missing_direction'] = 'left' if tree['default_left'] else 'right'
                node['missing_type'] = tree['missing_type']
                node['value'] = tree['internal_value']
                node['weight'] = tree['internal_weight']
                node['count'] = tree['internal_count']
            else:
                node['value'] = tree['leaf_value']
                if not _is_single_node_tree(tree):
                    node['weight'] = tree['leaf_weight']
                    node['count'] = tree['leaf_count']

            return node

3389
3390
3391
3392
3393
3394
3395
        def tree_dict_to_node_list(
            tree: Dict[str, Any],
            node_depth: int = 1,
            tree_index: Optional[int] = None,
            feature_names: Optional[List[str]] = None,
            parent_node: Optional[str] = None
        ) -> List[Dict[str, Any]]:
3396

3397
            node = create_node_record(tree=tree,
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
                                      node_depth=node_depth,
                                      tree_index=tree_index,
                                      feature_names=feature_names,
                                      parent_node=parent_node)

            res = [node]

            if _is_split_node(tree):
                # traverse the next level of the tree
                children = ['left_child', 'right_child']
                for child in children:
                    subtree_list = tree_dict_to_node_list(
3410
                        tree=tree[child],
3411
3412
3413
                        node_depth=node_depth + 1,
                        tree_index=tree_index,
                        feature_names=feature_names,
3414
3415
                        parent_node=node['node_index']
                    )
3416
3417
3418
3419
3420
3421
3422
3423
3424
                    # In tree format, "subtree_list" is a list of node records (dicts),
                    # and we add node to the list.
                    res.extend(subtree_list)
            return res

        model_dict = self.dump_model()
        feature_names = model_dict['feature_names']
        model_list = []
        for tree in model_dict['tree_info']:
3425
            model_list.extend(tree_dict_to_node_list(tree=tree['tree_structure'],
3426
3427
3428
                                                     tree_index=tree['tree_index'],
                                                     feature_names=feature_names))

3429
        return pd_DataFrame(model_list, columns=model_list[0].keys())
3430

3431
    def set_train_data_name(self, name: str) -> "Booster":
3432
3433
3434
3435
        """Set the name to the training Dataset.

        Parameters
        ----------
3436
        name : str
Nikita Titov's avatar
Nikita Titov committed
3437
3438
3439
3440
3441
3442
            Name for the training Dataset.

        Returns
        -------
        self : Booster
            Booster with set training Dataset name.
3443
        """
3444
        self._train_data_name = name
Nikita Titov's avatar
Nikita Titov committed
3445
        return self
wxchan's avatar
wxchan committed
3446

3447
    def add_valid(self, data: Dataset, name: str) -> "Booster":
3448
        """Add validation data.
wxchan's avatar
wxchan committed
3449
3450
3451
3452

        Parameters
        ----------
        data : Dataset
3453
            Validation data.
3454
        name : str
3455
            Name of validation data.
Nikita Titov's avatar
Nikita Titov committed
3456
3457
3458
3459
3460

        Returns
        -------
        self : Booster
            Booster with set validation data.
wxchan's avatar
wxchan committed
3461
        """
Guolin Ke's avatar
Guolin Ke committed
3462
        if not isinstance(data, Dataset):
3463
            raise TypeError(f'Validation data should be Dataset instance, met {type(data).__name__}')
Guolin Ke's avatar
Guolin Ke committed
3464
        if data._predictor is not self.__init_predictor:
3465
3466
            raise LightGBMError("Add validation data failed, "
                                "you should use same predictor for these data")
wxchan's avatar
wxchan committed
3467
        _safe_call(_LIB.LGBM_BoosterAddValidData(
3468
3469
            self._handle,
            data.construct()._handle))
wxchan's avatar
wxchan committed
3470
3471
3472
3473
3474
        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
3475
        return self
wxchan's avatar
wxchan committed
3476

3477
    def reset_parameter(self, params: Dict[str, Any]) -> "Booster":
3478
        """Reset parameters of Booster.
wxchan's avatar
wxchan committed
3479
3480
3481
3482

        Parameters
        ----------
        params : dict
3483
            New parameters for Booster.
Nikita Titov's avatar
Nikita Titov committed
3484
3485
3486
3487
3488

        Returns
        -------
        self : Booster
            Booster with new parameters.
wxchan's avatar
wxchan committed
3489
        """
3490
        params_str = _param_dict_to_str(params)
wxchan's avatar
wxchan committed
3491
3492
        if params_str:
            _safe_call(_LIB.LGBM_BoosterResetParameter(
3493
                self._handle,
3494
                _c_str(params_str)))
Guolin Ke's avatar
Guolin Ke committed
3495
        self.params.update(params)
Nikita Titov's avatar
Nikita Titov committed
3496
        return self
wxchan's avatar
wxchan committed
3497

3498
3499
3500
3501
3502
    def update(
        self,
        train_set: Optional[Dataset] = None,
        fobj: Optional[_LGBM_CustomObjectiveFunction] = None
    ) -> bool:
Nikita Titov's avatar
Nikita Titov committed
3503
        """Update Booster for one iteration.
3504

wxchan's avatar
wxchan committed
3505
3506
        Parameters
        ----------
3507
3508
3509
3510
        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
3511
            Customized objective function.
3512
3513
3514
            Should accept two parameters: preds, train_data,
            and return (grad, hess).

3515
                preds : numpy 1-D array or numpy 2-D array (for multi-class task)
3516
                    The predicted values.
3517
3518
                    Predicted values are returned before any transformation,
                    e.g. they are raw margin instead of probability of positive class for binary task.
3519
3520
                train_data : Dataset
                    The training dataset.
3521
                grad : numpy 1-D array or numpy 2-D array (for multi-class task)
3522
3523
                    The value of the first order derivative (gradient) of the loss
                    with respect to the elements of preds for each sample point.
3524
                hess : numpy 1-D array or numpy 2-D array (for multi-class task)
3525
3526
                    The value of the second order derivative (Hessian) of the loss
                    with respect to the elements of preds for each sample point.
wxchan's avatar
wxchan committed
3527

3528
            For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes],
3529
            and grad and hess should be returned in the same format.
3530

wxchan's avatar
wxchan committed
3531
3532
        Returns
        -------
3533
3534
        is_finished : bool
            Whether the update was successfully finished.
wxchan's avatar
wxchan committed
3535
        """
3536
        # need reset training data
3537
3538
3539
3540
3541
3542
        if train_set is None and self.train_set_version != self.train_set.version:
            train_set = self.train_set
            is_the_same_train_set = False
        else:
            is_the_same_train_set = train_set is self.train_set and self.train_set_version == train_set.version
        if train_set is not None and not is_the_same_train_set:
Guolin Ke's avatar
Guolin Ke committed
3543
            if not isinstance(train_set, Dataset):
3544
                raise TypeError(f'Training data should be Dataset instance, met {type(train_set).__name__}')
Guolin Ke's avatar
Guolin Ke committed
3545
            if train_set._predictor is not self.__init_predictor:
3546
3547
                raise LightGBMError("Replace training data failed, "
                                    "you should use same predictor for these data")
wxchan's avatar
wxchan committed
3548
3549
            self.train_set = train_set
            _safe_call(_LIB.LGBM_BoosterResetTrainingData(
3550
3551
                self._handle,
                self.train_set.construct()._handle))
wxchan's avatar
wxchan committed
3552
            self.__inner_predict_buffer[0] = None
3553
            self.train_set_version = self.train_set.version
wxchan's avatar
wxchan committed
3554
3555
        is_finished = ctypes.c_int(0)
        if fobj is None:
3556
            if self.__set_objective_to_none:
3557
                raise LightGBMError('Cannot update due to null objective function.')
wxchan's avatar
wxchan committed
3558
            _safe_call(_LIB.LGBM_BoosterUpdateOneIter(
3559
                self._handle,
wxchan's avatar
wxchan committed
3560
                ctypes.byref(is_finished)))
3561
            self.__is_predicted_cur_iter = [False for _ in range(self.__num_dataset)]
wxchan's avatar
wxchan committed
3562
3563
            return is_finished.value == 1
        else:
3564
            if not self.__set_objective_to_none:
Nikita Titov's avatar
Nikita Titov committed
3565
                self.reset_parameter({"objective": "none"}).__set_objective_to_none = True
wxchan's avatar
wxchan committed
3566
3567
3568
            grad, hess = fobj(self.__inner_predict(0), self.train_set)
            return self.__boost(grad, hess)

3569
3570
3571
3572
3573
    def __boost(
        self,
        grad: np.ndarray,
        hess: np.ndarray
    ) -> bool:
3574
        """Boost Booster for one iteration with customized gradient statistics.
Nikita Titov's avatar
Nikita Titov committed
3575

Nikita Titov's avatar
Nikita Titov committed
3576
3577
        .. note::

3578
3579
            Score is returned before any transformation,
            e.g. it is raw margin instead of probability of positive class for binary task.
3580
            For multi-class task, score are numpy 2-D array of shape = [n_samples, n_classes],
3581
            and grad and hess should be returned in the same format.
3582

wxchan's avatar
wxchan committed
3583
3584
        Parameters
        ----------
3585
        grad : numpy 1-D array or numpy 2-D array (for multi-class task)
3586
3587
            The value of the first order derivative (gradient) of the loss
            with respect to the elements of score for each sample point.
3588
        hess : numpy 1-D array or numpy 2-D array (for multi-class task)
3589
3590
            The value of the second order derivative (Hessian) of the loss
            with respect to the elements of score for each sample point.
wxchan's avatar
wxchan committed
3591
3592
3593

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
3594
3595
        is_finished : bool
            Whether the boost was successfully finished.
wxchan's avatar
wxchan committed
3596
        """
3597
3598
3599
        if self.__num_class > 1:
            grad = grad.ravel(order='F')
            hess = hess.ravel(order='F')
3600
3601
        grad = _list_to_1d_numpy(grad, dtype=np.float32, name='gradient')
        hess = _list_to_1d_numpy(hess, dtype=np.float32, name='hessian')
3602
3603
        assert grad.flags.c_contiguous
        assert hess.flags.c_contiguous
wxchan's avatar
wxchan committed
3604
        if len(grad) != len(hess):
3605
3606
            raise ValueError(f"Lengths of gradient ({len(grad)}) and Hessian ({len(hess)}) don't match")
        num_train_data = self.train_set.num_data()
3607
        if len(grad) != num_train_data * self.__num_class:
3608
3609
3610
            raise ValueError(
                f"Lengths of gradient ({len(grad)}) and Hessian ({len(hess)}) "
                f"don't match training data length ({num_train_data}) * "
3611
                f"number of models per one iteration ({self.__num_class})"
3612
            )
wxchan's avatar
wxchan committed
3613
3614
        is_finished = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterUpdateOneIterCustom(
3615
            self._handle,
wxchan's avatar
wxchan committed
3616
3617
3618
            grad.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            hess.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
            ctypes.byref(is_finished)))
3619
        self.__is_predicted_cur_iter = [False for _ in range(self.__num_dataset)]
wxchan's avatar
wxchan committed
3620
3621
        return is_finished.value == 1

3622
    def rollback_one_iter(self) -> "Booster":
Nikita Titov's avatar
Nikita Titov committed
3623
3624
3625
3626
3627
3628
3629
        """Rollback one iteration.

        Returns
        -------
        self : Booster
            Booster with rolled back one iteration.
        """
wxchan's avatar
wxchan committed
3630
        _safe_call(_LIB.LGBM_BoosterRollbackOneIter(
3631
            self._handle))
3632
        self.__is_predicted_cur_iter = [False for _ in range(self.__num_dataset)]
Nikita Titov's avatar
Nikita Titov committed
3633
        return self
wxchan's avatar
wxchan committed
3634

3635
    def current_iteration(self) -> int:
3636
3637
3638
3639
3640
3641
3642
        """Get the index of the current iteration.

        Returns
        -------
        cur_iter : int
            The index of the current iteration.
        """
Guolin Ke's avatar
Guolin Ke committed
3643
        out_cur_iter = ctypes.c_int(0)
wxchan's avatar
wxchan committed
3644
        _safe_call(_LIB.LGBM_BoosterGetCurrentIteration(
3645
            self._handle,
wxchan's avatar
wxchan committed
3646
3647
3648
            ctypes.byref(out_cur_iter)))
        return out_cur_iter.value

3649
    def num_model_per_iteration(self) -> int:
3650
3651
3652
3653
3654
3655
3656
3657
3658
        """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(
3659
            self._handle,
3660
3661
3662
            ctypes.byref(model_per_iter)))
        return model_per_iter.value

3663
    def num_trees(self) -> int:
3664
3665
3666
3667
3668
3669
3670
3671
3672
        """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(
3673
            self._handle,
3674
3675
3676
            ctypes.byref(num_trees)))
        return num_trees.value

3677
    def upper_bound(self) -> float:
3678
3679
3680
3681
        """Get upper bound value of a model.

        Returns
        -------
3682
        upper_bound : float
3683
3684
3685
3686
            Upper bound value of the model.
        """
        ret = ctypes.c_double(0)
        _safe_call(_LIB.LGBM_BoosterGetUpperBoundValue(
3687
            self._handle,
3688
3689
3690
            ctypes.byref(ret)))
        return ret.value

3691
    def lower_bound(self) -> float:
3692
3693
3694
3695
        """Get lower bound value of a model.

        Returns
        -------
3696
        lower_bound : float
3697
3698
3699
3700
            Lower bound value of the model.
        """
        ret = ctypes.c_double(0)
        _safe_call(_LIB.LGBM_BoosterGetLowerBoundValue(
3701
            self._handle,
3702
3703
3704
            ctypes.byref(ret)))
        return ret.value

3705
3706
3707
3708
3709
3710
    def eval(
        self,
        data: Dataset,
        name: str,
        feval: Optional[Union[_LGBM_CustomEvalFunction, List[_LGBM_CustomEvalFunction]]] = None
    ) -> List[_LGBM_BoosterEvalMethodResultType]:
3711
        """Evaluate for data.
wxchan's avatar
wxchan committed
3712
3713
3714

        Parameters
        ----------
3715
3716
        data : Dataset
            Data for the evaluating.
3717
        name : str
3718
            Name of the data.
3719
        feval : callable, list of callable, or None, optional (default=None)
3720
            Customized evaluation function.
3721
            Each evaluation function should accept two parameters: preds, eval_data,
3722
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
3723

3724
                preds : numpy 1-D array or numpy 2-D array (for multi-class task)
3725
                    The predicted values.
3726
                    For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes].
3727
                    If custom objective function is used, predicted values are returned before any transformation,
3728
                    e.g. they are raw margin instead of probability of positive class for binary task in this case.
3729
                eval_data : Dataset
3730
                    A ``Dataset`` to evaluate.
3731
                eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
3732
                    The name of evaluation function (without whitespace).
3733
3734
3735
3736
3737
                eval_result : float
                    The eval result.
                is_higher_better : bool
                    Is eval result higher better, e.g. AUC is ``is_higher_better``.

wxchan's avatar
wxchan committed
3738
3739
        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
3740
        result : list
3741
            List with (dataset_name, eval_name, eval_result, is_higher_better) tuples.
wxchan's avatar
wxchan committed
3742
        """
Guolin Ke's avatar
Guolin Ke committed
3743
3744
        if not isinstance(data, Dataset):
            raise TypeError("Can only eval for Dataset instance")
wxchan's avatar
wxchan committed
3745
3746
3747
3748
        data_idx = -1
        if data is self.train_set:
            data_idx = 0
        else:
3749
            for i in range(len(self.valid_sets)):
wxchan's avatar
wxchan committed
3750
3751
3752
                if data is self.valid_sets[i]:
                    data_idx = i + 1
                    break
3753
        # need to push new valid data
wxchan's avatar
wxchan committed
3754
3755
3756
3757
3758
3759
        if data_idx == -1:
            self.add_valid(data, name)
            data_idx = self.__num_dataset - 1

        return self.__inner_eval(name, data_idx, feval)

3760
3761
3762
3763
    def eval_train(
        self,
        feval: Optional[Union[_LGBM_CustomEvalFunction, List[_LGBM_CustomEvalFunction]]] = None
    ) -> List[_LGBM_BoosterEvalMethodResultType]:
3764
        """Evaluate for training data.
wxchan's avatar
wxchan committed
3765
3766
3767

        Parameters
        ----------
3768
        feval : callable, list of callable, or None, optional (default=None)
3769
            Customized evaluation function.
3770
            Each evaluation function should accept two parameters: preds, eval_data,
3771
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
3772

3773
                preds : numpy 1-D array or numpy 2-D array (for multi-class task)
3774
                    The predicted values.
3775
                    For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes].
3776
                    If custom objective function is used, predicted values are returned before any transformation,
3777
                    e.g. they are raw margin instead of probability of positive class for binary task in this case.
Akshita Dixit's avatar
Akshita Dixit committed
3778
                eval_data : Dataset
3779
                    The training dataset.
3780
                eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
3781
                    The name of evaluation function (without whitespace).
3782
3783
3784
3785
3786
                eval_result : float
                    The eval result.
                is_higher_better : bool
                    Is eval result higher better, e.g. AUC is ``is_higher_better``.

wxchan's avatar
wxchan committed
3787
3788
        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
3789
        result : list
3790
            List with (train_dataset_name, eval_name, eval_result, is_higher_better) tuples.
wxchan's avatar
wxchan committed
3791
        """
3792
        return self.__inner_eval(self._train_data_name, 0, feval)
wxchan's avatar
wxchan committed
3793

3794
3795
3796
3797
    def eval_valid(
        self,
        feval: Optional[Union[_LGBM_CustomEvalFunction, List[_LGBM_CustomEvalFunction]]] = None
    ) -> List[_LGBM_BoosterEvalMethodResultType]:
3798
        """Evaluate for validation data.
wxchan's avatar
wxchan committed
3799
3800
3801

        Parameters
        ----------
3802
        feval : callable, list of callable, or None, optional (default=None)
3803
            Customized evaluation function.
3804
            Each evaluation function should accept two parameters: preds, eval_data,
3805
            and return (eval_name, eval_result, is_higher_better) or list of such tuples.
3806

3807
                preds : numpy 1-D array or numpy 2-D array (for multi-class task)
3808
                    The predicted values.
3809
                    For multi-class task, preds are numpy 2-D array of shape = [n_samples, n_classes].
3810
                    If custom objective function is used, predicted values are returned before any transformation,
3811
                    e.g. they are raw margin instead of probability of positive class for binary task in this case.
Akshita Dixit's avatar
Akshita Dixit committed
3812
                eval_data : Dataset
3813
                    The validation dataset.
3814
                eval_name : str
Andrew Ziem's avatar
Andrew Ziem committed
3815
                    The name of evaluation function (without whitespace).
3816
3817
3818
3819
3820
                eval_result : float
                    The eval result.
                is_higher_better : bool
                    Is eval result higher better, e.g. AUC is ``is_higher_better``.

wxchan's avatar
wxchan committed
3821
3822
        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
3823
        result : list
3824
            List with (validation_dataset_name, eval_name, eval_result, is_higher_better) tuples.
wxchan's avatar
wxchan committed
3825
        """
3826
        return [item for i in range(1, self.__num_dataset)
wxchan's avatar
wxchan committed
3827
                for item in self.__inner_eval(self.name_valid_sets[i - 1], i, feval)]
wxchan's avatar
wxchan committed
3828

3829
3830
3831
3832
3833
3834
3835
    def save_model(
        self,
        filename: Union[str, Path],
        num_iteration: Optional[int] = None,
        start_iteration: int = 0,
        importance_type: str = 'split'
    ) -> "Booster":
3836
        """Save Booster to file.
wxchan's avatar
wxchan committed
3837
3838
3839

        Parameters
        ----------
3840
        filename : str or pathlib.Path
3841
            Filename to save Booster.
3842
3843
3844
3845
        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
3846
        start_iteration : int, optional (default=0)
3847
            Start index of the iteration that should be saved.
3848
        importance_type : str, optional (default="split")
3849
3850
3851
            What type of feature importance should be saved.
            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.
Nikita Titov's avatar
Nikita Titov committed
3852
3853
3854
3855
3856

        Returns
        -------
        self : Booster
            Returns self.
wxchan's avatar
wxchan committed
3857
        """
3858
        if num_iteration is None:
3859
            num_iteration = self.best_iteration
3860
        importance_type_int = _FEATURE_IMPORTANCE_TYPE_MAPPER[importance_type]
wxchan's avatar
wxchan committed
3861
        _safe_call(_LIB.LGBM_BoosterSaveModel(
3862
            self._handle,
3863
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
3864
            ctypes.c_int(num_iteration),
3865
            ctypes.c_int(importance_type_int),
3866
            _c_str(str(filename))))
3867
        _dump_pandas_categorical(self.pandas_categorical, filename)
Nikita Titov's avatar
Nikita Titov committed
3868
        return self
wxchan's avatar
wxchan committed
3869

3870
3871
3872
3873
3874
    def shuffle_models(
        self,
        start_iteration: int = 0,
        end_iteration: int = -1
    ) -> "Booster":
3875
        """Shuffle models.
Nikita Titov's avatar
Nikita Titov committed
3876

3877
3878
3879
        Parameters
        ----------
        start_iteration : int, optional (default=0)
3880
            The first iteration that will be shuffled.
3881
3882
        end_iteration : int, optional (default=-1)
            The last iteration that will be shuffled.
3883
            If <= 0, means the last available iteration.
3884

Nikita Titov's avatar
Nikita Titov committed
3885
3886
3887
3888
        Returns
        -------
        self : Booster
            Booster with shuffled models.
3889
        """
3890
        _safe_call(_LIB.LGBM_BoosterShuffleModels(
3891
            self._handle,
Guolin Ke's avatar
Guolin Ke committed
3892
3893
            ctypes.c_int(start_iteration),
            ctypes.c_int(end_iteration)))
Nikita Titov's avatar
Nikita Titov committed
3894
        return self
3895

3896
    def model_from_string(self, model_str: str) -> "Booster":
3897
3898
3899
3900
        """Load Booster from a string.

        Parameters
        ----------
3901
        model_str : str
3902
3903
3904
3905
            Model will be loaded from this string.

        Returns
        -------
Nikita Titov's avatar
Nikita Titov committed
3906
        self : Booster
3907
3908
            Loaded Booster object.
        """
3909
3910
        if self._handle is not None:
            _safe_call(_LIB.LGBM_BoosterFree(self._handle))
3911
        self._free_buffer()
3912
        self._handle = ctypes.c_void_p()
3913
3914
        out_num_iterations = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterLoadModelFromString(
3915
            _c_str(model_str),
3916
            ctypes.byref(out_num_iterations),
3917
            ctypes.byref(self._handle)))
3918
3919
        out_num_class = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetNumClasses(
3920
            self._handle,
3921
3922
            ctypes.byref(out_num_class)))
        self.__num_class = out_num_class.value
3923
        self.pandas_categorical = _load_pandas_categorical(model_str=model_str)
3924
3925
        return self

3926
3927
3928
3929
3930
3931
    def model_to_string(
        self,
        num_iteration: Optional[int] = None,
        start_iteration: int = 0,
        importance_type: str = 'split'
    ) -> str:
3932
        """Save Booster to string.
3933

3934
3935
3936
3937
3938
3939
        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
3940
        start_iteration : int, optional (default=0)
3941
            Start index of the iteration that should be saved.
3942
        importance_type : str, optional (default="split")
3943
3944
3945
            What type of feature importance should be saved.
            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.
3946
3947
3948

        Returns
        -------
3949
        str_repr : str
3950
3951
            String representation of Booster.
        """
3952
        if num_iteration is None:
3953
            num_iteration = self.best_iteration
3954
        importance_type_int = _FEATURE_IMPORTANCE_TYPE_MAPPER[importance_type]
3955
        buffer_len = 1 << 20
3956
        tmp_out_len = ctypes.c_int64(0)
3957
3958
3959
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterSaveModelToString(
3960
            self._handle,
3961
            ctypes.c_int(start_iteration),
3962
            ctypes.c_int(num_iteration),
3963
            ctypes.c_int(importance_type_int),
3964
            ctypes.c_int64(buffer_len),
3965
3966
3967
            ctypes.byref(tmp_out_len),
            ptr_string_buffer))
        actual_len = tmp_out_len.value
3968
        # if buffer length is not long enough, re-allocate a buffer
3969
3970
3971
3972
        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(
3973
                self._handle,
3974
                ctypes.c_int(start_iteration),
3975
                ctypes.c_int(num_iteration),
3976
                ctypes.c_int(importance_type_int),
3977
                ctypes.c_int64(actual_len),
3978
3979
                ctypes.byref(tmp_out_len),
                ptr_string_buffer))
3980
        ret = string_buffer.value.decode('utf-8')
3981
3982
        ret += _dump_pandas_categorical(self.pandas_categorical)
        return ret
3983

3984
3985
3986
3987
3988
3989
3990
    def dump_model(
        self,
        num_iteration: Optional[int] = None,
        start_iteration: int = 0,
        importance_type: str = 'split',
        object_hook: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None
    ) -> Dict[str, Any]:
Nikita Titov's avatar
Nikita Titov committed
3991
        """Dump Booster to JSON format.
wxchan's avatar
wxchan committed
3992

3993
3994
        Parameters
        ----------
3995
3996
3997
3998
        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
3999
        start_iteration : int, optional (default=0)
4000
            Start index of the iteration that should be dumped.
4001
        importance_type : str, optional (default="split")
4002
4003
4004
            What type of feature importance should be dumped.
            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.
4005
4006
4007
4008
4009
4010
4011
4012
4013
        object_hook : callable or None, optional (default=None)
            If not None, ``object_hook`` is a function called while parsing the json
            string returned by the C API. It may be used to alter the json, to store
            specific values while building the json structure. It avoids
            walking through the structure again. It saves a significant amount
            of time if the number of trees is huge.
            Signature is ``def object_hook(node: dict) -> dict``.
            None is equivalent to ``lambda node: node``.
            See documentation of ``json.loads()`` for further details.
4014

wxchan's avatar
wxchan committed
4015
4016
        Returns
        -------
4017
        json_repr : dict
Nikita Titov's avatar
Nikita Titov committed
4018
            JSON format of Booster.
wxchan's avatar
wxchan committed
4019
        """
4020
        if num_iteration is None:
4021
            num_iteration = self.best_iteration
4022
        importance_type_int = _FEATURE_IMPORTANCE_TYPE_MAPPER[importance_type]
wxchan's avatar
wxchan committed
4023
        buffer_len = 1 << 20
4024
        tmp_out_len = ctypes.c_int64(0)
wxchan's avatar
wxchan committed
4025
4026
4027
        string_buffer = ctypes.create_string_buffer(buffer_len)
        ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)])
        _safe_call(_LIB.LGBM_BoosterDumpModel(
4028
            self._handle,
4029
            ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
4030
            ctypes.c_int(num_iteration),
4031
            ctypes.c_int(importance_type_int),
4032
            ctypes.c_int64(buffer_len),
wxchan's avatar
wxchan committed
4033
            ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
4034
            ptr_string_buffer))
wxchan's avatar
wxchan committed
4035
        actual_len = tmp_out_len.value
4036
        # if buffer length is not long enough, reallocate a buffer
wxchan's avatar
wxchan committed
4037
4038
4039
4040
        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(
4041
                self._handle,
4042
                ctypes.c_int(start_iteration),
Guolin Ke's avatar
Guolin Ke committed
4043
                ctypes.c_int(num_iteration),
4044
                ctypes.c_int(importance_type_int),
4045
                ctypes.c_int64(actual_len),
wxchan's avatar
wxchan committed
4046
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
4047
                ptr_string_buffer))
4048
        ret = json.loads(string_buffer.value.decode('utf-8'), object_hook=object_hook)
4049
        ret['pandas_categorical'] = json.loads(json.dumps(self.pandas_categorical,
4050
                                                          default=_json_default_with_numpy))
4051
        return ret
wxchan's avatar
wxchan committed
4052

4053
4054
    def predict(
        self,
4055
        data: _LGBM_PredictDataType,
4056
4057
4058
4059
4060
4061
4062
4063
        start_iteration: int = 0,
        num_iteration: Optional[int] = None,
        raw_score: bool = False,
        pred_leaf: bool = False,
        pred_contrib: bool = False,
        data_has_header: bool = False,
        validate_features: bool = False,
        **kwargs: Any
4064
    ) -> Union[np.ndarray, scipy.sparse.spmatrix, List[scipy.sparse.spmatrix]]:
4065
        """Make a prediction.
wxchan's avatar
wxchan committed
4066
4067
4068

        Parameters
        ----------
4069
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame or scipy.sparse
4070
            Data source for prediction.
4071
            If str or pathlib.Path, it represents the path to a text file (CSV, TSV, or LibSVM).
4072
        start_iteration : int, optional (default=0)
4073
            Start index of the iteration to predict.
4074
            If <= 0, starts from the first iteration.
4075
        num_iteration : int or None, optional (default=None)
4076
4077
4078
4079
            Total number of iterations used in the prediction.
            If None, if the best iteration exists and start_iteration <= 0, the best iteration is used;
            otherwise, all iterations from ``start_iteration`` are used (no limits).
            If <= 0, all iterations from ``start_iteration`` are used (no limits).
4080
4081
4082
4083
        raw_score : bool, optional (default=False)
            Whether to predict raw scores.
        pred_leaf : bool, optional (default=False)
            Whether to predict leaf index.
4084
4085
        pred_contrib : bool, optional (default=False)
            Whether to predict feature contributions.
4086

Nikita Titov's avatar
Nikita Titov committed
4087
4088
4089
4090
4091
4092
4093
            .. note::

                If you want to get more explanations for your model's predictions using SHAP values,
                like SHAP interaction values,
                you can install the shap package (https://github.com/slundberg/shap).
                Note that unlike the shap package, with ``pred_contrib`` we return a matrix with an extra
                column, where the last column is the expected value.
4094

4095
4096
        data_has_header : bool, optional (default=False)
            Whether the data has header.
4097
            Used only if data is str.
4098
4099
4100
        validate_features : bool, optional (default=False)
            If True, ensure that the features used to predict match the ones used to train.
            Used only if data is pandas DataFrame.
4101
4102
        **kwargs
            Other parameters for the prediction.
wxchan's avatar
wxchan committed
4103
4104
4105

        Returns
        -------
4106
        result : numpy array, scipy.sparse or list of scipy.sparse
4107
            Prediction result.
4108
            Can be sparse or a list of sparse objects (each element represents predictions for one class) for feature contributions (when ``pred_contrib=True``).
wxchan's avatar
wxchan committed
4109
        """
4110
        predictor = self._to_predictor(pred_parameter=deepcopy(kwargs))
4111
        if num_iteration is None:
4112
            if start_iteration <= 0:
4113
4114
4115
                num_iteration = self.best_iteration
            else:
                num_iteration = -1
4116
4117
4118
4119
4120
4121
4122
4123
4124
4125
        return predictor.predict(
            data=data,
            start_iteration=start_iteration,
            num_iteration=num_iteration,
            raw_score=raw_score,
            pred_leaf=pred_leaf,
            pred_contrib=pred_contrib,
            data_has_header=data_has_header,
            validate_features=validate_features
        )
wxchan's avatar
wxchan committed
4126

4127
4128
    def refit(
        self,
4129
        data: _LGBM_TrainDataType,
4130
        label: _LGBM_LabelType,
4131
4132
        decay_rate: float = 0.9,
        reference: Optional[Dataset] = None,
4133
4134
4135
        weight: Optional[_LGBM_WeightType] = None,
        group: Optional[_LGBM_GroupType] = None,
        init_score: Optional[_LGBM_InitScoreType] = None,
4136
4137
        feature_name: _LGBM_FeatureNameConfiguration = 'auto',
        categorical_feature: _LGBM_CategoricalFeatureConfiguration = 'auto',
4138
4139
4140
        dataset_params: Optional[Dict[str, Any]] = None,
        free_raw_data: bool = True,
        validate_features: bool = False,
4141
        **kwargs
4142
    ) -> "Booster":
Guolin Ke's avatar
Guolin Ke committed
4143
4144
4145
4146
        """Refit the existing Booster by new data.

        Parameters
        ----------
4147
        data : str, pathlib.Path, numpy array, pandas DataFrame, H2O DataTable's Frame, scipy.sparse, Sequence, list of Sequence or list of numpy array
Guolin Ke's avatar
Guolin Ke committed
4148
            Data source for refit.
4149
            If str or pathlib.Path, it represents the path to a text file (CSV, TSV, or LibSVM).
4150
        label : list, numpy 1-D array or pandas Series / one-column DataFrame
Guolin Ke's avatar
Guolin Ke committed
4151
4152
            Label for refit.
        decay_rate : float, optional (default=0.9)
4153
4154
            Decay rate of refit,
            will use ``leaf_output = decay_rate * old_leaf_output + (1.0 - decay_rate) * new_leaf_output`` to refit trees.
4155
4156
        reference : Dataset or None, optional (default=None)
            Reference for ``data``.
4157
4158
4159

            .. versionadded:: 4.0.0

4160
        weight : list, numpy 1-D array, pandas Series or None, optional (default=None)
4161
            Weight for each ``data`` instance. Weights should be non-negative.
4162
4163
4164

            .. versionadded:: 4.0.0

4165
4166
4167
4168
4169
4170
        group : list, numpy 1-D array, pandas Series or None, optional (default=None)
            Group/query size for ``data``.
            Only used in the learning-to-rank task.
            sum(group) = n_samples.
            For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
            where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
4171
4172
4173

            .. versionadded:: 4.0.0

4174
4175
        init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), or None, optional (default=None)
            Init score for ``data``.
4176
4177
4178

            .. versionadded:: 4.0.0

4179
4180
4181
        feature_name : list of str, or 'auto', optional (default="auto")
            Feature names for ``data``.
            If 'auto' and data is pandas DataFrame, data columns names are used.
4182
4183
4184

            .. versionadded:: 4.0.0

4185
4186
4187
4188
4189
        categorical_feature : list of str or int, or 'auto', optional (default="auto")
            Categorical features for ``data``.
            If list of int, interpreted as indices.
            If list of str, interpreted as feature names (need to specify ``feature_name`` as well).
            If 'auto' and data is pandas DataFrame, pandas unordered categorical columns are used.
4190
            All values in categorical features will be cast to int32 and thus should be less than int32 max value (2147483647).
4191
4192
4193
            Large values could be memory consuming. Consider using consecutive integers starting from zero.
            All negative values in categorical features will be treated as missing values.
            The output cannot be monotonically constrained with respect to a categorical feature.
4194
            Floating point numbers in categorical features will be rounded towards 0.
4195
4196
4197

            .. versionadded:: 4.0.0

4198
4199
        dataset_params : dict or None, optional (default=None)
            Other parameters for Dataset ``data``.
4200
4201
4202

            .. versionadded:: 4.0.0

4203
4204
        free_raw_data : bool, optional (default=True)
            If True, raw data is freed after constructing inner Dataset for ``data``.
4205
4206
4207

            .. versionadded:: 4.0.0

4208
4209
4210
        validate_features : bool, optional (default=False)
            If True, ensure that the features used to refit the model match the original ones.
            Used only if data is pandas DataFrame.
4211
4212
4213

            .. versionadded:: 4.0.0

4214
4215
        **kwargs
            Other parameters for refit.
4216
            These parameters will be passed to ``predict`` method.
Guolin Ke's avatar
Guolin Ke committed
4217
4218
4219
4220
4221
4222

        Returns
        -------
        result : Booster
            Refitted Booster.
        """
4223
4224
        if self.__set_objective_to_none:
            raise LightGBMError('Cannot refit due to null objective function.')
4225
4226
        if dataset_params is None:
            dataset_params = {}
4227
        predictor = self._to_predictor(pred_parameter=deepcopy(kwargs))
4228
        leaf_preds: np.ndarray = predictor.predict(  # type: ignore[assignment]
4229
4230
4231
4232
4233
            data=data,
            start_iteration=-1,
            pred_leaf=True,
            validate_features=validate_features
        )
4234
        nrow, ncol = leaf_preds.shape
4235
        out_is_linear = ctypes.c_int(0)
4236
        _safe_call(_LIB.LGBM_BoosterGetLinear(
4237
            self._handle,
4238
            ctypes.byref(out_is_linear)))
Nikita Titov's avatar
Nikita Titov committed
4239
4240
4241
4242
4243
        new_params = _choose_param_value(
            main_param_name="linear_tree",
            params=self.params,
            default_value=None
        )
4244
        new_params["linear_tree"] = bool(out_is_linear.value)
4245
4246
4247
4248
4249
4250
4251
4252
4253
4254
4255
4256
4257
        new_params.update(dataset_params)
        train_set = Dataset(
            data=data,
            label=label,
            reference=reference,
            weight=weight,
            group=group,
            init_score=init_score,
            feature_name=feature_name,
            categorical_feature=categorical_feature,
            params=new_params,
            free_raw_data=free_raw_data,
        )
4258
        new_params['refit_decay_rate'] = decay_rate
4259
        new_booster = Booster(new_params, train_set)
Guolin Ke's avatar
Guolin Ke committed
4260
4261
        # Copy models
        _safe_call(_LIB.LGBM_BoosterMerge(
4262
4263
            new_booster._handle,
            predictor._handle))
Guolin Ke's avatar
Guolin Ke committed
4264
        leaf_preds = leaf_preds.reshape(-1)
4265
        ptr_data, _, _ = _c_int_array(leaf_preds)
Guolin Ke's avatar
Guolin Ke committed
4266
        _safe_call(_LIB.LGBM_BoosterRefit(
4267
            new_booster._handle,
Guolin Ke's avatar
Guolin Ke committed
4268
            ptr_data,
4269
4270
            ctypes.c_int32(nrow),
            ctypes.c_int32(ncol)))
4271
        new_booster._network = self._network
Guolin Ke's avatar
Guolin Ke committed
4272
4273
        return new_booster

4274
    def get_leaf_output(self, tree_id: int, leaf_id: int) -> float:
4275
4276
4277
4278
4279
4280
4281
4282
4283
4284
4285
4286
4287
4288
        """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.
        """
4289
4290
        ret = ctypes.c_double(0)
        _safe_call(_LIB.LGBM_BoosterGetLeafValue(
4291
            self._handle,
4292
4293
4294
4295
4296
            ctypes.c_int(tree_id),
            ctypes.c_int(leaf_id),
            ctypes.byref(ret)))
        return ret.value

4297
4298
4299
4300
4301
4302
4303
4304
    def set_leaf_output(
        self,
        tree_id: int,
        leaf_id: int,
        value: float,
    ) -> 'Booster':
        """Set the output of a leaf.

4305
4306
        .. versionadded:: 4.0.0

4307
4308
4309
4310
4311
4312
4313
4314
4315
4316
4317
4318
4319
4320
4321
4322
        Parameters
        ----------
        tree_id : int
            The index of the tree.
        leaf_id : int
            The index of the leaf in the tree.
        value : float
            Value to set as the output of the leaf.

        Returns
        -------
        self : Booster
            Booster with the leaf output set.
        """
        _safe_call(
            _LIB.LGBM_BoosterSetLeafValue(
4323
                self._handle,
4324
4325
4326
4327
4328
4329
4330
                ctypes.c_int(tree_id),
                ctypes.c_int(leaf_id),
                ctypes.c_double(value)
            )
        )
        return self

4331
4332
    def _to_predictor(
        self,
4333
        pred_parameter: Dict[str, Any]
4334
    ) -> _InnerPredictor:
4335
        """Convert to predictor."""
4336
        predictor = _InnerPredictor(booster_handle=self._handle, pred_parameter=pred_parameter)
4337
        predictor.pandas_categorical = self.pandas_categorical
wxchan's avatar
wxchan committed
4338
4339
        return predictor

4340
    def num_feature(self) -> int:
4341
4342
4343
4344
4345
4346
4347
        """Get number of features.

        Returns
        -------
        num_feature : int
            The number of features.
        """
4348
4349
        out_num_feature = ctypes.c_int(0)
        _safe_call(_LIB.LGBM_BoosterGetNumFeature(
4350
            self._handle,
4351
4352
4353
            ctypes.byref(out_num_feature)))
        return out_num_feature.value

4354
    def feature_name(self) -> List[str]:
4355
        """Get names of features.
wxchan's avatar
wxchan committed
4356
4357
4358

        Returns
        -------
4359
        result : list of str
4360
            List with names of features.
wxchan's avatar
wxchan committed
4361
        """
4362
        num_feature = self.num_feature()
4363
        # Get name of features
wxchan's avatar
wxchan committed
4364
        tmp_out_len = ctypes.c_int(0)
4365
4366
        reserved_string_buffer_size = 255
        required_string_buffer_size = ctypes.c_size_t(0)
4367
        string_buffers = [ctypes.create_string_buffer(reserved_string_buffer_size) for _ in range(num_feature)]
wxchan's avatar
wxchan committed
4368
4369
        ptr_string_buffers = (ctypes.c_char_p * num_feature)(*map(ctypes.addressof, string_buffers))
        _safe_call(_LIB.LGBM_BoosterGetFeatureNames(
4370
            self._handle,
4371
            ctypes.c_int(num_feature),
wxchan's avatar
wxchan committed
4372
            ctypes.byref(tmp_out_len),
4373
            ctypes.c_size_t(reserved_string_buffer_size),
4374
            ctypes.byref(required_string_buffer_size),
wxchan's avatar
wxchan committed
4375
4376
4377
            ptr_string_buffers))
        if num_feature != tmp_out_len.value:
            raise ValueError("Length of feature names doesn't equal with num_feature")
4378
4379
4380
4381
4382
4383
        actual_string_buffer_size = required_string_buffer_size.value
        # if buffer length is not long enough, reallocate buffers
        if reserved_string_buffer_size < actual_string_buffer_size:
            string_buffers = [ctypes.create_string_buffer(actual_string_buffer_size) for _ in range(num_feature)]
            ptr_string_buffers = (ctypes.c_char_p * num_feature)(*map(ctypes.addressof, string_buffers))
            _safe_call(_LIB.LGBM_BoosterGetFeatureNames(
4384
                self._handle,
4385
4386
4387
4388
4389
                ctypes.c_int(num_feature),
                ctypes.byref(tmp_out_len),
                ctypes.c_size_t(actual_string_buffer_size),
                ctypes.byref(required_string_buffer_size),
                ptr_string_buffers))
4390
        return [string_buffers[i].value.decode('utf-8') for i in range(num_feature)]
wxchan's avatar
wxchan committed
4391

4392
4393
4394
4395
4396
    def feature_importance(
        self,
        importance_type: str = 'split',
        iteration: Optional[int] = None
    ) -> np.ndarray:
4397
        """Get feature importances.
4398

4399
4400
        Parameters
        ----------
4401
        importance_type : str, optional (default="split")
4402
4403
4404
            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.
4405
4406
4407
4408
        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).
4409

4410
4411
        Returns
        -------
4412
4413
        result : numpy array
            Array with feature importances.
4414
        """
4415
4416
        if iteration is None:
            iteration = self.best_iteration
4417
        importance_type_int = _FEATURE_IMPORTANCE_TYPE_MAPPER[importance_type]
4418
        result = np.empty(self.num_feature(), dtype=np.float64)
4419
        _safe_call(_LIB.LGBM_BoosterFeatureImportance(
4420
            self._handle,
4421
4422
4423
            ctypes.c_int(iteration),
            ctypes.c_int(importance_type_int),
            result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
4424
        if importance_type_int == _C_API_FEATURE_IMPORTANCE_SPLIT:
4425
            return result.astype(np.int32)
4426
4427
        else:
            return result
4428

4429
4430
4431
4432
4433
4434
    def get_split_value_histogram(
        self,
        feature: Union[int, str],
        bins: Optional[Union[int, str]] = None,
        xgboost_style: bool = False
    ) -> Union[Tuple[np.ndarray, np.ndarray], np.ndarray, pd_DataFrame]:
4435
4436
4437
4438
        """Get split value histogram for the specified feature.

        Parameters
        ----------
4439
        feature : int or str
4440
4441
            The feature name or index the histogram is calculated for.
            If int, interpreted as index.
4442
            If str, interpreted as name.
4443

Nikita Titov's avatar
Nikita Titov committed
4444
4445
4446
            .. warning::

                Categorical features are not supported.
4447

4448
        bins : int, str or None, optional (default=None)
4449
4450
4451
            The maximum number of bins.
            If None, or int and > number of unique split values and ``xgboost_style=True``,
            the number of bins equals number of unique split values.
4452
            If str, it should be one from the list of the supported values by ``numpy.histogram()`` function.
4453
4454
4455
4456
4457
4458
4459
4460
4461
4462
4463
4464
4465
4466
        xgboost_style : bool, optional (default=False)
            Whether the returned result should be in the same form as it is in XGBoost.
            If False, the returned value is tuple of 2 numpy arrays as it is in ``numpy.histogram()`` function.
            If True, the returned value is matrix, in which the first column is the right edges of non-empty bins
            and the second one is the histogram values.

        Returns
        -------
        result_tuple : tuple of 2 numpy arrays
            If ``xgboost_style=False``, the values of the histogram of used splitting values for the specified feature
            and the bin edges.
        result_array_like : numpy array or pandas DataFrame (if pandas is installed)
            If ``xgboost_style=True``, the histogram of used splitting values for the specified feature.
        """
4467
        def add(root: Dict[str, Any]) -> None:
4468
4469
            """Recursively add thresholds."""
            if 'split_index' in root:  # non-leaf
4470
                if feature_names is not None and isinstance(feature, str):
4471
4472
4473
4474
                    split_feature = feature_names[root['split_feature']]
                else:
                    split_feature = root['split_feature']
                if split_feature == feature:
4475
                    if isinstance(root['threshold'], str):
4476
4477
4478
                        raise LightGBMError('Cannot compute split value histogram for the categorical feature')
                    else:
                        values.append(root['threshold'])
4479
4480
4481
4482
4483
4484
                add(root['left_child'])
                add(root['right_child'])

        model = self.dump_model()
        feature_names = model.get('feature_names')
        tree_infos = model['tree_info']
4485
        values: List[float] = []
4486
4487
4488
        for tree_info in tree_infos:
            add(tree_info['tree_structure'])

4489
        if bins is None or isinstance(bins, int) and xgboost_style:
4490
4491
4492
4493
4494
4495
4496
            n_unique = len(np.unique(values))
            bins = max(min(n_unique, bins) if bins is not None else n_unique, 1)
        hist, bin_edges = np.histogram(values, bins=bins)
        if xgboost_style:
            ret = np.column_stack((bin_edges[1:], hist))
            ret = ret[ret[:, 1] > 0]
            if PANDAS_INSTALLED:
4497
                return pd_DataFrame(ret, columns=['SplitValue', 'Count'])
4498
4499
4500
4501
4502
            else:
                return ret
        else:
            return hist, bin_edges

4503
4504
4505
4506
    def __inner_eval(
        self,
        data_name: str,
        data_idx: int,
4507
        feval: Optional[Union[_LGBM_CustomEvalFunction, List[_LGBM_CustomEvalFunction]]]
4508
    ) -> List[_LGBM_BoosterEvalMethodResultType]:
4509
        """Evaluate training or validation data."""
wxchan's avatar
wxchan committed
4510
        if data_idx >= self.__num_dataset:
4511
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
4512
4513
4514
        self.__get_eval_info()
        ret = []
        if self.__num_inner_eval > 0:
4515
            result = np.empty(self.__num_inner_eval, dtype=np.float64)
Guolin Ke's avatar
Guolin Ke committed
4516
            tmp_out_len = ctypes.c_int(0)
wxchan's avatar
wxchan committed
4517
            _safe_call(_LIB.LGBM_BoosterGetEval(
4518
                self._handle,
Guolin Ke's avatar
Guolin Ke committed
4519
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
4520
                ctypes.byref(tmp_out_len),
Guolin Ke's avatar
Guolin Ke committed
4521
                result.ctypes.data_as(ctypes.POINTER(ctypes.c_double))))
wxchan's avatar
wxchan committed
4522
            if tmp_out_len.value != self.__num_inner_eval:
4523
                raise ValueError("Wrong length of eval results")
4524
            for i in range(self.__num_inner_eval):
4525
4526
                ret.append((data_name, self.__name_inner_eval[i],
                            result[i], self.__higher_better_inner_eval[i]))
4527
4528
        if callable(feval):
            feval = [feval]
wxchan's avatar
wxchan committed
4529
4530
4531
4532
4533
        if feval is not None:
            if data_idx == 0:
                cur_data = self.train_set
            else:
                cur_data = self.valid_sets[data_idx - 1]
4534
4535
4536
4537
4538
4539
4540
4541
4542
            for eval_function in feval:
                if eval_function is None:
                    continue
                feval_ret = eval_function(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
wxchan's avatar
wxchan committed
4543
4544
4545
                    ret.append((data_name, eval_name, val, is_higher_better))
        return ret

4546
    def __inner_predict(self, data_idx: int) -> np.ndarray:
4547
        """Predict for training and validation dataset."""
wxchan's avatar
wxchan committed
4548
        if data_idx >= self.__num_dataset:
4549
            raise ValueError("Data_idx should be smaller than number of dataset")
wxchan's avatar
wxchan committed
4550
4551
4552
4553
4554
        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
4555
            self.__inner_predict_buffer[data_idx] = np.empty(n_preds, dtype=np.float64)
4556
        # avoid to predict many time in one iteration
wxchan's avatar
wxchan committed
4557
4558
        if not self.__is_predicted_cur_iter[data_idx]:
            tmp_out_len = ctypes.c_int64(0)
4559
            data_ptr = self.__inner_predict_buffer[data_idx].ctypes.data_as(ctypes.POINTER(ctypes.c_double))  # type: ignore[union-attr]
wxchan's avatar
wxchan committed
4560
            _safe_call(_LIB.LGBM_BoosterGetPredict(
4561
                self._handle,
Guolin Ke's avatar
Guolin Ke committed
4562
                ctypes.c_int(data_idx),
wxchan's avatar
wxchan committed
4563
4564
                ctypes.byref(tmp_out_len),
                data_ptr))
4565
            if tmp_out_len.value != len(self.__inner_predict_buffer[data_idx]):  # type: ignore[arg-type]
4566
                raise ValueError(f"Wrong length of predict results for data {data_idx}")
wxchan's avatar
wxchan committed
4567
            self.__is_predicted_cur_iter[data_idx] = True
4568
        result: np.ndarray = self.__inner_predict_buffer[data_idx]  # type: ignore[assignment]
4569
4570
4571
4572
        if self.__num_class > 1:
            num_data = result.size // self.__num_class
            result = result.reshape(num_data, self.__num_class, order='F')
        return result
wxchan's avatar
wxchan committed
4573

4574
    def __get_eval_info(self) -> None:
4575
        """Get inner evaluation count and names."""
wxchan's avatar
wxchan committed
4576
4577
        if self.__need_reload_eval_info:
            self.__need_reload_eval_info = False
Guolin Ke's avatar
Guolin Ke committed
4578
            out_num_eval = ctypes.c_int(0)
4579
            # Get num of inner evals
wxchan's avatar
wxchan committed
4580
            _safe_call(_LIB.LGBM_BoosterGetEvalCounts(
4581
                self._handle,
wxchan's avatar
wxchan committed
4582
4583
4584
                ctypes.byref(out_num_eval)))
            self.__num_inner_eval = out_num_eval.value
            if self.__num_inner_eval > 0:
4585
                # Get name of eval metrics
Guolin Ke's avatar
Guolin Ke committed
4586
                tmp_out_len = ctypes.c_int(0)
4587
4588
4589
                reserved_string_buffer_size = 255
                required_string_buffer_size = ctypes.c_size_t(0)
                string_buffers = [
4590
                    ctypes.create_string_buffer(reserved_string_buffer_size) for _ in range(self.__num_inner_eval)
4591
                ]
wxchan's avatar
wxchan committed
4592
                ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
wxchan's avatar
wxchan committed
4593
                _safe_call(_LIB.LGBM_BoosterGetEvalNames(
4594
                    self._handle,
4595
                    ctypes.c_int(self.__num_inner_eval),
wxchan's avatar
wxchan committed
4596
                    ctypes.byref(tmp_out_len),
4597
                    ctypes.c_size_t(reserved_string_buffer_size),
4598
                    ctypes.byref(required_string_buffer_size),
wxchan's avatar
wxchan committed
4599
4600
                    ptr_string_buffers))
                if self.__num_inner_eval != tmp_out_len.value:
4601
                    raise ValueError("Length of eval names doesn't equal with num_evals")
4602
4603
4604
4605
4606
4607
4608
4609
                actual_string_buffer_size = required_string_buffer_size.value
                # if buffer length is not long enough, reallocate buffers
                if reserved_string_buffer_size < actual_string_buffer_size:
                    string_buffers = [
                        ctypes.create_string_buffer(actual_string_buffer_size) for _ in range(self.__num_inner_eval)
                    ]
                    ptr_string_buffers = (ctypes.c_char_p * self.__num_inner_eval)(*map(ctypes.addressof, string_buffers))
                    _safe_call(_LIB.LGBM_BoosterGetEvalNames(
4610
                        self._handle,
4611
4612
4613
4614
4615
4616
4617
4618
4619
4620
4621
                        ctypes.c_int(self.__num_inner_eval),
                        ctypes.byref(tmp_out_len),
                        ctypes.c_size_t(actual_string_buffer_size),
                        ctypes.byref(required_string_buffer_size),
                        ptr_string_buffers))
                self.__name_inner_eval = [
                    string_buffers[i].value.decode('utf-8') for i in range(self.__num_inner_eval)
                ]
                self.__higher_better_inner_eval = [
                    name.startswith(('auc', 'ndcg@', 'map@', 'average_precision')) for name in self.__name_inner_eval
                ]