quantum-dev/numerov-3.py

61 lines
1.9 KiB
Python
Raw Normal View History

2024-04-23 19:40:39 -04:00
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()