在 matplotlib 中为 3d 箭袋图添加颜色

Adding colors to a 3d quiver plot in matplotlib

我想要在我的 3d 箭袋图中使用与颜色图相对应的颜色。绘图的二维版本有一个可选数组,用于将颜色映射到箭头。如何在 3d 版本中创建相同的效果?

3D 箭袋图是 1.4 中的一项全新功能,它(及其文档)的边缘可能仍然有点粗糙。在这种情况下,我们可以尝试使用 quiver 实现为 LineCollection 的事实,它(最终)继承自 ScalarMappable 这意味着它知道颜色图是什么并且返回的艺术家具有方法 set_array.

基于文档here

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

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

x, y, z = np.meshgrid(np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.8))

u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
w = (np.sqrt(2.0 / 3.0) * np.cos(np.pi * x) * np.cos(np.pi * y) *
     np.sin(np.pi * z))

q = ax.quiver(x, y, z, u, v, w, length=0.1, cmap='Reds', lw=2)
q.set_array(np.random.rand(np.prod(x.shape)))

plt.show()

但是,您会注意到 headsshaft 的颜色不同,这是由于实现细节的方式它实现了每个部分都被绘制成它自己的线。

直接使用 Norm 和颜色映射函数并将结果传递给 colors 可能是更好的方法。

可以 使用 colors 参数为每​​个箭头指定自定义颜色,尽管这样做的方式一点也不简单(至于 matplotlib 2.0.0)。我在 this issue the logic of how the quiver plot is actually drawn and a work-around to specify coloring. You can check this gist 中指定了一个简单示例,它可以生成与此类似的图表:

总而言之,以下是要遵循的步骤:

  1. 假设您有一个 3 元组(或 RGBA 的 4 元组)的列表(大小为 x*y*z),指定要绘制的每个矢量的 RGB 值(在 0~1 之间)。
  2. 过滤掉与长度为 0 的向量对应的 RGB(或 RGBA)元组,因为它们实际上不会被绘制。
  3. [color_1, color_2, ..., color_n] 成为您在第 2 步之后获得的列表,您应该指定 colors=[color_1, color_2, ..., color_n, color_1, color_1, color_2, color_2, ..., color_n, color_n] 因为实际上所有非零箭头的“-”部分(由 1 行组成)“- >" 将首先绘制,然后是 ">" 部分(由 2 行组成)。

希望对您有所帮助。

基于@tacaswell 和@sytrus 的回答,这里有一个为 3d 箭袋图着色的示例

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

# Make the grid
x, y, z = np.meshgrid(np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.8))

# Make the direction data for the arrows
u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
w = (np.sqrt(2.0 / 3.0) * np.cos(np.pi * x) * np.cos(np.pi * y) *
     np.sin(np.pi * z))

# Color by azimuthal angle
c = np.arctan2(v, u)
# Flatten and normalize
c = (c.ravel() - c.min()) / c.ptp()
# Repeat for each body line and two head lines
c = np.concatenate((c, np.repeat(c, 2)))
# Colormap
c = plt.cm.hsv(c)

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.quiver(x, y, z, u, v, w, colors=c, length=0.1, normalize=True)
plt.show()

扩展来自@slek120 的答案。我遇到了一个问题,即存在长度为零的向量。这些搞乱了箭头尖端颜色之间的对应关系。我的解决方案是给它们一个非零长度并使它们透明。出于某种我不明白的原因,简单地丢弃它们是行不通的。对最后一部分添加一个小改动,也可以包含一个颜色条。颜色条明确要求 q.set_array()。这会更改颜色,但 q.set_edgecolor(c); q.set_facecolor(c) 允许您插入自定义颜色图。

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

cmap = 'hsv'

# Make the grid
x, y, z = np.meshgrid(np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.2),
                      np.arange(-0.8, 1, 0.8))

# Make the direction data for the arrows
u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
w = np.sqrt(2.0 / 3.0) * np.cos(np.pi * x) * np.cos(np.pi * y) * np.sin(np.pi * z)

# check what happens if all values are zero
# no quivers are plotted, colors don't match anymore
u[:,2:4] = v[:,2:4] = w[:,2:4] = 0
# change values that are zero to something close to zero
uvw = np.vstack((u[np.newaxis],v[np.newaxis],w[np.newaxis]))
norm = np.linalg.norm(uvw, axis = 0)
max_norm = np.max(norm)
mask = norm == 0
min_norm = 0.3  # you want every arrow to be longer than this fraction of max_norm
# rescale vs for illustrative purposes, so small vectors become visible
# and zero vectors become nonzero so colors of the arrow shaft and head correspond. Later these are made transparent
uvw = uvw + min_norm * np.tile(mask[np.newaxis], (3, 1, 1, 1)) / max_norm
# recalculate norms so you don't divide by zero
norm = np.linalg.norm(uvw, axis=0)
uvw = min_norm * uvw / norm + (1 - min_norm) * uvw / max_norm
u, v, w = uvw

# Color by azimuthal angle
c = np.arctan2(v, u)
# Flatten and normalize
c = (c.ravel() - c.min()) / c.ptp()
# Adjust for missing quivers
# c = c[np.nonzero((u.ravel() != 0) * (v.ravel() != 0) * (w.ravel() != 0))]
# Repeat for each body line and two head lines
c = np.concatenate((c, np.repeat(c, 2)))
repeated_mask = np.concatenate((mask.ravel(), np.repeat(mask.ravel(), 2)))
# Colormap
c = getattr(plt.cm, cmap)(c)
# set zero values transparent, you made them nonzero not to mess up the tip colors
c[repeated_mask, 3] = 0.1

fig = plt.figure()
ax = fig.gca(projection='3d')
q = ax.quiver(x, y, z, u, v, w, cmap = cmap, length=0.1)
q.set_array(np.linspace(0,max_norm,10))
fig.colorbar(q)
q.set_edgecolor(c)
q.set_facecolor(c)

plt.show()