61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
import pylab as lab
|
|
import numpy as np
|
|
|
|
iterations = 60000 # iterations for approximation
|
|
|
|
step = 0.0001 # step size for x
|
|
step_sqrd = pow(step, 2) # square of step size
|
|
|
|
well_width = 2
|
|
|
|
# try values 1, 2, and 3 for various solutions
|
|
n = 1
|
|
epsilon = (n * np.pi) ** 2 / (well_width ** 2) # energy level, should be integer n+1/2 for good solutions
|
|
|
|
psi = 0.0 # initial value of wave function
|
|
potential = 0.0 # initial value of the potential energy function
|
|
pos = -1 * (iterations - 2) * step # initial value of the position
|
|
|
|
potential_past_2 = epsilon + pos - 2 * step # k_0, potential energy at two steps before current
|
|
potential_past_1 = epsilon + pos - step # k_1, potential energy at one step before current
|
|
amplitude = 0.1 # initial amplitude of wave function
|
|
psi_past_2 = 0 # y_0, wave function at two steps before current
|
|
psi_past_1 = amplitude # y_1, wave function at one step before current
|
|
|
|
x_out = [] # list to store x values for plotting
|
|
y_out = [] # list to store y values for plotting
|
|
|
|
count = -1 * iterations + 2 # counter for the loop
|
|
|
|
# Numerov integration loop
|
|
while count < iterations - 2:
|
|
count += 1
|
|
pos += step
|
|
# close enough to infinity :)
|
|
potential = 100000000
|
|
if abs(pos) < well_width / 2:
|
|
potential = epsilon # potential energy function
|
|
b = step_sqrd / 12 # constant used for Numerov
|
|
# Numerov method to calculate psi at current step
|
|
psi = ((2 * (1 - 5 * b * potential_past_1) * psi_past_1 - (1 + b * potential_past_2) * psi_past_2)
|
|
/ (1 + b * potential))
|
|
|
|
# Save for plotting
|
|
x_out.append(pos)
|
|
y_out.append(psi)
|
|
|
|
# Shift for next iteration
|
|
psi_past_2 = psi_past_1
|
|
psi_past_1 = psi
|
|
potential_past_2 = potential_past_1
|
|
potential_past_1 = potential
|
|
|
|
# Plot
|
|
lab.figure(1)
|
|
lab.plot(x_out, y_out, label=f'epsilon = {epsilon}')
|
|
lab.xlabel("x")
|
|
lab.ylabel("y")
|
|
lab.title("Schrodinger Eqn in Harmonic Potential")
|
|
lab.legend(loc=1)
|
|
lab.show()
|