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

7
8
import numpy as np

9
from .basic import Booster, _log_warning
10
from .compat import GRAPHVIZ_INSTALLED, MATPLOTLIB_INSTALLED
11
12
13
from .sklearn import LGBMModel


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


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


26
27
28
29
30
31
32
33
34
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',
35
    importance_type: str = 'auto',
36
37
38
39
40
41
42
43
    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:
44
    """Plot model's feature importances.
45
46
47

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

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

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

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

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

116
    tuples = sorted(zip(feature_name, importance), key=lambda x: x[1])
117
118
119
120
    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:]
121
    labels, values = zip(*tuples)
122
123

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

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

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

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

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

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

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


162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
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:
178
179
180
181
182
183
    """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.
184
    feature : int or str
185
186
        The feature name or index the histogram is plotted for.
        If int, interpreted as index.
187
188
        If str, interpreted as name.
    bins : int, str or None, optional (default=None)
189
190
        The maximum number of bins.
        If None, the number of bins equals number of unique split values.
191
        If str, it should be one from the list of the supported values by ``numpy.histogram()`` function.
192
193
194
195
196
197
198
199
200
    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()``.
201
    title : str or None, optional (default="Split value histogram for feature with @index/name@ @feature@")
202
203
204
205
206
        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
207
208
        or ``name`` word in case of ``str`` type ``feature`` parameter.
    xlabel : str or None, optional (default="Feature split value")
209
210
        X-axis title label.
        If None, title is disabled.
211
    ylabel : str or None, optional (default="Count")
212
213
214
215
        Y-axis title label.
        If None, title is disabled.
    figsize : tuple of 2 elements or None, optional (default=None)
        Figure size.
216
217
    dpi : int or None, optional (default=None)
        Resolution of the figure.
218
219
220
221
222
223
224
225
226
227
228
229
230
231
    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:
232
        raise ImportError('You must install matplotlib and restart your session to plot split value histogram.')
233
234
235
236
237
238
239
240
241

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

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

    if ax is None:
        if figsize is not None:
            _check_not_tuple_of_2_elements(figsize, 'figsize')
249
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268

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

    if xlim is not None:
        _check_not_tuple_of_2_elements(xlim, 'xlim')
    else:
        range_result = bins[-1] - bins[0]
        xlim = (bins[0] - range_result * 0.2, bins[-1] + range_result * 0.2)
    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))
269
        title = title.replace('@index/name@', ('name' if isinstance(feature, str) else 'index'))
270
271
272
273
274
275
276
277
278
        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


279
280
281
282
283
284
285
286
287
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',
288
    ylabel: Optional[str] = '@metric@',
289
290
291
292
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    grid: bool = True
) -> Any:
293
294
295
296
297
    """Plot one metric during training.

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

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

    if isinstance(booster, LGBMModel):
        eval_results = deepcopy(booster.evals_result_)
    elif isinstance(booster, dict):
        eval_results = deepcopy(booster)
    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:
355
            _check_not_tuple_of_2_elements(figsize, 'figsize')
356
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
357
358
359
360
361
362
363
364
365
366
367
368
369

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

    name = next(dataset_names)  # take one as sample
    metrics_for_one = eval_results[name]
    num_metric = len(metrics_for_one)
    if metric is None:
        if num_metric > 1:
370
            _log_warning("More than one metric available, picking one to plot.")
371
372
373
374
375
        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]
376
377
378
    num_iteration = len(results)
    max_result = max(results)
    min_result = min(results)
379
    x_ = range(num_iteration)
380
381
382
383
384
    ax.plot(x_, results, label=name)

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

    ax.legend(loc='best')

    if xlim is not None:
392
        _check_not_tuple_of_2_elements(xlim, 'xlim')
393
394
395
396
397
    else:
        xlim = (0, num_iteration)
    ax.set_xlim(xlim)

    if ylim is not None:
398
        _check_not_tuple_of_2_elements(ylim, 'ylim')
399
400
401
402
403
404
    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 ylabel == 'auto':
405
406
        _log_warning("'auto' value of 'ylabel' argument is deprecated and will be removed in a future release of LightGBM. "
                     "Use '@metric@' placeholder instead.")
407
        ylabel = '@metric@'
408
409
410
411
412
413

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


420
421
422
423
424
425
426
427
428
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,
    **kwargs: Any
) -> Any:
429
430
431
    """Convert specified tree to graphviz instance.

    See:
