如何使用 sklearn plot_tree 更改决策树 plot 的 colors？

Question

How to change colors in decision tree plot using sklearn.tree.plot_tree without using graphviz as in this question: Changing colors for decision tree plot created using 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)

Answer 1

许多 matplotlib 函数遵循颜色循环器来分配默认 colors，但这里似乎并不适用。

以下方法循环通过生成的艺术家和结构来分配颜色，具体取决于多数 class 和杂质 (gini)。 请注意，我们不能使用 alpha，因为透明背景会显示通常隐藏的箭头部分。

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

# generate some test data, separated by a diagonal line
X = np.random.rand(50, 2)
y = X[:, 0] - X[:, 1] > 0.2

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

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

colors = ['dodgerblue', 'limegreen']
ax1.plot(X[:, 0][y == 0], X[:, 1][y == 0], "bo")
ax1.plot(X[:, 0][y == 1], X[:, 1][y == 1], "g^")
xx, yy = np.meshgrid(np.linspace(0, 1, 100), np.linspace(0, 1, 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=["blue", "green"], 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)])
    artist.get_bbox_patch().set_facecolor((r + (1 - r) * impurity, g + (1 - g) * impurity, b + (1 - b) * impurity))

plt.tight_layout()
plt.show()