如何为带有嵌套循环的函数设置动画?

时间:2019-03-20 21:37:41

标签: python python-3.x matplotlib animation

我当前的代码将进行绘图,然后使用plt.draw命令更新该绘图。它还会使用plt.clf删除以前的更新。我想将其转换为动画并将其另存为视频。但是,所有在线示例都具有功能animate(i),但是我定义的功能已经遍历了所有内容。我认为有一个更简单的方法可以解决此问题,但我看不到它。我在下面尝试过。

def schrodinger_equation(self, total_time):
    count = 0
    for j in range(0, total_time, dt):

        time = j
        self.psi_xt = np.zeros((len(x),1),dtype=complex).reshape(len(x),1)

        for k in range(0, quantum_number):

            self.psi_xt[:,0] = self.psi_xt[:,0] + (self.Cn[0,k]
                * self.phi[:,k] * (np.exp((-1j * self.En[0,k] * time)/hbar)))

            count += 1

        style.use('seaborn-dark')

        plt.plot(x, np.real(self.psi_xt),'r',
            label='real' r'$\psi(x,t)$', linewidth = 0.75)
        plt.plot(x, np.imag(self.psi_xt),'b',
            label=r'$imag \psi(x,t)$', linewidth = 0.75)
        plt.plot(x, np.abs(self.psi_xt),'y',
            label=r'$|\psi(x,t)|$', linewidth = 0.75)

        x_min = min(self.x[:,0]-5)
        x_max = max(self.x[:,0]+5)
        psi_min = -A
        psi_max = A
        plt.xlim((x_min, x_max))
        plt.ylim((psi_min, psi_max))

        plt.legend(prop=dict(size=6))
        psi_x_line, = plt.plot([], [], c='r')
        V_x_line, = plt.plot([], [], c='k')
        left_wall_line = plt.axvline(0, c='k', linewidth=2)
        right_well_line = plt.axvline(x[-1], c='k', linewidth=2)

        plt.pause(0.01)
        plt.draw()
        plt.clf()
        plt.cla()

    print('The number of iterations: ' + str(count))

我的尝试

style.use('seaborn-dark')

fig = plt.figure()

ax = fig.add_subplot(111)

ax.legend(prop=dict(size=6))
psi_xt_real, = ax.plot([], [], c='r')
psi_xt_imag, = ax.plot([], [], c='b')
psi_xt_abs, = ax.plot([], [], c='y')
left_wall_line = ax.axvline(0, c='k', linewidth=2)
right_well_line = ax.axvline(x[-1], c='k', linewidth=2)

title = ax.set_title("")
ax.legend(prop=dict(size=12))
ax.set_xlabel('$x$')
ax.set_ylabel(r'$|\psi(x)|$')

def init():
    psi_xt_real.set_data([], [])
    psi_xt_imag.set_data([], [])
    psi_xt_abs.set_data([], [])
    return psi_xt_real,

def animate(i):

    psi_xt_real.set_data(Schrodinger.x, np.real(Schrodinger.psi_xt))
    psi_xt_imag.set_data(Schrodinger.x, np.imag(Schrodinger.psi_xt))
    psi_xt_abs.set_data(Schrodinger.x, np.abs(Schrodinger.psi_xt))

    return psi_xt_real,

ani = matplotlib.animation.FuncAnimation(fig,Schrodinger.schrodinger_equation, init_func=init, frames=10)
#ani.save('animation.gif', writer='imagemagick', fps=30)

plt.show()

1 个答案:

答案 0 :(得分:0)

我不合时宜地更改了数组的嵌套循环。因此,我实际上只是制作了第一个for循环,将时间控制为一个时间数组,并且该循环适用于动画情节。如果其他人遇到此问题,我将在下面提供解决方法。

style.use('seaborn-dark')

fig = plt.figure()

ax = fig.add_subplot(111)

ax.legend(prop=dict(size=6))
psi_xt_real, = ax.plot([], [], c='r')
psi_xt_imag, = ax.plot([], [], c='b')
psi_xt_abs, = ax.plot([], [], c='y')
left_wall_line = ax.axvline(0, c='k', linewidth=2)
right_well_line = ax.axvline(x[-1], c='k', linewidth=2)

x_min = min(Schrodinger.x[:,0]-5)
x_max = max(Schrodinger.x[:,0]+5)
psi_min = -A
psi_max = A
plt.xlim((x_min, x_max))
plt.ylim((psi_min, psi_max))

title = ax.set_title('')
ax.legend(prop=dict(size=12))
ax.set_xlabel('$x$')
ax.set_ylabel(r'$|\psi(x)|$')

######################################################################
# Animate and save the plot

def init():
    psi_xt_real.set_data([], [])
    psi_xt_imag.set_data([], [])
    psi_xt_abs.set_data([], [])

    title.set_text('')


def animate(i):

    time = np.linspace(0,total_time,dt).astype(complex)
    psi_xt = np.zeros((len(x),1),dtype=complex).reshape(len(x),1)

    for k in range(0, quantum_number):

        psi_xt[:,0] = psi_xt[:,0] + (Schrodinger.Cn[0,k] *
            Schrodinger.phi[:,k] * (np.exp((-1j *
                Schrodinger.En[0,k] * i)/hbar)))

        psi_xt_real.set_data(x, np.real(psi_xt))
        psi_xt_imag.set_data(x, np.imag(psi_xt))
        psi_xt_abs.set_data(x, np.abs(psi_xt))
        title.set_text('t = %1' %time)


ani = matplotlib.animation.FuncAnimation(fig, animate,
init_func=init, frames=100)
#ani.save('animation.gif', writer='imagemagick', fps=30)

plt.show()