Added specific x and y drag (x drag coefficient and cross sectional area

4 files changed, 41 insertions, 28 deletions

diff --git a/src/main.py b/src/main.py
index 0d70213..08c6229 100644
--- a/src/main.py
+++ b/src/main.py
@@ -26,8 +26,10 @@ def main(argv):
                         max_engine_gimbaling_angle=30,
                         fuel_type=methane_fuel,
                         fuel_mass=3600000,
-                        drag_coefficient=1.18,
-                        cross_sectional_area=(math.pi * (9**2))
+                        x_drag_coefficient=1.16,#https://www.sciencedirect.com/science/article/abs/pii/S002980181400167X
+                        x_cross_sectional_area=(69 * 9), #booster height: 69m, diameter:9m
+                        y_drag_coefficient=1.28,#https://www.grc.nasa.gov/www/k-12/rocket/shaped.html
+                        y_cross_sectional_area=(math.pi * (4.5**2)) #booster radius: 4.5m
                         )
 
     second_stage = Stage(name="starship",
@@ -37,8 +39,10 @@ def main(argv):
                         max_engine_gimbaling_angle=30,
                         fuel_type=methane_fuel,
                         fuel_mass=1200000,
-                        drag_coefficient=1.18,
-                        cross_sectional_area=(math.pi * (9**2))
+                        x_drag_coefficient=1.16,#https://www.sciencedirect.com/science/article/abs/pii/S002980181400167X
+                        x_cross_sectional_area=(49 * 9), #rocket height: 49m, diameter:9m
+                        y_drag_coefficient=0.8,#https://www.grc.nasa.gov/www/k-12/rocket/shaped.html
+                        y_cross_sectional_area=(math.pi * (4.5**2))#rocket radius: 4.5m
                         )
 
     rocket = Rocket(name="starship launch system", 
@@ -94,7 +98,6 @@ def main(argv):
             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
@@ -115,21 +118,21 @@ def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[S
                 return (0, 0, 0)
 
         #gradient for atmosphere
-        simulation_display.fill(get_color_for_height(simulation.y))
+        simulation_display.fill(get_color_for_height(simulation.rocket_altitude()))
 
         #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:
+        if simulation.rocket_altitude() > 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)
+        g = simulation.body.g(simulation.universe.G, simulation.rocket_altitude())
+        curr_thrust = simulation.rocket.current_stage().current_thrust(g, simulation.heading)
 
         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))
@@ -140,7 +143,7 @@ def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[S
         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("Altitude: {:.0f}m".format(simulation.rocket_altitude()), 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))
@@ -209,4 +212,4 @@ def handle_key_press(simulation, key):
             current_stage.gimbal += 1
 
 if __name__ == "__main__":
-    main(sys.argv)
\ No newline at end of file
+    main(sys.argv)
diff --git a/src/rocket.py b/src/rocket.py
index fb8ea7c..372e421 100644
--- a/src/rocket.py
+++ b/src/rocket.py
@@ -31,13 +31,21 @@ class Rocket():
             fuel_mass += stage.fuel_mass
         return fuel_mass
 
-    def s_cross_sectional_area(self):
-        return self.top_stage().cross_sectional_area
+    #maybe we should average this with the rest of the stages?
+    def rocket_x_drag_coefficient(self):
+        return self.top_stage().x_drag_coefficient
 
-    def s_drag_coefficient(self):
-        return self.top_stage().drag_coefficient
+    def rocket_x_cross_sectional_area(self):
+        total_cross_sectional_area = 0
+        for stage in self.stages:
+            total_cross_sectional_area += stage.x_cross_sectional_area
+        return total_cross_sectional_area
+
+    def rocket_y_drag_coefficient(self):
+        return self.top_stage().y_drag_coefficient
 
-    #TODO: IMPLEMENT rocket_x_drag_coefficient() that adds the x drag coefficient of all stages, same with cross sectional area
+    def rocket_y_cross_sectional_area(self):
+        return self.top_stage().y_cross_sectional_area
 
     def __str__(self):
         return "eue"
\ No newline at end of file
diff --git a/src/simulation.py b/src/simulation.py
index 106c5a1..8968da5 100644
--- a/src/simulation.py
+++ b/src/simulation.py
@@ -36,11 +36,14 @@ class Simulation():
             fuel_used = current_stage.fuel_mass
         current_stage.fuel_mass -= fuel_used
         print("Fuel remaining: " + str(current_stage.fuel_mass))
+
+        g = self.body.g(G=self.universe.G, height=self.rocket_altitude())
+        print("g: " + str(g))
                 
         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)
+            total_thrust = current_stage.current_thrust(g, self.heading)
             force_x = total_thrust[0]
             force_y = total_thrust[1]
         
@@ -51,27 +54,24 @@ class Simulation():
         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))
+        gravitational_force_y = g * self.rocket.total_mass()
+        print("Gravity Y: " + str(gravitational_force_y))
 
         #Remove gravity from force
-        force_y -= gravitational_force
+        force_y -= gravitational_force_y
 
         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_force_x = (1/2) * curr_atmospheric_density * (self.speed_x ** 2) * self.rocket.rocket_x_drag_coefficient() * self.rocket.rocket_x_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_force_y = (1/2) * curr_atmospheric_density * (self.speed_y ** 2) * self.rocket.rocket_y_drag_coefficient() * self.rocket.rocket_y_cross_sectional_area()
         #drag goes against speed
         if force_y < 0:
             drag_force_y *= -1
diff --git a/src/stage.py b/src/stage.py
index bb05eaf..e1f6b04 100644
--- a/src/stage.py
+++ b/src/stage.py
@@ -4,15 +4,17 @@ 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):
+    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, x_drag_coefficient: float, x_cross_sectional_area: float, y_drag_coefficient: float, y_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.x_drag_coefficient = x_drag_coefficient
+        self.x_cross_sectional_area = x_cross_sectional_area
+        self.y_drag_coefficient = y_drag_coefficient
+        self.y_cross_sectional_area = y_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)