432
      - https://graphviz.readthedocs.io/en/stable/api.html#digraph
433
    """
434
    if GRAPHVIZ_INSTALLED:
435
        from graphviz import Digraph
436
    else:
437
        raise ImportError('You must install graphviz and restart your session to plot tree.')
wxchan's avatar
wxchan committed
438

439
    def add(root, total_count, parent=None, decision=None):
440
        """Recursively add node or edge."""
wxchan's avatar
wxchan committed
441
        if 'split_index' in root:  # non-leaf
442
443
            l_dec = 'yes'
            r_dec = 'no'
444
            if root['decision_type'] == '<=':
445
446
                lte_symbol = "&#8804;"
                operator = lte_symbol
447
            elif root['decision_type'] == '==':
448
                operator = "="
449
450
            else:
                raise ValueError('Invalid decision type in tree model.')
451
            name = f"split{root['split_index']}"
452
            if feature_names is not None:
453
                label = f"<B>{feature_names[root['split_feature']]}</B> {operator}"
454
            else:
455
456
                label = f"feature <B>{root['split_feature']}</B> {operator} "
            label += f"<B>{_float2str(root['threshold'], precision)}</B>"
457
458
459
460
            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'}:
461
                        label += f"<br/>{_float2str(root[info], precision)} {output}"
462
                    elif info == 'internal_count':
463
                        label += f"<br/>{output}: {root[info]}"
464
                    elif info == "data_percentage":
465
                        label += f"<br/>{_float2str(root['internal_count'] / total_count * 100, 2)}% of data"
466
467
468
469
470
471
472
473
474

            fillcolor = "white"
            style = ""
            if constraints:
                if constraints[root['split_feature']] == 1:
                    fillcolor = "#ddffdd"  # light green
                if constraints[root['split_feature']] == -1:
                    fillcolor = "#ffdddd"  # light red
                style = "filled"
475
            label = f"<{label}>"
476
477
478
            graph.node(name, label=label, shape="rectangle", style=style, fillcolor=fillcolor)
            add(root['left_child'], total_count, name, l_dec)
            add(root['right_child'], total_count, name, r_dec)
wxchan's avatar
wxchan committed
479
        else:  # leaf
480
481
482
            name = f"leaf{root['leaf_index']}"
            label = f"leaf {root['leaf_index']}: "
            label += f"<B>{_float2str(root['leaf_value'], precision)}</B>"
483
            if 'leaf_weight' in show_info:
484
                label += f"<br/>{_float2str(root['leaf_weight'], precision)} weight"
485
            if 'leaf_count' in show_info:
486
                label += f"<br/>count: {root['leaf_count']}"
487
            if "data_percentage" in show_info:
488
489
                label += f"<br/>{_float2str(root['leaf_count'] / total_count * 100, 2)}% of data"
            label = f"<{label}>"
wxchan's avatar
wxchan committed
490
491
492
493
            graph.node(name, label=label)
        if parent is not None:
            graph.edge(parent, name, decision)

494
    graph = Digraph(**kwargs)
495
496
    rankdir = "LR" if orientation == "horizontal" else "TB"
    graph.attr("graph", nodesep="0.05", ranksep="0.3", rankdir=rankdir)
497
498
499
    if "internal_count" in tree_info['tree_structure']:
        add(tree_info['tree_structure'], tree_info['tree_structure']["internal_count"])
    else:
500
        raise Exception("Cannot plot trees with no split")
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519

    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")
520
521
522
    return graph


523
524
525
526
527
528
529
530
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',
    **kwargs: Any
) -> Any:
531
    """Create a digraph representation of specified tree.
