matplotlib里的动画

本文由Mathematrix译自由Jake Vanderplas撰写的Matplotlib Animation Tutorial

Matplotlib 1.1 版新添加了一些非常帅的用来制作动画的工具，你可以在Matplotlib的Example页看到一些非常棒的例子（OldExample, Example）。在这里我与大家分享下我使用这些工具的一些经验。

基本动画

动画工具的中心是matplotlib.animation.Animation基类，这个类提供了动画功能的基础。两个主要接口则分别是imeAnimation和FuncAnimation，你可以在这里查看API文档。下面我说说怎么使用我认为最实用的FuncAnimation。

首先我们用FuncAnimation做一个基本的动画，这个动画可以在屏幕上动态显示正弦函数的图像。

  1. “””  
  2. matplotlib animation example  
  3.  
  4. author: jake vanderplas  
  5. email: vanderplas@astro.washington.edu  
  6. website: http://jakevdp.github.com  
  7. license: bsd  
  8. please feel free to use and modify this, but keep the above information. thanks!  
  9. “””  
  10.   
  11. import numpy as np   
  12. from matplotlib import pyplot as plt   
  13. from matplotlib import animation   
  14.   
  15. # first set up the figure, the axis, and the plot element we want to animate   
  16. fig = plt.figure()   
  17. ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))   
  18. line, = ax.plot([], [], lw=2)   
  19.   
  20. # initialization function: plot the background of each frame   
  21. def init():   
  22.     line.set_data([], [])   
  23.     return line,   
  24.   
  25. # animation function.  this is called sequentially   
  26. def animate(i):   
  27.     x = np.linspace(0, 2, 1000)   
  28.     y = np.sin(2 * np.pi * (x - 0.01 * i))   
  29.     line.set_data(x, y)   
  30.     return line,   
  31.   
  32. # call the animator.  blit=true means only re-draw the parts that have changed.   
  33. anim = animation.funcanimation(fig, animate, init_func=init,   
  34.                                frames=200, interval=20, blit=true)   
  35.   
  36. # save the animation as an mp4.  this requires ffmpeg or mencoder to be  
  37. # installed.  the extra_args ensure that the x264 codec is used, so that  
  38. # the video can be embedded in html5.  you may need to adjust this for   
  39. # your system: for more information, see   
  40. # http://matplotlib.sourceforge.net/api/animation_api.html   
  41. anim.save(‘basic_animation.mp4′, fps=30, extra_args=['-vcodec', 'libx264'])   
  42.   
  43. plt.show()  

 

让我们一步一步来研究这个并看看下面将会发生什么。在引入几个必须的numpy和matplotlib库后，脚本开始设置这个plot:

  1. fig = plt.figure()   
  2. ax = plt.axes(xlim=(0, 2), ylim=(-2, 2))   
  3. line, = ax.plot([], [], lw=2)  

 

这里我们建立了一个figure窗口，然后建立了一个轴，接着建立了一个我们要在动画中不断被修改的line对象。注意我们现在仅仅是画了一个空白的线，我们将稍后添加数据

下面我们就创建动画发生时调用的函数了。Init()是我们的动画在在创建动画基础框架(base frame)时调用的函数。这里我们们用一个非常简单的对line什么都不做的函数。这个函数一定要返回line对象，这个很重要，因为这样就能告诉动画之后要更新的内容，也就是动作的内容是line。

  1. def init():   
  2. line.set_data([], [])   
  3. return line,  

 

下面是动画函数。这个函数需要一个参数i，就是帧数，并且一边位移一边根据i画出正弦函数图象。

  1. # animation function. this is called sequentially   
  2. def animate(i):   
  3. x = np.linspace(0, 2, 1000)   
  4. y = np.sin(2 * np.pi * (x - 0.01 * i))   
  5. line.set_data(x, y)   
  6. return line,  

 

再一次注意这里我们返回了我们修改过的对象，这将告诉动画框架哪一部分是要有动作的。

最后，我们创建动画对象

  1. anim = animation.funcanimation(fig, animate, init_func=init, frames=100, interval=20, blit=true)  

 

这个对象需要持续存在，所有我们要将它赋给一个变量，我们选择了一个100帧的动画（译者注：你上边的代码还是200帧，怎么到这儿就变成100帧了……，另外，这里也不一定一定要是个数字，可以是个generator 或iterable，详见API说明）并且帧与帧之间间隔20ms，blit是一个非常重要的关键字，它告诉动画只重绘修改的部分，结合上面保存的时间， blit=true会使动画显示得会非常非常快。

