coupled_sho_qm.py

"""
Simulating coupled 1D Quantum Harmonic Oscillators.

This is based on the following article:

Scott C. Johnson and Thomas D. Gutierrez, 
"Visualizing the phonon wave function",
American Journal of Physics 70, 227-237 (2002) 
https://doi.org/10.1119/1.1446858 

"""
from time import perf_counter
import numpy as np
import numpy.linalg as linalg
from scipy.special import gamma
from scipy.special import hermite
import matplotlib.pyplot as plt
import matplotlib.animation as animation


def sho_eigenstate(n, x, m=1.0, omega=1.0, hbar=1.0):
    """
    The analytical eigenstates and energy eigenvalues
    for the 1D harmonic oscillator can be found here:
    https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator
    #One-dimensional_harmonic_oscillator
    """
    return (1.0/np.sqrt(2**n*gamma(n+1))*
            (m*omega/(np.pi*hbar))**0.25*
            np.exp(-m*omega*x**2/(2.0*hbar))*
            hermite(n)(np.sqrt(m*omega/hbar)*x)
           )


K = 1.0
M = 1.0
N = 20
eq_matrix = np.zeros([N, N])
for i in range(N):
    eq_matrix[i, i] = 2.0*K/M
    if i-1 >= 0: eq_matrix[i, i-1] = -K/M
    if i+1 < N: eq_matrix[i, i+1] = -K/M
eq_matrix[0, N-1] = -K/M
eq_matrix[N-1, 0] = -K/M
# plt.imshow(eq_matrix)
# plt.show()
# plt.close()
e, v = linalg.eigh(eq_matrix)
e[0] = 0.0

L = 20.0
S = np.linspace(-L/2.0, L/2.0, 200)
X1, X2 = np.meshgrid(S, S)
coords = [X1, X2]
ns = [np.array([1.0]) for _ in range(N)]
n_basis = []
n_evals = []
for i in range(len(e)):
    omega = np.sqrt(e[i])
    n_basis.append(np.array([sho_eigenstate(n, S, omega=omega)
                             for n in range(60)]))
    n_evals.append(np.array([omega*(n) for n in range(60)]))

es = [np.exp(2.0j*5.0*np.pi*S/L), np.exp(-2.0j*5.0*np.pi*S/L)]
init_func = np.exp(-0.5*(S/L + 0.15)**2/0.06**2)
ns_2d = [3, 4]
for j, i in enumerate(ns_2d):
    basis = n_basis[i]
    coeffs = np.dot(basis, init_func*es[j])
    ns[i] = coeffs

states = []
for i, coeffs in enumerate(ns):
    states.append(np.dot(coeffs, n_basis[i][0:len(coeffs)]))
    n_factor = np.sqrt(np.sum(np.abs(states[i])**2))
    if n_factor != 0.0:
        ns[i] *= 1.0/n_factor

psi = np.outer(states[ns_2d[0]], states[ns_2d[1]])
psi *= 1.0/np.sqrt(np.sum(np.abs(psi)**2))
# psi = sho_eigenstate(1, sum([v[i]*coords[0] 
#                              for i in ns_2d]), 
#                              omega=np.sqrt(e[0]))
# psi *= sho_eigenstate(2, sum([v[i]*coords[1]
#                               for i in ns_2d]), 
#                               omega=np.sqrt(e[1]))
xn_exp = np.array([np.dot(S, np.abs(states[i])**2)
                   for i in range(len(e))])
# x2n_exp = np.array([np.dot(S**2, np.abs(states[i])**2)
#                     for i in range(len(e))])
# sigma2 = x2n_exp - xn_exp**2
# sigma = np.sqrt(np.abs(sigma2))
# xn_exp = xn_exp/np.sqrt(np.sum(xn_exp*xn_exp))

# for i, e_vals in enumerate(n_evals):
#     plt.scatter([i for _ in range(len(e_vals))], e_vals,
#                 marker='_')
# plt.title('Energy Levels')
# plt.xlabel('$q_i$')
# plt.ylabel('E')
# plt.xticks([i for i in range(len(n_evals)) if i%2 == 0])
# plt.show()
# plt.close()

# for i in range(3):
#     plt.plot(v.T[i], label='$q_{%d}$' % i)
# plt.title('Coupled Harmonic Oscillator Stationary Modes')
# plt.xlabel('$x_{i}$')
# plt.legend()
# plt.show()
# plt.close()

fig = plt.figure()
axes = fig.subplots(1, 2)
inv_v = np.linalg.inv(v).T
line, = axes[0].plot(np.arange(0, len(e), 1), np.dot(inv_v, xn_exp))
axes[0].set_title('<$x_i$>')
axes[0].set_xlabel('$i$')
axes[0].set_ylabel('$x_i$')
# axes[0].set_yticks([])
axes[0].set_ylim(S[0]/2.0, S[-1]/2.0)
# axes[0].set_xticks(np.arange(0, len(e), 1))
axes[0].grid()
axes[1].set_title('$q_{%d}q_{%d}$ plane' % (ns_2d[0], ns_2d[1]))
axes[1].set_xlabel('q')
im2 = axes[1].imshow(np.real(psi*0), cmap='gray', 
                     extent=[S[0], S[-1], S[0], S[-1]])
max_val = np.amax(np.abs(psi))
im = axes[1].imshow(np.angle(X1 + 1.0j*X2), 
                    alpha=np.abs(psi)/max_val, 
                    origin='lower',
                    interpolation='nearest', cmap='hsv',
                    extent=[S[0], S[-1], S[0], S[-1]])


t0 = perf_counter()
data = {'t': 0.0, 'frames': 0}
def animation_func(*arg):
    states = []
    data['t'] += 0.1
    t = data['t']
    for i, coeffs in enumerate(ns):
        e_val = n_evals[i]
        states.append(np.dot(coeffs*np.exp(-1.0j*e_val[0:len(coeffs)]*t), 
                             n_basis[i][0:len(coeffs)]))
    xn_exp = np.array([np.dot(np.conj(states[i]), 
                   S*states[i]) for i in range(len(e))])
    psi = np.outer(states[ns_2d[0]], states[ns_2d[1]])
    psi *= 1.0/np.sqrt(np.sum(np.abs(psi)**2))
    line.set_ydata(np.dot(inv_v, np.real(xn_exp)))
    # phase = np.prod(np.array([states[j+1][int(200*(x + L/2.0)/L)]
    #                            for j, x in enumerate(xn_exp[1:])]))
    # print(phase)
    # col = complex_to_colour(np.array([phase])/np.abs(phase))
    # col = [min(c[0], 1.0) for c in col]
    # line.set_color(np.abs(col))
    # alpha = np.abs(np.prod(np.array(
    #                         [states[j+1][int(200*(x + L/2.0)/L)]
    #                         for j, x in enumerate(xn_exp[1:])])))
    # line.set_alpha(np.real(alpha*N*N))
    im.set_alpha(np.abs(psi)/max_val)
    im.set_data(np.angle(psi))
    data['frames'] += 1
    return line, im2, im

ani = animation.FuncAnimation(fig, animation_func, blit=True, interval=1.0)
plt.show()
fps = 1.0/((perf_counter() - t0)/data['frames'])
print('fps: %d' % int(fps))