A simple rocket simulation in Python.

rocket.py

#! /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())