简体   繁体   English

使用 matplotlib 使用颜色条动画散布 plot

[英]Animate scatter plot with colorbar using matplotlib

I have spent a serious amount of time trying to animate scatter plot where the colour of the marker is defined by a number.我花了很多时间尝试为散点图 plot 设置动画,其中标记的颜色由数字定义。

Below is my attempt, which sort of works but not really as planned:以下是我的尝试,虽然效果不错,但并没有按计划进行:

  • After each animation step, the old points should be removed.在每个 animation 步骤之后,应删除旧点。 Instead, the new points are simply added to the points already on plot.相反,新点只是简单地添加到 plot 上已有的点。
  • The colorbar should also be updated at each step according to the values (just like time text is).颜色栏也应该根据值在每一步更新(就像时间文本一样)。

However, whatever I seem to do, produces a blank chart.但是,无论我做什么,都会产生一张空白图表。 Is this really the best I can do with python when it comes to animating scatter?在动画散布方面,这真的是我能用 python 做的最好的吗?

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

fig = plt.figure()
marker_size = 1
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def init():
    """ Initialize animation. """
    scat = ax.scatter(positions[0][0], positions[0][1], s = marker_size, c = solutions[0], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(scat)
    time_text.set_text('Time step = %d' % 0)

    return scat, time_text

def animate(i):
    """ Perform animation step. """
    scat = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    time_text.set_text('Time step = %d' % i)

    return scat, time_text

plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.grid(b=None)
plt.show()
ani = animation.FuncAnimation(fig, animate, interval=100, blit=True, repeat=True, init_func=init)

ani.save('animation.gif', writer='imagemagick', fps = 8)

I'm not sure if you were indicating this from your post, but I couldn't get your code to run as is.我不确定您是否在帖子中指出了这一点,但我无法让您的代码按原样运行。 However, I believe the main issue relates to the first point you mention: " After each animation step, the old points should be removed. " You do need to be explicit about this when drawing the animation.但是,我认为主要问题与您提到的第一点有关:“在每个 animation 步骤之后,应删除旧点。 ”在绘制 animation 时,您确实需要明确说明这一点。 Currently, your code is repeatedly creating a scatter for the same Axes .目前,您的代码正在重复为相同的Axes创建scatter Just as if you were to do this outside of an animation, this will result in multiple sets of data being drawn over each other.就像您要在 animation 之外执行此操作一样,这将导致多组数据相互重叠。

I have seen 2 major ways people do this: either using some set_... methods of the plot to update the data (seen here for scatter plots or here in general ) or clearing the Axes or Figure each iteration in order to plot new data.我已经看到人们这样做的两种主要方式:要么使用 plot 的一些set_...方法来更新数据(参见此处的散点图此处一般清除AxesFigure每次迭代以 plot 新数据. I find the latter easier/more universal (if lazier).我发现后者更容易/更普遍(如果更懒惰)。 Here is an approach for your example doing so (I've edited this code to remove calls to plt.grid and plt.label , as those were not functional):这是您的示例的一种方法(我已编辑此代码以删除对plt.gridplt.label的调用,因为它们不起作用):

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

fig, ax = plt.subplots()
marker_size = 5 #upped this to make points more visible

def animate(i):
    """ Perform animation step. """
    #important - the figure is cleared and new axes are added
    fig.clear()
    ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
    #the new axes must be re-formatted
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)
    ax.grid(b=None)
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    # and the elements for this frame are added
    ax.text(0.02, 0.95, 'Time step = %d' % i, transform=ax.transAxes)
    s = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(s)

ani = animation.FuncAnimation(fig, animate, interval=100, frames=range(time_steps))

ani.save('animation.gif', writer='pillow')

Producing the following GIF:生成以下 GIF:

在此处输入图像描述

Here, I use fig.clear() to clear the colorbar each frame;在这里,我使用fig.clear()来清除每一帧的colorbar条; otherwise, many of them will be drawn.否则,其中许多将被绘制。 This means you have to re-add the Axes and the formatting each time.这意味着您每次都必须重新添加Axes和格式。 In other cases, using ax.clear() can be fine and save the step of add_subplot .在其他情况下,使用ax.clear()可以很好并节省add_subplot的步骤。

There is another way to do this however, following here .然而,还有另一种方法可以做到这一点,如下此处 If you have the handle for the colorbar Axes , you can just clear them (rather than clearing the entire Figure ), similar to the scatter plot axes:如果您有颜色条Axes的句柄,则可以清除它们(而不是清除整个Figure ),类似于分散 plot 轴:

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

time_steps = 50
N_nodes = 100

positions = []
solutions = []
for i in range(time_steps):
    positions.append(np.random.rand(2, N_nodes))
    solutions.append(np.random.random(N_nodes))

# init the figure, so the colorbar can be initially placed somewhere
marker_size = 5
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, xlim=(0, 1), ylim=(0, 1))
s = ax.scatter(positions[0][0], positions[0][1], s = marker_size, c = solutions[0], cmap = "RdBu_r", marker = ".", edgecolor = None)
cb = fig.colorbar(s)

# get the axis for the colobar
cax = cb.ax

def animate(i):
    """ Perform animation step. """
    # clear both plotting axis and colorbar axis
    ax.clear()
    cax.cla()
    #the new axes must be re-formatted
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)
    ax.grid(b=None)
    ax.set_xlabel('x [m]')
    ax.set_ylabel('y [m]')
    # and the elements for this frame are added
    ax.text(0.02, 0.95, 'Time step = %d' % i, transform=ax.transAxes)
    s = ax.scatter(positions[i][0], positions[i][1], s = marker_size, c = solutions[i], cmap = "RdBu_r", marker = ".", edgecolor = None)
    fig.colorbar(s, cax=cax)

ani = animation.FuncAnimation(fig, animate, interval=100, frames=range(time_steps))

ani.save('animation2.gif', writer='pillow')

Producing the same figure.产生相同的图形。

声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM