ENH: refactor retrieve arrays cached properties and controller sim post processing

MateusStano · MateusStano · commit 830ea15445d9 · 2024-04-15T15:05:31.000+02:00
diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py
@@ -593,7 +593,6 @@ def __init__(
         if self.rail_length <= 0:
             raise ValueError("Rail length must be a positive value.")
         self.parachutes = self.rocket.parachutes[:]
-        self._controllers = self.rocket._controllers[:]
         self.inclination = inclination
         self.heading = heading
         self.max_time = max_time
@@ -607,6 +606,13 @@ def __init__(
         self.name = name
         self.equations_of_motion = equations_of_motion
 
+        # Controller initialization
+        self._controllers = self.rocket._controllers[:]
+        if self._controllers:
+            # reset controllable object to initial state (only airbrakes for now)
+            for air_brakes in self.rocket.air_brakes:
+                air_brakes.reset()
+
         # Flight initialization
         self.__init_post_process_variables()
         self.__init_solution_monitors()
@@ -1076,8 +1082,14 @@ def __init__(
                                             [self.t, parachute]
                                         )
 
+                    # If controlled flight, post process must be done on sim time
+                    if self._controllers:
+                        phase.derivative(self.t, self.y_sol, post_processing=True)
+
         self.t_final = self.t
         self._calculate_pressure_signal()
+        if self._controllers:
+            self.__cache_post_process_variables()
         if verbose:
             print("Simulation Completed at Time: {:3.4f} s".format(self.t))
 
@@ -1089,6 +1101,25 @@ def __init_post_process_variables(self):
         self._bearing = Function(0)
         self._latitude = Function(0)
         self._longitude = Function(0)
+        # Initialize state derivatives, force and atmospheric arrays
+        self.ax_list = []
+        self.ay_list = []
+        self.az_list = []
+        self.alpha1_list = []
+        self.alpha2_list = []
+        self.alpha3_list = []
+        self.R1_list = []
+        self.R2_list = []
+        self.R3_list = []
+        self.M1_list = []
+        self.M2_list = []
+        self.M3_list = []
+        self.pressure_list = []
+        self.density_list = []
+        self.dynamic_viscosity_list = []
+        self.speed_of_sound_list = []
+        self.wind_velocity_x_list = []
+        self.wind_velocity_y_list = []
 
     def __init_solution_monitors(self):
         # Initialize solution monitors
@@ -1181,10 +1212,28 @@ def __init_equations_of_motion(self):
         if self.equations_of_motion == "solid_propulsion":
             self.u_dot_generalized = self.u_dot
 
-    def __init_equations_of_motion(self):
-        """Initialize equations of motion."""
-        if self.equations_of_motion == "solid_propulsion":
-            self.u_dot_generalized = self.u_dot
+    def __cache_post_process_variables(self):
+        """Cache post-process variables for simulations with controllers."""
+        self.__retrieve_arrays = [
+            self.ax_list,
+            self.ay_list,
+            self.az_list,
+            self.alpha1_list,
+            self.alpha2_list,
+            self.alpha3_list,
+            self.R1_list,
+            self.R2_list,
+            self.R3_list,
+            self.M1_list,
+            self.M2_list,
+            self.M3_list,
+            self.pressure_list,
+            self.density_list,
+            self.dynamic_viscosity_list,
+            self.speed_of_sound_list,
+            self.wind_velocity_x_list,
+            self.wind_velocity_y_list,
+        ]
 
     @cached_property
     def effective_1rl(self):
@@ -1261,11 +1310,6 @@ def udot_rail1(self, t, u, post_processing=False):
             e0dot, e1dot, e2dot, e3dot, alpha1, alpha2, alpha3].
 
         """
-        # Check if post processing mode is on
-        if post_processing:
-            # Use u_dot post processing code
-            return self.u_dot_generalized(t, u, True)
-
         # Retrieve integration data
         x, y, z, vx, vy, vz, e0, e1, e2, e3, omega1, omega2, omega3 = u
 
@@ -1296,6 +1340,17 @@ def udot_rail1(self, t, u, post_processing=False):
         else:
             ax, ay, az = 0, 0, 0
 
+        if post_processing:
+            # Use u_dot post processing code for forces, moments and env data
+            self.u_dot_generalized(t, u, post_processing=True)
+            # Save feasible accelerations
+            self.ax_list[-1] = [t, ax]
+            self.ay_list[-1] = [t, ay]
+            self.az_list[-1] = [t, az]
+            self.alpha1_list[-1] = [t, 0]
+            self.alpha2_list[-1] = [t, 0]
+            self.alpha3_list[-1] = [t, 0]
+
         return [vx, vy, vz, ax, ay, az, 0, 0, 0, 0, 0, 0, 0]
 
     def udot_rail2(self, t, u, post_processing=False):
@@ -1585,6 +1640,13 @@ def u_dot(self, t, u, post_processing=False):
         ]
 
         if post_processing:
+            # Accelerations
+            self.ax_list.append([t, ax])
+            self.ay_list.append([t, ay])
+            self.az_list.append([t, az])
+            self.alpha1_list.append([t, alpha1])
+            self.alpha2_list.append([t, alpha2])
+            self.alpha3_list.append([t, alpha3])
             # Dynamics variables
             self.R1_list.append([t, R1])
             self.R2_list.append([t, R2])
@@ -1860,6 +1922,13 @@ def u_dot_generalized(self, t, u, post_processing=False):
         u_dot = [*r_dot, *v_dot, *e_dot, *w_dot]
 
         if post_processing:
+            # Accelerations
+            self.ax_list.append([t, v_dot[0]])
+            self.ay_list.append([t, v_dot[1]])
+            self.az_list.append([t, v_dot[2]])
+            self.alpha1_list.append([t, w_dot[0]])
+            self.alpha2_list.append([t, w_dot[1]])
+            self.alpha3_list.append([t, w_dot[2]])
             # Dynamics variables
             self.R1_list.append([t, R1])
             self.R2_list.append([t, R2])
@@ -1944,6 +2013,13 @@ def u_dot_parachute(self, t, u, post_processing=False):
         az = (Dz - 9.8 * mp) / (mp + ma)
 
         if post_processing:
+            # Accelerations
+            self.ax_list.append([t, ax])
+            self.ay_list.append([t, ay])
+            self.az_list.append([t, az])
+            self.alpha1_list.append([t, 0])
+            self.alpha2_list.append([t, 0])
+            self.alpha3_list.append([t, 0])
             # Dynamics variables
             self.R1_list.append([t, Dx])
             self.R2_list.append([t, Dy])
@@ -1952,13 +2028,20 @@ def u_dot_parachute(self, t, u, post_processing=False):
             self.M2_list.append([t, 0])
             self.M3_list.append([t, 0])
             # Atmospheric Conditions
-            self.wind_velocity_x_list.append([t, self.env.wind_velocity_x(z)])
-            self.wind_velocity_y_list.append([t, self.env.wind_velocity_y(z)])
-            self.density_list.append([t, self.env.density(z)])
-            self.dynamic_viscosity_list.append([t, self.env.dynamic_viscosity(z)])
-            self.pressure_list.append([t, self.env.pressure(z)])
-            self.speed_of_sound_list.append([t, self.env.speed_of_sound(z)])
-
+            self.wind_velocity_x_list.append(
+                [t, self.env.wind_velocity_x.get_value_opt(z)]
+            )
+            self.wind_velocity_y_list.append(
+                [t, self.env.wind_velocity_y.get_value_opt(z)]
+            )
+            self.density_list.append([t, self.env.density.get_value_opt(z)])
+            self.dynamic_viscosity_list.append(
+                [t, self.env.dynamic_viscosity.get_value_opt(z)]
+            )
+            self.pressure_list.append([t, self.env.pressure.get_value_opt(z)])
+            self.speed_of_sound_list.append(
+                [t, self.env.speed_of_sound.get_value_opt(z)]
+            )
         return [vx, vy, vz, ax, ay, az, 0, 0, 0, 0, 0, 0, 0]
 
     @cached_property
@@ -2790,13 +2873,64 @@ def longitude(self):
 
         return np.column_stack((self.time, longitude))
 
+    @cached_property
+    def __retrieve_arrays(self):
+        """post processing function to retrieve arrays from the integration
+        scheme and store them in lists for further analysis.
+
+        Returns
+        -------
+        temp_values: list
+            List containing the following arrays: ``ax`` , ``ay`` , ``az`` ,
+            ``alpha1`` , ``alpha2`` , ``alpha3`` , ``R1`` , ``R2`` , ``R3`` ,
+            ``M1`` , ``M2`` , ``M3`` , ``pressure`` , ``density`` ,
+            ``dynamic_viscosity`` , ``speed_of_sound`` , ``wind_velocity_x`` ,
+            ``wind_velocity_y``.
+        """
+        # Go through each time step and calculate forces and atmospheric values
+        # Get flight phases
+        for phase_index, phase in self.time_iterator(self.FlightPhases):
+            init_time = phase.t
+            final_time = self.FlightPhases[phase_index + 1].t
+            current_derivative = phase.derivative
+            # Call callback functions
+            for callback in phase.callbacks:
+                callback(self)
+            # Loop through time steps in flight phase
+            for step in self.solution:  # Can be optimized
+                if init_time < step[0] <= final_time or (
+                    init_time == self.t_initial and step[0] == self.t_initial
+                ):
+                    # Call derivatives in post processing mode
+                    current_derivative(step[0], step[1:], post_processing=True)
+
+        temp_values = [
+            self.ax_list,
+            self.ay_list,
+            self.az_list,
+            self.alpha1_list,
+            self.alpha2_list,
+            self.alpha3_list,
+            self.R1_list,
+            self.R2_list,
+            self.R3_list,
+            self.M1_list,
+            self.M2_list,
+            self.M3_list,
+            self.pressure_list,
+            self.density_list,
+            self.dynamic_viscosity_list,
+            self.speed_of_sound_list,
+            self.wind_velocity_x_list,
+            self.wind_velocity_y_list,
+        ]
+
+        return temp_values
+
     @cached_property
     def retrieve_acceleration_arrays(self):
         """Retrieve acceleration arrays from the integration scheme
 
-        Parameters
-        ----------
-
         Returns
         -------
         ax: list
@@ -2812,35 +2946,7 @@ def retrieve_acceleration_arrays(self):
         alpha3: list
             angular acceleration in z direction
         """
-        # Initialize acceleration arrays
-        ax, ay, az = [[0, 0]], [[0, 0]], [[0, 0]]
-        alpha1, alpha2, alpha3 = [[0, 0]], [[0, 0]], [[0, 0]]
-        # Go through each time step and calculate accelerations
-        # Get flight phases
-        for phase_index, phase in self.time_iterator(self.FlightPhases):
-            init_time = phase.t
-            final_time = self.FlightPhases[phase_index + 1].t
-            current_derivative = phase.derivative
-            # Call callback functions
-            for callback in phase.callbacks:
-                callback(self)
-            # Loop through time steps in flight phase
-            for step in self.solution:  # Can be optimized
-                if init_time < step[0] <= final_time:
-                    # Get derivatives
-                    u_dot = current_derivative(step[0], step[1:])
-                    # Get accelerations
-                    ax_value, ay_value, az_value = u_dot[3:6]
-                    alpha1_value, alpha2_value, alpha3_value = u_dot[10:]
-                    # Save accelerations
-                    ax.append([step[0], ax_value])
-                    ay.append([step[0], ay_value])
-                    az.append([step[0], az_value])
-                    alpha1.append([step[0], alpha1_value])
-                    alpha2.append([step[0], alpha2_value])
-                    alpha3.append([step[0], alpha3_value])
-
-        return ax, ay, az, alpha1, alpha2, alpha3
+        return self.__retrieve_arrays[:6]
 
     @cached_property
     def retrieve_temporary_values_arrays(self):
@@ -2877,54 +2983,7 @@ def retrieve_temporary_values_arrays(self):
         self.wind_velocity_y_list: list
             Wind velocity in y direction at each time step.
         """
-
-        # Initialize force and atmospheric arrays
-        self.R1_list = []
-        self.R2_list = []
-        self.R3_list = []
-        self.M1_list = []
-        self.M2_list = []
-        self.M3_list = []
-        self.pressure_list = []
-        self.density_list = []
-        self.dynamic_viscosity_list = []
-        self.speed_of_sound_list = []
-        self.wind_velocity_x_list = []
-        self.wind_velocity_y_list = []
-
-        # Go through each time step and calculate forces and atmospheric values
-        # Get flight phases
-        for phase_index, phase in self.time_iterator(self.FlightPhases):
-            init_time = phase.t
-            final_time = self.FlightPhases[phase_index + 1].t
-            current_derivative = phase.derivative
-            # Call callback functions
-            for callback in phase.callbacks:
-                callback(self)
-            # Loop through time steps in flight phase
-            for step in self.solution:  # Can be optimized
-                if init_time < step[0] <= final_time or (
-                    init_time == self.t_initial and step[0] == self.t_initial
-                ):
-                    # Call derivatives in post processing mode
-                    u_dot = current_derivative(step[0], step[1:], post_processing=True)
-
-        temporary_values = [
-            self.R1_list,
-            self.R2_list,
-            self.R3_list,
-            self.M1_list,
-            self.M2_list,
-            self.M3_list,
-            self.pressure_list,
-            self.density_list,
-            self.dynamic_viscosity_list,
-            self.speed_of_sound_list,
-            self.wind_velocity_x_list,
-            self.wind_velocity_y_list,
-        ]
-
-        return temporary_values
+        return self.__retrieve_arrays[6:]
 
     def get_controller_observed_variables(self):
         """Retrieve the observed variables related to air brakes from the
