提问者:小点点

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


如何使用sklearn更改决策树图中的颜色。树在不使用graphviz的情况下绘制树,如本问题中所述:更改使用export graphviz创建的决策树图的颜色?

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)


共1个答案

匿名用户

许多matplotlib函数跟随颜色循环器来分配默认颜色,但在这里似乎不适用。

下面的方法循环生成的艺术家和结构,根据大多数类别和杂质(基尼)分配颜色。请注意,我们不能使用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)
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
    artist.get_bbox_patch().set_facecolor((f + (1 - f) * r, f + (1 - f) * g, f + (1 - f) * b))

plt.tight_layout()
plt.show()