rocket_predictor.py

import numpy
import matplotlib.pyplot


x = []
y = []


gravitational_acceleration = 9.81 # m/s^2
time_step = 0.1 # s

mass = 1 # kg

class flight_path:

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

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



def calculate():

    motor_thrust = [50 for x in range(10)]

    current_altitude = 0
    current_time = 0
    current_velocity = 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_velocity = current_velocity + (time_step * -gravitational_acceleration)
        current_altitude = current_altitude + (current_velocity * time_step)

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

        

        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)