我是matplotlib
的新手,我试图绘制从头开始构建的决策树(不是使用sklearn
),因此它基本上是一个Node
对象,包含left
、right
和其他递归构建的标识变量。
这是我的计划:
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 + space
和y_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
1条答案
按热度按时间z9smfwbn1#
你最好使用Graphviz,因为它会为你处理空间问题。下载Graphviz和它的Python bindings,然后你就可以很容易地渲染图形了,如下所示: