plotting.py 29.4 KB
Newer Older
1
# coding: utf-8
2
"""Plotting library."""
3
import math
4
from copy import deepcopy
wxchan's avatar
wxchan committed
5
from io import BytesIO
6
from typing import Any, Dict, List, Optional, Tuple, Union
wxchan's avatar
wxchan committed
7

8
9
import numpy as np

10
11
from .basic import ZERO_THRESHOLD, Booster, _data_from_pandas, _is_zero, _log_warning, _MissingType
from .compat import GRAPHVIZ_INSTALLED, MATPLOTLIB_INSTALLED, pd_DataFrame
12
13
14
from .sklearn import LGBMModel


15
def _check_not_tuple_of_2_elements(obj: Any, obj_name: str = 'obj') -> None:
16
    """Check object is not tuple or does not have 2 elements."""
17
    if not isinstance(obj, tuple) or len(obj) != 2:
18
        raise TypeError(f"{obj_name} must be a tuple of 2 elements.")
wxchan's avatar
wxchan committed
19
20


21
def _float2str(value: float, precision: Optional[int] = None) -> str:
22
    return (f"{value:.{precision}f}"
23
            if precision is not None and not isinstance(value, str)
24
25
26
            else str(value))


27
28
29
30
31
32
33
34
35
def plot_importance(
    booster: Union[Booster, LGBMModel],
    ax=None,
    height: float = 0.2,
    xlim: Optional[Tuple[float, float]] = None,
    ylim: Optional[Tuple[float, float]] = None,
    title: Optional[str] = 'Feature importance',
    xlabel: Optional[str] = 'Feature importance',
    ylabel: Optional[str] = 'Features',
36
    importance_type: str = 'auto',
37
38
39
40
41
42
43
44
    max_num_features: Optional[int] = None,
    ignore_zero: bool = True,
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    grid: bool = True,
    precision: Optional[int] = 3,
    **kwargs: Any
) -> Any:
45
    """Plot model's feature importances.
46
47
48

    Parameters
    ----------
wxchan's avatar
wxchan committed
49
    booster : Booster or LGBMModel
50
51
52
53
54
55
56
57
58
59
        Booster or LGBMModel instance which feature importance should be plotted.
    ax : matplotlib.axes.Axes or None, optional (default=None)
        Target axes instance.
        If None, new figure and axes will be created.
    height : float, optional (default=0.2)
        Bar height, passed to ``ax.barh()``.
    xlim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.xlim()``.
    ylim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.ylim()``.
60
    title : str or None, optional (default="Feature importance")
61
62
        Axes title.
        If None, title is disabled.
63
    xlabel : str or None, optional (default="Feature importance")
64
65
        X-axis title label.
        If None, title is disabled.
66
        @importance_type@ placeholder can be used, and it will be replaced with the value of ``importance_type`` parameter.
67
    ylabel : str or None, optional (default="Features")
68
69
        Y-axis title label.
        If None, title is disabled.
70
    importance_type : str, optional (default="auto")
71
        How the importance is calculated.
72
        If "auto", if ``booster`` parameter is LGBMModel, ``booster.importance_type`` attribute is used; "split" otherwise.
73
74
75
        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.
    max_num_features : int or None, optional (default=None)
76
        Max number of top features displayed on plot.
77
78
79
80
81
        If None or <1, all features will be displayed.
    ignore_zero : bool, optional (default=True)
        Whether to ignore features with zero importance.
    figsize : tuple of 2 elements or None, optional (default=None)
        Figure size.
82
83
    dpi : int or None, optional (default=None)
        Resolution of the figure.
84
85
    grid : bool, optional (default=True)
        Whether to add a grid for axes.
86
    precision : int or None, optional (default=3)
87
        Used to restrict the display of floating point values to a certain precision.
88
    **kwargs
89
        Other parameters passed to ``ax.barh()``.
90
91
92

    Returns
    -------
93
94
    ax : matplotlib.axes.Axes
        The plot with model's feature importances.
95
    """
96
    if MATPLOTLIB_INSTALLED:
97
        import matplotlib.pyplot as plt
98
    else:
99
        raise ImportError('You must install matplotlib and restart your session to plot importance.')
100
101

    if isinstance(booster, LGBMModel):
102
103
        if importance_type == "auto":
            importance_type = booster.importance_type
wxchan's avatar
wxchan committed
104
        booster = booster.booster_
105
106
107
108
    elif isinstance(booster, Booster):
        if importance_type == "auto":
            importance_type = "split"
    else:
wxchan's avatar
wxchan committed
109
110
111
112
        raise TypeError('booster must be Booster or LGBMModel.')

    importance = booster.feature_importance(importance_type=importance_type)
    feature_name = booster.feature_name()
113
114

    if not len(importance):
115
        raise ValueError("Booster's feature_importance is empty.")
116

117
    tuples = sorted(zip(feature_name, importance), key=lambda x: x[1])
118
119
120
121
    if ignore_zero:
        tuples = [x for x in tuples if x[1] > 0]
    if max_num_features is not None and max_num_features > 0:
        tuples = tuples[-max_num_features:]
122
    labels, values = zip(*tuples)
123
124

    if ax is None:
125
        if figsize is not None:
126
            _check_not_tuple_of_2_elements(figsize, 'figsize')
127
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
128
129
130
131

    ylocs = np.arange(len(values))
    ax.barh(ylocs, values, align='center', height=height, **kwargs)

132
    for x, y in zip(values, ylocs):
133
134
135
        ax.text(x + 1, y,
                _float2str(x, precision) if importance_type == 'gain' else x,
                va='center')
136
137
138
139
140

    ax.set_yticks(ylocs)
    ax.set_yticklabels(labels)

    if xlim is not None:
141
        _check_not_tuple_of_2_elements(xlim, 'xlim')
142
143
144
145
146
    else:
        xlim = (0, max(values) * 1.1)
    ax.set_xlim(xlim)

    if ylim is not None:
147
        _check_not_tuple_of_2_elements(ylim, 'ylim')
148
149
150
151
152
153
154
    else:
        ylim = (-1, len(values))
    ax.set_ylim(ylim)

    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
155
        xlabel = xlabel.replace('@importance_type@', importance_type)
156
157
158
159
160
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    ax.grid(grid)
    return ax
wxchan's avatar
wxchan committed
161
162