532

533
534
535
536
537
538
539
540
541
542
    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
543
544
545
546
    .. note::

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

548
549
    Parameters
    ----------
550
    booster : Booster or LGBMModel
551
        Booster or LGBMModel instance to be converted.
552
553
    tree_index : int, optional (default=0)
        The index of a target tree to convert.
554
    show_info : list of str, or None, optional (default=None)
555
        What information should be shown in nodes.
556
557
558
559
560
561
562
563

            - ``'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
            - ``'leaf_weight'`` : total weight (sum of hessian) of all observations that fall into this leaf node
            - ``'data_percentage'`` : percentage of training data that fall into this node
564
    precision : int or None, optional (default=3)
565
        Used to restrict the display of floating point values to a certain precision.
566
    orientation : str, optional (default='horizontal')
567
568
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
569
    **kwargs
570
571
        Other parameters passed to ``Digraph`` constructor.
        Check https://graphviz.readthedocs.io/en/stable/api.html#digraph for the full list of supported parameters.
572
573
574

    Returns
    -------
575
576
    graph : graphviz.Digraph
        The digraph representation of specified tree.
577
578
579
580
581
582
583
584
585
586
587
588
589
    """
    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

590
591
    monotone_constraints = model.get('monotone_constraints', None)

592
593
594
595
596
597
598
599
    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 = []

600
601
    graph = _to_graphviz(tree_info, show_info, feature_names, precision,
                         orientation, monotone_constraints, **kwargs)
602

wxchan's avatar
wxchan committed
603
604
605
    return graph


606
607
608
609
610
611
612
613
614
615
616
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',
    **kwargs: Any
) -> Any:
wxchan's avatar
wxchan committed
617
618
    """Plot specified tree.

619
620
621
622
623
624
625
626
627
628
    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
629
630
631
632
    .. 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.
633

wxchan's avatar
wxchan committed
634
635
    Parameters
    ----------
636
637
638
639
640
641
642
643
    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
644
        Figure size.
645
646
    dpi : int or None, optional (default=None)
        Resolution of the figure.
647
    show_info : list of str, or None, optional (default=None)
648
        What information should be shown in nodes.
649
650
651
652
653
654
655
656

            - ``'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
            - ``'leaf_weight'`` : total weight (sum of hessian) of all observations that fall into this leaf node
            - ``'data_percentage'`` : percentage of training data that fall into this node
657
    precision : int or None, optional (default=3)
658
        Used to restrict the display of floating point values to a certain precision.
659
    orientation : str, optional (default='horizontal')
660
661
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
662
    **kwargs
663
664
        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
665
666
667

    Returns
    -------
668
669
    ax : matplotlib.axes.Axes
        The plot with single tree.
wxchan's avatar
wxchan committed
670
    """
671
    if MATPLOTLIB_INSTALLED:
wxchan's avatar
wxchan committed
672
        import matplotlib.image as image
673
        import matplotlib.pyplot as plt
674
    else:
675
        raise ImportError('You must install matplotlib and restart your session to plot tree.')
wxchan's avatar
wxchan committed
676
677

    if ax is None:
678
        if figsize is not None:
679
            _check_not_tuple_of_2_elements(figsize, 'figsize')
680
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
wxchan's avatar
wxchan committed
681

682
    graph = create_tree_digraph(booster=booster, tree_index=tree_index,
683
684
                                show_info=show_info, precision=precision,
                                orientation=orientation, **kwargs)
wxchan's avatar
wxchan committed
685
686

    s = BytesIO()
687
    s.write(graph.pipe(format='png'))
wxchan's avatar
wxchan committed
688
689
690
691
692
693
    s.seek(0)
    img = image.imread(s)

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