如何用ax.scatter和ax.plot在NumPy的3D图中正确连接同组点？

问题描述 #

需要绘制包含多条平行线的3D图，已通过ax.scatter成功绘制三组不同颜色的点，但使用ax.plot连接同组点时得到不符合预期的结果。

原始正确散点代码 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

ax = plt.gca(projection='3d')

ax.scatter(0, 0, 100, color='red')
ax.scatter(0,1,128, color='red')
ax.scatter(0,2,30, color='red')
ax.scatter(0,3,15, color='red')

ax.scatter(1, 0, 100, color='blue')
ax.scatter(1, 1, 45, color='blue')
ax.scatter(1,2,13, color='blue')
ax.scatter(1,3,6, color='blue')

ax.scatter(2, 0, 100, color='green')
ax.scatter(2, 1, 55, color='green')
ax.scatter(2, 2, 27, color='green')
ax.scatter(2, 3, 26, color='green')

plt.show()

错误的连线尝试代码 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

ax = plt.gca(projection='3d')

p1,p2,p3,p4 = [0,0,100],[0,1,128],[0,2,30],[0,3,15]
ax.scatter(p1,p2,p3,p4, c='r',s=10)
ax.plot(p1,p2,p3,p4, color='r')

plt.show()

错误原因 #

matplotlib的3D绘图函数ax.scatter和ax.plot要求传入三个一维数组，分别对应所有点的x坐标、y坐标、z坐标，而不是把每个点作为单独的列表传入。错误代码中把四个点的[x,y,z]列表分别传给scatter和plot，导致参数维度不匹配，无法正确识别坐标数据。

解决方法 #

将每组点的x、y、z坐标分别提取为独立数组，再传入scatter和plot函数：

方式一：直接定义坐标数组 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 创建3D坐标轴
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

# 第一组红色点与连线
red_x = [0, 0, 0, 0]
red_y = [0, 1, 2, 3]
red_z = [100, 128, 30, 15]
ax.scatter(red_x, red_y, red_z, color='red', s=10)
ax.plot(red_x, red_y, red_z, color='red')

# 第二组蓝色点与连线
blue_x = [1, 1, 1, 1]
blue_y = [0, 1, 2, 3]
blue_z = [100, 45, 13, 6]
ax.scatter(blue_x, blue_y, blue_z, color='blue', s=10)
ax.plot(blue_x, blue_y, blue_z, color='blue')

# 第三组绿色点与连线
green_x = [2, 2, 2, 2]
green_y = [0, 1, 2, 3]
green_z = [100, 55, 27, 26]
ax.scatter(green_x, green_y, green_z, color='green', s=10)
ax.plot(green_x, green_y, green_z, color='green')

plt.show()

方式二：批量处理多组数据 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(projection='3d')

# 定义三组数据，每组是n个点的(x,y,z)
groups = [
    ([[0,0,100], [0,1,128], [0,2,30], [0,3,15]], 'red'),
    ([[1,0,100], [1,1,45], [1,2,13], [1,3,6]], 'blue'),
    ([[2,0,100], [2,1,55], [2,2,27], [2,3,26]], 'green')
]

for points, color in groups:
    # 拆分x、y、z坐标数组
    x = [p[0] for p in points]
    y = [p[1] for p in points]
    z = [p[2] for p in points]
    ax.scatter(x, y, z, color=color, s=10)
    ax.plot(x, y, z, color=color)

plt.show()

两种方式都能正确将同组点用对应颜色的线段连接，生成符合预期的3D多组平行线图。

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