Simulating the re-entry of China's rocket body with Newton's Second Law

reentry.py

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def air_density(altitude):
    """
    Air density altitude profile.
    :param altitude: (m)
    :return: density (kg/m3)
    """
    # https://en.wikipedia.org/wiki/Barometric_formula#Density_equations
    return 1.225 * np.exp(-altitude * 1.189e-4)


def drag_force(altitude, speed, sphere_radius):
    """
    Drag force exerted on a sphere with drag_coefficient=0.1
    :param altitude: (m)
    :param speed: (m/s)
    :param sphere_radius: (m)
    :return: drag_force (N)
    """
    # https://en.wikipedia.org/wiki/Drag_(physics)
    return 0.157 * air_density(altitude) * (speed ** 2) * (sphere_radius ** 2)


def gravity_force(mass, distance):
    """
    Gravity exerted on an object around the Earth.
    :param mass: (kg)
    :param distance: (m)
    :return: gravity (N)
    """
    # https://en.wikipedia.org/wiki/Newton%27s_law_of_universal_gravitation#Modern_form
    return mu * mass / (distance ** 2)


# Initial Settings
T0 = pd.to_datetime('2021-05-06T12')  # Start date
DT = 1.  # Time resolution (s)
R = 6.371e6  # Earth radius (m)
mu = 3.986e14  # Earth gravity parameter = GM (m3/s2)

# Body parameters
r = 10.  # Object radius (m)
m = 18000.  # Object mass (kg)
ra = 267. * 1000.  # Apogee altitude(m)
rp = 156. * 1000.  # Perigee altitude (m)

va = np.sqrt(2 * mu * (R + rp) / ((R + ra) * (2 * R + ra + rp)))  # Velocity at apogee http://www.braeunig.us/space/
position = np.array([R + ra, 0.])
velocity = np.array([0., va])

log = {k: list() for k in ['t', 'x', 'y', 'V_x', 'V_y', 'g_force']}
dist = np.linalg.norm(position)
vel_norm = np.linalg.norm(velocity)
t = 0.
while t < 7. * 24. * 3600. and dist > R:
    # Calculate forces
    drag_norm = drag_force(altitude=dist - R, speed=vel_norm, sphere_radius=r)
    drag = drag_norm * (-velocity / vel_norm)

    gravity_norm = gravity_force(mass=m, distance=dist)
    gravity = gravity_norm * (-position / dist)

    # Newton's Second Law
    acceleration = (drag + gravity) / m
    velocity += acceleration * DT
    position += velocity * DT
    t += DT

    vel_norm = np.linalg.norm(velocity)
    dist = np.linalg.norm(position)

    log['t'].append(t)
    log['x'].append(position[0])
    log['y'].append(position[1])
    log['V_x'].append(velocity[0])
    log['V_y'].append(velocity[1])
    log['g_force'].append(drag_norm / (m * 9.807))

df = pd.DataFrame(data=log)
df.t = T0 + pd.to_timedelta(df.t, 's')
print("Impact date : {}".format(df.t.iloc[-1]))
df['distance'] = np.sqrt((df.x ** 2 + df.y ** 2))
df['velocity'] = np.sqrt((df.V_x ** 2 + df.V_y ** 2))
df['longitude'] = np.arctan2(df.y, df.x)
df['direction'] = np.arctan2(df.V_y, df.V_x)
df['sin_rv'] = np.sin(df.direction - df.longitude)

# Compute Apogee and Perigee http://www.braeunig.us/space/
D = (2 * mu) ** 2 - (2. * df.distance * df.sin_rv * df.velocity) ** 2 * (2 * mu / df.distance - df.velocity ** 2)
df['apogee'] = (2 * mu + np.sqrt(D)) / (2. * (2 * mu / df.distance - df.velocity ** 2))
df['perigee'] = (2 * mu - np.sqrt(D)) / (2. * (2 * mu / df.distance - df.velocity ** 2))

# https://en.wikipedia.org/wiki/Specific_orbital_energy
df['energy'] = -mu / (df.apogee + df.perigee)

# Convert to altitude in km
df['altitude'] = (df.distance - R) / 1000.
df.perigee = (df.perigee - R) / 1000.
df.apogee = (df.apogee - R) / 1000.

# Remove orbital data when suborbital
suborbital = df.perigee <= 0.
df.loc[suborbital, ['perigee', 'apogee', 'energy']] = float('nan')

# Remove past data
df = df.loc[df.t > pd.to_datetime('2021-05-07T14')]

# PLOTTING
for zoom_end in [False, True]:
    # Trajectory
    _, ax_pos = plt.subplots(1, figsize=(10, 10))
    ax_pos.set_title('Trajectory')
    ax_pos.add_patch(plt.Circle((0., 0.), R, color='xkcd:black'))
    ax_pos.plot(df.x, df.y)
    ax_pos.axis('equal')

    if zoom_end:
        zoom = 3e5
        ax_pos.set_xlim(df.x.iloc[-1] - zoom, df.x.iloc[-1] + zoom)
        ax_pos.set_ylim(df.y.iloc[-1] - zoom, df.y.iloc[-1] + zoom)
        plt.savefig('traj_zoomed.png')
    else:
        plt.savefig('traj.png')
        pass

    # noinspection PyTypeChecker
    # Flight stats
    _, (ax_alt, ax_vel, ax_force, ax_energy) = plt.subplots(4, 1, sharex=True, figsize=(5, 15))

    ax_alt.set_title('Altitude (km)')
    ax_alt.xaxis.set_tick_params(labelbottom=True)
    ax_alt.plot(df.t, df.altitude)
    ax_alt.plot(df.t, df.apogee, 'g', alpha=0.1)
    ax_alt.plot(df.t, df.perigee, 'g', alpha=0.1)

    ax_vel.set_title('Velocity (m/s)')
    ax_vel.xaxis.set_tick_params(labelbottom=True)
    ax_vel.plot(df.t, df.velocity)

    ax_force.set_title('g-force')
    ax_force.xaxis.set_tick_params(labelbottom=True)
    ax_force.plot(df.t, df.g_force)

    ax_energy.set_title('Orbital energy (m2/s2)')
    ax_energy.xaxis.set_tick_params(labelbottom=True)
    ax_energy.plot(df.t, df.energy)

    if zoom_end:
        ax_energy.set_xlim(df.t.iloc[-1] - pd.to_timedelta(10, 'm'), df.t.iloc[-1] + pd.to_timedelta(5, 'm'))
        plt.savefig('data_zoomed.png')
    else:
        plt.savefig('data.png')
        pass

# plt.show()