python - 如何使用 matplotlib 动态更改 scatter3d 动画中的点颜色？

Question

我正在尝试使用matplotlib绘制几个3DScatter的动画。我成功地画出了所有的点，但我在颜色方面遇到了困难。即使我调用函数 set_color(..) 也没有任何改变。这是我目前正在做的事情，to_plot是一个大小为total的数组，其中包含(5120, 3) float 元素和颜色 是一个大小为 total 的数组，其中包含 (5120,) 元素(等于“r”或“b”):

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

total = 10

num_whatever = 100    # old = 5120

to_plot = [np.random.rand(num_whatever, 3) for i in range(total)]
colors = [['r' if i%2==0 else 'b' for i in range(num_whatever)] for i in range(total)]
red_patch = mpatches.Patch(color='red', label='Men')
blue_patch = mpatches.Patch(color='blue', label='Women')

fig = plt.figure()
ax3d = Axes3D(fig)
scat3D = ax3d.scatter([],[],[], s=10)
ttl = ax3d.text2D(0.05, 0.95, "", transform=ax3d.transAxes)

def update_plot(i):
    print i, to_plot[i].shape
    ttl.set_text('PCA on 3 components at step = {}'.format(i*20))
    scat3D._offsets3d = np.transpose(to_plot[i])
    scat3D.set_color(colors[i])
    return scat3D,

def init():
    scat3D.set_offsets([[],[],[]])
    ax3d.set_xlim(-1.,2.)
    ax3d.set_ylim(-0.5,0.7)
    ax3d.set_zlim(-1.,0.75)
    plt.style.use('ggplot')
    plt.legend(handles=[red_patch, blue_patch])

ani = animation.FuncAnimation(fig, update_plot, init_func=init, blit=False, interval=100, frames=xrange(total))

# ani.save(os.path.join(config.workdir, 'gif', 'bins','anim.gif'), writer="imagemagick")

plt.plot()

Answer 1

散点图是 Path3DCollection 。它可以有一个与之关联的颜色图，以便它的点根据颜色数组进行着色。

因此，您可以通过 scat3D.set_array(colors[i]) 向散点图提供数值列表，其中colors[i] = [0,1,0,...,1,0,1] 。然后根据使用的颜色图映射这些值。对于蓝色/红色，这很简单，因为已经存在颜色图 "bwr"从蓝色到红色。

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

total = 10

num_whatever = 100

to_plot = [np.random.rand(num_whatever, 3) for i in range(total)]
colors = [np.tile([0,1],num_whatever//2) for i in range(total)]
red_patch = mpatches.Patch(color='red', label='Men')
blue_patch = mpatches.Patch(color='blue', label='Women')

fig = plt.figure()
ax3d = Axes3D(fig)
scat3D = ax3d.scatter([],[],[], s=10, cmap="bwr", vmin=0, vmax=1)
scat3D.set_cmap("bwr") # cmap argument above is ignored, so set it manually
ttl = ax3d.text2D(0.05, 0.95, "", transform=ax3d.transAxes)

def update_plot(i):
    print i, to_plot[i].shape
    ttl.set_text('PCA on 3 components at step = {}'.format(i*20))
    scat3D._offsets3d = np.transpose(to_plot[i])
    scat3D.set_array(colors[i])
    return scat3D,

def init():
    scat3D.set_offsets([[],[],[]])
    plt.style.use('ggplot')
    plt.legend(handles=[red_patch, blue_patch])

ani = animation.FuncAnimation(fig, update_plot, init_func=init, 
                              blit=False, interval=100, frames=xrange(total))

ani.save("ani.gif", writer="imagemagick")

plt.show()