simulate and animate a rocket launch to orbit in python

main.py
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
98 chars
4 lines

main.py
G = 6.67430e-11  # gravitational constant
M = 5.97e24  # mass of the earth
R = 6.38e6  # radius of the earth
h0 = R + 500e3  # initial altitude of the rocket
v0 = 0  # initial velocity of the rocket
t0 = 0  # initial time
dt = 1  # time step
242 chars
8 lines

main.py
def acceleration(r):
    a = -G*M*r / np.linalg.norm(r)**3
    return a

def velocity(v, a, dt):
    v = v + a*dt
    return v

def position(r, v, dt):
    r = r + v*dt
    return r
182 chars
12 lines

main.py
fig, ax = plt.subplots()
25 chars
2 lines

main.py
def update(frame):
    # calculate acceleration, velocity, and position at each time step
    a = acceleration(r)
    v = velocity(v, a, dt)
    r = position(r, v, dt)
    
    # update rocket position on plot
    line.set_data(r[0], r[1])
    
    return line,
262 chars
11 lines

main.py
r = np.array([h0, 0])
v = np.array([v0, np.sqrt(G*M/h0)])  # circular orbit velocity
    
line, = ax.plot(r[0], r[1], 'o', color='red')
136 chars
5 lines

main.py
ax.set_xlim(-2*R, 2*R)
ax.set_ylim(-2*R, 2*R)
ax.set_aspect('equal')
ax.set_title('Rocket Launch to Orbit')

ani = FuncAnimation(fig, update, frames=np.arange(0,1000), interval=50, blit=True)
plt.show()
203 chars
8 lines