[Python] Process simulation of the wave function of particle(electron) in a box
Assume there is an electron in a 1-dimension infinite potential well, and its initial state is the wave pocket whose center is at x0 = L/2:
In this expression, L = 10−8m, σ = 10−10m, k = 5×1010m−1. Set other physics constants as h = 1.0546 × 10−34J · s，Me = 9.109 × 10−31kg.
Then, set the time step as ht = 10−18s and the coordinate step as hx = L/N. N = 1000 refers to the length L being cut into 1000 equivalent portions.
Here is the code:
import numpy as np
import matplotlib.pyplot as plt
#set parameters
L = 10.0**(-8)
sgm = 10.0**(-10)
k = 5.0*10**10
hbar = 1.0546*10**(-34)
Me = 9.109*10**(-31)
x0 = L/2
N = 1000
ht = 10.0**(-18)
hx = L/N
a1 = 1 + (1j*hbar*ht)/(2*Me*hx**2)
a2 = -(1j*hbar*ht)/(4*Me*hx**2)
b1 = 1 - (1j*hbar*ht)/(2*Me*hx**2)
b2 = (1j*hbar*ht)/(4*Me*hx**2)
#define initial state function
x = np.arange(0,L+hx,hx)
t = np.arange(0,2000*ht+ht,ht)
psi0 = np.exp(-(x-x0)**2/(2*sgm**2))*np.exp(1j*k*x)
#the matrix of the wave functions at different t
phi = np.zeros((len(psi0), len(t)),dtype=complex)
phi[:,0]=psi0
#set transfer matrix
A = np.zeros((len(psi0), len(psi0)),dtype=complex)
B = np.zeros((len(psi0), len(psi0)),dtype=complex)
for n in range(len(psi0)):
if n==0:
A[n,0]=a1
A[n,1]=a2
B[n,0]=b1
B[n,1]=b2
elif n==len(psi0)-1:
A[n,n]=a1
A[n,n-1]=a2
B[n,n]=b1
B[n,n-1]=b2
else:
A[n,n-1]=a2
A[n,n]=a1
A[n,n+1]=a2
B[n,n-1]=b2
B[n,n]=b1
B[n,n+1]=b2
#compute the wave functions
for n in range(1,len(t)):
phi[:,n]=np.linalg.inv(A) @ B @ phi[:,n-1]
#drive curves
for m in [0, 300, 500, 800, 2000]:
prob = phi[:,m]*np.conjugate(phi[:,m]) #compute the probability distribution
lbl='t='+str(m)
plt.figure(1,figsize=(15,3),dpi=300)
plt.plot(x,np.real(phi[:,m]),label=lbl)
plt.xlabel('x')
plt.ylabel('real')
plt.legend()
plt.grid()
plt.figure(2,figsize=(15,3),dpi=300)
plt.plot(x,np.imag(phi[:,m]),label=lbl)
plt.xlabel('x')
plt.ylabel('image')
plt.legend()
plt.grid()
plt.figure(3,figsize=(15,3),dpi=300)
plt.plot(x,np.sqrt(prob),label=lbl)
plt.xlabel('x')
plt.ylabel('probability')
plt.legend()
plt.grid()
Output:
The running time is kinda long. Be patient.