如何在matplotlib 3D绘图中避免循环绘制独立数据集？

问题场景 #

把独立3D数据集的坐标分别存在X、Y、Z数组的列中，直接调用ax.plot(X, Y, Z)会报错，但循环遍历每列的方式能正常绘图，想找不用循环的实现方法。

代码示例 #

import numpy as np
import matplotlib.pyplot as plt

X = np.array([[1, 2 ,5], [6, 1, 7]])
Y = np.array([[4, 6, 9], [7, 1, 5]])
Z = np.array([[5, 7, 4], [2, 1, 6]])

fig1, (ax) = plt.subplots(subplot_kw={"projection": "3d"})

# 报错代码
ax.plot(X, Y, Z)
# 可行的循环代码
for i in range(X.shape[1]):
    ax.plot(X[:, i], Y[:, i], Z[:, i])

报错信息（翻译后） #

ValueError: 操作数无法一起广播，重映射形状为[原始->重映射]: (6,) 和请求的形状(2,)
...
AttributeError: 'Line3D'对象没有属性'_verts3d'

无需循环的解决方案 #

利用Matplotlib遇到NaN会中断曲线绘制的特性，把每条独立曲线的坐标用NaN分隔后，一次性传入ax.plot就能实现多曲线绘制，不用循环。

方法1：拼接列并添加NaN分隔符

import numpy as np
import matplotlib.pyplot as plt

X = np.array([[1, 2 ,5], [6, 1, 7]])
Y = np.array([[4, 6, 9], [7, 1, 5]])
Z = np.array([[5, 7, 4], [2, 1, 6]])

fig1, ax = plt.subplots(subplot_kw={"projection": "3d"})

# 为每列数据末尾添加NaN，再拼接成一维数组
x_data = np.hstack([X[:, col] for col in range(X.shape[1])] + [np.full(X.shape[0], np.nan)])
y_data = np.hstack([Y[:, col] for col in range(Y.shape[1])] + [np.full(Y.shape[0], np.nan)])
z_data = np.hstack([Z[:, col] for col in range(Z.shape[1])] + [np.full(Z.shape[0], np.nan)])

ax.plot(x_data, y_data, z_data)
plt.show()

方法2：转置数组后插入NaN行再扁平化

这种写法更简洁，批量处理数组：

import numpy as np
import matplotlib.pyplot as plt

X = np.array([[1, 2 ,5], [6, 1, 7]])
Y = np.array([[4, 6, 9], [7, 1, 5]])
Z = np.array([[5, 7, 4], [2, 1, 6]])

fig1, ax = plt.subplots(subplot_kw={"projection": "3d"})

# 转置数组，在每两条曲线之间插入NaN行，再扁平化
X_flat = np.insert(X.T, np.arange(1, X.shape[1]), np.nan, axis=0).flatten()
Y_flat = np.insert(Y.T, np.arange(1, Y.shape[1]), np.nan, axis=0).flatten()
Z_flat = np.insert(Z.T, np.arange(1, Z.shape[1]), np.nan, axis=0).flatten()

ax.plot(X_flat, Y_flat, Z_flat)
plt.show()

原理说明 #

Matplotlib的plot函数绘制时，碰到NaN就会结束当前曲线，后续数据会作为新曲线开始绘制。把所有独立曲线的坐标用NaN分隔成一个一维数组，一次调用就能画出多条曲线，效果和循环遍历列完全一致。

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