def intsch(psi, n, x, h): # integrate Sch eqn for n steps
psix = [psi[0]]
for i in range(n):
psi = ode.RK45n(sch, psi, x, h)
x += h
psix.append(psi[0])
return psix, psi # return WF and last point
def go():
xrk4, vrk4 = 5, -4
xrk45, vrk45 = 5, -4
t, h = 0.0, 0.01
tt = []
xrk4_array, vrk4_array = [], []
xrk45_array, vrk45_array = [], []
while t < 4:
xrk4_array.append(xrk4)
vrk4_array.append(vrk4)
xrk4, vrk4 = ode.rk4(electron_ode, [xrk4, vrk4], t, h)
xrk45_array.append(xrk45)
vrk45_array.append(vrk45)
xrk45, vrk45 = ode.RK45n(electron_ode, [xrk45, vrk45], t, h)
tt.append(t)
t += h
plt.figure(1)
plt.plot(tt, xrk4_array, '.', label='RK4')
plt.plot(tt, xrk45_array, '.', label='RK45')
plt.xlabel('Time (s)')
plt.ylabel('x (m)')
plt.legend(loc='best')
plt.show()
plt.figure(2)
plt.plot(tt, vrk4_array, '.', label='RK4')
plt.plot(tt, vrk45_array, '.', label='RK45')
plt.xlabel('Time (s)')
plt.ylabel('v (m/s)')
plt.legend(loc='best')
plt.show()
plt.figure(3)
plt.plot(xrk4_array, vrk4_array, '.', label='RK4')
plt.plot(xrk45_array, vrk45_array, '.', label='RK45')
plt.xlabel('x (m)')
plt.ylabel('v (m/s)')
plt.legend(loc='best')
plt.show()
accel += Q[i] * (Y[0] - loc[i]) / (vp.mag(Y[0] - loc[i]))**3
return [Y[1], q * accel] # list for non-vectorized solver
a, b, w = 1., 0.5, 0.125 # rink size, goal width
q, qcor, qmid, qcen = -1.0, 1.0, -2., 2. # Qs: puck, cor., mid, cen.
Q = [qcor, qmid, qcor, qcor, qmid, qcor, qcen] # charges, locations
loc = [[-a, b], [0, b], [a, b], [a, -b], [0, -b], [-a, -b], [0, 0]]
scene = vp.display(title='Electric hockey', background=(.2, .5, 1))
puck = vp.sphere(pos=(-a, 0, 0), radius=0.05, make_trail=True) # trail
rink = vp.curve(pos=loc[:-1] + [loc[0]], radius=0.01) # closed curve
goalL = vp.curve(pos=[(-a, w, 0), (-a, -w, 0)], color=(0, 1, 0), radius=.02)
goalR = vp.curve(pos=[(a, w, 0), (a, -w, 0)], color=(0, 1, 0), radius=.02)
for i in range(len(loc)): # charges, red if Q>0, blue if Q<0
color = (1, 0, 0) if Q[i] > 0 else (0, 0, 1)
vp.sphere(pos=loc[i], radius=0.05, color=color)
v, theta = input('enter speed, theta; eg, 2.2, 19:') # try 2.2, 18.5
v, theta = min(4, v), max(1, theta) * np.pi / 180. # check valid input
Y = np.array([[-a, 0], [v * np.cos(theta), v * np.sin(theta)]])
while True:
vp.rate(200)
Y = ode.RK45n(hockey, Y, t=0., h=0.002)
x, y = Y[0][0], Y[0][1]
if (abs(x) > a or abs(y) > b):
txt = 'Goal!' if (x > 0 and abs(y) < w) else 'Miss!'
vp.label(pos=(x, y + .2), text=txt, box=False)
break
puck.pos = Y[0]
t = 0
h = 0.5
y_euler, y_rk2, y_rk4, y_rk45n = 1000, 1000, 1000, 1000
tt, exsol = [], []
nuclear_euler, nuclear_rk2, nuclear_rk4, nuclear_rk45n = [], [], [], []
abserr_euler, abserr_rk2, abserr_rk4, abserr_rk45n = [], [], [], []
clf_euler, clf_rk2, clf_rk4, clf_rk45n = [], [], [], []
while t <= 8:
yy_euler = ode.Euler(decay, [y_euler], t, h)
yy_rk2 = ode.rk2(decay, [y_rk2], t, h)
yy_rk4 = ode.rk4(decay, [y_rk4], t, h)
yy_rk45n = ode.RK45n(decay, [y_rk45n], t, h)
exs = exact_sol(t)
tt.append(t)
exsol.append(exs)
nuclear_euler.append(yy_euler[0])
nuclear_rk2.append(yy_rk2[0])
nuclear_rk4.append(yy_rk4[0])
nuclear_rk45n.append(yy_rk45n[0])
abserr_euler.append(abs(yy_euler[0] - exs))
abserr_rk2.append(abs(yy_rk2[0] - exs))
abserr_rk4.append(abs(yy_rk4[0] - exs))
abserr_rk45n.append(abs(yy_rk45n[0] - exs))
clf_euler.append(t * (1000 - yy_euler[0]))
return [psi[1], 2 * (V(x) - E) * psi[0]]
# initialization and animation setup
a, V0 = 4.0, 4. # well width, depth
R, N = 4 * a, 200 # limit, intervals
xa = np.linspace(-R, R, 2 * N + 1) # grid
h, z = xa[1] - xa[0], np.zeros(2 * N + 1) # step size
E, dE, dpsi, psix = -V0, 0.001, 1.0, np.zeros(2 * N + 1)
scene = vp.display(background=(.2, .5, 1), range=1.5 * a)
wf = vpm.line(xa, psix, z, vp.color.red, .05)
pot = vpm.line(xa, .5 * np.vectorize(V)(xa), z, (1, 1, 1), .04) # pot. V
info = vp.label(pos=(0, -0.6 * a, 0), box=False, height=20)
while (E < 0.0):
psi, x = [.0, .1], -R
for i in range(N): # WF for x <=0
psi = ode.RK45n(sch, psi, x, h)
x += h
psix[i + 1] = psi[0]
psix[N + 1:] = psix[N - 1::-1] # WF for x > 0 by reflection
if (dpsi * psi[1] < 0.): # dpsi/dx changes sign
info.text = 'Energy found, E=%5.4f' % (E - dE / 2)
vpm.pause(scene) # any key to continue
else:
info.text = 'E=%5.3f' % (E)
wf.move(xa, 2 * psix / max(psix), z), vpm.wait(scene), vp.rate(2000)
dpsi = psi[1] # old dpsi/dx at E
E += dE