# Visualizing time-dependent wavefunctions

In my previous post, I presented a method of visualizing wavefunctions that are inherently complex-valued using a single plot that shows both the probability density and phase but frozen in time. Here, I complete this visualization by animating the plot.

The wavefunction shown above is that of a particle in a 1D box of length $L$, which is in equal superposition of the ground $\phi_1$ and the first excited state $\phi_2$, i.e. $\Psi(x,t) = \frac{1}{\sqrt{2}} \phi_1 e^{i\omega_1 t} + \frac{1}{\sqrt{2}}\phi_2e^{i\omega_2 t}$.

Two things to consider regarding this plotting method:

1. The phase is plotted first using imshow.
2. The probability density is plotted next as a “white” fill_between between the probability density and the top of the plot.
3. When we animate the whole figure, we basically reiterate the two plotting steps above.

Animating fill_between is tricky. Interestingly, the workaround is quite simple. Appending the pointer to the plot and animating it works like a charm.

Here’s the code:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from numpy import pi as pi
import matplotlib.animation as animation

def psi(x,t):
w1 = 1 # eigenfrequency
w2 = 2 # eigenfrequency
psi = np.sin(1*pi*x)*np.exp(1j*w1*t) + np.sin(2*pi*x)*np.exp(1j*w2*t) # wavefunction
return psi

def plotComplexFunction1D(x,psi):
psi = np.conj(psi*np.exp(1j*pi)) # rotate zero angle to x-axis
psipsi = np.abs(psi)**2 # probability density
y = np.linspace(psipsi.min(),psipsi.max(),100) # range of values for y
h = np.angle(np.conj(psi)) # takes the argument of the complex number
z = np.tile(h, (y.size, 1)) # creates 2D image with phase along x
# create the background colormap for the phase
xylims = [x.min(),x.max(),y.min(),y.max()]
a=ax.get_ylim()
imax=plt.imshow(z,cmap='hsv',extent=xylims,aspect='auto', animated=True)
imax.set_clim(vmin=-pi, vmax=pi)
# fill the region above the curve with white
pfill = plt.fill_between(x, psipsi, y2=max(a), color='w', animated=True) # plots the probability densit
return imax, pfill

def animate(x,t):
ims = []
for i in range(4*15):
t += np.pi / 15.
im, pfill = plotComplexFunction1D(x,psi(x,t))
ims.append([im, pfill])
ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True,
repeat_delay=0)
return ani

t = 0
x = np.linspace(0,1,100)
fig, ax = plt.subplots()
ax.set_ylim([0, 3.5])
plt.ylabel('$|\Psi(x)|^2$')
plt.xlabel('x')
ax.spines['top'].set_visible(False)
ax.set_yticks([])
plt.xticks([0,0.5,1],[0,'L/2','L'])

ani = animate(x,t)

# save animation
# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']