Raketti Trajectory

import math
import numpy
import matplotlib.pyplot as plt

m = 0.25
g = 9.8
P = 15
T = 2
rho = 1.2
S = math.pi * 0.02**2
S_p = math.pi * 0.17**2
C = 0.5
C_p = 1.3
T_p = 8

def a1(v):
    return (P - m*g - numpy.sign(v) * C*rho*S*v**2/2) / m

def a2(v):
    return (-m*g- numpy.sign(v)* C*rho*S*v**2/2) / m

def a3(v):
    return (-m*g - numpy.sign(v)* C_p*rho*S_p*v**2/2) / m

h = [0]
v = [0]
t = 0
dt = 1e-3

while t < T:
    h.append(h[-1] + v[-1] * dt + a1(v[-1]) * dt ** 2 / 2)
    v.append(v[-1]+ a1(v[-1]) * dt)
    t = t + dt

while h[-1] > 0:
    if t < T_p:
        h.append(h[-1] + v[-1] * dt + a2(v[-1]) * dt ** 2 / 2)
        v.append(v[-1] + a2(v[-1]) * dt)
    else:
        h.append(h[-1] + v[-1] * dt + a3(v[-1]) * dt ** 2 / 2)
        v.append(v[-1] + a3(v[-1]) * dt)
    t = t + dt


plt.figure()
plt.plot([dt*i for i in range(len(h))], h)
plt.show()