diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py index 97ffbc88c..ca896cc48 100644 --- a/rocketpy/simulation/flight.py +++ b/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") diff --git a/tests/test_flight.py b/tests/test_flight.py index 212fad932..de66d2902 100644 --- a/tests/test_flight.py +++ b/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",