Matplotlib:使用ArtistAnimation动画plot_surface

时间:2019-02-12 09:23:28

标签: python python-3.x numpy matplotlib animation

我正在尝试在Matplotlib 3中制作plot_surface的动画。我看到了一个如何使用FuncAnimation的示例,但是我知道函数调用很慢,最好先保存所有帧然后使用ArtistAnimation对其进行动画处理。

如何将以下代码转换为可与ArtistAnimation一起使用(而不是按照此处的步骤每次绘制表面):

#!/usr/bin/env python
"""
Solving 2D Allen-Cahn Eq using pseudo-spectral with Implicit/Explicit
u_t= epsilon(u_{xx}+u_{yy}) + u - u^3
where u-u^3 is treated explicitly and u_{xx} and u_{yy} is treated implicitly
BC = Periodic
IC=v=sin(2*pi*x)+0.5*cos(4*pi*y)
"""

import math
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import time

plt.ion()

# Setup the grid
N = 64; h = 1.0/N;
x = [h*i for i in range(1,N+1)]
y = [h*i for i in range(1,N+1)]
dt = 0.05
xx,yy = (np.mat(A) for A in (np.meshgrid(x,y)))

# Initial Conditions
u = np.array(np.sin(2*math.pi*xx) + 0.5*np.cos(4*math.pi*yy), dtype=float)

epsilon = 0.01

# (ik) and (ik)^2 vectors in x and y direction
I = complex(0,1)
k_x = 2*np.pi*np.array([I*n for n in np.concatenate((np.arange(0,N/2),np.array([0]),np.arange(-N/2+1,0)))])
k_y = k_x

kxx = np.zeros((N,N), dtype=complex)
kyy = np.zeros((N,N), dtype=complex)
for i in range(N):
    for j in range(N):
        kxx[i,j] = k_x[i]**2
        kyy[i,j] = k_y[j]**2

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(xx, yy, u,rstride=1, cstride=1, cmap=cm.jet,
        linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.xlabel('x')
plt.ylabel('y')
plt.show()

v_hat = np.zeros((N,N), dtype=complex)
v_hat = np.fft.fft2(u)

for n in range(100):
    # calculate nonlinear term in real space
    v_nl = np.array(u**3, dtype=complex)
# FFT for nonlinear and linear terms
    v_nl = np.fft.fft2(v_nl)
    v_hat = (v_hat*(1+1/dt) - v_nl)
    v_hat=v_hat/(1/dt - (kxx+kyy)*epsilon) # Implicit/Explicit timestepping
    u = np.real(np.fft.ifft2(v_hat))
    # Remove old plot before drawing
    ax.collections.remove(surf)
    surf = ax.plot_surface(xx, yy, u,rstride=1, cstride=1, cmap=cm.jet,
         linewidth=0, antialiased=False)
    plt.draw()
plt.show()

0 个答案:

没有答案