diff options
Diffstat (limited to 'src/main.py')
| -rw-r--r-- | src/main.py | 212 |
1 files changed, 212 insertions, 0 deletions
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 |