%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt, mpld3
from matplotlib.animation import FuncAnimation
from JSAnimation import IPython_display
import scipy, scipy.linalg, scipy.sparse, scipy.sparse.linalg, scipy.fftpack
plt.rcParams['figure.figsize'] = 12, 8
mpld3.enable_notebook()
Initial wave packet:
x_0 = -10.
k = 5.
a = 1
def wavepacket(x, t=0):
return (np.pi * a**2 / 2)**(-1./4) \
* np.sqrt(a**2 / (a**2 + 2j*t)) \
* np.exp(1j*(k*x - k**2 * t / 2)) \
* np.exp(-(x - x_0 - k*t)**2 / (a**2 + 2j*t))
(Discretized) Hamiltonian for given potential:
def hamiltonian(V):
diag = V(xs) + 1. / dx**2
offdiag = -0.5 / dx**2 * np.ones(x_steps)
return scipy.sparse.dia_matrix(([diag, offdiag, offdiag], [0, -1, 1]), shape=(x_steps, x_steps))
Video recording:
def display_video(next, V=None):
fig, ax1 = plt.subplots()
if V:
ax2 = ax1.twinx()
ax2.set_aspect('equal')
ax2.plot(xs, V(xs), '--', color='gray')
for tl in ax2.get_yticklabels(): tl.set_color('gray')
line_real, = ax1.plot([], [], 'r-', label="Re \Psi")
line_imag, = ax1.plot([], [], 'g-', label="Im \Psi")
line_abs_sqr, = ax1.plot([], [], 'b-', lw=2, label="|\Psi|^2")
ax1.set_xlim(x_min, x_max)
ax1.set_ylim(-1, 1)
ax1.set_xlabel("x")
ax1.axvline(0, ls="--", color="gray")
ax1.legend()
def update(t):
psi = next(dt, t)
line_real.set_data(xs, psi.real)
line_imag.set_data(xs, psi.imag)
line_abs_sqr.set_data(xs, np.abs(psi)**2)
anim = FuncAnimation(fig, update, ts, interval=30, blit=True)
#anim.save(filename, fps=1/dt)
#anim.save(filename, fps=1/dt, writer="ffmpeg", codec="libx264")
#anim.save(filename, fps=1/dt, writer="ffmpeg", codec="libvpx")
#plt.close(fig)
return anim
Discretization parameters:
x_min = -50
x_max = 50
x_steps = 5000
xs, dx = np.linspace(x_min, x_max, x_steps, retstep=True)
t_min = 0
t_max = 15
t_steps = 100 # 1000
ts, dt = np.linspace(t_min, t_max, t_steps, retstep=True)
V_free = np.vectorize(lambda x: 0)
H_free = hamiltonian(V_free)
Exact Solution:
display_video(lambda dt, t: wavepacket(xs, t))