1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

from math import pi, cos, sin, atan2, sqrt


def project_ballistic_trajectory(
    x=0,
    y=100,
    degree=30.0,
    init_vel=100,
    air_density=.1,
    gravity=9.81,
    mass=0.1,
    cs_area=0.1,
    drag_coef=0.09,
    start_time=0.,
    step=0.01,
):
    if step > 0.01:
        step = 0.01

    theta = pi / (180 / degree)
    vel_x = init_vel * cos(theta)
    vel_y = init_vel * sin(theta)
    distance_list = [x]
    height_list = [y]

    def drag(vel, theta):
        force = 0.5 * air_density * cs_area * drag_coef * (vel ** 2)
        return force * cos(theta), force * sin(theta)

    while y > 0 or vel_y > 0:
        force_x, force_y = drag(init_vel, theta)
        acc_x = -force_x / mass
        acc_y = -force_y / mass - gravity
        x = x + vel_x * step + 0.5 * acc_x * step * step
        y = y + vel_y * step + 0.5 * acc_y * step * step
        vel_x = vel_x + acc_x * step
        vel_y = vel_y + acc_y * step
        init_vel = sqrt(vel_x * vel_x + vel_y * vel_y)
        theta = atan2(vel_y, vel_x)

        distance_list.append(x)
        height_list.append(y)
        start_time = start_time + step

    ft = height_list[-2] / (height_list[-2] - height_list[-1])  # fractional time to last point
    distance_list[-1] = distance_list[-2] + (distance_list[-1] - distance_list[-2]) * ft
    height_list[-1] = 0.
    return [(i, height_list[c], 0) for c, i in enumerate(distance_list)]