Merge pull request #441 from RocketPy-Team/enh/get-solution-at-time

ENH: new Flight.get_solution_at_time() method

Author: Gui-FernandesBR

rocketpy/simulation/flight.py

@@ -1934,6 +1934,37 @@ def time(self):
         """Returns time array from solution."""
         return self.solution_array[:, 0]
 
+    def get_solution_at_time(self, t, atol=1e-3):
+        """Returns the solution state vector at a given time. If the time is
+        not found in the solution, the closest time is used and a warning is
+        raised.
+
+        Parameters
+        ----------
+        t : float
+            Time in seconds.
+        atol : float, optional
+            Absolute tolerance for time comparison. Default is 1e-3. If the
+            difference between the time and the closest time in the solution
+            is greater than this value, a warning is raised.
+
+        Returns
+        -------
+        solution : np.array
+            Solution state at time t. The array follows the format of the
+            solution array, with the first column being time like this:
+            [t, x, y, z, vx, vy, vz, e0, e1, e2, e3, omega1, omega2, omega3].
+
+        """
+        time_index = find_closest(self.time, t)
+        if abs(self.time[time_index] - t) > atol:
+            warnings.warn(
+                f"Time {t} not found in solution. Closest time is "
+                f"{self.time[time_index]}. Using closest time.",
+                UserWarning,
+            )
+        return self.solution_array[time_index, :]
+
     # Process first type of outputs - state vector
     # Transform solution array into Functions
     @funcify_method("Time (s)", "X (m)", "spline", "constant")

tests/test_flight.py

@@ -141,6 +141,48 @@ def test_initial_solution(mock_show, example_env, calisto_robust):
     assert test_flight.all_info() == None
 
 
+def test_get_solution_at_time(flight_calisto):
+    """Test the get_solution_at_time method of the Flight class. This test
+    simply calls the method at the initial and final time and checks if the
+    returned values are correct.
+
+    Parameters
+    ----------
+    flight_calisto : rocketpy.Flight
+        Flight object to be tested. See the conftest.py file for more info
+        regarding this pytest fixture.
+    """
+    assert np.allclose(
+        flight_calisto.get_solution_at_time(0),
+        np.array([0, 0, 0, 0, 0, 0, 0, 0.99904822, -0.04361939, 0, 0, 0, 0, 0]),
+        rtol=1e-05,
+        atol=1e-08,
+    )
+    assert np.allclose(
+        flight_calisto.get_solution_at_time(flight_calisto.t_final),
+        np.array(
+            [
+                48.4313533,
+                0.0,
+                985.7665845,
+                -0.00000229951048,
+                0.0,
+                11.2223284,
+                -341.028803,
+                0.999048222,
+                -0.0436193874,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+            ]
+        ),
+        rtol=1e-02,
+        atol=5e-03,
+    )
+
+
 @pytest.mark.parametrize("wind_u, wind_v", [(0, 10), (0, -10), (10, 0), (-10, 0)])
 @pytest.mark.parametrize(
     "static_margin, max_time",