diff --git a/rocketpy/stochastic/stochastic_flight.py b/rocketpy/stochastic/stochastic_flight.py index 0fde38c3a..ecac053a3 100644 --- a/rocketpy/stochastic/stochastic_flight.py +++ b/rocketpy/stochastic/stochastic_flight.py @@ -32,6 +32,9 @@ class StochasticFlight(StochasticModel): time_overshoot : bool If False, the simulation will run at the time step defined by the controller sampling rate. Be aware that this will make the simulation run much slower. + max_time : int, float + The maximum time of the flight simulation. If the flight simulation + reaches this time, it will terminate. This attribute can not be randomized. """ def __init__( @@ -43,6 +46,7 @@ def __init__( initial_solution=None, terminate_on_apogee=None, time_overshoot=None, + max_time=None, ): """Initializes the Stochastic Flight class. @@ -70,6 +74,9 @@ def __init__( time_overshoot : bool If False, the simulation will run at the time step defined by the controller sampling rate. Be aware that this will make the simulation run much slower. + max_time : int, float + The maximum time of the flight simulation. If the flight simulation + reaches this time, it will terminate. This attribute can not be randomized. """ if terminate_on_apogee is not None: assert isinstance(terminate_on_apogee, bool), ( @@ -78,6 +85,9 @@ def __init__( if time_overshoot is not None: if not isinstance(time_overshoot, bool): raise TypeError("`time_overshoot` must be a boolean") + if max_time is not None: + if not isinstance(max_time, (int, float)): + raise TypeError("`max_time` must be a number") super().__init__( flight, rail_length=rail_length, @@ -87,6 +97,10 @@ def __init__( self.initial_solution = initial_solution self.terminate_on_apogee = terminate_on_apogee + if max_time is None: + self.max_time = flight.max_time + else: + self.max_time = max_time if time_overshoot is None: self.time_overshoot = flight.time_overshoot else: @@ -135,12 +149,21 @@ def create_object(self): generated_dict = next(self.dict_generator()) # TODO: maybe we should use generated_dict["rail_length"] instead return Flight( + rocket=self.obj.rocket, environment=self.obj.env, rail_length=self._randomize_rail_length(), - rocket=self.obj.rocket, inclination=generated_dict["inclination"], heading=generated_dict["heading"], initial_solution=self.initial_solution, terminate_on_apogee=self.terminate_on_apogee, + max_time=self.max_time, + max_time_step=self.obj.max_time_step, + min_time_step=self.obj.min_time_step, + rtol=self.obj.rtol, + atol=self.obj.atol, time_overshoot=self.time_overshoot, + name=self.obj.name, + equations_of_motion=self.obj.equations_of_motion, + ode_solver=self.obj.ode_solver, + simulation_mode=self.obj.simulation_mode, ) diff --git a/tests/unit/stochastic/test_stochastic_flight.py b/tests/unit/stochastic/test_stochastic_flight.py index aeb71b906..e03917475 100644 --- a/tests/unit/stochastic/test_stochastic_flight.py +++ b/tests/unit/stochastic/test_stochastic_flight.py @@ -1,6 +1,47 @@ from rocketpy.simulation.flight import Flight +from rocketpy.stochastic import StochasticFlight def test_stochastic_flight_create_object(stochastic_flight): obj = stochastic_flight.create_object() assert isinstance(obj, Flight) + + +def test_stochastic_flight_inherited_attributes(calisto_robust, example_spaceport_env): + flight = Flight( + rocket=calisto_robust, + environment=example_spaceport_env, + rail_length=5.2, + max_time_step=10, + min_time_step=0.01, + rtol=1e-4, + atol=1e-6, + name="FlightName", + equations_of_motion="solid_propulsion", + ode_solver="BDF", + simulation_mode="3 DOF", + ) + stochastic_flight = StochasticFlight(flight=flight) + + obj = stochastic_flight.create_object() + assert flight.max_time_step == obj.max_time_step + assert flight.min_time_step == obj.min_time_step + assert flight.rtol == obj.rtol + assert flight.atol == obj.atol + assert flight.name == obj.name + assert flight.equations_of_motion == obj.equations_of_motion + assert flight.ode_solver == obj.ode_solver + assert flight.simulation_mode == obj.simulation_mode + + +def test_stochastic_flight_optional_attributes(flight_calisto_robust): + stochastic_flight = StochasticFlight( + flight=flight_calisto_robust, + terminate_on_apogee=True, + time_overshoot=True, + max_time=987.6, + ) + obj = stochastic_flight.create_object() + assert obj.terminate_on_apogee is True + assert obj.time_overshoot is True + assert obj.max_time == 987.6