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):
    # define your wavefunction here
    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()]
    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,
    return ani

t = 0
x = np.linspace(0,1,100)
fig, ax = plt.subplots()
ax.set_ylim([0, 3.5])

ani = animate(x,t)

# save animation
# Set up formatting for the movie files
Writer = animation.writers['ffmpeg']
writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)'im.mp4', writer=writer)

Note that this code saves the animation to an mp4 file. I converted the video to gif using

Have fun!

[Featured image is an analog oscilloscope I used in one of my research visits to the US. This is one of the old instruments available in the lab. These oldies rock!]


Leave a Reply

Fill in your details below or click an icon to log in: Logo

You are commenting using your account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s