更改使用导出graphviz创建的决策树图的颜色

时间:2017-03-19 19:26:16

标签: scikit-learn visualization graphviz cart

我正在使用scikit的回归树函数和graphviz来生成一些决策树的精彩,易于解释的视觉效果:

dot_data = tree.export_graphviz(Run.reg, out_file=None, 
                         feature_names=Xvar,  
                         filled=True, rounded=True,  
                         special_characters=True) 
graph = pydotplus.graph_from_dot_data(dot_data)
graph.write_png('CART.png')
graph.write_svg("CART.svg")

enter image description here

这完全可以运行,但如果可能的话我想改变颜色方案?该图表示CO 2 通量,因此我想将负值设为绿色且为正棕色。我可以导出为svg而不是手动更改所有内容,但是当我这样做时,文本与盒子不完全对齐,因此手动更改颜色并修复所有文本为我的工作流程添加了一个非常繁琐的步骤,我真的很喜欢避免! enter image description here

另外,我已经看到一些树,其中连接节点的线的长度与分裂所解释的%方差成比例。如果可能的话,我也希望能够做到这一点吗?

1 个答案:

答案 0 :(得分:7)

  • 您可以通过graph.get_edge_list()
  • 获取所有边的列表
  • 每个源节点应该有两个目标节点,具有较低索引的节点被评估为True,较高的索引为False
  • 可以通过set_fillcolor()
  • 分配颜色

enter image description here

import pydotplus
from sklearn.datasets import load_iris
from sklearn import tree
import collections

clf = tree.DecisionTreeClassifier(random_state=42)
iris = load_iris()

clf = clf.fit(iris.data, iris.target)

dot_data = tree.export_graphviz(clf,
                                feature_names=iris.feature_names,
                                out_file=None,
                                filled=True,
                                rounded=True)
graph = pydotplus.graph_from_dot_data(dot_data)

colors = ('brown', 'forestgreen')
edges = collections.defaultdict(list)

for edge in graph.get_edge_list():
    edges[edge.get_source()].append(int(edge.get_destination()))

for edge in edges:
    edges[edge].sort()    
    for i in range(2):
        dest = graph.get_node(str(edges[edge][i]))[0]
        dest.set_fillcolor(colors[i])

graph.write_png('tree.png')
  

另外,我已经看到一些连接线的长度   节点与拆分解释的%变量成比例。 ID   如果可能的话,也喜欢能够做到这一点!?

你可以玩set_weight()set_len(),但这有点棘手,需要一些摆弄才能把它弄好但是这里有一些代码可以让你开始。

for edge in edges:
    edges[edge].sort()
    src = graph.get_node(edge)[0]
    total_weight = int(src.get_attributes()['label'].split('samples = ')[1].split('<br/>')[0])
    for i in range(2):
        dest = graph.get_node(str(edges[edge][i]))[0]
        weight = int(dest.get_attributes()['label'].split('samples = ')[1].split('<br/>')[0])
        graph.get_edge(edge, str(edges[edge][0]))[0].set_weight((1 - weight / total_weight) * 100)
        graph.get_edge(edge, str(edges[edge][0]))[0].set_len(weight / total_weight)
        graph.get_edge(edge, str(edges[edge][0]))[0].set_minlen(weight / total_weight)