ENH: Refactor Flight class to improve time node handling and sensor/controllers (#843)

* ENH: Refactor Flight class to improve time node handling and sensor/controller processing

TST: add tests for overshootable controllers

TST: fix slow tests

fix merge conflicts

* fix tests

fix slow tests

update CHANGELOG

Update rocketpy/simulation/flight.py

fix lint

Author: Gui-FernandesBR

CHANGELOG.md

@@ -49,7 +49,7 @@ Attention: The newest changes should be on top -->
 
 ### Changed
 
--
+- ENH: Refactor Flight class to improve time node handling and sensor/controllers [#843](https://github.com/RocketPy-Team/RocketPy/pull/843)
 
 ### Fixed
 

rocketpy/plots/flight_plots.py

@@ -720,9 +720,16 @@ def energy_data(self, *, filename=None):  # pylint: disable=too-many-statements
         ax2.grid()
 
         ax3 = plt.subplot(413)
+        # Handle both array-based and callable-based Functions
+        thrust_power = self.flight.thrust_power
+        if callable(thrust_power.source):
+            # For callable sources, discretize based on speed
+            thrust_power = thrust_power.set_discrete_based_on_model(
+                self.flight.speed, mutate_self=False
+            )
         ax3.plot(
-            self.flight.thrust_power[:, 0],
-            self.flight.thrust_power[:, 1],
+            thrust_power[:, 0],
+            thrust_power[:, 1],
             label="|Thrust Power|",
         )
         ax3.set_xlim(0, self.flight.rocket.motor.burn_out_time)
@@ -734,9 +741,16 @@ def energy_data(self, *, filename=None):  # pylint: disable=too-many-statements
         ax3.grid()
 
         ax4 = plt.subplot(414)
+        # Handle both array-based and callable-based Functions
+        drag_power = self.flight.drag_power
+        if callable(drag_power.source):
+            # For callable sources, discretize based on speed
+            drag_power = drag_power.set_discrete_based_on_model(
+                self.flight.speed, mutate_self=False
+            )
         ax4.plot(
-            self.flight.drag_power[:, 0],
-            -self.flight.drag_power[:, 1],
+            drag_power[:, 0],
+            -drag_power[:, 1],
             label="|Drag Power|",
         )
         ax4.set_xlim(

rocketpy/simulation/flight.py

@@ -76,6 +76,7 @@ def pytest_collection_modifyitems(config, items):
             item.add_marker(skip_slow)
 
 
+# TODO: move this to Environment fixtures when possible
 @pytest.fixture
 def merra2_file_path(tmp_path):  # pylint: disable=too-many-statements
     """

tests/conftest.py

@@ -246,6 +246,22 @@ def flight_calisto_air_brakes(calisto_air_brakes_clamp_on, example_plain_env):
     )
 
 
+@pytest.fixture
+def flight_calisto_air_brakes_time_overshoot(
+    calisto_air_brakes_clamp_on, example_plain_env
+):
+    """Same as flight_calisto_air_brakes but with time_overshoot=True."""
+    return Flight(
+        rocket=calisto_air_brakes_clamp_on,
+        environment=example_plain_env,
+        rail_length=5.2,
+        inclination=85,
+        heading=0,
+        time_overshoot=True,
+        terminate_on_apogee=True,
+    )
+
+
 @pytest.fixture
 def flight_calisto_with_sensors(calisto_with_sensors, example_plain_env):
     """A rocketpy.Flight object of the Calisto rocket. This uses the calisto

tests/fixtures/flight/flight_fixtures.py

@@ -160,9 +160,8 @@ def test_nam_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=un
 @pytest.mark.slow
 @patch("matplotlib.pyplot.show")
 def test_rap_atmosphere(mock_show, example_spaceport_env):  # pylint: disable=unused-argument
-    today = date.today()
     now = datetime.now(timezone.utc)
-    example_spaceport_env.set_date((today.year, today.month, today.day, now.hour))
+    example_spaceport_env.set_date((now.year, now.month, now.day, now.hour))
     example_spaceport_env.set_atmospheric_model(type="Forecast", file="RAP")
     assert example_spaceport_env.all_info() is None
 
@@ -281,7 +280,7 @@ def test_merra2_full_specification_compliance(merra2_file_path, example_plain_en
     # Input: 9806.65 m2/s2
     # Expected: 1000.0 m
     print(f"Calculated Elevation: {env.elevation} m")
-    assert abs(env.elevation - 1000.0) < 1e-6, (
+    assert abs(env.elevation - 1000.0) < 1e-4, (
         f"Failed to convert PHIS (m2/s2) to meters. Got {env.elevation}, expected 1000.0"
     )
 

tests/integration/environment/test_environment.py

@@ -102,7 +102,7 @@ def test_liquid_motor_flight(mock_show, flight_calisto_liquid_modded):  # pylint
 
 @pytest.mark.slow
 @patch("matplotlib.pyplot.show")
-def test_time_overshoot(mock_show, calisto_robust, example_spaceport_env):  # pylint: disable=unused-argument
+def test_time_overshoot_false(mock_show, calisto_robust, example_spaceport_env):  # pylint: disable=unused-argument
     """Test the time_overshoot parameter of the Flight class. This basically
     calls the all_info() method for a simulation without time_overshoot and
     checks if it returns None. It is not testing if the values are correct,
@@ -293,6 +293,19 @@ def test_air_brakes_flight(mock_show, flight_calisto_air_brakes):  # pylint: dis
     assert air_brakes.prints.all() is None
 
 
+@patch("matplotlib.pyplot.show")
+def test_air_brakes_flight_with_overshoot(
+    mock_show, flight_calisto_air_brakes_time_overshoot
+):  # pylint: disable=unused-argument
+    """
+    Same as test_air_brakes_flight but with time_overshoot=True.
+    """
+    test_flight = flight_calisto_air_brakes_time_overshoot
+    air_brakes = test_flight.rocket.air_brakes[0]
+    assert air_brakes.plots.all() is None
+    assert air_brakes.prints.all() is None
+
+
 @patch("matplotlib.pyplot.show")
 def test_initial_solution(mock_show, example_plain_env, calisto_robust):  # pylint: disable=unused-argument
     """Tests the initial_solution option of the Flight class. This test simply

tests/integration/simulation/test_flight.py

@@ -74,10 +74,10 @@ def test_compare_flights(mock_show, mock_figure_show, calisto, example_plain_env
     )
 
     calisto.set_rail_buttons(-0.5, 0.2)
-    inclinations = [60, 70, 80, 90]
-    headings = [0, 45, 90, 180]
+    inclinations = [60, 90]
+    headings = [0, 180]
     flights = []
-    # Create (4 * 4) = 16 different flights to be compared
+    # Create (2 * 2) = 4 different flights to be compared
     for heading in headings:
         for inclination in inclinations:
             flight = Flight(