163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def plot_split_value_histogram(
    booster: Union[Booster, LGBMModel],
    feature: Union[int, str],
    bins: Union[int, str, None] = None,
    ax=None,
    width_coef: float = 0.8,
    xlim: Optional[Tuple[float, float]] = None,
    ylim: Optional[Tuple[float, float]] = None,
    title: Optional[str] = 'Split value histogram for feature with @index/name@ @feature@',
    xlabel: Optional[str] = 'Feature split value',
    ylabel: Optional[str] = 'Count',
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    grid: bool = True,
    **kwargs: Any
) -> Any:
179
180
181
182
183
184
    """Plot split value histogram for the specified feature of the model.

    Parameters
    ----------
    booster : Booster or LGBMModel
        Booster or LGBMModel instance of which feature split value histogram should be plotted.
185
    feature : int or str
186
187
        The feature name or index the histogram is plotted for.
        If int, interpreted as index.
188
189
        If str, interpreted as name.
    bins : int, str or None, optional (default=None)
190
191
        The maximum number of bins.
        If None, the number of bins equals number of unique split values.
192
        If str, it should be one from the list of the supported values by ``numpy.histogram()`` function.
193
194
195
196
197
198
199
200
201
    ax : matplotlib.axes.Axes or None, optional (default=None)
        Target axes instance.
        If None, new figure and axes will be created.
    width_coef : float, optional (default=0.8)
        Coefficient for histogram bar width.
    xlim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.xlim()``.
    ylim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.ylim()``.
202
    title : str or None, optional (default="Split value histogram for feature with @index/name@ @feature@")
203
204
205
206
207
        Axes title.
        If None, title is disabled.
        @feature@ placeholder can be used, and it will be replaced with the value of ``feature`` parameter.
        @index/name@ placeholder can be used,
        and it will be replaced with ``index`` word in case of ``int`` type ``feature`` parameter
208
209
        or ``name`` word in case of ``str`` type ``feature`` parameter.
    xlabel : str or None, optional (default="Feature split value")
210
211
        X-axis title label.
        If None, title is disabled.
212
    ylabel : str or None, optional (default="Count")
213
214
215
216
        Y-axis title label.
        If None, title is disabled.
    figsize : tuple of 2 elements or None, optional (default=None)
        Figure size.
217
218
    dpi : int or None, optional (default=None)
        Resolution of the figure.
219
220
221
222
223
224
225
226
227
228
229
230
231
232
    grid : bool, optional (default=True)
        Whether to add a grid for axes.
    **kwargs
        Other parameters passed to ``ax.bar()``.

    Returns
    -------
    ax : matplotlib.axes.Axes
        The plot with specified model's feature split value histogram.
    """
    if MATPLOTLIB_INSTALLED:
        import matplotlib.pyplot as plt
        from matplotlib.ticker import MaxNLocator
    else:
233
        raise ImportError('You must install matplotlib and restart your session to plot split value histogram.')
234
235
236
237
238
239

    if isinstance(booster, LGBMModel):
        booster = booster.booster_
    elif not isinstance(booster, Booster):
        raise TypeError('booster must be Booster or LGBMModel.')

240
    hist, split_bins = booster.get_split_value_histogram(feature=feature, bins=bins, xgboost_style=False)
241
242
    if np.count_nonzero(hist) == 0:
        raise ValueError('Cannot plot split value histogram, '
243
                         f'because feature {feature} was not used in splitting')
244
245
    width = width_coef * (split_bins[1] - split_bins[0])
    centred = (split_bins[:-1] + split_bins[1:]) / 2
246
247
248
249

    if ax is None:
        if figsize is not None:
            _check_not_tuple_of_2_elements(figsize, 'figsize')
250
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
251
252
253
254
255
256

    ax.bar(centred, hist, align='center', width=width, **kwargs)

    if xlim is not None:
        _check_not_tuple_of_2_elements(xlim, 'xlim')
    else:
257
258
        range_result = split_bins[-1] - split_bins[0]
        xlim = (split_bins[0] - range_result * 0.2, split_bins[-1] + range_result * 0.2)
259
260
261
262
263
264
265
266
267
268
269
    ax.set_xlim(xlim)

    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    if ylim is not None:
        _check_not_tuple_of_2_elements(ylim, 'ylim')
    else:
        ylim = (0, max(hist) * 1.1)
    ax.set_ylim(ylim)

    if title is not None:
        title = title.replace('@feature@', str(feature))
270
        title = title.replace('@index/name@', ('name' if isinstance(feature, str) else 'index'))
271
272
273
274
275
276
277
278
279
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    ax.grid(grid)
    return ax


280
281
282
283
284
285
286
287
288
def plot_metric(
    booster: Union[Dict, LGBMModel],
    metric: Optional[str] = None,
    dataset_names: Optional[List[str]] = None,
    ax=None,
    xlim: Optional[Tuple[float, float]] = None,
    ylim: Optional[Tuple[float, float]] = None,
    title: Optional[str] = 'Metric during training',
    xlabel: Optional[str] = 'Iterations',
289
    ylabel: Optional[str] = '@metric@',
290
291
292
293
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    grid: bool = True
) -> Any:
294
295
296
297
298
    """Plot one metric during training.

    Parameters
    ----------
    booster : dict or LGBMModel
299
        Dictionary returned from ``lightgbm.train()`` or LGBMModel instance.
300
    metric : str or None, optional (default=None)
301
302
        The metric name to plot.
        Only one metric supported because different metrics have various scales.
303
        If None, first metric picked from dictionary (according to hashcode).
304
    dataset_names : list of str, or None, optional (default=None)
305
306
307
308
309
310
311
312
313
        List of the dataset names which are used to calculate metric to plot.
        If None, all datasets are used.
    ax : matplotlib.axes.Axes or None, optional (default=None)
        Target axes instance.
        If None, new figure and axes will be created.
    xlim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.xlim()``.
    ylim : tuple of 2 elements or None, optional (default=None)
        Tuple passed to ``ax.ylim()``.
314
    title : str or None, optional (default="Metric during training")
315
316
        Axes title.
        If None, title is disabled.
317
    xlabel : str or None, optional (default="Iterations")
318
319
        X-axis title label.
        If None, title is disabled.
320
    ylabel : str or None, optional (default="@metric@")
321
322
323
        Y-axis title label.
        If 'auto', metric name is used.
        If None, title is disabled.
324
        @metric@ placeholder can be used, and it will be replaced with metric name.
325
326
    figsize : tuple of 2 elements or None, optional (default=None)
        Figure size.
327
328
    dpi : int or None, optional (default=None)
        Resolution of the figure.
329
330
    grid : bool, optional (default=True)
        Whether to add a grid for axes.
331
332
333

    Returns
    -------
334
335
    ax : matplotlib.axes.Axes
        The plot with metric's history over the training.
336
    """
337
    if MATPLOTLIB_INSTALLED:
338
        import matplotlib.pyplot as plt
339
    else:
340
        raise ImportError('You must install matplotlib and restart your session to plot metric.')
341
342
343
344
345

    if isinstance(booster, LGBMModel):
        eval_results = deepcopy(booster.evals_result_)
    elif isinstance(booster, dict):
        eval_results = deepcopy(booster)
346
347
    elif isinstance(booster, Booster):
        raise TypeError("booster must be dict or LGBMModel. To use plot_metric with Booster type, first record the metrics using record_evaluation callback then pass that to plot_metric as argument `booster`")
348
349
350
351
352
353
354
355
356
357
    else:
        raise TypeError('booster must be dict or LGBMModel.')

    num_data = len(eval_results)

    if not num_data:
        raise ValueError('eval results cannot be empty.')

    if ax is None:
        if figsize is not None:
358
            _check_not_tuple_of_2_elements(figsize, 'figsize')
359
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
360
361

    if dataset_names is None:
362
        dataset_names_iter = iter(eval_results.keys())
363
364
365
    elif not isinstance(dataset_names, (list, tuple, set)) or not dataset_names:
        raise ValueError('dataset_names should be iterable and cannot be empty')
    else:
366
        dataset_names_iter = iter(dataset_names)
367

368
    name = next(dataset_names_iter)  # take one as sample
369
370
371
372
    metrics_for_one = eval_results[name]
    num_metric = len(metrics_for_one)
    if metric is None:
        if num_metric > 1:
373
            _log_warning("More than one metric available, picking one to plot.")
374
375
376
377
378
        metric, results = metrics_for_one.popitem()
    else:
        if metric not in metrics_for_one:
            raise KeyError('No given metric in eval results.')
        results = metrics_for_one[metric]
379
380
381
    num_iteration = len(results)
    max_result = max(results)
    min_result = min(results)
382
    x_ = range(num_iteration)
383
384
    ax.plot(x_, results, label=name)

385
    for name in dataset_names_iter:
386
387
        metrics_for_one = eval_results[name]
        results = metrics_for_one[metric]
388
389
        max_result = max(max(results), max_result)
        min_result = min(min(results), min_result)
390
391
392
393
394
        ax.plot(x_, results, label=name)

    ax.legend(loc='best')

    if xlim is not None:
395
        _check_not_tuple_of_2_elements(xlim, 'xlim')
396
397
398
399
400
    else:
        xlim = (0, num_iteration)
    ax.set_xlim(xlim)

    if ylim is not None:
401
        _check_not_tuple_of_2_elements(ylim, 'ylim')
402
403
404
405
406
407
408
409
410
411
    else:
        range_result = max_result - min_result
        ylim = (min_result - range_result * 0.2, max_result + range_result * 0.2)
    ax.set_ylim(ylim)

    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
412
        ylabel = ylabel.replace('@metric@', metric)
413
414
415
416
417
        ax.set_ylabel(ylabel)
    ax.grid(grid)
    return ax


418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
def _determine_direction_for_numeric_split(
    fval: float,
    threshold: float,
    missing_type_str: str,
    default_left: bool,
) -> str:
    missing_type = _MissingType(missing_type_str)
    if math.isnan(fval) and missing_type != _MissingType.NAN:
        fval = 0.0
    if ((missing_type == _MissingType.ZERO and _is_zero(fval))
            or (missing_type == _MissingType.NAN and math.isnan(fval))):
        direction = 'left' if default_left else 'right'
    else:
        direction = 'left' if fval <= threshold else 'right'
    return direction


def _determine_direction_for_categorical_split(fval: float, thresholds: str) -> str:
    if math.isnan(fval) or int(fval) < 0:
        return 'right'
    int_thresholds = {int(t) for t in thresholds.split('||')}
    return 'left' if int(fval) in int_thresholds else 'right'


442
443
444
445
446
447
448
def _to_graphviz(
    tree_info: Dict[str, Any],
    show_info: List[str],
    feature_names: Union[List[str], None],
    precision: Optional[int] = 3,
    orientation: str = 'horizontal',
    constraints: Optional[List[int]] = None,
449
    example_case: Optional[Union[np.ndarray, pd_DataFrame]] = None,
450
451
    **kwargs: Any
) -> Any:
452
453
454
    """Convert specified tree to graphviz instance.

    See:
455
      - https://graphviz.readthedocs.io/en/stable/api.html#digraph
456
    """
457
    if GRAPHVIZ_INSTALLED:
458
        from graphviz import Digraph
459
    else:
460
        raise ImportError('You must install graphviz and restart your session to plot tree.')
wxchan's avatar
wxchan committed
461

462
    def add(root, total_count, parent=None, decision=None, highlight=False):
463
        """Recursively add node or edge."""
464
465
466
467
468
469
470
471
        fillcolor = 'white'
        style = ''
        if highlight:
            color = 'blue'
            penwidth = '3'
        else:
            color = 'black'
            penwidth = '1'
wxchan's avatar
wxchan committed
472
        if 'split_index' in root:  # non-leaf
473
            shape = "rectangle"
474
475
            l_dec = 'yes'
            r_dec = 'no'
476
            if root['decision_type'] == '<=':
477
                operator = "&#8804;"
478
            elif root['decision_type'] == '==':
479
                operator = "="
480
481
            else:
                raise ValueError('Invalid decision type in tree model.')
482
            name = f"split{root['split_index']}"
483
            split_feature = root['split_feature']
484
            if feature_names is not None:
485
                label = f"<B>{feature_names[split_feature]}</B> {operator}"
486
            else:
487
488
489
490
491
492
493
494
495
                label = f"feature <B>{split_feature}</B> {operator} "
            direction = None
            if example_case is not None:
                if root['decision_type'] == '==':
                    direction = _determine_direction_for_categorical_split(example_case[split_feature], root['threshold'])
                else:
                    direction = _determine_direction_for_numeric_split(
                        example_case[split_feature], root['threshold'], root['missing_type'], root['default_left']
                    )
496
            label += f"<B>{_float2str(root['threshold'], precision)}</B>"
