66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
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
|
|
|
|
fig, ax = plt.subplots()
|
|
line, = ax.plot([], [], label=f'epsilon = {epsilon}')
|
|
|
|
ax.set_xlim(-1, 1)
|
|
ax.set_ylim(-1, 1)
|
|
ax.set_xlabel("x")
|
|
ax.set_ylabel("y")
|
|
ax.set_title("Schrodinger Eqn in Harmonic Potential")
|
|
ax.legend(loc=1)
|
|
|
|
|
|
def init():
|
|
line.set_data([], [])
|
|
return line,
|
|
|
|
|
|
def update(frame):
|
|
global pos, potential, psi, potential_past_2, potential_past_1, psi_past_2, psi_past_1
|
|
for i in range(frame * 1000, (frame * 1000) + 1000):
|
|
print(i)
|
|
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
|
|
|
|
line.set_data(x_out, y_out)
|
|
return line,
|
|
|
|
|
|
ani = animation.FuncAnimation(fig, update, frames=np.arange(0, (iterations - 2)), init_func=init, blit=True)
|
|
plt.show()
|