ENH: Add optional max time to stochastic_flight and create stochastic_flight objects with attributes in base flight (#1070)

* Add max_time option to stochastic_flight

* Update: use default attributes in flight object for stochastic flights

* Add tests to check attributes

* Refactor implementation

Author: zuorenchen

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,
         )

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