tauplot = np.empty(nstep)
xplot, yplot, zplot = np.empty(nstep), np.empty(nstep), np.empty(nstep)
for istep in range(nstep):
#* Record values for plotting
x, y, z = state[0], state[1], state[2]
tplot[istep] = time
tauplot[istep] = tau
xplot[istep] = x
yplot[istep] = y
zplot[istep] = z
if (istep + 1) % 50 < 1:
print('Finished ', istep, ' steps out of ', nstep)
#* Find new state using adaptive Runge-Kutta
[state, time, tau] = rka(state, time, tau, err, lorzrk, param)
#* Print max and min time step returned by rka
tauMax = np.max(tauplot[1:nstep])
tauMin = np.min(tauplot[1:nstep])
print('Adaptive time step: Max = ', tauMax, ' Min = ', tauMin)
#* Graph the time series x(t)
plt.plot(tplot, xplot, '-')
plt.xlabel('Time')
plt.ylabel('x(t)')
plt.title('Lorenz model time series')
plt.show()
#* Graph the x,y,z phase space trajectory
# Mark the location of the three steady states
accel = -GM * r / np.linalg.norm(r)**3
r = r + tau * v # Euler step
v = v + tau * accel
time = time + tau
elif NumericalMethod == 2:
accel = -GM * r / np.linalg.norm(r)**3
v = v + tau * accel
r = r + tau * v # Euler-Cromer step
time = time + tau
elif NumericalMethod == 3:
state = rk4(state, time, tau, gravrk, GM)
r = np.array([state[0], state[1]]) # 4th order Runge-Kutta
v = np.array([state[2], state[3]])
time = time + tau
else:
[state, time, tau] = rka(state, time, tau, adaptErr, gravrk, GM)
r = np.array([state[0], state[1]]) # Adaptive Runge-Kutta
v = np.array([state[2], state[3]])
#* Graph the trajectory of the comet.
ax = plt.subplot(111, projection='polar') # Use polar plot for graphing orbit
ax.plot(thplot, rplot, '+')
ax.set_title('Distance (AU)')
ax.grid(True)
plt.show()
#* Graph the energy of the comet versus time.
totalE = kinetic + potential # Total energy
plt.plot(tplot, kinetic, '-.', tplot, potential, '--', tplot, totalE, '-')
plt.legend(['Kinetic', 'Potential', 'Total'])
plt.xlabel('Time (yr)')