Square WellΒΆ

[1]:
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()
<Figure size 640x480 with 3 Axes>
../_images/sq_well.gif