Started preparation to implement proper altitude calculations and gravity

3 files changed, 15 insertions, 7 deletions

diff --git a/main.py b/main.py
index a4fad42..0d70213 100644
--- a/main.py
+++ b/main.py
@@ -94,6 +94,7 @@ 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
@@ -106,7 +107,7 @@ def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[S
                 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))
diff --git a/rocket.py b/rocket.py
index 962678b..fb8ea7c 100644
--- a/rocket.py
+++ b/rocket.py
@@ -37,5 +37,7 @@ class Rocket():
     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
index 46a7c57..106c5a1 100644
--- a/simulation.py
+++ b/simulation.py
@@ -40,7 +40,7 @@ class Simulation():
         force_x = 0
         force_y = 0
         if fuel_used > 0:
-            total_thrust = current_stage.current_thrust(self.body.g(self.universe.G, self.y), self.heading)
+            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]
         
@@ -51,25 +51,27 @@ class Simulation():
         print("ROCKET TOTAL MASS: " + str(self.rocket.total_mass()))
 
         #calculate downwards force by drag and gravity
-        print("g: " + str(self.body.g(G=self.universe.G, height=self.y)))
+        g = self.body.g(G=self.universe.G, height=self.rocket_altitude())
+        print("g: " + str(g))
 
-        gravitational_force = self.body.g(G=self.universe.G, height=self.y) * self.rocket.total_mass()
+        gravitational_force = g * self.rocket.total_mass()
         print("Gravity: " + str(gravitational_force))
 
         #Remove gravity from force
         force_y -= gravitational_force
 
-        print("Atmosphere density: " + str(self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.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) * self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.y)) * (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.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) * self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.y)) * (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.s_drag_coefficient() * self.rocket.s_cross_sectional_area()
         #drag goes against speed
         if force_y < 0:
             drag_force_y *= -1
@@ -119,6 +121,9 @@ class Simulation():
         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)