用pydot绘制决策树 [英] Plotting a decision tree with pydot
问题描述
我已经训练了如下决定tree
(Python词典).现在,我尝试使用 pydot 对其进行绘制.在定义树的每个节点(pydot图)时,我为其指定了唯一的(且冗长的)名称和简短的标签.
I have trained a decision tree
(Python dictionary) as below. Now I am trying to plot it using pydot. In defining each node of the tree (pydot graph), I appoint it a unique (and verbose) name and a brief label.
我的问题是,在写成.png的结果图中,我看到的是冗长的node names
而不是node labels
.
My problem is that in the resulting figure that I get by writing to a .png, I see the verbosenode names
and not the node labels
.
I have followed the answer by @Martijn Pieters here. I do not know what I am missing, any ideas?
import pydot
tree= {'salary': {'41k-45k': 'junior', '46k-50k': {'department': {'marketing': 'senior', 'sales': 'senior', 'systems': 'junior'}}, '36k-40k': 'senior', '26k-30k': 'junior', '31k-35k': 'junior', '66k-70k': 'senior'}}
def walk_dictionaryv2(graph, dictionary, parent_node=None):
'''
Recursive plotting function for the decision tree stored as a dictionary
'''
for k in dictionary.keys():
if parent_node is not None:
from_name = parent_node.get_name().replace("\"", "") + '_' + str(k)
from_label = str(k)
node_from = pydot.Node(from_name, label=from_label)
graph.add_edge( pydot.Edge(parent_node, node_from) )
if isinstance(dictionary[k], dict): # if interim node
walk_dictionaryv2(graph, dictionary[k], node_from)
else: # if leaf node
to_name = str(k) + '_' + str(dictionary[k]) # unique name
to_label = str(dictionary[k])
node_to = pydot.Node(to_name, label=to_label, shape='box')
graph.add_edge(pydot.Edge(node_from, node_to))
#node_from.set_name(to_name)
else:
from_name = str(k)
from_label = str(k)
node_from = pydot.Node(from_name, label=from_label)
walk_dictionaryv2(graph, dictionary[k], node_from)
def plot_tree(tree, name):
# first you create a new graph, you do that with pydot.Dot()
graph = pydot.Dot(graph_type='graph')
walk_dictionaryv2(graph, tree)
graph.write_png(name+'.png')
plot_tree(tree,'name')
这是我在上面的代码中得到的(不需要的)输出:
This is the (undesired) output I get with the code above:
推荐答案
您需要将创建的节点显式添加到图中:
You need to explicitly add the nodes you create to the graph:
node_from = pydot.Node(from_name, label=from_label)
graph.add_node(node_from)
和
node_to = pydot.Node(to_name, label=to_label, shape='box')
graph.add_node(node_to)
否则,渲染器将看不到名称. graph.add_node()
将节点元数据包含在生成的.dot
文件中.
otherwise the renderer won't see the names. graph.add_node()
includes the node metadata in the generated .dot
file.
添加了这些graph.add_node()
行,结果是:
With those graph.add_node()
lines added, the result is:
这篇关于用pydot绘制决策树的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!