From 9ebc639fb90a878e6b65800c689b89750b607b33 Mon Sep 17 00:00:00 2001 From: Baitinq Date: Thu, 10 Feb 2022 23:39:43 +0000 Subject: Structured source files into src folder --- atmosphere.py | 14 ---- body.py | 21 ------ engine.py | 14 ---- fuel.py | 3 - main.py | 212 ------------------------------------------------------ rocket.py | 43 ----------- simulation.py | 133 ---------------------------------- src/atmosphere.py | 14 ++++ src/body.py | 21 ++++++ src/engine.py | 14 ++++ src/fuel.py | 3 + src/main.py | 212 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/rocket.py | 43 +++++++++++ src/simulation.py | 133 ++++++++++++++++++++++++++++++++++ src/stage.py | 49 +++++++++++++ src/universe.py | 7 ++ stage.py | 49 ------------- universe.py | 7 -- 18 files changed, 496 insertions(+), 496 deletions(-) delete mode 100644 atmosphere.py delete mode 100644 body.py delete mode 100644 engine.py delete mode 100644 fuel.py delete mode 100644 main.py delete mode 100644 rocket.py delete mode 100644 simulation.py create mode 100644 src/atmosphere.py create mode 100644 src/body.py create mode 100644 src/engine.py create mode 100644 src/fuel.py create mode 100644 src/main.py create mode 100644 src/rocket.py create mode 100644 src/simulation.py create mode 100644 src/stage.py create mode 100644 src/universe.py delete mode 100644 stage.py delete mode 100644 universe.py diff --git a/atmosphere.py b/atmosphere.py deleted file mode 100644 index 48f9bd1..0000000 --- a/atmosphere.py +++ /dev/null @@ -1,14 +0,0 @@ -import math - -class Atmosphere(): - def __init__(self, avg_sea_level_pressure: int, molar_mass_air: float, standard_temp: float): - self.avg_sea_level_pressure = avg_sea_level_pressure - self.molar_mass_air = molar_mass_air - self.standard_temp = standard_temp - - #https://math24.net/barometric-formula.html - def density_at_height(self, height: int, g: float) -> None: - R = 8.3144598 #universal gas constant - pressure = self.avg_sea_level_pressure * math.e ** (-(self.molar_mass_air * g * height)/(R * self.standard_temp)) - density = pressure / (R * 10000) - return density \ No newline at end of file diff --git a/body.py b/body.py deleted file mode 100644 index 3be2c5e..0000000 --- a/body.py +++ /dev/null @@ -1,21 +0,0 @@ -import math - -from atmosphere import Atmosphere - -class Body(): - def __init__(self, name: str, density: int, radius: int, atmosphere: type[Atmosphere]): - self.name = name - self.density = density - self.radius = radius - self.atmosphere = atmosphere - - def mass(self): - body_volume = (4/3) * math.pi * (self.radius**3) - return body_volume * self.density * 1000 - - def g(self, G: float, height: int): - return (G * self.mass()) / ((self.radius + height) ** 2) - - - def __str__(self): - return "uwu" \ No newline at end of file diff --git a/engine.py b/engine.py deleted file mode 100644 index f1da646..0000000 --- a/engine.py +++ /dev/null @@ -1,14 +0,0 @@ -import fuel - -class Engine(): - def __init__(self, name: str, isp: int, max_flow_rate: int): - self.name = name - self.max_flow_rate = max_flow_rate - self.isp = isp - - def thrust(self, throttle: int, g: float): - #https://www.grc.nasa.gov/www/k-12/airplane/specimp.html - return self.flow_rate(throttle) * self.isp * g - - def flow_rate(self, throttle: int): - return self.max_flow_rate * (throttle / 100) \ No newline at end of file diff --git a/fuel.py b/fuel.py deleted file mode 100644 index caca9a9..0000000 --- a/fuel.py +++ /dev/null @@ -1,3 +0,0 @@ -class Fuel(): - def __init__(self, name: str): - self.name = name \ No newline at end of file diff --git a/main.py b/main.py deleted file mode 100644 index 0d70213..0000000 --- a/main.py +++ /dev/null @@ -1,212 +0,0 @@ -import sys -import math -from random import randint - -from engine import Engine -from fuel import Fuel -from stage import Stage -from rocket import Rocket -from atmosphere import Atmosphere -from body import Body -from universe import Universe -from simulation import Simulation - -import pygame -from pygame.locals import * - -def main(argv): - raptor_engine = Engine(name="raptor", isp=360, max_flow_rate=931) #https://en.wikipedia.org/wiki/SpaceX_Raptor - methane_fuel = Fuel(name="methane") #TODO: more - - #https://en.wikipedia.org/wiki/SpaceX_Starship - first_stage = Stage(name="superheavy booster", - stage_mass=180000, - engine=raptor_engine, - engine_number=33, - max_engine_gimbaling_angle=30, - fuel_type=methane_fuel, - fuel_mass=3600000, - drag_coefficient=1.18, - cross_sectional_area=(math.pi * (9**2)) - ) - - second_stage = Stage(name="starship", - stage_mass=80000, - engine=raptor_engine, - engine_number=6, - max_engine_gimbaling_angle=30, - fuel_type=methane_fuel, - fuel_mass=1200000, - drag_coefficient=1.18, - cross_sectional_area=(math.pi * (9**2)) - ) - - rocket = Rocket(name="starship launch system", - stages=[first_stage, second_stage], - payload_mass=100 - ) - - body = Body(name="earth", - density=5.51, - radius=6371000, - atmosphere=Atmosphere( - avg_sea_level_pressure=101325, - molar_mass_air=0.02896, - standard_temp=288.15 - ) - ) - - universe = Universe(name="conventional", - G=6.67E-11 - ) - - simulation = Simulation(universe, body, rocket) - simulation.rocket.current_stage().engines_on = True - - pygame.init() - pygame.display.set_caption("OSLS - Overly Simple Launch Simulator") - clock = pygame.time.Clock() - - SCREEN_WIDTH = 1024 - SCREEN_HEIGHT = 720 - - simulation_display = pygame.display.set_mode((SCREEN_WIDTH,SCREEN_HEIGHT)) - paused = False - while True: - if not paused: - draw_simulation(simulation_display, simulation) - pygame.display.update() - - for event in pygame.event.get(): - if event.type == pygame.QUIT: - pygame.quit() - quit() - elif event.type == pygame.KEYDOWN: - if event.key == pygame.K_q: - quit() - elif event.key == pygame.K_SPACE: - paused = not paused - else: - handle_key_press(simulation, event.key) - - delta = clock.tick(60) / 1000 #60fps #are we using delta in the simulation tick everywhere needed? - if not paused: #tick with pause messes up delta TODO: TODOODODODODODOOD TODO - print("delta: " + str(delta)) - simulation.tick(delta=delta) - - #TODO: IMPLEMENT rocket_x_drag_coefficient() that adds the x drag coefficient of all stages, same with cross sectional area - #TODO: draw body sprite, rocket sprite, clouds sprites, etc. - #TODO: implement height properly (body radius) + actually implement body - #TODO: do max load on rocket so it blows up - #TODO: allow multilanguage api for landing algorithms etc - -def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[Simulation]) -> None: - #draw background - def linear_gradient(start_color, end_color, length, value_at): - return [ - int(start_color[j] + (float(value_at)/(length-1))*(end_color[j]-start_color[j])) - for j in range(3) - ] - - def get_color_for_height(height: float) -> (int, int, int): - if height < 70000: - return linear_gradient((31,118,194), (0, 0, 0), 70000, int(height)) - else: - return (0, 0, 0) - - #gradient for atmosphere - simulation_display.fill(get_color_for_height(simulation.y)) - - #draw clouds and stars - #draw clouds (we need continuity TODO) - #if simulation.y < 20000 and randint(0, 100) < 5: - # pygame.draw.circle(simulation_display, (255, 255, 255), (randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), 30) - #draw stars - if simulation.y > 30000: - for _ in range(100): - simulation_display.set_at((randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), (255, 255, 255)) - #draw stats text - font = pygame.font.SysFont("Comic Sans MS", 30) - - curr_thrust = simulation.rocket.current_stage().current_thrust(simulation.body.g(simulation.universe.G, simulation.y), simulation.heading) - g = simulation.body.g(simulation.universe.G, simulation.y) - - simulation_display.blit(font.render("Simulation time: {:.0f}s".format(simulation.time), False, (255, 255, 255)),(0,0)) - simulation_display.blit(font.render("X: {:.0f}m".format(simulation.x), False, (255, 255, 255)),(0,40)) - simulation_display.blit(font.render("Y: {:.0f}m".format(simulation.y), False, (255, 255, 255)),(0,80)) - simulation_display.blit(font.render("Speed x: {:.0f}m/s".format(simulation.speed_x), False, (255, 255, 255)),(0,120)) - simulation_display.blit(font.render("Speed y: {:.0f}m/s".format(simulation.speed_y), False, (255, 255, 255)),(0,160)) - simulation_display.blit(font.render("Acceleration x: {:.2f}m/s2".format(simulation.acceleration_x), False, (255, 255, 255)),(0,200)) - simulation_display.blit(font.render("Acceleration y: {:.2f}m/s2".format(simulation.acceleration_y), False, (255, 255, 255)),(0,240)) - simulation_display.blit(font.render("Thrust x: {:.0f}N".format(simulation.rocket.current_stage().current_thrust(g, simulation.heading)[0]), False, (255, 255, 255)),(0,280)) - simulation_display.blit(font.render("Thrust y: {:.0f}N".format(simulation.rocket.current_stage().current_thrust(g, simulation.heading)[1]), False, (255, 255, 255)),(0,320)) - simulation_display.blit(font.render("Altitude: {:.0f}m".format(simulation.y), False, (255, 255, 255)),(0,360)) - simulation_display.blit(font.render("Fuel in stage: {:.0f}kg".format(simulation.rocket.current_stage().fuel_mass), False, (255, 255, 255)),(0,400)) - simulation_display.blit(font.render("Stage mass: {:.0f}kg".format(simulation.rocket.current_stage().total_mass()), False, (255, 255, 255)),(0,440)) - simulation_display.blit(font.render("Rocket mass: {:.0f}kg".format(simulation.rocket.total_mass()), False, (255, 255, 255)),(0,480)) - simulation_display.blit(font.render("Stage number: {:.0f}".format(simulation.rocket.stages_spent), False, (255, 255, 255)),(0,520)) - simulation_display.blit(font.render("Throttle: {:.0f}%".format(simulation.rocket.current_stage().throttle), False, (255, 255, 255)),(0,560)) - simulation_display.blit(font.render("Gimbal: {:.2f}deg".format(simulation.rocket.current_stage().gimbal), False, (255, 255, 255)),(0,600)) - simulation_display.blit(font.render("Heading: {:.2f}deg".format(simulation.heading), False, (255, 255, 255)),(0,640)) - - #draw rocket - first_stage_height = 90 #TODO - first_stage_width = 60 - - def calculate_rocket_y_based_on_y_speed_accel(display_height: int, rocket_height: int, speed_y: float, accel_y: float) -> int: - top = display_height / 5 - (rocket_height / 2) #in the case we are accelerating positively - bottom = display_height - (top * 2) - - return bottom - - def calculate_rocket_x_based_on_x_speed_accel(display_width: int, rocket_width: int, speed_x: float, accel_x: float) -> int: - return display_width / 2 - (rocket_width / 2) - - rocket_x = calculate_rocket_x_based_on_x_speed_accel(simulation_display.get_width(), first_stage_width, None, None) - rocket_y = calculate_rocket_y_based_on_y_speed_accel(simulation_display.get_height(), first_stage_height, simulation.speed_y, simulation.acceleration_y) - - rocket_color = (244, 67, 54) - - flame_radius = 10 - flame_color = (255, 125, 100) - - #TODO: Rotate rocket with heading - i = simulation.rocket.stages_spent - stage_height = first_stage_height / (i + 1) - stage_y = rocket_y + first_stage_height - stage_height - for _ in simulation.rocket.stages: - stage_width = first_stage_width / (i + 1) - stage_x = rocket_x + i * (stage_width / 2) - pygame.draw.rect(simulation_display, rocket_color, pygame.Rect(stage_x, stage_y, stage_width, stage_height)) - stage_y -= stage_height / 2 - stage_height /= 2 - i += 1 - - #draw flame - if simulation.rocket.current_stage().engines_on and simulation.rocket.current_stage().fuel_mass > 0: - pygame.draw.circle(simulation_display, flame_color, (rocket_x + (first_stage_width / 2), rocket_y + first_stage_height + flame_radius), flame_radius) - -def handle_key_press(simulation, key): - if key == pygame.K_x: - simulation.rocket.current_stage().engines_on = not simulation.rocket.current_stage().engines_on - elif key == pygame.K_z: - simulation.rocket.perform_stage_separation(True) - elif key == pygame.K_DOWN: - current_stage = simulation.rocket.current_stage() - if current_stage.throttle > 0: - current_stage.throttle -= 1 - elif key == pygame.K_UP: - current_stage = simulation.rocket.current_stage() - if current_stage.throttle < 100: - current_stage.throttle += 1 - elif key == pygame.K_LEFT: - current_stage = simulation.rocket.current_stage() - if current_stage.gimbal > 0 - current_stage.max_engine_gimbaling_angle: - current_stage.gimbal -= 1 - elif key == pygame.K_RIGHT: - current_stage = simulation.rocket.current_stage() - if current_stage.gimbal < 0 + current_stage.max_engine_gimbaling_angle: - current_stage.gimbal += 1 - -if __name__ == "__main__": - main(sys.argv) \ No newline at end of file diff --git a/rocket.py b/rocket.py deleted file mode 100644 index fb8ea7c..0000000 --- a/rocket.py +++ /dev/null @@ -1,43 +0,0 @@ -from stage import Stage - -class Rocket(): - def __init__(self, name: str, stages: [type[Stage]], payload_mass: int): - self.name = name - self.stages = stages - self.stages_spent = 0 - self.payload_mass = payload_mass - - def current_stage(self) -> type[Stage]: - return self.stages[0] - - def top_stage(self) -> type[Stage]: - return self.stages[len(self.stages) - 1] #TODO: drag coef and cross sectional area of top stage - - def perform_stage_separation(self, engines_on: bool): - if len(self.stages) > 1: - self.stages.pop(0) - self.stages_spent += 1 - self.current_stage().engines_on = engines_on - - def total_mass(self): - total_mass = self.payload_mass - for stage in self.stages: - total_mass += stage.total_mass() - return total_mass - - def total_fuel(self): - fuel_mass = 0 - for stage in self.stages: - fuel_mass += stage.fuel_mass - return fuel_mass - - def s_cross_sectional_area(self): - return self.top_stage().cross_sectional_area - - def s_drag_coefficient(self): - return self.top_stage().drag_coefficient - - #TODO: IMPLEMENT rocket_x_drag_coefficient() that adds the x drag coefficient of all stages, same with cross sectional area - - def __str__(self): - return "eue" \ No newline at end of file diff --git a/simulation.py b/simulation.py deleted file mode 100644 index 106c5a1..0000000 --- a/simulation.py +++ /dev/null @@ -1,133 +0,0 @@ -import math -from dataclasses import dataclass - -from universe import Universe -from body import Body -from rocket import Rocket - -@dataclass -class Simulation_Snapshot: - universe: type[Universe] - body: type[Body] - rocket: type[Rocket] - -class Simulation(): - def __init__(self, universe: type[Universe], body: type[Body], rocket: type[Rocket]): - self.ticks = 0 - self.time = 0 - self.universe = universe - self.body = body - self.rocket = rocket - self.x = 0#TODO - self.y = 0 #TODO: we need to make it so there is height() to calc height based on x and y - self.speed_x = 0 - self.speed_y = 0 - self.acceleration_x = 0 - self.acceleration_y = 0 - - self.heading = 0 - - #simulation logic - def tick(self, delta: int) -> None: - current_stage = self.rocket.current_stage() - #calculate upwards force by fuel - fuel_used = current_stage.total_fuel_used(delta) - if current_stage.fuel_mass < fuel_used: - fuel_used = current_stage.fuel_mass - current_stage.fuel_mass -= fuel_used - print("Fuel remaining: " + str(current_stage.fuel_mass)) - - force_x = 0 - force_y = 0 - if fuel_used > 0: - total_thrust = current_stage.current_thrust(self.body.g(self.universe.G, self.rocket_altitude()), self.heading) - force_x = total_thrust[0] - force_y = total_thrust[1] - - print("Thrust X: " + str(force_x)) - print("Thrust Y: " + str(force_y)) - - print("BODY MASS: " + str(self.body.mass())) - print("ROCKET TOTAL MASS: " + str(self.rocket.total_mass())) - - #calculate downwards force by drag and gravity - g = self.body.g(G=self.universe.G, height=self.rocket_altitude()) - print("g: " + str(g)) - - gravitational_force = g * self.rocket.total_mass() - print("Gravity: " + str(gravitational_force)) - - #Remove gravity from force - force_y -= gravitational_force - - curr_atmospheric_density = self.body.atmosphere.density_at_height(self.rocket_altitude(), g) - print("Atmosphere density: " + str(curr_atmospheric_density)) - - #TODO: cross sectional area and drag coef for x should b different - drag_force_x = (1/2) * curr_atmospheric_density * (self.speed_x ** 2) * self.rocket.s_drag_coefficient() * self.rocket.s_cross_sectional_area() - #drag goes against speed - if force_x < 0: - drag_force_x *= -1 - print("Drag X: " + str(drag_force_x)) - - #https://www.grc.nasa.gov/www/k-12/airplane/drageq.html - drag_force_y = (1/2) * curr_atmospheric_density * (self.speed_y ** 2) * self.rocket.s_drag_coefficient() * self.rocket.s_cross_sectional_area() - #drag goes against speed - if force_y < 0: - drag_force_y *= -1 - print("Drag Y: " + str(drag_force_y)) - - #remove drag - force_x -= drag_force_x - force_y -= drag_force_y - - print("Total Force X: " + str(force_x)) - print("Total Force Y: " + str(force_y)) - - self.acceleration_x = force_x / self.rocket.total_mass() - self.acceleration_y = force_y / self.rocket.total_mass() - print("Acceleration x: " + str(self.acceleration_x)) - print("Acceleration y: " + str(self.acceleration_y)) - - self.speed_x = self.speed_x + (self.acceleration_x * delta) - self.speed_y = self.speed_y + (self.acceleration_y * delta) - - print("Speed x: " + str(self.speed_x)) - print("Speed y: " + str(self.speed_y)) - - #TODO: WELL CALCULATED? (angle well?) - ref_vec = (0, 1) - acc_vec = (self.speed_x, self.speed_y) - dot = (acc_vec[0] * ref_vec[0]) + (acc_vec[1] * ref_vec[1]) - det = (acc_vec[0] * ref_vec[1]) - (acc_vec[1] * ref_vec[0]) - self.heading = math.degrees(math.atan2(det, dot)) - print("Heading: " + str(self.heading)) - - #update position based on velocity and delta - self.x += self.speed_x * delta - - #in future u should be able to go negative y (y and height are different) - self.y += self.speed_y * delta - if self.y < 0: - self.y = 0 - self.speed_y = 0 - - print("X: " + str(self.x)) - print("Y: " + str(self.y)) - - print("Total Simulation Time: " + str(self.time)) - print("") - - self.ticks += 1 - self.time += delta - - def rocket_altitude(self): - return self.y #TODO: take into account body and allow for 360 height - - def snapshot(self) -> Simulation_Snapshot: - return Simulation_Snapshot(self.universe, self.body, self.rocket) - - def str_snapshot(self) -> str: - return str(self.universe) + "\n" + \ - str(self.body) + "\n" + \ - str(self.rocket) \ No newline at end of file diff --git a/src/atmosphere.py b/src/atmosphere.py new file mode 100644 index 0000000..48f9bd1 --- /dev/null +++ b/src/atmosphere.py @@ -0,0 +1,14 @@ +import math + +class Atmosphere(): + def __init__(self, avg_sea_level_pressure: int, molar_mass_air: float, standard_temp: float): + self.avg_sea_level_pressure = avg_sea_level_pressure + self.molar_mass_air = molar_mass_air + self.standard_temp = standard_temp + + #https://math24.net/barometric-formula.html + def density_at_height(self, height: int, g: float) -> None: + R = 8.3144598 #universal gas constant + pressure = self.avg_sea_level_pressure * math.e ** (-(self.molar_mass_air * g * height)/(R * self.standard_temp)) + density = pressure / (R * 10000) + return density \ No newline at end of file diff --git a/src/body.py b/src/body.py new file mode 100644 index 0000000..3be2c5e --- /dev/null +++ b/src/body.py @@ -0,0 +1,21 @@ +import math + +from atmosphere import Atmosphere + +class Body(): + def __init__(self, name: str, density: int, radius: int, atmosphere: type[Atmosphere]): + self.name = name + self.density = density + self.radius = radius + self.atmosphere = atmosphere + + def mass(self): + body_volume = (4/3) * math.pi * (self.radius**3) + return body_volume * self.density * 1000 + + def g(self, G: float, height: int): + return (G * self.mass()) / ((self.radius + height) ** 2) + + + def __str__(self): + return "uwu" \ No newline at end of file diff --git a/src/engine.py b/src/engine.py new file mode 100644 index 0000000..f1da646 --- /dev/null +++ b/src/engine.py @@ -0,0 +1,14 @@ +import fuel + +class Engine(): + def __init__(self, name: str, isp: int, max_flow_rate: int): + self.name = name + self.max_flow_rate = max_flow_rate + self.isp = isp + + def thrust(self, throttle: int, g: float): + #https://www.grc.nasa.gov/www/k-12/airplane/specimp.html + return self.flow_rate(throttle) * self.isp * g + + def flow_rate(self, throttle: int): + return self.max_flow_rate * (throttle / 100) \ No newline at end of file diff --git a/src/fuel.py b/src/fuel.py new file mode 100644 index 0000000..caca9a9 --- /dev/null +++ b/src/fuel.py @@ -0,0 +1,3 @@ +class Fuel(): + def __init__(self, name: str): + self.name = name \ No newline at end of file diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..0d70213 --- /dev/null +++ b/src/main.py @@ -0,0 +1,212 @@ +import sys +import math +from random import randint + +from engine import Engine +from fuel import Fuel +from stage import Stage +from rocket import Rocket +from atmosphere import Atmosphere +from body import Body +from universe import Universe +from simulation import Simulation + +import pygame +from pygame.locals import * + +def main(argv): + raptor_engine = Engine(name="raptor", isp=360, max_flow_rate=931) #https://en.wikipedia.org/wiki/SpaceX_Raptor + methane_fuel = Fuel(name="methane") #TODO: more + + #https://en.wikipedia.org/wiki/SpaceX_Starship + first_stage = Stage(name="superheavy booster", + stage_mass=180000, + engine=raptor_engine, + engine_number=33, + max_engine_gimbaling_angle=30, + fuel_type=methane_fuel, + fuel_mass=3600000, + drag_coefficient=1.18, + cross_sectional_area=(math.pi * (9**2)) + ) + + second_stage = Stage(name="starship", + stage_mass=80000, + engine=raptor_engine, + engine_number=6, + max_engine_gimbaling_angle=30, + fuel_type=methane_fuel, + fuel_mass=1200000, + drag_coefficient=1.18, + cross_sectional_area=(math.pi * (9**2)) + ) + + rocket = Rocket(name="starship launch system", + stages=[first_stage, second_stage], + payload_mass=100 + ) + + body = Body(name="earth", + density=5.51, + radius=6371000, + atmosphere=Atmosphere( + avg_sea_level_pressure=101325, + molar_mass_air=0.02896, + standard_temp=288.15 + ) + ) + + universe = Universe(name="conventional", + G=6.67E-11 + ) + + simulation = Simulation(universe, body, rocket) + simulation.rocket.current_stage().engines_on = True + + pygame.init() + pygame.display.set_caption("OSLS - Overly Simple Launch Simulator") + clock = pygame.time.Clock() + + SCREEN_WIDTH = 1024 + SCREEN_HEIGHT = 720 + + simulation_display = pygame.display.set_mode((SCREEN_WIDTH,SCREEN_HEIGHT)) + paused = False + while True: + if not paused: + draw_simulation(simulation_display, simulation) + pygame.display.update() + + for event in pygame.event.get(): + if event.type == pygame.QUIT: + pygame.quit() + quit() + elif event.type == pygame.KEYDOWN: + if event.key == pygame.K_q: + quit() + elif event.key == pygame.K_SPACE: + paused = not paused + else: + handle_key_press(simulation, event.key) + + delta = clock.tick(60) / 1000 #60fps #are we using delta in the simulation tick everywhere needed? + if not paused: #tick with pause messes up delta TODO: TODOODODODODODOOD TODO + print("delta: " + str(delta)) + simulation.tick(delta=delta) + + #TODO: IMPLEMENT rocket_x_drag_coefficient() that adds the x drag coefficient of all stages, same with cross sectional area + #TODO: draw body sprite, rocket sprite, clouds sprites, etc. + #TODO: implement height properly (body radius) + actually implement body + #TODO: do max load on rocket so it blows up + #TODO: allow multilanguage api for landing algorithms etc + +def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[Simulation]) -> None: + #draw background + def linear_gradient(start_color, end_color, length, value_at): + return [ + int(start_color[j] + (float(value_at)/(length-1))*(end_color[j]-start_color[j])) + for j in range(3) + ] + + def get_color_for_height(height: float) -> (int, int, int): + if height < 70000: + return linear_gradient((31,118,194), (0, 0, 0), 70000, int(height)) + else: + return (0, 0, 0) + + #gradient for atmosphere + simulation_display.fill(get_color_for_height(simulation.y)) + + #draw clouds and stars + #draw clouds (we need continuity TODO) + #if simulation.y < 20000 and randint(0, 100) < 5: + # pygame.draw.circle(simulation_display, (255, 255, 255), (randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), 30) + #draw stars + if simulation.y > 30000: + for _ in range(100): + simulation_display.set_at((randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), (255, 255, 255)) + #draw stats text + font = pygame.font.SysFont("Comic Sans MS", 30) + + curr_thrust = simulation.rocket.current_stage().current_thrust(simulation.body.g(simulation.universe.G, simulation.y), simulation.heading) + g = simulation.body.g(simulation.universe.G, simulation.y) + + simulation_display.blit(font.render("Simulation time: {:.0f}s".format(simulation.time), False, (255, 255, 255)),(0,0)) + simulation_display.blit(font.render("X: {:.0f}m".format(simulation.x), False, (255, 255, 255)),(0,40)) + simulation_display.blit(font.render("Y: {:.0f}m".format(simulation.y), False, (255, 255, 255)),(0,80)) + simulation_display.blit(font.render("Speed x: {:.0f}m/s".format(simulation.speed_x), False, (255, 255, 255)),(0,120)) + simulation_display.blit(font.render("Speed y: {:.0f}m/s".format(simulation.speed_y), False, (255, 255, 255)),(0,160)) + simulation_display.blit(font.render("Acceleration x: {:.2f}m/s2".format(simulation.acceleration_x), False, (255, 255, 255)),(0,200)) + simulation_display.blit(font.render("Acceleration y: {:.2f}m/s2".format(simulation.acceleration_y), False, (255, 255, 255)),(0,240)) + simulation_display.blit(font.render("Thrust x: {:.0f}N".format(simulation.rocket.current_stage().current_thrust(g, simulation.heading)[0]), False, (255, 255, 255)),(0,280)) + simulation_display.blit(font.render("Thrust y: {:.0f}N".format(simulation.rocket.current_stage().current_thrust(g, simulation.heading)[1]), False, (255, 255, 255)),(0,320)) + simulation_display.blit(font.render("Altitude: {:.0f}m".format(simulation.y), False, (255, 255, 255)),(0,360)) + simulation_display.blit(font.render("Fuel in stage: {:.0f}kg".format(simulation.rocket.current_stage().fuel_mass), False, (255, 255, 255)),(0,400)) + simulation_display.blit(font.render("Stage mass: {:.0f}kg".format(simulation.rocket.current_stage().total_mass()), False, (255, 255, 255)),(0,440)) + simulation_display.blit(font.render("Rocket mass: {:.0f}kg".format(simulation.rocket.total_mass()), False, (255, 255, 255)),(0,480)) + simulation_display.blit(font.render("Stage number: {:.0f}".format(simulation.rocket.stages_spent), False, (255, 255, 255)),(0,520)) + simulation_display.blit(font.render("Throttle: {:.0f}%".format(simulation.rocket.current_stage().throttle), False, (255, 255, 255)),(0,560)) + simulation_display.blit(font.render("Gimbal: {:.2f}deg".format(simulation.rocket.current_stage().gimbal), False, (255, 255, 255)),(0,600)) + simulation_display.blit(font.render("Heading: {:.2f}deg".format(simulation.heading), False, (255, 255, 255)),(0,640)) + + #draw rocket + first_stage_height = 90 #TODO + first_stage_width = 60 + + def calculate_rocket_y_based_on_y_speed_accel(display_height: int, rocket_height: int, speed_y: float, accel_y: float) -> int: + top = display_height / 5 - (rocket_height / 2) #in the case we are accelerating positively + bottom = display_height - (top * 2) + + return bottom + + def calculate_rocket_x_based_on_x_speed_accel(display_width: int, rocket_width: int, speed_x: float, accel_x: float) -> int: + return display_width / 2 - (rocket_width / 2) + + rocket_x = calculate_rocket_x_based_on_x_speed_accel(simulation_display.get_width(), first_stage_width, None, None) + rocket_y = calculate_rocket_y_based_on_y_speed_accel(simulation_display.get_height(), first_stage_height, simulation.speed_y, simulation.acceleration_y) + + rocket_color = (244, 67, 54) + + flame_radius = 10 + flame_color = (255, 125, 100) + + #TODO: Rotate rocket with heading + i = simulation.rocket.stages_spent + stage_height = first_stage_height / (i + 1) + stage_y = rocket_y + first_stage_height - stage_height + for _ in simulation.rocket.stages: + stage_width = first_stage_width / (i + 1) + stage_x = rocket_x + i * (stage_width / 2) + pygame.draw.rect(simulation_display, rocket_color, pygame.Rect(stage_x, stage_y, stage_width, stage_height)) + stage_y -= stage_height / 2 + stage_height /= 2 + i += 1 + + #draw flame + if simulation.rocket.current_stage().engines_on and simulation.rocket.current_stage().fuel_mass > 0: + pygame.draw.circle(simulation_display, flame_color, (rocket_x + (first_stage_width / 2), rocket_y + first_stage_height + flame_radius), flame_radius) + +def handle_key_press(simulation, key): + if key == pygame.K_x: + simulation.rocket.current_stage().engines_on = not simulation.rocket.current_stage().engines_on + elif key == pygame.K_z: + simulation.rocket.perform_stage_separation(True) + elif key == pygame.K_DOWN: + current_stage = simulation.rocket.current_stage() + if current_stage.throttle > 0: + current_stage.throttle -= 1 + elif key == pygame.K_UP: + current_stage = simulation.rocket.current_stage() + if current_stage.throttle < 100: + current_stage.throttle += 1 + elif key == pygame.K_LEFT: + current_stage = simulation.rocket.current_stage() + if current_stage.gimbal > 0 - current_stage.max_engine_gimbaling_angle: + current_stage.gimbal -= 1 + elif key == pygame.K_RIGHT: + current_stage = simulation.rocket.current_stage() + if current_stage.gimbal < 0 + current_stage.max_engine_gimbaling_angle: + current_stage.gimbal += 1 + +if __name__ == "__main__": + main(sys.argv) \ No newline at end of file diff --git a/src/rocket.py b/src/rocket.py new file mode 100644 index 0000000..fb8ea7c --- /dev/null +++ b/src/rocket.py @@ -0,0 +1,43 @@ +from stage import Stage + +class Rocket(): + def __init__(self, name: str, stages: [type[Stage]], payload_mass: int): + self.name = name + self.stages = stages + self.stages_spent = 0 + self.payload_mass = payload_mass + + def current_stage(self) -> type[Stage]: + return self.stages[0] + + def top_stage(self) -> type[Stage]: + return self.stages[len(self.stages) - 1] #TODO: drag coef and cross sectional area of top stage + + def perform_stage_separation(self, engines_on: bool): + if len(self.stages) > 1: + self.stages.pop(0) + self.stages_spent += 1 + self.current_stage().engines_on = engines_on + + def total_mass(self): + total_mass = self.payload_mass + for stage in self.stages: + total_mass += stage.total_mass() + return total_mass + + def total_fuel(self): + fuel_mass = 0 + for stage in self.stages: + fuel_mass += stage.fuel_mass + return fuel_mass + + def s_cross_sectional_area(self): + return self.top_stage().cross_sectional_area + + def s_drag_coefficient(self): + return self.top_stage().drag_coefficient + + #TODO: IMPLEMENT rocket_x_drag_coefficient() that adds the x drag coefficient of all stages, same with cross sectional area + + def __str__(self): + return "eue" \ No newline at end of file diff --git a/src/simulation.py b/src/simulation.py new file mode 100644 index 0000000..106c5a1 --- /dev/null +++ b/src/simulation.py @@ -0,0 +1,133 @@ +import math +from dataclasses import dataclass + +from universe import Universe +from body import Body +from rocket import Rocket + +@dataclass +class Simulation_Snapshot: + universe: type[Universe] + body: type[Body] + rocket: type[Rocket] + +class Simulation(): + def __init__(self, universe: type[Universe], body: type[Body], rocket: type[Rocket]): + self.ticks = 0 + self.time = 0 + self.universe = universe + self.body = body + self.rocket = rocket + self.x = 0#TODO + self.y = 0 #TODO: we need to make it so there is height() to calc height based on x and y + self.speed_x = 0 + self.speed_y = 0 + self.acceleration_x = 0 + self.acceleration_y = 0 + + self.heading = 0 + + #simulation logic + def tick(self, delta: int) -> None: + current_stage = self.rocket.current_stage() + #calculate upwards force by fuel + fuel_used = current_stage.total_fuel_used(delta) + if current_stage.fuel_mass < fuel_used: + fuel_used = current_stage.fuel_mass + current_stage.fuel_mass -= fuel_used + print("Fuel remaining: " + str(current_stage.fuel_mass)) + + force_x = 0 + force_y = 0 + if fuel_used > 0: + total_thrust = current_stage.current_thrust(self.body.g(self.universe.G, self.rocket_altitude()), self.heading) + force_x = total_thrust[0] + force_y = total_thrust[1] + + print("Thrust X: " + str(force_x)) + print("Thrust Y: " + str(force_y)) + + print("BODY MASS: " + str(self.body.mass())) + print("ROCKET TOTAL MASS: " + str(self.rocket.total_mass())) + + #calculate downwards force by drag and gravity + g = self.body.g(G=self.universe.G, height=self.rocket_altitude()) + print("g: " + str(g)) + + gravitational_force = g * self.rocket.total_mass() + print("Gravity: " + str(gravitational_force)) + + #Remove gravity from force + force_y -= gravitational_force + + curr_atmospheric_density = self.body.atmosphere.density_at_height(self.rocket_altitude(), g) + print("Atmosphere density: " + str(curr_atmospheric_density)) + + #TODO: cross sectional area and drag coef for x should b different + drag_force_x = (1/2) * curr_atmospheric_density * (self.speed_x ** 2) * self.rocket.s_drag_coefficient() * self.rocket.s_cross_sectional_area() + #drag goes against speed + if force_x < 0: + drag_force_x *= -1 + print("Drag X: " + str(drag_force_x)) + + #https://www.grc.nasa.gov/www/k-12/airplane/drageq.html + drag_force_y = (1/2) * curr_atmospheric_density * (self.speed_y ** 2) * self.rocket.s_drag_coefficient() * self.rocket.s_cross_sectional_area() + #drag goes against speed + if force_y < 0: + drag_force_y *= -1 + print("Drag Y: " + str(drag_force_y)) + + #remove drag + force_x -= drag_force_x + force_y -= drag_force_y + + print("Total Force X: " + str(force_x)) + print("Total Force Y: " + str(force_y)) + + self.acceleration_x = force_x / self.rocket.total_mass() + self.acceleration_y = force_y / self.rocket.total_mass() + print("Acceleration x: " + str(self.acceleration_x)) + print("Acceleration y: " + str(self.acceleration_y)) + + self.speed_x = self.speed_x + (self.acceleration_x * delta) + self.speed_y = self.speed_y + (self.acceleration_y * delta) + + print("Speed x: " + str(self.speed_x)) + print("Speed y: " + str(self.speed_y)) + + #TODO: WELL CALCULATED? (angle well?) + ref_vec = (0, 1) + acc_vec = (self.speed_x, self.speed_y) + dot = (acc_vec[0] * ref_vec[0]) + (acc_vec[1] * ref_vec[1]) + det = (acc_vec[0] * ref_vec[1]) - (acc_vec[1] * ref_vec[0]) + self.heading = math.degrees(math.atan2(det, dot)) + print("Heading: " + str(self.heading)) + + #update position based on velocity and delta + self.x += self.speed_x * delta + + #in future u should be able to go negative y (y and height are different) + self.y += self.speed_y * delta + if self.y < 0: + self.y = 0 + self.speed_y = 0 + + print("X: " + str(self.x)) + print("Y: " + str(self.y)) + + print("Total Simulation Time: " + str(self.time)) + print("") + + self.ticks += 1 + self.time += delta + + def rocket_altitude(self): + return self.y #TODO: take into account body and allow for 360 height + + def snapshot(self) -> Simulation_Snapshot: + return Simulation_Snapshot(self.universe, self.body, self.rocket) + + def str_snapshot(self) -> str: + return str(self.universe) + "\n" + \ + str(self.body) + "\n" + \ + str(self.rocket) \ No newline at end of file diff --git a/src/stage.py b/src/stage.py new file mode 100644 index 0000000..bb05eaf --- /dev/null +++ b/src/stage.py @@ -0,0 +1,49 @@ +import math + +from engine import Engine +from fuel import Fuel + +class Stage(): + def __init__(self, name: str, stage_mass: int, engine: type[Engine], engine_number: int, max_engine_gimbaling_angle: int, fuel_type: type[Fuel], fuel_mass: int, drag_coefficient: float, cross_sectional_area: float): + self.name = name + self.stage_mass = stage_mass + self.engine = engine + self.engine_number = engine_number + self.fuel_type = fuel_type + self.fuel_mass = fuel_mass + self.drag_coefficient = drag_coefficient + self.cross_sectional_area = cross_sectional_area + + self.max_engine_gimbaling_angle = max_engine_gimbaling_angle + self.gimbal = 0 #one thing is gimbal another is rocket angle (TODO TOODODODODODODOD) + self.throttle = 100 + self.engines_on = False + + def total_mass(self): + return (self.stage_mass + self.fuel_mass) + + def current_thrust(self, g: float, heading: int) -> (float, float): + if(self.engines_on and self.fuel_mass > 0): + total_thrust = self.engine.thrust(self.throttle, g) * self.engine_number + #gimbal and heading components + thrust_x = (math.sin(math.radians(self.gimbal + heading)) * total_thrust) + thrust_y = (math.cos(math.radians(self.gimbal + heading)) * total_thrust) + + return (thrust_x, thrust_y) + else: + return (0, 0) + + def total_fuel_used(self, delta: int): + if(self.engines_on): + return self.engine.flow_rate(self.throttle) * self.engine_number * delta + else: + return 0 + + def convert_y_component_to_total_with_gimbal(self, y_value): + return math.fabs(y_value / math.cos(math.radians(self.gimbal))) + + #total drag coefficient is just the upper stage + #engines on is just the lower stage + #thrust is just the lower stage + #fuel used is just lower stage + #when stage separation lower stage jetissoned \ No newline at end of file diff --git a/src/universe.py b/src/universe.py new file mode 100644 index 0000000..a8279ec --- /dev/null +++ b/src/universe.py @@ -0,0 +1,7 @@ +class Universe(): + def __init__(self, name: str, G: float): + self.name = name + self.G = G + + def __str__(self): + return "kuu" \ No newline at end of file diff --git a/stage.py b/stage.py deleted file mode 100644 index bb05eaf..0000000 --- a/stage.py +++ /dev/null @@ -1,49 +0,0 @@ -import math - -from engine import Engine -from fuel import Fuel - -class Stage(): - def __init__(self, name: str, stage_mass: int, engine: type[Engine], engine_number: int, max_engine_gimbaling_angle: int, fuel_type: type[Fuel], fuel_mass: int, drag_coefficient: float, cross_sectional_area: float): - self.name = name - self.stage_mass = stage_mass - self.engine = engine - self.engine_number = engine_number - self.fuel_type = fuel_type - self.fuel_mass = fuel_mass - self.drag_coefficient = drag_coefficient - self.cross_sectional_area = cross_sectional_area - - self.max_engine_gimbaling_angle = max_engine_gimbaling_angle - self.gimbal = 0 #one thing is gimbal another is rocket angle (TODO TOODODODODODODOD) - self.throttle = 100 - self.engines_on = False - - def total_mass(self): - return (self.stage_mass + self.fuel_mass) - - def current_thrust(self, g: float, heading: int) -> (float, float): - if(self.engines_on and self.fuel_mass > 0): - total_thrust = self.engine.thrust(self.throttle, g) * self.engine_number - #gimbal and heading components - thrust_x = (math.sin(math.radians(self.gimbal + heading)) * total_thrust) - thrust_y = (math.cos(math.radians(self.gimbal + heading)) * total_thrust) - - return (thrust_x, thrust_y) - else: - return (0, 0) - - def total_fuel_used(self, delta: int): - if(self.engines_on): - return self.engine.flow_rate(self.throttle) * self.engine_number * delta - else: - return 0 - - def convert_y_component_to_total_with_gimbal(self, y_value): - return math.fabs(y_value / math.cos(math.radians(self.gimbal))) - - #total drag coefficient is just the upper stage - #engines on is just the lower stage - #thrust is just the lower stage - #fuel used is just lower stage - #when stage separation lower stage jetissoned \ No newline at end of file diff --git a/universe.py b/universe.py deleted file mode 100644 index a8279ec..0000000 --- a/universe.py +++ /dev/null @@ -1,7 +0,0 @@ -class Universe(): - def __init__(self, name: str, G: float): - self.name = name - self.G = G - - def __str__(self): - return "kuu" \ No newline at end of file -- cgit 1.4.1