import numpy as np from numpy import sin, cos, exp, sqrt, abs, pi import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation, PillowWriter # Time evolution of wave function in Infinite square well. # Steps: # 1 - decide on initial state # 2 - construct suitable CONB of eigenstates -- should be large enough # 3 - find coefficients c_n, check if basis was large enough # 4 - for all t values, construct linear combination with time-dependent weights def overlap(f, g): """ calculate the overlap between two complex functions using trapeze method """ return dx*(np.sum(f.conj()*g) - 0.5*(f[0]*g[0] + f[-1]*g[-1])) # units: use nm for length, eV for energy, hbar for action. length_unit, energy_unit, action_unit = 1e-9, 1.602e-19, 1.05e-34 mass_unit = energy_unit *time_unit**2 /length_unit**2 print("mass unit = %.3g kg"%mass_unit) time_unit = action_unit/energy_unit print("time unit = %.3g s"%time_unit) # electron mass: m = 9.11e-31/mass_unit # potential width: a = 200 x_resolution = 1000 x = np.linspace(0, a, x_resolution) dx = x[1]-x[0] # --- Step 1: initial state wavepacket_width = 4 x0 = a/4. k0 = 2. psi = exp(-(x-x0)**2/(4*wavepacket_width**2)) +0.j psi *= exp(1.j * k0*x) v = k0/m print("expected speed = %.3g"%v) x1 = a/2. psi += exp(-(x-x1)**2/(4*wavepacket_width**2)) +0.j NN = overlap(psi, psi) psi /= sqrt(NN) # --- Step 2: construct basis n_max = 10 psi_n, E_n = [], [] for n in range(n_max): k = (n+1)*np.pi/a psi_n.append(sqrt(2./a)*sin(k*x)+0.j) E_n.append(k*k/(2.*m)) # --- Step 3: Calculate c_n's, expand basis if needed c_n = [] sum_c2 = 0 for n in range(n_max): c_n.append(overlap(psi_n[n], psi)) sum_c2 += abs(c_n[-1])**2 while sum_c2<0.999: n_max += 1 k = (n_max)*np.pi/a psi_n.append(sqrt(2./a)*sin(k*x)+0.j) E_n.append(k*k/(2.*m)) c_n.append(overlap(psi_n[-1], psi)) sum_c2 += abs(c_n[-1])**2 print("n_max = %d sum_c2 = %.5g"%(n_max, sum_c2)) c_n = np.array(c_n) psi_n = np.array(psi_n) E_n = np.array(E_n) # --- Step 4: For all t, create the linear combinations psit = [] time_resolution = 200 T = 2*a/v #T = 1000 print("T = %.3g"%T + " time units = %.3g"%(T*time_unit*1e15)+ " fs") tt = np.linspace(0, T, time_resolution) for t in tt: #psi0 = np.zeros_like(psi) #for n in range(n_max): # psi0 += c_n[n]*exp(-1.j*E_n[n]*t) * psi_n[n] #psit.append(psi0) psit.append(np.sum(c_n*exp(-1.j*E_n*t)*psi_n.T, axis=1)) psit = np.array(psit) plt.imshow(abs(psit)**2, origin='lower', extent=(0, a, 0, T*time_unit*1e15), aspect="auto") plt.xlabel("position [nm]") plt.ylabel("time [fs]") plt.title("probability distribution") plt.colorbar() plt.show() num_of_frames = 150 fig, ax = plt.subplots(figsize=(14, 5)) def animate(i): print(".", end="") ax.clear() nnt = int(time_resolution*i/num_of_frames) ax.fill_between(x, abs(psit[nnt]), 0, color="lightgray") ax.plot(x, psit[nnt].real, color="darkblue") ax.plot(x, psit[nnt].imag, color="red") ax.set_ylim(-2*max(psi.real), 2*max(psi.real)) ax.set_xlim(-0.5, a+0.5) ax.set_title(fr'\$t = \$ {T*i/num_of_frames*time_unit*1e15:.2g} \$fs\$') ax.set_ylabel("wavefunction") ax.set_xlabel("position [nm]") plt.tight_layout() return # create FuncAnimation object and save the animation into gif ani = FuncAnimation(fig, animate, interval=10, repeat=True, frames=num_of_frames) ani.save("squarewell.gif", dpi=300, writer=PillowWriter(fps=24))