Aluode/PerceptionLabPortable
0
1"""
2This module defines export functions for decision trees.
3"""
4
5# Authors: The scikit-learn developers
6# SPDX-License-Identifier: BSD-3-Clause
7
8from collections.abc import Iterable
9from io import StringIO
10from numbers import Integral
11
12import numpy as np
13
14from ..base import is_classifier
15from ..utils._param_validation import HasMethods, Interval, StrOptions, validate_params
16from ..utils.validation import check_array, check_is_fitted
17from . import DecisionTreeClassifier, DecisionTreeRegressor, _criterion, _tree
18from ._reingold_tilford import Tree, buchheim
19
20
21def _color_brew(n):
22 """Generate n colors with equally spaced hues.
23
24 Parameters
25 ----------
26 n : int
27 The number of colors required.
28
29 Returns
30 -------
31 color_list : list, length n
32 List of n tuples of form (R, G, B) being the components of each color.
33 """
34 color_list = []
35
36 # Initialize saturation & value; calculate chroma & value shift
37 s, v = 0.75, 0.9
38 c = s * v
39 m = v - c
40
41 for h in np.arange(25, 385, 360.0 / n).astype(int):
42 # Calculate some intermediate values
43 h_bar = h / 60.0
44 x = c * (1 - abs((h_bar % 2) - 1))
45 # Initialize RGB with same hue & chroma as our color
46 rgb = [
47 (c, x, 0),
48 (x, c, 0),
49 (0, c, x),
50 (0, x, c),
51 (x, 0, c),
52 (c, 0, x),
53 (c, x, 0),
54 ]
55 r, g, b = rgb[int(h_bar)]
56 # Shift the initial RGB values to match value and store
57 rgb = [(int(255 * (r + m))), (int(255 * (g + m))), (int(255 * (b + m)))]
58 color_list.append(rgb)
59
60 return color_list
61
62
63class Sentinel:
64 def __repr__(self):
65 return '"tree.dot"'
66
67
68SENTINEL = Sentinel()
69
70
71@validate_params(
72 {
73 "decision_tree": [DecisionTreeClassifier, DecisionTreeRegressor],
74 "max_depth": [Interval(Integral, 0, None, closed="left"), None],
75 "feature_names": ["array-like", None],
76 "class_names": ["array-like", "boolean", None],
77 "label": [StrOptions({"all", "root", "none"})],
78 "filled": ["boolean"],
79 "impurity": ["boolean"],
80 "node_ids": ["boolean"],
81 "proportion": ["boolean"],
82 "rounded": ["boolean"],
83 "precision": [Interval(Integral, 0, None, closed="left"), None],
84 "ax": "no_validation", # delegate validation to matplotlib
85 "fontsize": [Interval(Integral, 0, None, closed="left"), None],
86 },
87 prefer_skip_nested_validation=True,
88)
89def plot_tree(
90 decision_tree,
91 *,
92 max_depth=None,
93 feature_names=None,
94 class_names=None,
95 label="all",
96 filled=False,
97 impurity=True,
98 node_ids=False,
99 proportion=False,
100 rounded=False,
101 precision=3,
102 ax=None,
103 fontsize=None,
104):
105 """Plot a decision tree.
106
107 The sample counts that are shown are weighted with any sample_weights that
108 might be present.
109
110 The visualization is fit automatically to the size of the axis.
111 Use the ``figsize`` or ``dpi`` arguments of ``plt.figure`` to control
112 the size of the rendering.
113
114 Read more in the :ref:`User Guide <tree>`.
115
116 .. versionadded:: 0.21
117
118 Parameters
119 ----------
120 decision_tree : decision tree regressor or classifier
121 The decision tree to be plotted.
122
123 max_depth : int, default=None
124 The maximum depth of the representation. If None, the tree is fully
125 generated.
126
127 feature_names : array-like of str, default=None
128 Names of each of the features.
129 If None, generic names will be used ("x[0]", "x[1]", ...).
130
131 class_names : array-like of str or True, default=None
132 Names of each of the target classes in ascending numerical order.
133 Only relevant for classification and not supported for multi-output.
134 If ``True``, shows a symbolic representation of the class name.
135
136 label : {'all', 'root', 'none'}, default='all'
137 Whether to show informative labels for impurity, etc.
138 Options include 'all' to show at every node, 'root' to show only at
139 the top root node, or 'none' to not show at any node.
140
141 filled : bool, default=False
142 When set to ``True``, paint nodes to indicate majority class for
143 classification, extremity of values for regression, or purity of node
144 for multi-output.
145
146 impurity : bool, default=True
147 When set to ``True``, show the impurity at each node.
148
149 node_ids : bool, default=False
150 When set to ``True``, show the ID number on each node.
151
152 proportion : bool, default=False
153 When set to ``True``, change the display of 'values' and/or 'samples'
154 to be proportions and percentages respectively.
155
156 rounded : bool, default=False
157 When set to ``True``, draw node boxes with rounded corners and use
158 Helvetica fonts instead of Times-Roman.
159
160 precision : int, default=3
161 Number of digits of precision for floating point in the values of
162 impurity, threshold and value attributes of each node.
163
164 ax : matplotlib axis, default=None
165 Axes to plot to. If None, use current axis. Any previous content
166 is cleared.
167
168 fontsize : int, default=None
169 Size of text font. If None, determined automatically to fit figure.
170
171 Returns
172 -------
173 annotations : list of artists
174 List containing the artists for the annotation boxes making up the
175 tree.
176
177 Examples
178 --------
179 >>> from sklearn.datasets import load_iris
180 >>> from sklearn import tree
181
182 >>> clf = tree.DecisionTreeClassifier(random_state=0)
183 >>> iris = load_iris()
184
185 >>> clf = clf.fit(iris.data, iris.target)
186 >>> tree.plot_tree(clf)
187 [...]
188 """
189
190 check_is_fitted(decision_tree)
191
192 exporter = _MPLTreeExporter(
193 max_depth=max_depth,
194 feature_names=feature_names,
195 class_names=class_names,
196 label=label,
197 filled=filled,
198 impurity=impurity,
199 node_ids=node_ids,
200 proportion=proportion,
201 rounded=rounded,
202 precision=precision,
203 fontsize=fontsize,
204 )
205 return exporter.export(decision_tree, ax=ax)
206
207
208class _BaseTreeExporter:
209 def __init__(
210 self,
211 max_depth=None,
212 feature_names=None,
213 class_names=None,
214 label="all",
215 filled=False,
216 impurity=True,
217 node_ids=False,
218 proportion=False,
219 rounded=False,
220 precision=3,
221 fontsize=None,
222 ):
223 self.max_depth = max_depth
224 self.feature_names = feature_names
225 self.class_names = class_names
226 self.label = label
227 self.filled = filled
228 self.impurity = impurity
229 self.node_ids = node_ids
230 self.proportion = proportion
231 self.rounded = rounded
232 self.precision = precision
233 self.fontsize = fontsize
234
235 def get_color(self, value):
236 # Find the appropriate color & intensity for a node
237 if self.colors["bounds"] is None:
238 # Classification tree
239 color = list(self.colors["rgb"][np.argmax(value)])
240 sorted_values = sorted(value, reverse=True)
241 if len(sorted_values) == 1:
242 alpha = 0.0
243 else:
244 alpha = (sorted_values[0] - sorted_values[1]) / (1 - sorted_values[1])
245 else:
246 # Regression tree or multi-output
247 color = list(self.colors["rgb"][0])
248 alpha = (value - self.colors["bounds"][0]) / (
249 self.colors["bounds"][1] - self.colors["bounds"][0]
250 )
251 # compute the color as alpha against white
252 color = [int(round(alpha * c + (1 - alpha) * 255, 0)) for c in color]
253 # Return html color code in #RRGGBB format
254 return "#%2x%2x%2x" % tuple(color)
255
256 def get_fill_color(self, tree, node_id):
257 # Fetch appropriate color for node
258 if "rgb" not in self.colors:
259 # Initialize colors and bounds if required
260 self.colors["rgb"] = _color_brew(tree.n_classes[0])
261 if tree.n_outputs != 1:
262 # Find max and min impurities for multi-output
263 # The next line uses -max(impurity) instead of min(-impurity)
264 # and -min(impurity) instead of max(-impurity) on purpose, in
265 # order to avoid what looks like an issue with SIMD on non
266 # memory aligned arrays on 32bit OS. For more details see
267 # https://github.com/scikit-learn/scikit-learn/issues/27506.
268 self.colors["bounds"] = (-np.max(tree.impurity), -np.min(tree.impurity))
269 elif tree.n_classes[0] == 1 and len(np.unique(tree.value)) != 1:
270 # Find max and min values in leaf nodes for regression
271 self.colors["bounds"] = (np.min(tree.value), np.max(tree.value))
272 if tree.n_outputs == 1:
273 node_val = tree.value[node_id][0, :]
274 if (
275 tree.n_classes[0] == 1
276 and isinstance(node_val, Iterable)
277 and self.colors["bounds"] is not None
278 ):
279 # Unpack the float only for the regression tree case.
280 # Classification tree requires an Iterable in `get_color`.
281 node_val = node_val.item()
282 else:
283 # If multi-output color node by impurity
284 node_val = -tree.impurity[node_id]
285 return self.get_color(node_val)
286
287 def node_to_str(self, tree, node_id, criterion):
288 # Generate the node content string
289 if tree.n_outputs == 1:
290 value = tree.value[node_id][0, :]
291 else:
292 value = tree.value[node_id]
293
294 # Should labels be shown?
295 labels = (self.label == "root" and node_id == 0) or self.label == "all"
296
297 characters = self.characters
298 node_string = characters[-1]
299
300 # Write node ID
301 if self.node_ids:
302 if labels:
303 node_string += "node "
304 node_string += characters[0] + str(node_id) + characters[4]
305
306 # Write decision criteria
307 if tree.children_left[node_id] != _tree.TREE_LEAF:
308 # Always write node decision criteria, except for leaves
309 if self.feature_names is not None:
310 feature = self.feature_names[tree.feature[node_id]]
311 feature = self.str_escape(feature)
312 else:
313 feature = "x%s%s%s" % (
314 characters[1],
315 tree.feature[node_id],
316 characters[2],
317 )
318 node_string += "%s %s %s%s" % (
319 feature,
320 characters[3],
321 round(tree.threshold[node_id], self.precision),
322 characters[4],
323 )
324
325 # Write impurity
326 if self.impurity:
327 if isinstance(criterion, _criterion.FriedmanMSE):
328 criterion = "friedman_mse"
329 elif isinstance(criterion, _criterion.MSE) or criterion == "squared_error":
330 criterion = "squared_error"
331 elif not isinstance(criterion, str):
332 criterion = "impurity"
333 if labels:
334 node_string += "%s = " % criterion
335 node_string += (
336 str(round(tree.impurity[node_id], self.precision)) + characters[4]
337 )
338
339 # Write node sample count
340 if labels:
341 node_string += "samples = "
342 if self.proportion:
343 percent = (
344 100.0 * tree.n_node_samples[node_id] / float(tree.n_node_samples[0])
345 )
346 node_string += str(round(percent, 1)) + "%" + characters[4]
347 else:
348 node_string += str(tree.n_node_samples[node_id]) + characters[4]
349
350 # Write node class distribution / regression value
351 if not self.proportion and tree.n_classes[0] != 1:
352 # For classification this will show the proportion of samples
353 value = value * tree.weighted_n_node_samples[node_id]
354 if labels:
355 node_string += "value = "
356 if tree.n_classes[0] == 1:
357 # Regression
358 value_text = np.around(value, self.precision)
359 elif self.proportion:
360 # Classification
361 value_text = np.around(value, self.precision)
362 elif np.all(np.equal(np.mod(value, 1), 0)):
363 # Classification without floating-point weights
364 value_text = value.astype(int)
365 else:
366 # Classification with floating-point weights
367 value_text = np.around(value, self.precision)
368 # Strip whitespace
369 value_text = str(value_text.astype("S32")).replace("b'", "'")
370 value_text = value_text.replace("' '", ", ").replace("'", "")
371 if tree.n_classes[0] == 1 and tree.n_outputs == 1:
372 value_text = value_text.replace("[", "").replace("]", "")
373 value_text = value_text.replace("\n ", characters[4])
374 node_string += value_text + characters[4]
375
376 # Write node majority class
377 if (
378 self.class_names is not None
379 and tree.n_classes[0] != 1
380 and tree.n_outputs == 1
381 ):
382 # Only done for single-output classification trees
383 if labels:
384 node_string += "class = "
385 if self.class_names is not True:
386 class_name = self.class_names[np.argmax(value)]
387 class_name = self.str_escape(class_name)
388 else:
389 class_name = "y%s%s%s" % (
390 characters[1],
391 np.argmax(value),
392 characters[2],
393 )
394 node_string += class_name
395
396 # Clean up any trailing newlines
397 if node_string.endswith(characters[4]):
398 node_string = node_string[: -len(characters[4])]
399
400 return node_string + characters[5]
401
402 def str_escape(self, string):
403 return string
404
405
406class _DOTTreeExporter(_BaseTreeExporter):
407 def __init__(
408 self,
409 out_file=SENTINEL,
410 max_depth=None,
411 feature_names=None,
412 class_names=None,
413 label="all",
414 filled=False,
415 leaves_parallel=False,
416 impurity=True,
417 node_ids=False,
418 proportion=False,
419 rotate=False,
420 rounded=False,
421 special_characters=False,
422 precision=3,
423 fontname="helvetica",
424 ):
425 super().__init__(
426 max_depth=max_depth,
427 feature_names=feature_names,
428 class_names=class_names,
429 label=label,
430 filled=filled,
431 impurity=impurity,
432 node_ids=node_ids,
433 proportion=proportion,
434 rounded=rounded,
435 precision=precision,
436 )
437 self.leaves_parallel = leaves_parallel
438 self.out_file = out_file
439 self.special_characters = special_characters
440 self.fontname = fontname
441 self.rotate = rotate
442
443 # PostScript compatibility for special characters
444 if special_characters:
445 self.characters = ["#", "<SUB>", "</SUB>", "≤", "<br/>", ">", "<"]
446 else:
447 self.characters = ["#", "[", "]", "<=", "\\n", '"', '"']
448
449 # The depth of each node for plotting with 'leaf' option
450 self.ranks = {"leaves": []}
451 # The colors to render each node with
452 self.colors = {"bounds": None}
453
454 def export(self, decision_tree):
455 # Check length of feature_names before getting into the tree node
456 # Raise error if length of feature_names does not match
457 # n_features_in_ in the decision_tree
458 if self.feature_names is not None:
459 if len(self.feature_names) != decision_tree.n_features_in_:
460 raise ValueError(
461 "Length of feature_names, %d does not match number of features, %d"
462 % (len(self.feature_names), decision_tree.n_features_in_)
463 )
464 # each part writes to out_file
465 self.head()
466 # Now recurse the tree and add node & edge attributes
467 if isinstance(decision_tree, _tree.Tree):
468 self.recurse(decision_tree, 0, criterion="impurity")
469 else:
470 self.recurse(decision_tree.tree_, 0, criterion=decision_tree.criterion)
471
472 self.tail()
473
474 def tail(self):
475 # If required, draw leaf nodes at same depth as each other
476 if self.leaves_parallel:
477 for rank in sorted(self.ranks):
478 self.out_file.write(
479 "{rank=same ; " + "; ".join(r for r in self.ranks[rank]) + "} ;\n"
480 )
481 self.out_file.write("}")
482
483 def head(self):
484 self.out_file.write("digraph Tree {\n")
485
486 # Specify node aesthetics
487 self.out_file.write("node [shape=box")
488 rounded_filled = []
489 if self.filled:
490 rounded_filled.append("filled")
491 if self.rounded:
492 rounded_filled.append("rounded")
493 if len(rounded_filled) > 0:
494 self.out_file.write(
495 ', style="%s", color="black"' % ", ".join(rounded_filled)
496 )
497
498 self.out_file.write(', fontname="%s"' % self.fontname)
499 self.out_file.write("] ;\n")
500
501 # Specify graph & edge aesthetics
502 if self.leaves_parallel:
503 self.out_file.write("graph [ranksep=equally, splines=polyline] ;\n")
504
505 self.out_file.write('edge [fontname="%s"] ;\n' % self.fontname)
506
507 if self.rotate:
508 self.out_file.write("rankdir=LR ;\n")
509
510 def recurse(self, tree, node_id, criterion, parent=None, depth=0):
511 if node_id == _tree.TREE_LEAF:
512 raise ValueError("Invalid node_id %s" % _tree.TREE_LEAF)
513
514 left_child = tree.children_left[node_id]
515 right_child = tree.children_right[node_id]
516
517 # Add node with description
518 if self.max_depth is None or depth <= self.max_depth:
519 # Collect ranks for 'leaf' option in plot_options
520 if left_child == _tree.TREE_LEAF:
521 self.ranks["leaves"].append(str(node_id))
522 elif str(depth) not in self.ranks:
523 self.ranks[str(depth)] = [str(node_id)]
524 else:
525 self.ranks[str(depth)].append(str(node_id))
526
527 self.out_file.write(
528 "%d [label=%s" % (node_id, self.node_to_str(tree, node_id, criterion))
529 )
530
531 if self.filled:
532 self.out_file.write(
533 ', fillcolor="%s"' % self.get_fill_color(tree, node_id)
534 )
535 self.out_file.write("] ;\n")
536
537 if parent is not None:
538 # Add edge to parent
539 self.out_file.write("%d -> %d" % (parent, node_id))
540 if parent == 0:
541 # Draw True/False labels if parent is root node
542 angles = np.array([45, -45]) * ((self.rotate - 0.5) * -2)
543 self.out_file.write(" [labeldistance=2.5, labelangle=")
544 if node_id == 1:
545 self.out_file.write('%d, headlabel="True"]' % angles[0])
546 else:
547 self.out_file.write('%d, headlabel="False"]' % angles[1])
548 self.out_file.write(" ;\n")
549
550 if left_child != _tree.TREE_LEAF:
551 self.recurse(
552 tree,
553 left_child,
554 criterion=criterion,
555 parent=node_id,
556 depth=depth + 1,
557 )
558 self.recurse(
559 tree,
560 right_child,
561 criterion=criterion,
562 parent=node_id,
563 depth=depth + 1,
564 )
565
566 else:
567 self.ranks["leaves"].append(str(node_id))
568
569 self.out_file.write('%d [label="(...)"' % node_id)
570 if self.filled:
571 # color cropped nodes grey
572 self.out_file.write(', fillcolor="#C0C0C0"')
573 self.out_file.write("] ;\n" % node_id)
574
575 if parent is not None:
576 # Add edge to parent
577 self.out_file.write("%d -> %d ;\n" % (parent, node_id))
578
579 def str_escape(self, string):
580 # override default escaping for graphviz
581 return string.replace('"', r"\"")
582
583
584class _MPLTreeExporter(_BaseTreeExporter):
585 def __init__(
586 self,
587 max_depth=None,
588 feature_names=None,
589 class_names=None,
590 label="all",
591 filled=False,
592 impurity=True,
593 node_ids=False,
594 proportion=False,
595 rounded=False,
596 precision=3,
597 fontsize=None,
598 ):
599 super().__init__(
600 max_depth=max_depth,
601 feature_names=feature_names,
602 class_names=class_names,
603 label=label,
604 filled=filled,
605 impurity=impurity,
606 node_ids=node_ids,
607 proportion=proportion,
608 rounded=rounded,
609 precision=precision,
610 )
611 self.fontsize = fontsize
612
613 # The depth of each node for plotting with 'leaf' option
614 self.ranks = {"leaves": []}
615 # The colors to render each node with
616 self.colors = {"bounds": None}
617
618 self.characters = ["#", "[", "]", "<=", "\n", "", ""]
619 self.bbox_args = dict()
620 if self.rounded:
621 self.bbox_args["boxstyle"] = "round"
622
623 self.arrow_args = dict(arrowstyle="<-")
624
625 def _make_tree(self, node_id, et, criterion, depth=0):
626 # traverses _tree.Tree recursively, builds intermediate
627 # "_reingold_tilford.Tree" object
628 name = self.node_to_str(et, node_id, criterion=criterion)
629 if et.children_left[node_id] != _tree.TREE_LEAF and (
630 self.max_depth is None or depth <= self.max_depth
631 ):
632 children = [
633 self._make_tree(
634 et.children_left[node_id], et, criterion, depth=depth + 1
635 ),
636 self._make_tree(
637 et.children_right[node_id], et, criterion, depth=depth + 1
638 ),
639 ]
640 else:
641 return Tree(name, node_id)
642 return Tree(name, node_id, *children)
643
644 def export(self, decision_tree, ax=None):
645 import matplotlib.pyplot as plt
646 from matplotlib.text import Annotation
647
648 if ax is None:
649 ax = plt.gca()
650 ax.clear()
651 ax.set_axis_off()
652 my_tree = self._make_tree(0, decision_tree.tree_, decision_tree.criterion)
653 draw_tree = buchheim(my_tree)
654
655 # important to make sure we're still
656 # inside the axis after drawing the box
657 # this makes sense because the width of a box
658 # is about the same as the distance between boxes
659 max_x, max_y = draw_tree.max_extents() + 1
660 ax_width = ax.get_window_extent().width
661 ax_height = ax.get_window_extent().height
662
663 scale_x = ax_width / max_x
664 scale_y = ax_height / max_y
665 self.recurse(draw_tree, decision_tree.tree_, ax, max_x, max_y)
666
667 anns = [ann for ann in ax.get_children() if isinstance(ann, Annotation)]
668
669 # update sizes of all bboxes
670 renderer = ax.figure.canvas.get_renderer()
671
672 for ann in anns:
673 ann.update_bbox_position_size(renderer)
674
675 if self.fontsize is None:
676 # get figure to data transform
677 # adjust fontsize to avoid overlap
678 # get max box width and height
679 extents = [
680 bbox_patch.get_window_extent()
681 for ann in anns
682 if (bbox_patch := ann.get_bbox_patch()) is not None
683 ]
684 max_width = max([extent.width for extent in extents])
685 max_height = max([extent.height for extent in extents])
686 # width should be around scale_x in axis coordinates
687 size = anns[0].get_fontsize() * min(
688 scale_x / max_width, scale_y / max_height
689 )
690 for ann in anns:
691 ann.set_fontsize(size)
692
693 return anns
694
695 def recurse(self, node, tree, ax, max_x, max_y, depth=0):
696 import matplotlib.pyplot as plt
697
698 # kwargs for annotations without a bounding box
699 common_kwargs = dict(
700 zorder=100 - 10 * depth,
701 xycoords="axes fraction",
702 )
703 if self.fontsize is not None:
704 common_kwargs["fontsize"] = self.fontsize
705
706 # kwargs for annotations with a bounding box
707 kwargs = dict(
708 ha="center",
709 va="center",
710 bbox=self.bbox_args.copy(),
711 arrowprops=self.arrow_args.copy(),
712 **common_kwargs,
713 )
714 kwargs["arrowprops"]["edgecolor"] = plt.rcParams["text.color"]
715
716 # offset things by .5 to center them in plot
717 xy = ((node.x + 0.5) / max_x, (max_y - node.y - 0.5) / max_y)
718
719 if self.max_depth is None or depth <= self.max_depth:
720 if self.filled:
721 kwargs["bbox"]["fc"] = self.get_fill_color(tree, node.tree.node_id)
722 else:
723 kwargs["bbox"]["fc"] = ax.get_facecolor()
724
725 if node.parent is None:
726 # root
727 ax.annotate(node.tree.label, xy, **kwargs)
728 else:
729 xy_parent = (
730 (node.parent.x + 0.5) / max_x,
731 (max_y - node.parent.y - 0.5) / max_y,
732 )
733 ax.annotate(node.tree.label, xy_parent, xy, **kwargs)
734
735 # Draw True/False labels if parent is root node
736 if node.parent.parent is None:
737 # Adjust the position for the text to be slightly above the arrow
738 text_pos = (
739 (xy_parent[0] + xy[0]) / 2,
740 (xy_parent[1] + xy[1]) / 2,
741 )
742 # Annotate the arrow with the edge label to indicate the child
743 # where the sample-split condition is satisfied
744 if node.parent.left() == node:
745 label_text, label_ha = ("True ", "right")
746 else:
747 label_text, label_ha = (" False", "left")
748 ax.annotate(label_text, text_pos, ha=label_ha, **common_kwargs)
749 for child in node.children:
750 self.recurse(child, tree, ax, max_x, max_y, depth=depth + 1)
751
752 else:
753 xy_parent = (
754 (node.parent.x + 0.5) / max_x,
755 (max_y - node.parent.y - 0.5) / max_y,
756 )
757 kwargs["bbox"]["fc"] = "grey"
758 ax.annotate("\n (...) \n", xy_parent, xy, **kwargs)
759
760
761@validate_params(
762 {
763 "decision_tree": "no_validation",
764 "out_file": [str, None, HasMethods("write")],
765 "max_depth": [Interval(Integral, 0, None, closed="left"), None],
766 "feature_names": ["array-like", None],
767 "class_names": ["array-like", "boolean", None],
768 "label": [StrOptions({"all", "root", "none"})],
769 "filled": ["boolean"],
770 "leaves_parallel": ["boolean"],
771 "impurity": ["boolean"],
772 "node_ids": ["boolean"],
773 "proportion": ["boolean"],
774 "rotate": ["boolean"],
775 "rounded": ["boolean"],
776 "special_characters": ["boolean"],
777 "precision": [Interval(Integral, 0, None, closed="left"), None],
778 "fontname": [str],
779 },
780 prefer_skip_nested_validation=True,
781)
782def export_graphviz(
783 decision_tree,
784 out_file=None,
785 *,
786 max_depth=None,
787 feature_names=None,
788 class_names=None,
789 label="all",
790 filled=False,
791 leaves_parallel=False,
792 impurity=True,
793 node_ids=False,
794 proportion=False,
795 rotate=False,
796 rounded=False,
797 special_characters=False,
798 precision=3,
799 fontname="helvetica",
800):
801 """Export a decision tree in DOT format.
802
803 This function generates a GraphViz representation of the decision tree,
804 which is then written into `out_file`. Once exported, graphical renderings
805 can be generated using, for example::
806
807 $ dot -Tps tree.dot -o tree.ps (PostScript format)
808 $ dot -Tpng tree.dot -o tree.png (PNG format)
809
810 The sample counts that are shown are weighted with any sample_weights that
811 might be present.
812
813 Read more in the :ref:`User Guide <tree>`.
814
815 Parameters
816 ----------
817 decision_tree : object
818 The decision tree estimator to be exported to GraphViz.
819
820 out_file : object or str, default=None
821 Handle or name of the output file. If ``None``, the result is
822 returned as a string.
823
824 .. versionchanged:: 0.20
825 Default of out_file changed from "tree.dot" to None.
826
827 max_depth : int, default=None
828 The maximum depth of the representation. If None, the tree is fully
829 generated.
830
831 feature_names : array-like of shape (n_features,), default=None
832 An array containing the feature names.
833 If None, generic names will be used ("x[0]", "x[1]", ...).
834
835 class_names : array-like of shape (n_classes,) or bool, default=None
836 Names of each of the target classes in ascending numerical order.
837 Only relevant for classification and not supported for multi-output.
838 If ``True``, shows a symbolic representation of the class name.
839
840 label : {'all', 'root', 'none'}, default='all'
841 Whether to show informative labels for impurity, etc.
842 Options include 'all' to show at every node, 'root' to show only at
843 the top root node, or 'none' to not show at any node.
844
845 filled : bool, default=False
846 When set to ``True``, paint nodes to indicate majority class for
847 classification, extremity of values for regression, or purity of node
848 for multi-output.
849
850 leaves_parallel : bool, default=False
851 When set to ``True``, draw all leaf nodes at the bottom of the tree.
852
853 impurity : bool, default=True
854 When set to ``True``, show the impurity at each node.
855
856 node_ids : bool, default=False
857 When set to ``True``, show the ID number on each node.
858
859 proportion : bool, default=False
860 When set to ``True``, change the display of 'values' and/or 'samples'
861 to be proportions and percentages respectively.
862
863 rotate : bool, default=False
864 When set to ``True``, orient tree left to right rather than top-down.
865
866 rounded : bool, default=False
867 When set to ``True``, draw node boxes with rounded corners.
868
869 special_characters : bool, default=False
870 When set to ``False``, ignore special characters for PostScript
871 compatibility.
872
873 precision : int, default=3
874 Number of digits of precision for floating point in the values of
875 impurity, threshold and value attributes of each node.
876
877 fontname : str, default='helvetica'
878 Name of font used to render text.
879
880 Returns
881 -------
882 dot_data : str
883 String representation of the input tree in GraphViz dot format.
884 Only returned if ``out_file`` is None.
885
886 .. versionadded:: 0.18
887
888 Examples
889 --------
890 >>> from sklearn.datasets import load_iris
891 >>> from sklearn import tree
892
893 >>> clf = tree.DecisionTreeClassifier()
894 >>> iris = load_iris()
895
896 >>> clf = clf.fit(iris.data, iris.target)
897 >>> tree.export_graphviz(clf)
898 'digraph Tree {...
899 """
900 if feature_names is not None:
901 feature_names = check_array(
902 feature_names, ensure_2d=False, dtype=None, ensure_min_samples=0
903 )
904 if class_names is not None and not isinstance(class_names, bool):
905 class_names = check_array(
906 class_names, ensure_2d=False, dtype=None, ensure_min_samples=0
907 )
908
909 check_is_fitted(decision_tree)
910 own_file = False
911 return_string = False
912 try:
913 if isinstance(out_file, str):
914 out_file = open(out_file, "w", encoding="utf-8")
915 own_file = True
916
917 if out_file is None:
918 return_string = True
919 out_file = StringIO()
920
921 exporter = _DOTTreeExporter(
922 out_file=out_file,
923 max_depth=max_depth,
924 feature_names=feature_names,
925 class_names=class_names,
926 label=label,
927 filled=filled,
928 leaves_parallel=leaves_parallel,
929 impurity=impurity,
930 node_ids=node_ids,
931 proportion=proportion,
932 rotate=rotate,
933 rounded=rounded,
934 special_characters=special_characters,
935 precision=precision,
936 fontname=fontname,
937 )
938 exporter.export(decision_tree)
939
940 if return_string:
941 return exporter.out_file.getvalue()
942
943 finally:
944 if own_file:
945 out_file.close()
946
947
948def _compute_depth(tree, node):
949 """
950 Returns the depth of the subtree rooted in node.
951 """
952
953 def compute_depth_(
954 current_node, current_depth, children_left, children_right, depths
955 ):
956 depths += [current_depth]
957 left = children_left[current_node]
958 right = children_right[current_node]
959 if left != -1 and right != -1:
960 compute_depth_(
961 left, current_depth + 1, children_left, children_right, depths
962 )
963 compute_depth_(
964 right, current_depth + 1, children_left, children_right, depths
965 )
966
967 depths = []
968 compute_depth_(node, 1, tree.children_left, tree.children_right, depths)
969 return max(depths)
970
971
972@validate_params(
973 {
974 "decision_tree": [DecisionTreeClassifier, DecisionTreeRegressor],
975 "feature_names": ["array-like", None],
976 "class_names": ["array-like", None],
977 "max_depth": [Interval(Integral, 0, None, closed="left"), None],
978 "spacing": [Interval(Integral, 1, None, closed="left"), None],
979 "decimals": [Interval(Integral, 0, None, closed="left"), None],
980 "show_weights": ["boolean"],
981 },
982 prefer_skip_nested_validation=True,
983)
984def export_text(
985 decision_tree,
986 *,
987 feature_names=None,
988 class_names=None,
989 max_depth=10,
990 spacing=3,
991 decimals=2,
992 show_weights=False,
993):
994 """Build a text report showing the rules of a decision tree.
995
996 Note that backwards compatibility may not be supported.
997
998 Parameters
999 ----------
1000 decision_tree : object
1001 The decision tree estimator to be exported.
1002 It can be an instance of
1003 DecisionTreeClassifier or DecisionTreeRegressor.
1004
1005 feature_names : array-like of shape (n_features,), default=None
1006 An array containing the feature names.
1007 If None generic names will be used ("feature_0", "feature_1", ...).
1008
1009 class_names : array-like of shape (n_classes,), default=None
1010 Names of each of the target classes in ascending numerical order.
1011 Only relevant for classification and not supported for multi-output.
1012
1013 - if `None`, the class names are delegated to `decision_tree.classes_`;
1014 - otherwise, `class_names` will be used as class names instead of
1015 `decision_tree.classes_`. The length of `class_names` must match
1016 the length of `decision_tree.classes_`.
1017
1018 .. versionadded:: 1.3
1019
1020 max_depth : int, default=10
1021 Only the first max_depth levels of the tree are exported.
1022 Truncated branches will be marked with "...".
1023
1024 spacing : int, default=3
1025 Number of spaces between edges. The higher it is, the wider the result.
1026
1027 decimals : int, default=2
1028 Number of decimal digits to display.
1029
1030 show_weights : bool, default=False
1031 If true the classification weights will be exported on each leaf.
1032 The classification weights are the number of samples each class.
1033
1034 Returns
1035 -------
1036 report : str
1037 Text summary of all the rules in the decision tree.
1038
1039 Examples
1040 --------
1041
1042 >>> from sklearn.datasets import load_iris
1043 >>> from sklearn.tree import DecisionTreeClassifier
1044 >>> from sklearn.tree import export_text
1045 >>> iris = load_iris()
1046 >>> X = iris['data']
1047 >>> y = iris['target']
1048 >>> decision_tree = DecisionTreeClassifier(random_state=0, max_depth=2)
1049 >>> decision_tree = decision_tree.fit(X, y)
1050 >>> r = export_text(decision_tree, feature_names=iris['feature_names'])
1051 >>> print(r)
1052 |--- petal width (cm) <= 0.80
1053 | |--- class: 0
1054 |--- petal width (cm) > 0.80
1055 | |--- petal width (cm) <= 1.75
1056 | | |--- class: 1
1057 | |--- petal width (cm) > 1.75
1058 | | |--- class: 2
1059 """
1060 if feature_names is not None:
1061 feature_names = check_array(
1062 feature_names, ensure_2d=False, dtype=None, ensure_min_samples=0
1063 )
1064 if class_names is not None:
1065 class_names = check_array(
1066 class_names, ensure_2d=False, dtype=None, ensure_min_samples=0
1067 )
1068
1069 check_is_fitted(decision_tree)
1070 tree_ = decision_tree.tree_
1071 if is_classifier(decision_tree):
1072 if class_names is None:
1073 class_names = decision_tree.classes_
1074 elif len(class_names) != len(decision_tree.classes_):
1075 raise ValueError(
1076 "When `class_names` is an array, it should contain as"
1077 " many items as `decision_tree.classes_`. Got"
1078 f" {len(class_names)} while the tree was fitted with"
1079 f" {len(decision_tree.classes_)} classes."
1080 )
1081 right_child_fmt = "{} {} <= {}\n"
1082 left_child_fmt = "{} {} > {}\n"
1083 truncation_fmt = "{} {}\n"
1084
1085 if feature_names is not None and len(feature_names) != tree_.n_features:
1086 raise ValueError(
1087 "feature_names must contain %d elements, got %d"
1088 % (tree_.n_features, len(feature_names))
1089 )
1090
1091 if isinstance(decision_tree, DecisionTreeClassifier):
1092 value_fmt = "{}{} weights: {}\n"
1093 if not show_weights:
1094 value_fmt = "{}{}{}\n"
1095 else:
1096 value_fmt = "{}{} value: {}\n"
1097
1098 if feature_names is not None:
1099 feature_names_ = [
1100 feature_names[i] if i != _tree.TREE_UNDEFINED else None
1101 for i in tree_.feature
1102 ]
1103 else:
1104 feature_names_ = ["feature_{}".format(i) for i in tree_.feature]
1105
1106 export_text.report = ""
1107
1108 def _add_leaf(value, weighted_n_node_samples, class_name, indent):
1109 val = ""
1110 if isinstance(decision_tree, DecisionTreeClassifier):
1111 if show_weights:
1112 val = [
1113 "{1:.{0}f}, ".format(decimals, v * weighted_n_node_samples)
1114 for v in value
1115 ]
1116 val = "[" + "".join(val)[:-2] + "]"
1117 weighted_n_node_samples
1118 val += " class: " + str(class_name)
1119 else:
1120 val = ["{1:.{0}f}, ".format(decimals, v) for v in value]
1121 val = "[" + "".join(val)[:-2] + "]"
1122 export_text.report += value_fmt.format(indent, "", val)
1123
1124 def print_tree_recurse(node, depth):
1125 indent = ("|" + (" " * spacing)) * depth
1126 indent = indent[:-spacing] + "-" * spacing
1127
1128 value = None
1129 if tree_.n_outputs == 1:
1130 value = tree_.value[node][0]
1131 else:
1132 value = tree_.value[node].T[0]
1133 class_name = np.argmax(value)
1134
1135 if tree_.n_classes[0] != 1 and tree_.n_outputs == 1:
1136 class_name = class_names[class_name]
1137
1138 weighted_n_node_samples = tree_.weighted_n_node_samples[node]
1139
1140 if depth <= max_depth + 1:
1141 info_fmt = ""
1142 info_fmt_left = info_fmt
1143 info_fmt_right = info_fmt
1144
1145 if tree_.feature[node] != _tree.TREE_UNDEFINED:
1146 name = feature_names_[node]
1147 threshold = tree_.threshold[node]
1148 threshold = "{1:.{0}f}".format(decimals, threshold)
1149 export_text.report += right_child_fmt.format(indent, name, threshold)
1150 export_text.report += info_fmt_left
1151 print_tree_recurse(tree_.children_left[node], depth + 1)
1152
1153 export_text.report += left_child_fmt.format(indent, name, threshold)
1154 export_text.report += info_fmt_right
1155 print_tree_recurse(tree_.children_right[node], depth + 1)
1156 else: # leaf
1157 _add_leaf(value, weighted_n_node_samples, class_name, indent)
1158 else:
1159 subtree_depth = _compute_depth(tree_, node)
1160 if subtree_depth == 1:
1161 _add_leaf(value, weighted_n_node_samples, class_name, indent)
1162 else:
1163 trunc_report = "truncated branch of depth %d" % subtree_depth
1164 export_text.report += truncation_fmt.format(indent, trunc_report)
1165
1166 print_tree_recurse(0, 1)
1167 return export_text.report
1168