如何使用 sklearn plot_tree 更改决策树图的颜色？

Question

How to change colors in decision tree plot using sklearn.tree.plot_tree without using graphviz as in this question: ?

plt.figure(figsize=[21, 6])
ax1 = plt.subplot(121)
ax2 = plt.subplot(122)

ax1.plot(X[:, 0][y == 0], X[:, 1][y == 0], "bo")
ax1.plot(X[:, 0][y == 1], X[:, 1][y == 1], "g^")
ax1.contourf(xx, yy, pred.reshape(xx.shape), cmap=matplotlib.colors.ListedColormap(['b', 'g']), alpha=0.25)
ax1.set_title(title)

plot_tree(tree_clf, feature_names=["X", "y"], class_names=["blue", "green"], filled=True, rounded=True)

Answer 1

许多 matplotlib 函数遵循 color cycler 分配默认颜色，但这似乎不适用于此处。

以下方法遍历生成的注释文本 (artists) 和 clf tree structure 以根据多数 class 和杂质 (gini) 分配颜色。请注意，我们不能使用 alpha，因为透明背景会显示通常隐藏的部分箭头。

from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap, to_rgb
import numpy as np
from sklearn import tree

X = np.random.rand(50, 2) * np.r_[100, 50]
y = X[:, 0] - X[:, 1] > 20

clf = tree.DecisionTreeClassifier(random_state=2021)
clf = clf.fit(X, y)

fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=[21, 6])

colors = ['crimson', 'dodgerblue']
ax1.plot(X[:, 0][y == 0], X[:, 1][y == 0], "o", color=colors[0])
ax1.plot(X[:, 0][y == 1], X[:, 1][y == 1], "^", color=colors[1])
xx, yy = np.meshgrid(np.linspace(X[:, 0].min(), X[:, 0].max(), 100), np.linspace(X[:, 1].min(), X[:, 1].max(), 100))
pred = clf.predict(np.c_[(xx.ravel(), yy.ravel())])
ax1.contourf(xx, yy, pred.reshape(xx.shape), cmap=ListedColormap(colors), alpha=0.25)

# ax2.set_prop_cycle(mpl.cycler(color=colors)) # doesn't seem to work

artists = tree.plot_tree(clf, feature_names=["X", "y"], class_names=colors,
                         filled=True, rounded=True, ax=ax2)
for artist, impurity, value in zip(artists, clf.tree_.impurity, clf.tree_.value):
    # let the max value decide the color; whiten the color depending on impurity (gini)
    r, g, b = to_rgb(colors[np.argmax(value)])
    f = impurity * 2 # for N colors: f = impurity * N/(N-1) if N>1 else 0
    artist.get_bbox_patch().set_facecolor((f + (1-f)*r, f + (1-f)*g, f + (1-f)*b))
    artist.get_bbox_patch().set_edgecolor('black')

plt.tight_layout()
plt.show()

如何使用 sklearn plot_tree 更改决策树图的颜色？

How to change colors for decision tree plot using sklearn plot_tree?

python

matplotlib

scikit-learn