projectile motion with air resistance in python

1# -*- coding: utf-8 -*-
2
3import matplotlib.pyplot as plt
4import numpy as np
5import math
6import scipy.constants as const
7
8g = const.g     #gravitation constant
9dt = 1e-3       #integration time step (delta t)
10v0 = 40         #initial speed at t=0
11
12angle = math.pi / 4     #launch angle in radians
13time = np.arange(0,10, dt)      #create time axis
14
15gamm = 0.005    #gamma (used to compute f, below)
16h = 100         #height (used to compute f, below)
17
18def traj_fr(angle, v0):             #function that computes trajectory for some launch angle & velocity
19    vx0 = math.cos(angle)*v0        #compute x components of starting velocity
20    vy0 = math.sin(angle)*v0        #compute y components of starting velocity
21    x = np.zeros(len(time))         #initialise x array
22    y = np.zeros(len(time))         #initialise y array
23  
24    x[0],y[0] = 0,0     #initial position at t=0s, ie motion starts at (0,0)
25    x[1],y[1] = x[0] + vx0*(2*dt), y[0]+vy0*(2*dt)  #calculating 2nd elements of x & y based on init velocity
26  
27    i=1
28    while y[i]>=0:  #loop continuous until y becomes <0, ie projectile hits ground
29        f = 0.5 * gamm * (h - y[i]) * dt        #intermediate 'function'; used in calculating x & y vals below
30        x[i+1] = ((2*x[i]-x[i-1]) + (f * x[i-1])) / (1 + f)                 #numerical integration to find x[i+1]...
31        y[i+1] = ((2*y[i]-y[i-1]) + (f * y[i-1]) - g*(dt**2) ) / (1 + f)      # ...& y[i+1]
32        i = i+1     #increment i for next iteration
33  
34    x = x[0:i+1]        #truncate x array
35    y = y[0:i+1]        #truncate y array
36    return x, y, (dt*i), x[i]       #return x, y, flight time, range of projectile
37
38x,y,duration,distance = traj_fr(angle,v0)   #define variables for output of traj_fr function
39
40print 'Distance: ' , distance
41print 'Duration: ' , duration
42
43n = 5
44angles = np.linspace(0, math.pi/2, n)   #generate array of n angles
45print 'Angles: ' , angles
46maxrange = np.zeros(n)                  #generate array of n elements to take range from traj_fr
47
48
49for i in range(n):      #loop to run angles through traj_fr function & populate maxrange array with distance results
50    x,y,duration,maxrange[i] = traj_fr(angles[i], v0)
51
52angles = angles / 2 / math.pi * 360       #convert radians to degrees
53print 'Launch Angles: ', angles
54
55print 'Optimum Angle: ', angles[np.where(maxrange==np.max(maxrange))]
56
57plt.plot(x,y)       #quick plot of x vs y to check trajectory
58plt.xlabel('x')
59plt.ylabel('y')