matplotlib 用pyplot手工绘制决策树

k97glaaz  于 2022-11-15  发布在  其他
关注(0)|答案(1)|浏览(200)

我是matplotlib的新手,我试图绘制从头开始构建的决策树(不是使用sklearn),因此它基本上是一个Node对象,包含leftright和其他递归构建的标识变量。
这是我的计划:

def plot_tree(node, x_axis=0, y_axis=10, space=5):
    if node.label is not None:
        ax.text(x_axis, y_axis, node.label[0],
                bbox=dict(boxstyle='round', facecolor='green', edgecolor='g'), ha='center', va='center')

    else:
        ax.text(x_axis, y_axis, f'{node.value:.2f}\nidx:{node.feature_idx}',
                     bbox=dict(boxstyle='round', facecolor='red', edgecolor='r'), ha='center', va='center')

        # x2, y2, w2, h2 = t2.get_tightbbox(fig.canvas.get_renderer()).bounds
        # plt.annotate(' ', xy=(x2 + w2, y2 + h2), xytext=(x_axis, y_axis), xycoords='figure points',
        #              arrowprops=dict(arrowstyle="<|-,head_length=1,head_width=0.5", lw=2, color='b'))

        plot_tree(node.left, x_axis + space, y_axis + space)
        plot_tree(node.right, x_axis + space, y_axis - space)

if __name__ == '__main__':
    node = root.load_tree()
    fig, ax = plt.subplots(1, 1)
    ax.axis('off')
    ax.set_aspect('equal')
    ax.autoscale_view()
    ax.set_xlim(0, 30)
    ax.set_ylim(-10, 30)
    plt.tick_params(axis='both', labelsize=0, length=0)
    plot_tree(node)

结果是:

我知道y轴会发生碰撞,因为y_axis + spacey_axis - space,但是我不知道如何让它在间距上保持对称,并且不出现这个,你们可以看到箭头被注解掉了,因为它们本身就是一团乱麻,这个库非常丰富,弄清楚它有点困难。

**编辑:**这是树的打印表示:

split is at feature:  27  and value  0.14235  and depth is:  1
     split is at feature:  20  and value  17.615000000000002  and depth is:  2
         label is:  B and depth is:  3
         split is at feature:  8  and value  0.15165  and depth is:  3
             label is:  B and depth is:  4
             label is:  M and depth is:  4
     split is at feature:  13  and value  13.93  and depth is:  2
         label is:  B and depth is:  3
         label is:  M and depth is:  3
z9smfwbn

z9smfwbn1#

你最好使用Graphviz,因为它会为你处理空间问题。下载Graphviz和它的Python bindings,然后你就可以很容易地渲染图形了,如下所示:

dot = graphviz.Digraph(comment="A graph", format="svg")
dot.node('A', 'King Arthur')
dot.node('B', 'Sir Bedevere the Wise')
dot.node('C', 'Sir Lancelot the Brave')
dot.edge('A', 'B')
dot.edge('A', 'C')
dot.render('digraph.gv', view=True)

相关问题