497
498
499
500
            for info in ['split_gain', 'internal_value', 'internal_weight', "internal_count", "data_percentage"]:
                if info in show_info:
                    output = info.split('_')[-1]
                    if info in {'split_gain', 'internal_value', 'internal_weight'}:
501
                        label += f"<br/>{_float2str(root[info], precision)} {output}"
502
                    elif info == 'internal_count':
503
                        label += f"<br/>{output}: {root[info]}"
504
                    elif info == "data_percentage":
505
                        label += f"<br/>{_float2str(root['internal_count'] / total_count * 100, 2)}% of data"
506
507
508
509
510
511
512

            if constraints:
                if constraints[root['split_feature']] == 1:
                    fillcolor = "#ddffdd"  # light green
                if constraints[root['split_feature']] == -1:
                    fillcolor = "#ffdddd"  # light red
                style = "filled"
513
            label = f"<{label}>"
514
515
            add(root['left_child'], total_count, name, l_dec, highlight and direction == "left")
            add(root['right_child'], total_count, name, r_dec, highlight and direction == "right")
wxchan's avatar
wxchan committed
516
        else:  # leaf
517
            shape = "ellipse"
518
519
520
            name = f"leaf{root['leaf_index']}"
            label = f"leaf {root['leaf_index']}: "
            label += f"<B>{_float2str(root['leaf_value'], precision)}</B>"
521
            if 'leaf_weight' in show_info:
522
                label += f"<br/>{_float2str(root['leaf_weight'], precision)} weight"
523
            if 'leaf_count' in show_info:
524
                label += f"<br/>count: {root['leaf_count']}"
525
            if "data_percentage" in show_info:
526
527
                label += f"<br/>{_float2str(root['leaf_count'] / total_count * 100, 2)}% of data"
            label = f"<{label}>"
528
        graph.node(name, label=label, shape=shape, style=style, fillcolor=fillcolor, color=color, penwidth=penwidth)
wxchan's avatar
wxchan committed
529
        if parent is not None:
530
            graph.edge(parent, name, decision, color=color, penwidth=penwidth)
wxchan's avatar
wxchan committed
531

532
    graph = Digraph(**kwargs)
533
534
    rankdir = "LR" if orientation == "horizontal" else "TB"
    graph.attr("graph", nodesep="0.05", ranksep="0.3", rankdir=rankdir)
535
    if "internal_count" in tree_info['tree_structure']:
536
        add(tree_info['tree_structure'], tree_info['tree_structure']["internal_count"], highlight=example_case is not None)
537
    else:
538
        raise Exception("Cannot plot trees with no split")
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557

    if constraints:
        # "#ddffdd" is light green, "#ffdddd" is light red
        legend = """<
            <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">
             <TR>
              <TD COLSPAN="2"><B>Monotone constraints</B></TD>
             </TR>
             <TR>
              <TD>Increasing</TD>
              <TD BGCOLOR="#ddffdd"></TD>
             </TR>
             <TR>
              <TD>Decreasing</TD>
              <TD BGCOLOR="#ffdddd"></TD>
             </TR>
            </TABLE>
           >"""
        graph.node("legend", label=legend, shape="rectangle", color="white")
558
559
560
    return graph


