65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
import pylab as pl
|
|
import numpy as np
|
|
import matplotlib.animation as animation
|
|
|
|
iterations = 60000 # iterations for approximation
|
|
|
|
step = 0.0001 # step size for x
|
|
step_sqrd = pow(step, 2) # square of step size
|
|
|
|
epsilon = 2.5 # 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
|
|
|
|
fig, ax = pl.subplots()
|
|
ax.set_xlim(-6, 6)
|
|
ax.set_ylim(-1, 1)
|
|
line, = ax.plot(x_out, y_out, label=f'epsilon = {epsilon}')
|
|
|
|
|
|
def update(frame):
|
|
global pos, potential_past_1, psi_past_1, potential_past_2, psi_past_2, ax
|
|
# Numerov integration loop
|
|
# count += 1
|
|
for i in range(0, 1000, 1):
|
|
pos += step
|
|
potential = 2 * epsilon - pow(pos, 2) # 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
|
|
|
|
ax.set_xlim(min(x_out), max(x_out))
|
|
ax.set_ylim(min(y_out), max(y_out))
|
|
line.set_data(x_out, y_out)
|
|
return line,
|
|
|
|
|
|
# Plot
|
|
ani = animation.FuncAnimation(fig, update, frames=np.arange((-1 * iterations + 2) // 1000, (iterations - 2) // 1000),
|
|
blit=True, repeat=False)
|
|
pl.show()
|