Как добавить цветовую карту в анимацию matplotlib?

Question

Предположим, я пытаюсь визуализировать положение объекта, используя его положения по осям X и Y, и используя некоторую другую переменную, 'Z', в качестве карты цветов.Этот упрощенный пример ниже иллюстрирует, как я сейчас делаю это.

import numpy as np
from matplotlib import pyplot as plt

dataX = np.linspace(-50,50,1000)
dataY = np.linspace(-50,50,1000)
dataZ = np.linspace(-50,50,1000)

plt.scatter(dataX, dataY, c=dataZ, cmap='winter', edgecolors='none')
plt.colorbar()
plt.show()

и результат:

Вместо этого я хочу добавить живую анимациюпросто показывает статическое изображение, но я изо всех сил пытаюсь добавить к нему цветовую карту.Приведенный ниже код показывает, как я это делаю без карты цветов.

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import Tkinter
import tkMessageBox

def restart():
    root = Tkinter.Tk()
    root.withdraw()
    result = tkMessageBox.askyesno("Restart", "Would you like to restart the animation?")
    if result:
        ani.frame_seq = ani.new_frame_seq() 
        ani.event_source.start()
    else:
        plt.close()

dataX = np.linspace(-50,50,1000)
dataY = np.linspace(-50,50,1000)
dataZ = np.linspace(-50,50,1000)


def function(num, dataX,dataY, line):
    line.set_data(dataX[..., :num],dataY[..., :num])
    if num == dataX.size :
        restart()
    return line,

fig = plt.figure()

l, = plt.plot([], [], 'ro', markeredgewidth=0.0)
limitsX = [min(dataX)-100,max(dataX)+100]
limitsY = [min(dataY)-100, max(dataY)+100]
plt.xlim(limitsX[0],limitsX[1] )
plt.ylim(limitsY[0],limitsY[1] )
plt.xlabel('x')
plt.ylabel('y')
plt.title('test')

ani = animation.FuncAnimation(fig, function, (dataX.size+1), fargs=(dataX,dataY,l),
                              interval=10, blit = True, repeat = False)

plt.show()

Повторно задаем вопрос: как добавить карту цветов в мою анимацию?

Answer 1

Этот ответ основывается на моем ответе на подобную проблему , хотя я чувствовал, что конкретный случай обработки PathCollection (как возвращено scatter()) требует нового ответа.

Вот как я бы подошел к проблеме.Хитрость заключается в том, чтобы создать отдельную статическую цветную полосу на втором топоре на рисунке.Затем при обновлении свойств PathCollection используйте эту цветовую полосу и нормализацию для обновления цвета точек.

dataX = np.linspace(-50,50,1000)
dataY = np.linspace(-50,50,1000)
dataZ = np.linspace(-50,50,1000)


def animate(num):
    data = np.hstack((dataX[:num,np.newaxis], dataY[:num, np.newaxis]))
    art.set_offsets(data)
    art.set_color(cmap(norm(dataZ[:num]))) # update colors using the colorbar and its normalization defined below
    return art,


fig,[ax,cax] = plt.subplots(1,2, gridspec_kw={"width_ratios":[50,1]})
# Set the colormap and norm to correspond to the data for which
# the colorbar will be used.
cmap = matplotlib.cm.winter
norm = matplotlib.colors.Normalize(vmin=-50, vmax=50)

cb1 = matplotlib.colorbar.ColorbarBase(cax, cmap=cmap,
                                norm=norm,
                                orientation='vertical')

ax.set_xlim(-60,60)
ax.set_ylim(-60,60)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('test')
art = ax.scatter([],[],c=[])

ani = animation.FuncAnimation(fig, animate,interval=2, blit=True, repeat=True)

plt.show()