Как оживить функцию с помощью вложенного цикла?

Question

Мой текущий код будет построен, а затем обновит график с помощью команды 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()

Answer 1

Я несвоевременно только что изменил вложенный цикл для массива. Следовательно, я по сути только что сделал первый цикл 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()