561
562
563
564
565
566
def create_tree_digraph(
    booster: Union[Booster, LGBMModel],
    tree_index: int = 0,
    show_info: Optional[List[str]] = None,
    precision: Optional[int] = 3,
    orientation: str = 'horizontal',
567
    example_case: Optional[Union[np.ndarray, pd_DataFrame]] = None,
568
569
    **kwargs: Any
) -> Any:
570
    """Create a digraph representation of specified tree.
571

572
573
574
575
576
577
578
579
580
581
    Each node in the graph represents a node in the tree.

    Non-leaf nodes have labels like ``Column_10 <= 875.9``, which means
    "this node splits on the feature named "Column_10", with threshold 875.9".

    Leaf nodes have labels like ``leaf 2: 0.422``, which means "this node is a
    leaf node, and the predicted value for records that fall into this node
    is 0.422". The number (``2``) is an internal unique identifier and doesn't
    have any special meaning.

Nikita Titov's avatar
Nikita Titov committed
582
583
584
585
    .. note::

        For more information please visit
        https://graphviz.readthedocs.io/en/stable/api.html#digraph.
586

587
588
    Parameters
    ----------
589
    booster : Booster or LGBMModel
590
        Booster or LGBMModel instance to be converted.
591
592
    tree_index : int, optional (default=0)
        The index of a target tree to convert.
593
    show_info : list of str, or None, optional (default=None)
594
        What information should be shown in nodes.
595
596
597
598
599
600

            - ``'split_gain'`` : gain from adding this split to the model
            - ``'internal_value'`` : raw predicted value that would be produced by this node if it was a leaf node
            - ``'internal_count'`` : number of records from the training data that fall into this non-leaf node
            - ``'internal_weight'`` : total weight of all nodes that fall into this non-leaf node
            - ``'leaf_count'`` : number of records from the training data that fall into this leaf node
601
            - ``'leaf_weight'`` : total weight (sum of Hessian) of all observations that fall into this leaf node
602
            - ``'data_percentage'`` : percentage of training data that fall into this node
603
    precision : int or None, optional (default=3)
604
        Used to restrict the display of floating point values to a certain precision.
605
    orientation : str, optional (default='horizontal')
606
607
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
608
609
610
    example_case : numpy 2-D array, pandas DataFrame or None, optional (default=None)
        Single row with the same structure as the training data.
        If not None, the plot will highlight the path that sample takes through the tree.
611
    **kwargs
612
613
        Other parameters passed to ``Digraph`` constructor.
        Check https://graphviz.readthedocs.io/en/stable/api.html#digraph for the full list of supported parameters.
614
615
616

    Returns
    -------
617
618
    graph : graphviz.Digraph
        The digraph representation of specified tree.
619
620
621
622
623
624
625
626
627
628
629
630
631
    """
    if isinstance(booster, LGBMModel):
        booster = booster.booster_
    elif not isinstance(booster, Booster):
        raise TypeError('booster must be Booster or LGBMModel.')

    model = booster.dump_model()
    tree_infos = model['tree_info']
    if 'feature_names' in model:
        feature_names = model['feature_names']
    else:
        feature_names = None

632
633
    monotone_constraints = model.get('monotone_constraints', None)

634
635
636
637
638
639
640
641
    if tree_index < len(tree_infos):
        tree_info = tree_infos[tree_index]
    else:
        raise IndexError('tree_index is out of range.')

    if show_info is None:
        show_info = []

642
643
644
645
646
647
648
649
650
    if example_case is not None:
        if not isinstance(example_case, (np.ndarray, pd_DataFrame)) or example_case.ndim != 2:
            raise ValueError('example_case must be a numpy 2-D array or a pandas DataFrame')
        if example_case.shape[0] != 1:
            raise ValueError('example_case must have a single row.')
        if isinstance(example_case, pd_DataFrame):
            example_case = _data_from_pandas(example_case, None, None, booster.pandas_categorical)[0]
        example_case = example_case[0]

651
    graph = _to_graphviz(tree_info, show_info, feature_names, precision,
652
                         orientation, monotone_constraints, example_case=example_case, **kwargs)
653

wxchan's avatar
wxchan committed
654
655
656
    return graph