我们用一个可选的保存命令来结束，并且用一个show()函数让这个图像显示出来，下面是运行的结果：

这个生成和保存动画的框架非常的强大且有着良好的灵活性：如果我们在动画中添加一些物理原理，那么可能性将是无止境的。下面是我试验过的两个物理动画。

二级摆

第一个例子讲给大家展示二级摆。这个例子是通过提前计算好摆在10s内的位置，然后用动画的形式播放出来。当我刚开始看到这个的时候还在怀疑python能不能把这个图形跑得足够快使人们看起来球就像是在非一样。代码是这样的：

  1. “””  
  2. general numerical solver for the 1d time-dependent schrodinger’s equation.  
  3.  
  4. adapted from code at http://matplotlib.sourceforge.net/examples/animation/double_pendulum_animated.py  
  5.  
  6. double pendulum formula translated from the c code at  
  7. http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c  
  8.  
  9. author: jake vanderplas  
  10. email: vanderplas@astro.washington.edu  
  11. website: http://jakevdp.github.com  
  12. license: bsd  
  13. please feel free to use and modify this, but keep the above information. thanks!  
  14. “””  
  15.   
  16. from numpy import sin, cos   
  17. import numpy as np   
  18. import matplotlib.pyplot as plt   
  19. import scipy.integrate as integrate   
  20. import matplotlib.animation as animation   
  21.   
  22. class doublependulum:   
  23.     “””double pendulum class
  24.  
  25.     init_state is [theta1, omega1, theta2, omega2] in degrees,  
  26.     where theta1, omega1 is the angular position and velocity of the first  
  27.     pendulum arm, and theta2, omega2 is that of the second pendulum arm  
  28.     “””  
  29.     def __init__(self,   
  30.                  init_state = [120, 0, -20, 0],   
  31.                  l1=1.0,  # length of pendulum 1 in m   
  32.                  l2=1.0,  # length of pendulum 2 in m   
  33.                  m1=1.0,  # mass of pendulum 1 in kg   
  34.                  m2=1.0,  # mass of pendulum 2 in kg   
  35.                  g=9.8,  # acceleration due to gravity, in m/s^2   
  36.                  origin=(0, 0)):    
  37.         self.init_state = np.asarray(init_state, dtype=’float‘)   
  38.         self.params = (l1, l2, m1, m2, g)   
  39.         self.origin = origin   
  40.         self.time_elapsed = 0   
  41.   
  42.         self.state = self.init_state * np.pi / 180.   
  43.        
  44.     def position(self):   
  45.         “””compute the current x,y positions of the pendulum arms”””  
  46.         (l1, l2, m1, m2, g) = self.params   
  47.   
  48.         x = np.cumsum([self.origin[0],   
  49.                        l1 * sin(self.state[0]),   
  50.                        l2 * sin(self.state[2])])   
  51.         y = np.cumsum([self.origin[1],   
  52.                        -l1 * cos(self.state[0]),   
  53.                        -l2 * cos(self.state[2])])   
  54.         return (x, y)   
  55.   
  56.     def energy(self):   
  57.         “””compute the energy of the current state”””  
  58.         (l1, l2, m1, m2, g) = self.params   
  59.   
  60.         x = np.cumsum([l1 * sin(self.state[0]),   
  61.                        l2 * sin(self.state[2])])   
  62.         y = np.cumsum([-l1 * cos(self.state[0]),   
  63.                        -l2 * cos(self.state[2])])   
  64.         vx = np.cumsum([l1 * self.state[1] * cos(self.state[0]),   
  65.                         l2 * self.state[3] * cos(self.state[2])])   
  66.         vy = np.cumsum([l1 * self.state[1] * sin(self.state[0]),   
  67.                         l2 * self.state[3] * sin(self.state[2])])   
  68.   
  69.         u = g * (m1 * y[0] + m2 * y[1])   
  70.         k = 0.5 * (m1 * np.dot(vx, vx) + m2 * np.dot(vy, vy))   
  71.   
  72.         return u + k   
  73.   
  74.     def dstate_dt(self, state, t):   
  75.         “””compute the derivative of the given state”””  
  76.         (m1, m2, l1, l2, g) = self.params   
  77.   
  78.         dydx = np.zeros_like(state)   
  79.         dydx[0] = state[1]   
  80.         dydx[2] = state[3]   
  81.   
  82.         cos_delta = cos(state[2] - state[0])   
  83.         sin_delta = sin(state[2] - state[0])   
  84.   
  85.         den1 = (m1 + m2) * l1 - m2 * l1 * cos_delta * cos_delta   
  86.         dydx[1] = (m2 * l1 * state[1] * state[1] * sin_delta * cos_delta   
  87.                    + m2 * g * sin(state[2]) * cos_delta   
  88.                    + m2 * l2 * state[3] * state[3] * sin_delta   
  89.                    - (m1 + m2) * g * sin(state[0])) / den1   
  90.   
  91.         den2 = (l2 / l1) * den1   
  92.         dydx[3] = (-m2 * l2 * state[3] * state[3] * sin_delta * cos_delta   
  93.                    + (m1 + m2) * g * sin(state[0]) * cos_delta   
  94.                    - (m1 + m2) * l1 * state[1] * state[1] * sin_delta   
  95.                    - (m1 + m2) * g * sin(state[2])) / den2   
  96.            
  97.         return dydx   
  98.   
  99.     def step(self, dt):   
  100.         “””execute one time step of length dt and update state”””  
  101.         self.state = integrate.odeint(self.dstate_dt, self.state, [0, dt])[1]   
  102.         self.time_elapsed += dt   
  103.   
  104. #————————————————————   
  105. # set up initial state and global variables   
  106. pendulum = doublependulum([180., 0.0, -20., 0.0])   
  107. dt = 1./30 # 30 fps   
  108.   
  109. #————————————————————   
  110. # set up figure and animation   
  111. fig = plt.figure()   
  112. ax = fig.add_subplot(111, aspect=’equal’, autoscale_on=false,   
  113.                      xlim=(-2, 2), ylim=(-2, 2))   
  114. ax.grid()   
  115.   
  116. line, = ax.plot([], [], ‘o-’, lw=2)   
  117. time_text = ax.text(0.02, 0.95, ”, transform=ax.transaxes)   
  118. energy_text = ax.text(0.02, 0.90, ”, transform=ax.transaxes)   
  119.   
  120. def init():   
  121.     “””initialize animation”””  
  122.     line.set_data([], [])   
  123.     time_text.set_text(”)   
  124.     energy_text.set_text(”)   
  125.     return line, time_text, energy_text   
  126.   
  127. def animate(i):   
  128.     “””perform animation step”””  
  129.     global pendulum, dt   
  130.     pendulum.step(dt)   
  131.        
  132.     line.set_data(*pendulum.position())   
  133.     time_text.set_text(‘time = %.1f’ % pendulum.time_elapsed)   
  134.     energy_text.set_text(‘energy = %.3f j’ % pendulum.energy())   
  135.     return line, time_text, energy_text   
  136.   
  137. # choose the interval based on dt and the time to animate one step   
  138. from time import time  
  139. t0 = time()   
  140. animate(0)   
  141. t1 = time()   
  142. interval = 1000 * dt - (t1 - t0)   
  143.   
  144. ani = animation.funcanimation(fig, animate, frames=300,   
  145.                               interval=interval, blit=true, init_func=init)   
  146.   
  147. # save the animation as an mp4.  this requires ffmpeg or mencoder to be  
  148. # installed.  the extra_args ensure that the x264 codec is used, so that  
  149. # the video can be embedded in html5.  you may need to adjust this for   
  150. # your system: for more information, see   
  151. # http://matplotlib.sourceforge.net/api/animation_api.html   
  152. #ani.save(‘double_pendulum.mp4′, fps=30, extra_args=['-vcodec', 'libx264'])   
  153.   
  154. plt.show()  

 

