如何在Matplotlib 3D堆叠立方体图中添加花括号标签

问题说明 #

我正在为报告制作一张3D可视化图，展示沿轴向堆叠的立方体，希望在每个堆叠层级的侧边添加花括号{并标注对应区域。已经完成了立方体生成代码，但不知道如何给3D图添加花括号。

原代码如下：

import matplotlib.pyplot as plt 
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def draw_cube(ax, origin, size, height, label, dotted = False):
    # Vertices of a cube
    x, y, z = origin
    vertices = [
        [x, y, z],
        [x + size, y, z],
        [x + size, y + size, z],
        [x, y + size, z],
        [x, y, z + height],
        [x + size, y, z + height],
        [x + size, y + size, z + height],
        [x, y + size, z + height]
    ]
    
    # Define the 6 faces of the cube
    faces = [
        [vertices[0], vertices[1], vertices[2], vertices[3]],
        [vertices[4], vertices[5], vertices[6], vertices[7]], 
        [vertices[0], vertices[1], vertices[5], vertices[4]], 
        [vertices[2], vertices[3], vertices[7], vertices[6]], 
        [vertices[1], vertices[2], vertices[6], vertices[5]],
        [vertices[4], vertices[7], vertices[3], vertices[0]]
    ]
    linestyle = 'dotted' if dotted else '-'
    # Draw the faces
    ax.add_collection3d(Poly3DCollection(faces, linewidths=1, edgecolors='k', facecolor = 'b',alpha=.25, linestyle = linestyle))


def main():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d', elev= 10, azim = 15)

    # cubes for section 1
    draw_cube(ax, origin=(0, 0, 0), size=1, height = 1, dotted = False, label= "$U_{j}$")

    # cubes for section 2
    draw_cube(ax, origin=(0, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")

    #cubes for section 3
    draw_cube(ax, origin=(0, 0, 4), size=1, height = 2, dotted = False, label= "$U_{j}$")

    # Set the aspect ratio to be equal
    ax.set_aspect('auto')

    # Set labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    ax.axis('off')
    # Set limits
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1.5])
    ax.set_zlim([0, 4.1])

    # Display the plot
    plt.show()
if __name__ == "__main__":
    main()

解决方案 #

Matplotlib没有原生支持3D花括号，但可以利用Annotation工具绘制花括号，并通过3D坐标转2D屏幕坐标的方式，将花括号精准定位到层级侧边。以下是修改后的代码：

import matplotlib.pyplot as plt 
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def draw_cube(ax, origin, size, height, label, dotted = False):
    # Vertices of a cube
    x, y, z = origin
    vertices = [
        [x, y, z],
        [x + size, y, z],
        [x + size, y + size, z],
        [x, y + size, z],
        [x, y, z + height],
        [x + size, y, z + height],
        [x + size, y + size, z + height],
        [x, y + size, z + height]
    ]
    
    # Define the 6 faces of the cube
    faces = [
        [vertices[0], vertices[1], vertices[2], vertices[3]],
        [vertices[4], vertices[5], vertices[6], vertices[7]], 
        [vertices[0], vertices[1], vertices[5], vertices[4]], 
        [vertices[2], vertices[3], vertices[7], vertices[6]], 
        [vertices[1], vertices[2], vertices[6], vertices[5]],
        [vertices[4], vertices[7], vertices[3], vertices[0]]
    ]
    linestyle = 'dotted' if dotted else '-'
    # Draw the faces
    ax.add_collection3d(Poly3DCollection(faces, linewidths=1, edgecolors='k', facecolor = 'b',alpha=.25, linestyle = linestyle))

def add_brace(ax, z_start, z_end, label, y_pos=1.6):
    # 将3D坐标转换为屏幕坐标，再转回Axes坐标系
    p1 = ax.transData.transform((0, y_pos, z_start))
    p2 = ax.transData.transform((0, y_pos, z_end))
    
    p1_ax = ax.transAxes.inverted().transform(p1)
    p2_ax = ax.transAxes.inverted().transform(p2)
    
    # 绘制左侧花括号
    ann = ax.annotate("",
                     xy=(p1_ax[0], p1_ax[1]), xycoords='axes fraction',
                     xytext=(p2_ax[0], p2_ax[1]), textcoords='axes fraction',
                     arrowprops=dict(arrowstyle="|-{", color='black', linewidth=2),
                     annotation_clip=False)
    
    # 添加标注文本，位置在花括号左侧
    text_x = p1_ax[0] - 0.05
    text_y = (p1_ax[1] + p2_ax[1])/2
    ax.text(text_x, text_y, label, transform=ax.transAxes, fontsize=12, ha='right', va='center')

def main():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d', elev= 10, azim = 15)

    # cubes for section 1
    draw_cube(ax, origin=(0, 0, 0), size=1, height = 1, dotted = False, label= "$U_{j}$")

    # cubes for section 2
    draw_cube(ax, origin=(0, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")
    draw_cube(ax, origin=(0.5, 0.5, 1), size=0.5, height = 3, dotted = False, label = "$U_{j+1}$")

    #cubes for section 3
    draw_cube(ax, origin=(0, 0, 4), size=1, height = 2, dotted = False, label= "$U_{j}$")

    # 添加三个层级的花括号标注
    add_brace(ax, z_start=0, z_end=1, label="$U_{j}$")
    add_brace(ax, z_start=1, z_end=4, label="$U_{j+1}$")
    add_brace(ax, z_start=4, z_end=6, label="$U_{j}$")

    # Set the aspect ratio to be equal
    ax.set_aspect('auto')

    # Set labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    ax.axis('off')
    # 调整z轴范围以显示完整的第三层
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1.5])
    ax.set_zlim([0, 6.1])

    # Display the plot
    plt.show()
if __name__ == "__main__":
    main()

关键说明 #

• add_brace函数通过坐标转换适配3D视角，确保花括号位置与堆叠层级对齐。
• 可通过调整y_pos、text_x等参数微调花括号和文本的位置，适配不同可视化需求。
• 调整了z轴范围，保证第三层立方体和对应花括号完整显示。

内容的提问来源于stack exchange，提问作者user23622626