drag.py

import numpy
import matplotlib.pyplot


x = []
y = []


gravitational_acceleration = 9.81 # m/s^2
time_step = 0.1 # s
air_density_at_sea_level = 1.22
frontal_area = 0.002
coefficient_of_drag = 0.25

mass = 1 # kg

class flight_path:

    def __init__(self):
        self.time = []
        self.altitude = []
        self.velocity = []
        self.drag = []

    def add_rocket_position(self, time, altitude, velocity, drag):
        self.time.append(time)
        self.altitude.append(altitude)
        self.velocity.append(velocity)
        self.drag.append(drag)



def calculate():

    motor_thrust = [50,50,50,50,50,50,50,50,50,50,50,50,50,50,50,50,50,50,50,50]

    current_altitude = 0
    current_time = 0
    current_velocity = 0
    current_drag = 0

    current_flightpath = flight_path()

    counter = 0

    while current_altitude > 0 or counter == 0:
        current_time = current_time + time_step
        if counter < len(motor_thrust):
            current_velocity = current_velocity + ((motor_thrust[counter] / mass) * time_step)
        
        current_drag = (air_density_at_sea_level / 2) * (current_velocity ** 2) * coefficient_of_drag * frontal_area
        current_velocity = current_velocity + (time_step * (-gravitational_acceleration + (current_drag / mass)))
        current_altitude = current_altitude + (current_velocity * time_step)

        if current_altitude > 0:
            print current_time, current_altitude, current_velocity, current_drag
            current_flightpath.add_rocket_position(current_time, current_altitude, current_velocity, current_drag)

        

        counter += 1

    return current_flightpath


def plot(flightpath):
    time = numpy.asarray(flightpath.time)
    altitude = numpy.asarray(flightpath.altitude)
    axes_height = matplotlib.pyplot.subplot(211)
    matplotlib.pyplot.plot(time, altitude)

    axes_height.set_ylabel('Height in m')
    axes_height.set_xlabel('Time in s')

    matplotlib.pyplot.show()




flightpath = calculate()
plot(flightpath)