yv = yv[idx] A = vstack([ones(size(ya)), log(abs(yv))]) [knot, kk2] = dot(inv(dot(A, A.T)), dot(A, log(ya + 9.8))) kk1 = exp(knot) def f4(t, y): return array([y[1], -9.8 + kk1 * (abs(y[1])) ** kk2]) drag_string = "Drag model:$F_d(v) = k_1 |v|^{{k_2}},\quad k_1 = {:.4f}, \quad k_2 = {:.4f}$".format(kk1, kk2) titles = ["Linear Drag", "Quadratic Drag", "Two Term Quadratic Drag", drag_string] titles.reverse() for f in (f1, f2, f3, f4): integrator = RungeKutta(f, [y[0], yy[0]], data[0][0], data[0][-1], 0.001) (mt, my) = integrator.integrate() idx = find(my[0] > 0) mt = mt[idx] my = my[:, idx] f = figure() subplot(121) plot(t, y, ".") plot(mt, my[0]) xlabel("Time (sec)") ylabel("Displacement (meters)") subplot(122) plot(tt, yy, ".") plot(mt, my[1]) xlabel("Time (sec)") ylabel("Velocity (meters)") f.text(0.5, 0.95, titles.pop(), horizontalalignment="center", verticalalignment="top")
a2 = -GM/r2**3*x2 + m1*GM/r21**3*(x1-x2) derivatives = array([v1,a1,v2,a2]).T.flatten(1) return derivatives scipy_result = odeint(two_body,xinit,times) scipy_result = scipy_result.T atoll=10**-13 rtoll=10**-14 high_prec_result = odeint(two_body,xinit,times,atol=atoll,rtol=rtoll) high_prec_result = high_prec_result.T ######## Do RK4 HERE ############ f = lambda t,x: two_body(x,t) rk4 = RungeKutta(f,xinit,0.,years,years/num_samples) (t,rk4_result) = rk4.integrate() rk4_result = rk4_result[:,:-2] ################################# iterates = ((scipy_result,"Scipy odeint default tolerances"), (rk4_result,"Runge Kutta 4th Order"), (high_prec_result,"odeint rtol={} atol={}".format(rtoll,atoll))) for x,t in iterates: figure() r1 = sqrt(x[0]**2 + x[1]**2) r2 = sqrt(x[4]**2 + x[5]**2) r21 = sqrt( (x[4]-x[0])**2 + (x[5]-x[1])**2) E = .5*(m1*(x[2]**2 + x[3]**2) + m2*(x[6]**2+x[7]**2)) - GM*( m1/r1 + m2/r2 + m1*m2/r21 ) deltaE = (E-E[0])/E[0]