这里我们建了一个类来保存二级摆的状态（包括摆臂的角度和角速度）并且提供了一些函数来计算这些动作。动画函数（animation functions）和上面是一样的，但是我们让更新函数（update function）要比以前稍稍复杂一些：现在我们不知要改变点的位置，还要让左上角显示两行字去追踪时间和能量状态（如果我们的数学公式没错的话能量应该是个定值）。下面的视频只录制了10s，但是通过观看运行后的脚本你可以看到二级摆不断混乱地（chaotically）摆动直到你笔记本跑完最后一点电力：

盒子里的粒子

另外一个我做的动画是一些粒子在重力的作用下在一个盒子里不断地发生弹性碰撞。注意，碰撞都是弹性的，每个里都具有一定的能量和平面上的动量，而且每当粒子碰到了盒子的内壁上都会栩栩如生地反弹到另一个方向并且所有粒子都受重力影响会向下落：

  1. “””  
  2. animation of elastic collisions with gravity  
  3.  
  4. author: jake vanderplas  
  5. email: vanderplas@astro.washington.edu  
  6. website: http://jakevdp.github.com  
  7. license: bsd  
  8. please feel free to use and modify this, but keep the above information. thanks!  
  9. “””  
  10. import numpy as np   
  11. from scipy.spatial.distance import pdist, squareform   
  12.   
  13. import matplotlib.pyplot as plt   
  14. import scipy.integrate as integrate   
  15. import matplotlib.animation as animation   
  16.   
  17. class particlebox:   
  18.     “””orbits class
  19.       
  20.     init_state is an [n x 4] array, where n is the number of particles:  
  21.        [[x1, y1, vx1, vy1],  
  22.         [x2, y2, vx2, vy2],  
  23.         …               ]  
  24.  
  25.     bounds is the size of the box: [xmin, xmax, ymin, ymax]  
  26.     “””  
  27.     def __init__(self,   
  28.                  init_state = [[1, 0, 0, -1],   
  29.                                [-0.5, 0.5, 0.5, 0.5],   
  30.                                [-0.5, -0.5, -0.5, 0.5]],   
  31.                  bounds = [-2, 2, -2, 2],   
  32.                  size = 0.04,   
  33.                  m = 0.05,   
  34.                  g = 9.8):   
  35.         self.init_state = np.asarray(init_state, dtype=float)   
  36.         self.m = m * np.ones(self.init_state.shape[0])   
  37.         self.size = size   
  38.         self.state = self.init_state.copycopy()   
  39.         self.time_elapsed = 0   
  40.         self.bounds = bounds   
  41.         self.g = g   
  42.   
  43.     def step(self, dt):   
  44.         “””step once by dt seconds”””  
  45.         self.time_elapsed += dt   
  46.            
  47.         # update positions   
  48.         self.state[:, :2] += dt * self.state[:, 2:]   
  49.   
  50.         # find pairs of particles undergoing a collision   
  51.         d = squareform(pdist(self.state[:, :2]))   
  52.         ind1, ind2 = np.where(d < 2 * self.size)   
  53.         unique = (ind1 < ind2)   
  54.         ind1 = ind1[unique]   
  55.         ind2 = ind2[unique]   
  56.   
  57.         # update velocities of colliding pairs   
  58.         for i1, i2 in zip(ind1, ind2):   
  59.             # mass   
  60.             m1 = self.m[i1]   
  61.             m2 = self.m[i2]   
  62.   
  63.             # location vector   
  64.             r1 = self.state[i1, :2]   
  65.             r2 = self.state[i2, :2]   
  66.   
  67.             # velocity vector   
  68.             v1 = self.state[i1, 2:]   
  69.             v2 = self.state[i2, 2:]   
  70.   
  71.             # relative location & velocity vectors   
  72.             r_rel = r1 - r2   
  73.             v_rel = v1 - v2   
  74.   
  75.             # momentum vector of the center of mass   
  76.             v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)   
  77.   
  78.             # collisions of spheres reflect v_rel over r_rel   
  79.             rr_rel = np.dot(r_rel, r_rel)   
  80.             vr_rel = np.dot(v_rel, r_rel)   
  81.             v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel   
  82.   
  83.             # assign new velocities   
  84.             self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)   
  85.             self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2)    
  86.   
  87.         # check for crossing boundary   
  88.         crossed_x1 = (self.state[:, 0] < self.bounds[0] + self.size)   
  89.         crossed_x2 = (self.state[:, 0] > self.bounds[1] - self.size)   
  90.         crossed_y1 = (self.state[:, 1] < self.bounds[2] + self.size)   
  91.         crossed_y2 = (self.state[:, 1] > self.bounds[3] - self.size)   
  92.   
  93.         self.state[crossed_x1, 0] = self.bounds[0] + self.size   
  94.         self.state[crossed_x2, 0] = self.bounds[1] - self.size   
  95.   
  96.         self.state[crossed_y1, 1] = self.bounds[2] + self.size   
  97.         self.state[crossed_y2, 1] = self.bounds[3] - self.size   
  98.   
  99.         self.state[crossed_x1 | crossed_x2, 2] *= -1   
  100.         self.state[crossed_y1 | crossed_y2, 3] *= -1   
  101.   
  102.         # add gravity   
  103.         self.state[:, 3] -= self.m * self.g * dt   
  104.   
  105.   
  106. #————————————————————   
  107. # set up initial state   
  108. np.random.seed(0)   
  109. init_state = -0.5 + np.random.random((50, 4))   
  110. init_state[:, :2] *= 3.9   
  111.   
  112. box = particlebox(init_state, size=0.04)   
  113. dt = 1. / 30 # 30fps   
  114.   
  115.   
  116. #————————————————————   
  117. # set up figure and animation   
  118. fig = plt.figure()   
  119. fig.subplots_adjust(left=0, right=1, bottom=0, top=1)   
  120. ax = fig.add_subplot(111, aspect=’equal’, autoscale_on=false,   
  121.                      xlim=(-3.2, 3.2), ylim=(-2.4, 2.4))   
  122.   
  123. # particles holds the locations of the particles   
  124. particles, = ax.plot([], [], ‘bo’, ms=6)   
  125.   
  126. # rect is the box edge   
  127. rect = plt.rectangle(box.bounds[::2],   
  128.                      box.bounds[1] - box.bounds[0],   
  129.                      box.bounds[3] - box.bounds[2],   
  130.                      ec=’none’, lw=2, fc=’none’)   
  131. ax.add_patch(rect)   
  132.   
  133. def init():   
  134.     “””initialize animation”””  
  135.     global box, rect   
  136.     particles.set_data([], [])   
  137.     rect.set_edgecolor(‘none’)   
  138.     return particles, rect   
  139.   
  140. def animate(i):   
  141.     “””perform animation step”””  
  142.     global box, rect, dt, ax, fig   
  143.     box.step(dt)   
  144.   
  145.     ms = int(fig.dpi * 2 * box.size * fig.get_figwidth()   
  146.              / np.diff(ax.get_xbound())[0])   
  147.        
  148.     # update pieces of the animation   
  149.     rect.set_edgecolor(‘k’)   
  150.     particles.set_data(box.state[:, 0], box.state[:, 1])   
  151.     particles.set_markersize(ms)   
  152.     return particles, rect   
  153.   
  154. ani = animation.funcanimation(fig, animate, frames=600,   
  155.                               interval=10, blit=true, init_func=init)   
  156.   
  157.   
  158. # save the animation as an mp4.  this requires ffmpeg or mencoder to be  
  159. # installed.  the extra_args ensure that the x264 codec is used, so that  
  160. # the video can be embedded in html5.  you may need to adjust this for   
  161. # your system: for more information, see   
  162. # http://matplotlib.sourceforge.net/api/animation_api.html   
  163. #ani.save(‘particle_box.mp4′, fps=30, extra_args=['-vcodec', 'libx264'])  
  164.   
  165. plt.show()  

 

如果你稍学过一些物理学知识，就应该不会对里面的公式太陌生，而且最终运行结果是令人入迷的。我在飞机上把代码写了出来，然后坐着看了它差不多10分钟。

这仅仅是个开始：往这些动画里加些别的元素是个很有趣的锻炼，像计算温度和气压来说明合理的油价（译者注：原话为like computation of the temperature and pressure to demonstrate the ideal gas law，考虑到美国油价以为加仑计量单位，原作者是名物理学的博士，所以这里应该是调侃油价计算应该仔细考虑温度和气压）或者是画速度分布的实时曲线来观察它能否接近麦克斯韦分布（译者注：麦克斯韦-玻尔兹曼分布，统计力学中的一个概率分布）。这个为可视化物理学演示打开了太多的可能……

总结

Matplotlib的动画模块可以说是对本来就很出色的Matplotlib一个非常好的改进。我认为我认为我已经用这些工具抓到了一些浅显的可能……你们还有什么能做出很酷的动画的电子么？

修改：在这篇博文里，我展示了如何做一个简单的量子力学系统的演示动画。

posted @ 2016-04-19 13:28  StevenLuke  阅读(229)  评论(0)  编辑  收藏  举报