如何为3D箭袋图（Quiver Plot）的箭头设置不同颜色？

核心逻辑 #

Matplotlib的3D箭袋图ax.quiver会自动将输入的三维坐标数组（x/y/z/u/v/w）展平为一维来逐个绘制箭头，因此要给每个箭头单独指定字符串类型的颜色，需要将三维的颜色数组也同步展平为与箭头数量匹配的一维列表，才能实现一一对应。

实现步骤 #

1. 确保你的三维颜色数组colour的维度和x/y/z/u/v/w完全一致（比如都是(N,M,L)的形状）。
2. 将三维颜色数组展平为一维列表，每个元素是'red'/'blue'这类合法的颜色字符串。
3. 把展平后的颜色列表传入quiver的color参数即可。

示例代码 #

import matplotlib.pyplot as plt
import numpy as np

# 模拟三维坐标与方向向量数据
N, M, L = 2, 2, 2
x, y, z = np.meshgrid(np.arange(N), np.arange(M), np.arange(L))
u = np.ones_like(x)
v = np.ones_like(y)
w = np.ones_like(z)

# 构造三维颜色数组：指定不同位置的箭头颜色
colour = np.empty((N, M, L), dtype=str)
colour[0, :, :] = 'red'
colour[1, :, :] = 'blue'

# 将三维颜色数组展平为一维
flat_colours = colour.flatten()

# 绘制3D箭袋图
fig = plt.figure(figsize=(15,15))
ax = fig.add_subplot(projection='3d')
ax.quiver(x, y, z, u, v, w, color=flat_colours, length=0.1, normalize=True)

plt.show()

补充说明 #

如果你的colour不是numpy数组，而是嵌套列表结构，可以用列表推导式手动展平：

flat_colours = [c for layer in colour for row in layer for c in row]

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