Time evolution of QM infinite square-well amplitudes & probs

#! /usr/bin/env python

from __future__ import print_function, division

## Parse args
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("COEFFS", help="list of coefficients (generally complex, e.g. 1+2j) for the inf-well eigenstates",
                nargs="*", type=complex, default=[1., 3., 0., 8., 2., 1.])
ap.add_argument("-o", "--out", dest="MPGOUT", help="write to this MPEG file rather than showing interactively", default=None)
args = ap.parse_args()


## Normalize coeffs and make separate u^+- coeff arrays
CS = args.COEFFS
#CS = [10., 0., 5., 0., 2.]
#CS = [0., 3., 0., 8., 0., 1.]
#CS = [1., 3., 0., 8., 2., 1.]
import numpy as np
Z = np.sqrt((np.abs(CS)**2).sum())
cpluses = np.array(CS[::2]).reshape((1,1,-1))
cminuses = np.array(CS[1::2]).reshape((1,1,-1))


## Make arrays of t and x values (could move to command-line args, ...)
A, NT, NX = 5, 500, 200
ts = np.linspace(0, 10, NT).reshape((-1,1,1))
xs = np.linspace(-A, A, NX).reshape((1,-1,1))
#print(ts,ks,sep="\n")


## Compute all terms
M = 1; HBAR = 1
npluses = np.array(range(1, cpluses.size+1)).reshape((1,1,-1))
nminuses = np.array(range(1, cminuses.size+1)).reshape((1,1,-1))
#print(npluses, npluses.shape)
#print(nminuses, nminuses.shape)
cupluses = cpluses * np.cos((2*npluses-1) * np.pi * xs / 2 / A) / np.sqrt(A) * \
           np.exp(1j / 2 / M * ((2*npluses-1) * np.pi * HBAR / 2 / A)**2 * ts / HBAR)
cuminuses = cminuses * np.sin((2*nminuses) * np.pi * xs / 2 / A) / np.sqrt(A) * \
            np.exp(1j / 2 / M * ((2*nminuses) * np.pi * HBAR / 2 / A)**2 * ts / HBAR)


## Sum over eigenstates
psis = cupluses.sum(axis=2) + cuminuses.sum(axis=2)
#print(psis.shape)


## Render the first frame
import matplotlib.pyplot as plt
import matplotlib.animation as animation
fig, ax = plt.subplots(figsize=(14,8))
line, = ax.plot(xs.reshape(-1), np.absolute(psis[0,:]), linewidth=2)
line_r, = ax.plot(xs.reshape(-1), np.real(psis[0,:]), linestyle="--", color="lightgray")
line_i, = ax.plot(xs.reshape(-1), np.imag(psis[0,:]), linestyle=":", color="lightgray")
ax.set_xlim(-A,A)
#ax.set_ylim(-8,12)


## Set up the animation callbacks for updating the lines
def init():  # only required for blitting to give a clean slate.
    line.set_ydata([np.nan] * psis.shape[1])
    line_r.set_ydata([np.nan] * psis.shape[1])
    line_i.set_ydata([np.nan] * psis.shape[1])
    return line, line_r, line_i
def animate(i):
    #print(np.absolute(psis[i,:]), np.real(psis[i,:]), np.imag(psis[i,:]), sep="\n")
    line.set_ydata(np.absolute(psis[i,:]))
    line_r.set_ydata(np.real(psis[i,:]))
    line_i.set_ydata(np.imag(psis[i,:]))
    return line, line_r, line_i
anim = animation.FuncAnimation(fig, animate, frames=NT, init_func=init, interval=20, blit=True, save_count=50)


## Display/save
if not args.MPGOUT:
    plt.show()
    #HTML(anim.to_html5_video())
else:
    Writer = animation.writers['ffmpeg']
    writer = Writer(fps=24, metadata=dict(artist='Andy Buckley'), bitrate=1800)
    anim.save(args.MPGOUT, writer=writer)