Added omega (launch angle). same results as scipy ode package.

exp/RungeKutta.py

@@ -8,6 +8,7 @@ Created on Sat Jul 05 09:54:39 2014
 from __future__ import division
 import matplotlib.pyplot as plt
 from scipy.integrate import ode
+import math
 
 
 class BallEuler(object):
@@ -52,44 +53,65 @@ def rk4(y, x, dx, f):
     return y + (k1 + 2*k2 + 2*k3 + k4) / 6
 
 def fy(y, t):
-    """ returns velocity of ball at time t. 
+    """ returns velocity of ball at time t. """
     return -9.8*t
 
 
 def fx(y, t):
-    """ returns velocity of ball. Need to set vx.vel prior to first call
+    """ returns velocity of ball. Need to set vx.vel prior to first call"""
     return fx.vel
+    
+    
+    
+def fx2(x,t):
+    return fx2.vel
+    
+def fy2(y,t):
+    return fy2.vel - 9.8*t
         
 class BallRungeKutta(object):
-    def __init__(self, y=100., vel=10.):
-        self.x = 0.
+    def __init__(self, x=0, y=100., vel=10., omega = 0.0):
+        
+        self.omega = math.radians(omega)
+        self.vx = math.cos(self.omega) * vel
+        self.vy = math.sin(self.omega) * vel
+        self.x = x
         self.y = y
-        self.vel = vel
-        self.y_vel = 0.0
         self.t = 0
         
-        fx.vel = vel
+        fx2.vel = self.vx
+        fy2.vel = self.vy
 
 
-    def step2 (self, dt):
-        self.x = rk4 (self.x, self.t, dt, fx)
-        self.y = rk4 (self.y, self.t, dt, fy)
+    def step (self, dt):
+        self.x = rk4 (self.x, self.t, dt, fx2)
+        self.y = rk4 (self.y, self.t, dt, fy2)
         self.t += dt
 
         print self.x, self.y
 
 
+    
+
 if __name__ == "__main__":
 
     dt = 1./30
     y0 = 15.
     vel = 100.
+    omega = 10.
+    vel_y = math.sin(math.radians(omega)) * vel
+    print vel_y
+    
+    def f(t,y):
+        #print t,y
+        return vel_y-9.8*t
+        
     be = BallEuler (y=y0, vel=vel)
-    brk = BallRungeKutta (y=y0, vel=vel)
+    brk = BallRungeKutta (y=y0, vel=vel, omega=omega)
 
 
     solver = ode(f).set_integrator('dopri5')
-    solver.set_initial_value(y0, 0)
+    solver.set_initial_value(y0)
     t = 0
     y = y0
 
@@ -100,11 +122,12 @@ if __name__ == "__main__":
 
 
     while brk.y >= 0:
-        brk.step2 (dt)
-
-        y = solver.integrate(t+dt)
         t += dt
-        #print brk.x, y[0]
+        brk.step (dt)
+
+        y = solver.integrate(t)
+
 
         plt.scatter (brk.x, brk.y, color='blue', marker='v')
         plt.scatter (brk.x, y[0], color='green', marker='+')
+        plt.axis('equal')