plotting.py 26.7 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

    if isinstance(booster, LGBMModel):
        eval_results = deepcopy(booster.evals_result_)
    elif isinstance(booster, dict):
        eval_results = deepcopy(booster)
345
346
    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`")
347
348
349
350
351
352
353
354
355
356
    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:
357
            _check_not_tuple_of_2_elements(figsize, 'figsize')
358
        _, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
359
360
361
362
363
364
365
366
367
368
369
370
371

    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:
372
            _log_warning("More than one metric available, picking one to plot.")
373
374
375
376
377
        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]
378
379
380
    num_iteration = len(results)
    max_result = max(results)
    min_result = min(results)
381
    x_ = range(num_iteration)
382
383
384
385
386
    ax.plot(x_, results, label=name)

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

    ax.legend(loc='best')

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

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

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


422
423
424
425
426
427
428
429
430
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:
431
432
433
    """Convert specified tree to graphviz instance.

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

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

            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"
477
            label = f"<{label}>"
478
479
480
            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
481
        else:  # leaf
482
483
484
            name = f"leaf{root['leaf_index']}"
            label = f"leaf {root['leaf_index']}: "
            label += f"<B>{_float2str(root['leaf_value'], precision)}</B>"
485
            if 'leaf_weight' in show_info:
486
                label += f"<br/>{_float2str(root['leaf_weight'], precision)} weight"
487
            if 'leaf_count' in show_info:
488
                label += f"<br/>count: {root['leaf_count']}"
489
            if "data_percentage" in show_info:
490
491
                label += f"<br/>{_float2str(root['leaf_count'] / total_count * 100, 2)}% of data"
            label = f"<{label}>"
wxchan's avatar
wxchan committed
492
493
494
495
            graph.node(name, label=label)
        if parent is not None:
            graph.edge(parent, name, decision)

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

    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")
522
523
524
    return graph


525
526
527
528
529
530
531
532
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:
533
    """Create a digraph representation of specified tree.
534

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

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

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

            - ``'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
566
    precision : int or None, optional (default=3)
567
        Used to restrict the display of floating point values to a certain precision.
568
    orientation : str, optional (default='horizontal')
569
570
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
571
    **kwargs
572
573
        Other parameters passed to ``Digraph`` constructor.
        Check https://graphviz.readthedocs.io/en/stable/api.html#digraph for the full list of supported parameters.
574
575
576

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

592
593
    monotone_constraints = model.get('monotone_constraints', None)

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

602
603
    graph = _to_graphviz(tree_info, show_info, feature_names, precision,
                         orientation, monotone_constraints, **kwargs)
604

wxchan's avatar
wxchan committed
605
606
607
    return graph


608
609
610
611
612
613
614
615
616
617
618
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
619
620
    """Plot specified tree.

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

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

            - ``'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
659
    precision : int or None, optional (default=3)
660
        Used to restrict the display of floating point values to a certain precision.
661
    orientation : str, optional (default='horizontal')
662
663
        Orientation of the tree.
        Can be 'horizontal' or 'vertical'.
664
    **kwargs
665
666
        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
667
668
669

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

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

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

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

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