_export.py 38 KB

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