import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import PillowWriter
import animatplot as amp
def psi(x, t):
return (2**-.5*np.exp(t*1j)*np.sin(np.pi*x)
+ .5*np.exp(t*4j)*np.sin(2*np.pi*x)
+ .5*np.exp(t*9j)*np.sin(3*np.pi*x))
x = np.linspace(0, 1, 20)
t = np.linspace(0, 10, 20)
X, T = np.meshgrid(x, t)
Y1 = psi(X, T).real
Y2 = psi(X, T).imag
timeline = amp.Timeline(t, 's', 24)
ax = plt.axes(xlim=[0, 1], ylim=[-2, 2])
block1 = amp.blocks.Line(X, Y1, ax=ax)
block2 = amp.blocks.Line(X, Y2, ax=ax)
anim = amp.Animation([block1, block2], timeline)
# Your standard matplotlib stuff
plt.title(r'Particle in a Box: $|\Psi\rangle = \frac{1}{\sqrt{2}}'
r'|E_1\rangle + \frac{1}{2}|E_2\rangle + \frac{1}{2}|E_3\rangle$',
y=1.03)
plt.xlabel('position')
plt.ylabel(r'$\Psi$')
plt.legend(['Real', 'Imaginary'])
anim.toggle()
anim.timeline_slider()
anim.save('sq_well.gif', writer=PillowWriter(fps=5))
plt.show()