用 pydot 绘制决策树

Posted 2023-03-12

【中文标题】用 pydot 绘制决策树

【英文标题】：Plotting a decision tree with pydot

【发布时间】：2014-08-30 16:44:26

【问题描述】：

我已经训练了一个决策tree（Python 字典），如下所示。现在我正在尝试使用pydot 绘制它。在定义树的每个节点（pydot 图）时，我为其指定了一个唯一（且详细）的名称和一个简短的标签。

我的问题是，在我通过写入 .png 得到的结果图中，我看到的是 verbosenode names 而不是 node labels。

我已关注@Martijn Pieters here 的回答。我不知道我错过了什么，有什么想法吗？

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')

这是我使用上面的代码得到的（不需要的）输出：

【参考方案1】：

您需要将您创建的节点显式添加到图表中：

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 文件中包含节点元数据。

加上那些graph.add_node() 行，结果是：

【讨论】：

谢谢，你已经看过代码了，你觉得递归函数的实现是对的吗？

@Zhubarb：我只是粗略地看了一眼，但从这里看起来不错。您无需拨打.keys()； for k in dictionary: 就够了。

【参考方案2】：

如果有人想要使用边缘标签的版本（显示决策树的传统方式）

import pydot
import uuid

def generate_unique_node():
    """ Generate a unique node label."""
    return str(uuid.uuid1())

def create_node(graph, label, shape='oval'):
    node = pydot.Node(generate_unique_node(), label=label, shape=shape)
    graph.add_node(node)
    return node

def create_edge(graph, node_parent, node_child, label):
    link = pydot.Edge(node_parent, node_child, label=label)
    graph.add_edge(link)
    return link

def walk_tree(graph, dictionary, prev_node=None):
    """ Recursive construction of a decision tree stored as a dictionary """
    for parent, child in dictionary.items():
        # root
        if not prev_node: 
            root = create_node(graph, parent)
            walk_tree(graph, child, root)
            continue
            
        # node
        if isinstance(child, dict):
            for p, c in child.items():
                n = create_node(graph, p)
                create_edge(graph, prev_node, n, str(parent))
                walk_tree(graph, c, n)
    
        # leaf
        else: 
            leaf = create_node(graph, str(child), shape='box')
            create_edge(graph, prev_node, leaf, str(parent))

def plot_tree(dictionary, filename="DecisionTree.png"):
    graph = pydot.Dot(graph_type='graph')
    walk_tree(graph, tree)
    graph.write_png(filename)
        

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'
plot_tree(tree)