657
658
659
660
661
662
663
664
665
def plot_tree(
    booster: Union[Booster, LGBMModel],
    ax=None,
    tree_index: int = 0,
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    show_info: Optional[List[str]] = None,
    precision: Optional[int] = 3,
    orientation: str = 'horizontal',
666
    example_case: Optional[Union[np.ndarray, pd_DataFrame]] = None,
667
668
    **kwargs: Any
) -> Any:
wxchan's avatar
wxchan committed
669
670
    """Plot specified tree.

671
672
673
674
675
676
677
678
679
680
    Each node in the graph represents a node in the tree.

    Non-leaf nodes have labels like ``Column_10 <= 875.9``, which means
    "this node splits on the feature named "Column_10", with threshold 875.9".

    Leaf nodes have labels like ``leaf 2: 0.422``, which means "this node is a
    leaf node, and the predicted value for records that fall into this node
    is 0.422". The number (``2``) is an internal unique identifier and doesn't
    have any special meaning.

Nikita Titov's avatar
Nikita Titov committed
681
682
683
684
    .. note::

        It is preferable to use ``create_tree_digraph()`` because of its lossless quality
        and returned objects can be also rendered and displayed directly inside a Jupyter notebook.
685

wxchan's avatar
wxchan committed
686
687
    Parameters
    ----------
688
689
690
691
692
693
694
695
    booster : Booster or LGBMModel
        Booster or LGBMModel instance to be plotted.
    ax : matplotlib.axes.Axes or None, optional (default=None)
        Target axes instance.
        If None, new figure and axes will be created.
    tree_index : int, optional (default=0)
        The index of a target tree to plot.
    figsize : tuple of 2 elements or None, optional (default=None)
wxchan's avatar
wxchan committed
696
        Figure size.
697
698
    dpi : int or None, optional (default=None)
        Resolution of the figure.
699
    show_info : list of str, or None, optional (default=None)
700
        What information should be shown in nodes.
701
702
703
704
705
706

            - ``'split_gain'`` : gain from adding this split to the model
            - ``'internal_value'`` : raw predicted value that would be produced by this node if it was a leaf node
            - ``'internal_count'`` : number of records from the training data that fall into this non-leaf node
            - ``'internal_weight'`` : total weight of all nodes that fall into this non-leaf node
            - ``'leaf_count'`` : number of records from the training data that fall into this leaf node
707
            - ``'leaf_weight'`` : total weight (sum of Hessian) of all observations that fall into this leaf node
708
            - ``'data_percentage'`` : percentage of training data that fall into this node
709
    precision : int or None, optional (default=3)
710
        Used to restrict the display of floating point values to a certain precision.
711
    orientation : str, optional (default='horizontal')
712
713
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
714
715
716
    example_case : numpy 2-D array, pandas DataFrame or None, optional (default=None)
        Single row with the same structure as the training data.
        If not None, the plot will highlight the path that sample takes through the tree.
717
    **kwargs
718
719
        Other parameters passed to ``Digraph`` constructor.
        Check https://graphviz.readthedocs.io/en/stable/api.html#digraph for the full list of supported parameters.
wxchan's avatar
wxchan committed
720
721
722

    Returns
    -------
723
724
    ax : matplotlib.axes.Axes
        The plot with single tree.
wxchan's avatar
wxchan committed
725
    """
726
    if MATPLOTLIB_INSTALLED:
wxchan's avatar
wxchan committed
727
        import matplotlib.image as image
728
        import matplotlib.pyplot as plt
729
    else:
730
        raise ImportError('You must install matplotlib and restart your session to plot tree.')
wxchan's avatar
wxchan committed
731
732

    if ax is None:
733
        if figsize is not None:
734
            _check_not_tuple_of_2_elements(figsize, 'figsize')
735
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
wxchan's avatar
wxchan committed
736

737
    graph = create_tree_digraph(booster=booster, tree_index=tree_index,
738
                                show_info=show_info, precision=precision,
739
                                orientation=orientation, example_case=example_case, **kwargs)
wxchan's avatar
wxchan committed
740
741

    s = BytesIO()
742
    s.write(graph.pipe(format='png'))
wxchan's avatar
wxchan committed
743
744
745
746
747
748
    s.seek(0)
    img = image.imread(s)

    ax.imshow(img)
    ax.axis('off')
    return ax