%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()

C:\Users\Michael\Anaconda\lib\site-packages\numpy\oldnumeric\__init__.py:11: ModuleDeprecationWarning: The oldnumeric module will be dropped in Numpy 1.9
  warnings.warn(_msg, ModuleDeprecationWarning)

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

1. "No" Wall

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))

– +

Once Loop Reflect

It is apparent from the analytical solution that the time-evolved wavepacket is again a wavepacket: More precisely, the probability amplitude follows a Gaussian distribution with center $x_0 + k_0 t$ and spread $\sqrt{(a/2)^2 + (t/a)^2}$. The probability current is given by $\mathrm{Im} \Psi^* \nabla \Psi$. It can be interpreted as $v(x) |\Psi(x)|^2$, with $v(x)$ the "local velocity", and hence its integral gives the average velocity, which for our wave packet is equal to $k_0$. The local kinetic energy can also be readily computed. Here is a snapshot:

def prob_current(psi):
    psi_mean = (psi[1:] + psi[:-1]) / 2. 
    psi_prime = np.diff(psi).conj() / dx
    return -(psi_mean * psi_prime).imag

def velocity(psi):
    return sum(prob_current(psi) * dx)

t = 2.5
xs_mid = (xs[1:]+xs[:-1])/2

plt.ylim(bottom=0)
plt.plot(xs, np.abs(wavepacket(xs, t))**2, color="blue", lw=2, label="$|\Psi(x)|^2$")
plt.plot(xs_mid, prob_current(wavepacket(xs, t)), color="green", label="$j(x)$")
#plt.plot(xs_mid, prob_current(wavepacket(xs, t)) / np.abs(wavepacket(xs_mid, t))**2, color="red", label="$v(x)$")
plt.xlabel('$x$')
plt.grid()
plt.title('v = %f (should be %f)' % (velocity(wavepacket(xs, t)), k))
plt.legend()
plt.show()

Spectral time evolution:

class SpectralEvolution:
    def __init__(self, H, initial_psi):
        self.H = H.todense()
        self.energies, self.eigenstates = scipy.linalg.eigh(self.H)
        self.initial_psi = np.dot(self.eigenstates.T.conj(), initial_psi(xs))   # to get a unit vector, we should multiply by sqrt(dx)
    
    def __call__(self, dt, t):
        psi = np.exp(-1j*self.energies*t) * self.initial_psi
        return np.dot(self.eigenstates, psi)

display_video(SpectralEvolution(H_free, wavepacket))

– +

Once Loop Reflect

Forward Euler scheme:

class EulerEvolution:
    def __init__(self, H, initial_psi):
        self.H = H # np.array(H.todense())
        self.current_psi = initial_psi(xs).reshape(-1, 1)
    
    def __call__(self, dt, t):
        A = scipy.sparse.identity(x_steps) - 1j*dt*self.H
        psi = A.dot(self.current_psi)
        self.current_psi = psi / np.linalg.norm(psi) * np.linalg.norm(self.current_psi)
        return self.current_psi

display_video(EulerEvolution(H_free, wavepacket))

– +

Once Loop Reflect

Forward/Backward stable, unitary scheme:

class StableEvolution:
    def __init__(self, H, initial_psi):
        self.H = H
        self.current_psi = initial_psi(xs)
    
    def __call__(self, dt, t):
        A = scipy.sparse.identity(x_steps) - 0.5j*dt*self.H
        tmp = A.dot(self.current_psi)
        self.current_psi = scipy.sparse.linalg.spsolve(A.conj(), tmp)
        return self.current_psi

display_video(StableEvolution(H_free, wavepacket))

– +

Once Loop Reflect

Split operator method (assumes periodic b.c., hence only valid as long as the wave function does not exceed x_min or x_max):

class SplitOperatorMethod:
    def __init__(self, V, initial_psi):
        self.V = V(xs)
        self.current_psi = initial_psi(xs)
        freqs = scipy.fftpack.fftfreq(x_steps)
        self.neg_laplace = (2*np.pi*(x_steps-1)/(x_max-x_min)*freqs)**2

    def __call__(self, dt, t):
        psi = np.exp(-0.5j*dt*self.V) * self.current_psi
        psi = scipy.fftpack.fft(psi)
        psi = np.exp(-1j*dt*self.neg_laplace/2) * psi
        psi = scipy.fftpack.ifft(psi)
        self.current_psi = np.exp(-0.5j*dt*self.V) * psi
        return self.current_psi

display_video(SplitOperatorMethod(V_free, wavepacket))

– +

Once Loop Reflect

2. Infinite Wall

The tridiagonal Hamiltonian for the kinetic part that we set up above automatically implements vanishing boundary conditions at $x_{\min}$ and $x_{\max}$. Thus we merely set $x_{\max} = 0$ (cf. discussion in the exercise class, also for the analytical solution):

x_min = -45
x_max = 0
x_steps = 1000
xs, dx = np.linspace(x_min, x_max, x_steps, retstep=True)

V_free = np.vectorize(lambda x: 0)
H_free = hamiltonian(V_free)

display_video(SpectralEvolution(H_free, wavepacket))

– +

Once Loop Reflect

3. Tilted Wall

Re-set simulation range:

x_min = -20
x_max = 50
x_steps = 500
xs, dx = np.linspace(x_min, x_max, x_steps, retstep=True)

# # fs13 scenario:
# x_min = -15
# x_max = 85
# x_steps = 1000
# xs, dx = np.linspace(x_min, x_max, x_steps, retstep=True)

# t_min = 0
# t_max = 15
# t_steps = 300
# ts, dt = np.linspace(t_min, t_max, t_steps, retstep=True)

# x_0 = -10
# a = 0.5
# k = 10.

Potential:

def V_wall(theta_deg):
    m = np.tan(theta_deg*2*np.pi/360)
    return np.vectorize(lambda x: max(.0, m*x))

def H_wall(theta_deg):
    return hamiltonian(V_wall(theta_deg))

Time evolution of the wave function and of the average velocity:

def velocities(theta_deg):
    evol = SpectralEvolution(H_wall(theta_deg), wavepacket)
    vs = [velocity(wavepacket(xs, 0))]
    for t in ts[1:]:
        psi = evol(dt, t)
        vs.append(velocity(psi))
    return vs

plt.xlabel("t")
plt.ylabel("v")
for theta_deg in [30, 45, 60]:
    vs = velocities(theta_deg)
    plt.plot(ts, vs, label="theta = %d" % theta_deg)
plt.legend()
plt.axhline(0, color="gray")
plt.show()

display_video(SpectralEvolution(H_wall(60), wavepacket), V_wall(60))

C:\Users\Michael\Anaconda\lib\site-packages\mpld3\mplexporter\exporter.py:82: UserWarning: Blended transforms not yet supported. Zoom behavior may not work as expected.
  warnings.warn("Blended transforms not yet supported. "

– +

Once Loop Reflect