Skip to content

Instantly share code, notes, and snippets.

@Sonictherocketman
Created March 10, 2023 08:59
Show Gist options
  • Save Sonictherocketman/2b3836cb2bc2ca7fdb5283deb67d20c6 to your computer and use it in GitHub Desktop.
Save Sonictherocketman/2b3836cb2bc2ca7fdb5283deb67d20c6 to your computer and use it in GitHub Desktop.
A simple rocket simulation in Python.
#! /usr/bin/env python3
""" rocket.py -- Draw simulated rocket launches with turtle graphics
Sample Parameters:
via https://en.wikipedia.org/wiki/Falcon_9
$ python3.10 rocket.py -i 282 -a 90 -r 60 -p 950 -f 3500 -b 162
author: Brian Schrader
"""
import argparse
from base64 import b16encode
from datetime import datetime, timedelta
import decimal
import io
import logging
import logging.config
import math
import turtle
from tkinter import ALL, EventType
from random import randrange
import sys
logging.config.dictConfig({
'version': 1,
'handlers': {
'console': {
'class': 'logging.StreamHandler',
'level': 'INFO',
},
},
'root': {
'level': 'INFO',
'handlers': ['console']
},
})
logger = logging.getLogger(__name__)
screen = canvas = None
GRAVITY = 9.81
class Rocket:
max_altitude = None
max_distance = None
launched_at = None
current_time = None
burn_seconds = 0
tick_increment = timedelta(seconds=1)
Isp = 0
burn_rate = 0
payload_mass = 0
fuel_mass = 0
altitude = 0
distance = 0
angle = 0
ticks = 0
v_x = v_y = 0
def __init__(self, **kwargs):
self.__dict__.update(**kwargs)
def _increment_timer(self):
self.current_time += self.tick_increment
self.flight_seconds = (self.current_time - self.launched_at).seconds
def _update_color(self):
if self.is_powered:
turtle.pencolor('green')
elif self.angle > 0:
turtle.pencolor('blue')
else:
turtle.pencolor('red')
def _move(self):
tick_increment = float(self.tick_increment.seconds)
if self.is_powered:
F = self.Isp * self.burn_rate * GRAVITY
else:
F = 0
Ft_x = F * math.cos(self.angle_rad)
Ft_y = F * math.sin(self.angle_rad)
mass = self.payload_mass + self.fuel_mass
at_x = Ft_x / mass
at_y = Ft_y / mass
aw_x = 0
aw_y = GRAVITY
ad_x = 0 # TODO: Add drag calculations
ad_y = 0
self.a_x = at_x - aw_x - ad_x
self.a_y = at_y - aw_y - ad_y
self.dv_x = self.a_x * tick_increment
self.dv_y = self.a_y * tick_increment
self.v_x += self.dv_x
self.v_y += self.dv_y
self.d_x = self.v_x * tick_increment
self.d_y = self.v_y * tick_increment
if self.d_y == 0:
self.angle = math.degrees(90)
elif self.d_x == 0:
self.angle = math.degrees(0)
else:
self.angle = math.degrees(math.atan(self.d_y / self.d_x))
distance = math.sqrt(math.pow(self.d_x, 2) + math.pow(self.d_y, 2))
self._update_color()
turtle.setheading(self.angle)
turtle.forward(distance)
if int(self.d_x) == 0:
# Override the d_x so we can see the change on the chart.
x, y = self.position()
turtle.setposition(x+100, y)
def launch(self, angle=85):
self.current_time = self.launched_at = datetime.utcnow()
self.starting_position = self.position()
self.landed = False
self.ticks = 0
self.angle = angle
self.burn_rate = self.fuel_mass / self.burn_seconds
turtle.setheading(angle)
def tick(self):
self.ticks += 1
_, y_initial = self.starting_position
x, y = self.position()
if y < y_initial:
self.landed = True
else:
self._increment_timer()
self._move()
self.fuel_mass -= self.burn_rate * self.tick_increment.seconds
self.fuel_mass = max(self.fuel_mass, 0)
return not self.landed
def position(self):
return turtle.position()
@property
def angle_rad(self):
return math.radians(self.angle)
@property
def is_powered(self):
return self.fuel_mass > 0
def setup_turtle():
global screen
global canvas
screen = turtle.getscreen()
canvas = screen.getcanvas()
turtle.color('black', 'yellow')
turtle.speed('fastest')
turtle.resizemode('user')
turtle.mode('world')
turtle.tracer(0, 0)
def recenter_screen(l, r, t, b, padding=100):
turtle.update()
screen.setworldcoordinates(l-padding, b-padding, r+padding, t+padding)
def make_interactive(l, r, t, b, padding=100):
turtle.penup()
def do_zoom(event):
x = canvas.canvasx(event.x)
y = canvas.canvasy(event.y)
factor = 1.01 ** event.delta
canvas.scale(ALL, x, y, factor, factor)
canvas.bind("<MouseWheel>", do_zoom)
canvas.bind('<ButtonPress-1>', lambda event: canvas.scan_mark(event.x, event.y))
canvas.bind("<B1-Motion>", lambda event: canvas.scan_dragto(event.x, event.y, gain=1))
width, height = int(abs(l) + abs(r)), int(abs(t) + abs(b))
centerx, centery = int((l + r) / 2), int((t + b) / 2)
turtle.setpos(centerx, centery)
def take_picture(fname, l, r, t, b):
turtle.update()
width, height = int(abs(l) + abs(r)), int(abs(t) + abs(b))
canvas.postscript(file=fname)
def shutdown_turtle(close=False):
if close:
turtle.bye()
else:
turtle.done()
def parse_args():
parser = argparse.ArgumentParser(
description='Simulated rocket launches.'
)
parser.add_argument(
'-o', '--output',
type=str,
help=(
'Where to save the resultant image. If no value is provided, then no '
'image is saved once the process completes. It will still be displayed. '
'pdraw generates .eps files which can be opened in PDF viewing apps.'
),
)
parser.add_argument(
'-a', '--angle',
type=int,
default=89,
help=(
'The launch angle. This is measured in degrees '
'up to 180.'
),
)
parser.add_argument(
'-b', '--burn-time',
type=int,
default=3,
help=(
'The amount of time to burn.'
),
)
parser.add_argument(
'-i', '--impulse',
type=float,
default=10,
help=(
'The amount of thrust.'
),
)
parser.add_argument(
'-p', '--payload-mass',
type=float,
default=10,
help=(
'The amount of payload.'
),
)
parser.add_argument(
'-f', '--fuel-mass',
type=float,
default=10,
help=(
'The amount of fuel.'
),
)
parser.add_argument(
'-r', '--refresh-rate',
type=int,
default=100,
help=(
'The number of iterations to perform before refreshing the screen. '
'A value of 1 refreshes after each turn.'
),
)
parser.add_argument(
'-q', '--quiet',
action='store_false',
default=True,
dest='verbose',
help='Do not display progress indicator and status messages.',
)
parser.add_argument(
'-c', '--close',
action='store_true',
default=False,
help='Close the turtle window when drawing is finished.',
)
return parser.parse_args()
def main(args):
setup_turtle()
l, r, t, b = 0, 0, 0, 0
max_x = max_y = 0
rocket = Rocket(
burn_seconds=args.burn_time,
Isp=args.impulse,
payload_mass=args.payload_mass,
fuel_mass=args.fuel_mass,
)
try:
rocket.launch(args.angle)
while rocket.tick():
# Record the position
x, y = rocket.position()
l = min(x, l)
r = max(x, r)
t = max(y, t)
b = min(y, b)
# Square the coordinates (it looks better w/ a 1:1 ratio)
l = b = min(l, b)
r = t = max(r, t)
max_x = max(int(x), max_x)
max_y = max(int(y), max_y)
# Update UI
if rocket.ticks % args.refresh_rate == 0:
recenter_screen(l, r, t, b)
max_x_km, max_y_km = max_x // 1000, max_y // 1000
duration = rocket.flight_seconds
logger.info(
f'Flight Statistics: {max_x_km=}km, {max_y_km=}km, {duration=}s'
)
except ValueError:
logger.warning(
'Could not convert text to useful integers. Did you mean to use --encode?'
)
return
width, height = int(abs(l) + abs(r)), int(abs(t) + abs(b))
recenter_screen(l, r, t, b)
if args.verbose: logger.info('Finished drawing')
if args.output:
if args.verbose: logger.info('Saving drawing...')
take_picture(args.output, l, r, t, b)
if args.verbose: logger.info('Interactive mode enabled')
make_interactive(l, r, t, b)
shutdown_turtle(close=args.close)
if __name__ == '__main__':
main(parse_args())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment