Structured source files into src folder

9 files changed, 496 insertions, 0 deletions

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