Изменить цвет графика рассеяния во время анимации

Question

Я пытаюсь создать анимацию точечного графика, где цвет каждой точки сетки - это единственное, что меняется.

Я использую его следующим образом:

fig = setAnimation(x_grid, y_grid, psi_init)
ani = anime.FuncAnimation(fig, animate, frames=range(int(num_iter)),
                              fargs=[x_grid, y_grid, psi_prop_mat_time],
                              interval=50, blit=True)
plt.show()

У меня есть следующее:

wave_packet_xy = None

def setAnimation(x_grid, y_grid, psi_init):
    global wave_packet_xy

    fig = plt.figure()
    # time domain
    # ax1 = fig.add_subplot(111, autoscale_on=False)
    # ax1.grid()
    wave_packet_xy = plt.scatter(x_grid, y_grid, c=np.abs(psi_init), cmap="jet")
    return fig


def animate(frame, x_grid, y_grid, psi_prop_mat_time):
    # wave_packet_xy.set_offsets(x_grid, y_grid)
    wave_packet_xy.set_array(np.abs(psi_prop_mat_time[frame, :, :]).flatten())
    # wave_packet_xy.set_facecolor(np.abs(psi_prop_mat_time[frame, :, :]).flatten())
    return wave_packet_xy

Функция setAnimation работает и выдает следующее изображение:

В то время как анимация завершается неудачно и дает следующий результат:

Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 388, in process
    proxy(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 228, in __call__
    return mtd(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1081, in _start
    self._init_draw()
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1792, in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1819, in _draw_frame
    for a in self._drawn_artists:
TypeError: 'PathCollection' object is not iterable
Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 388, in process
    proxy(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 228, in __call__
    return mtd(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1368, in _handle_resize
    self._init_draw()
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1792, in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1819, in _draw_frame
    for a in self._drawn_artists:
TypeError: 'PathCollection' object is not iterable

Если я делаю следующее:

def setAnimation(x_grid, y_grid, psi_init):
    global wave_packet_xy

    fig = plt.figure()
    # time domain
    ax1 = fig.add_subplot(111, autoscale_on=False)
    # ax1.grid()
    wave_packet_xy = ax1.scatter(x_grid, y_grid, c=np.abs(psi_init), cmap="jet")
    return fig


def animate(frame, x_grid, y_grid, psi_prop_mat_time):
    # wave_packet_xy.set_offsets(x_grid, y_grid)
    wave_packet_xy.set_array(np.abs(psi_prop_mat_time[frame, :, :]).flatten())
    # wave_packet_xy.set_facecolor(np.abs(psi_prop_mat_time[frame, :, :]).flatten())
    return wave_packet_xy

Я получаю следующую ошибку:

Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 388, in process
    proxy(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 228, in __call__
    return mtd(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1081, in _start
    self._init_draw()
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1792, in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1819, in _draw_frame
    for a in self._drawn_artists:
TypeError: 'PathCollection' object is not iterable
Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 388, in process
    proxy(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/cbook/__init__.py", line 228, in __call__
    return mtd(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1368, in _handle_resize
    self._init_draw()
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1792, in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
  File "/usr/lib/python3/dist-packages/matplotlib/animation.py", line 1819, in _draw_frame
    for a in self._drawn_artists:
TypeError: 'PathCollection' object is not iterable

Не могу показаться обновить симуляцию правильно. Я попробовал несколько вещей, но, похоже, не работает.

Буду признателен за помощь.