Merge pull request #125 from nasa/feature/UAV_updates

Updates to UAV model

src/progpy/loading/controllers.py

@@ -95,7 +95,7 @@ def __init__(self, x_ref, vehicle, **kwargs):
         self.n_inputs = len(self.inputs) - 1   # number of inputs (minus one to remove mission_complete)
         self.ref_traj = x_ref                  # reference state to follow during simulation (x_ref, y_ref, z_ref, phi_ref, theta_ref, psi_ref, ...)
         self.ss_input = vehicle.parameters['steadystate_input']
-        self.vehicle_max_thrust = vehicle.dynamics['max_thrust']
+        self.vehicle_max_thrust = vehicle.parameters['dynamics']['max_thrust']
 
         # Default control parameters
         # --------------------------------
@@ -122,7 +122,7 @@ def __call__(self, t, x=None):
         if x is None:
             x_k = np.zeros((self.n_states, 1))
         else:
-            x_k = np.array([x.matrix[ii][0] for ii in range(len(x.matrix)-2)])
+            x_k = np.array([x.matrix[ii][0] for ii in range(self.n_states)])
 
         # Identify reference state (desired state) at t
         t_k = np.round(t + self.dt/2.0, 1)  # current time step
@@ -221,7 +221,7 @@ def __init__(self, x_ref, vehicle, **kwargs):
         self.n_inputs = len(self.inputs) - 1  # number of inputs (minus one to remove mission_complete)
         self.ref_traj = x_ref                 # reference state to follow during simulation (x_ref, y_ref, z_ref, phi_ref, theta_ref, psi_ref, ...)
         self.ss_input = vehicle.parameters['steadystate_input']
-        self.vehicle_max_thrust = vehicle.dynamics['max_thrust']
+        self.vehicle_max_thrust = vehicle.parameters['dynamics']['max_thrust']
 
         self.outputs = vehicle.outputs[:3]    # output variables of the system to be controlled (x, y, z only)
         self.n_outputs = 3                    # number of outputs

src/progpy/models/aircraft_model/small_rotorcraft.py

@@ -148,28 +148,28 @@ def __init__(self, **kwargs):
         if not isinstance(self.parameters['vehicle_model'], str):
             raise TypeError("Vehicle model must be defined as a string.")
         if self.parameters['vehicle_model'].lower() == 'djis1000':
-            self.mass, self.geom, self.dynamics = vehicles.DJIS1000(self.parameters['vehicle_payload'], self.parameters['gravity'])
+            self.parameters['mass'], self.parameters['geom'], self.parameters['dynamics'] = vehicles.DJIS1000(self.parameters['vehicle_payload'], self.parameters['gravity'])
         elif self.parameters['vehicle_model'].lower() == 'tarot18':
-            self.mass, self.geom, self.dynamics = vehicles.TAROT18(self.parameters['vehicle_payload'], self.parameters['gravity'])
+            self.parameters['mass'], self.parameters['geom'], self.parameters['dynamics'] = vehicles.TAROT18(self.parameters['vehicle_payload'], self.parameters['gravity'])
         else:
             raise ValueError("Specified vehicle type is not supported. Only 'tarot18' and 'djis1000' are currently supported.")
 
         # Steady-state input value: hover, [weight, 0, 0, 0]
         if self.parameters['steadystate_input'] is None:
-            self.parameters['steadystate_input'] = self.mass['total'] * self.parameters['gravity']
+            self.parameters['steadystate_input'] = self.parameters['mass']['total'] * self.parameters['gravity']
 
         # Introduction of Aerodynamic effects:
-        self.aero = {
+        self.parameters['aero'] = {
           'drag': aero.DragModel(
-            bodyarea=self.dynamics['aero']['ad'],
-            Cd=self.dynamics['aero']['cd'],
+            bodyarea=self.parameters['dynamics']['aero']['ad'],
+            Cd=self.parameters['dynamics']['aero']['cd'],
             air_density=self.parameters['air_density']),
           'lift': None}
 
     def dx(self, x, u):
         # Extract useful values
-        m = self.mass['total']  # vehicle mass
-        Ixx, Iyy, Izz = self.mass['Ixx'], self.mass['Iyy'], self.mass['Izz']  # vehicle inertia
+        m = self.parameters['mass']['total']  # vehicle mass
+        Ixx, Iyy, Izz = self.parameters['mass']['Ixx'], self.parameters['mass']['Iyy'], self.parameters['mass']['Izz']  # vehicle inertia
 
         # Input vector
         T = u['T']  # Thrust (along body z)
@@ -202,7 +202,7 @@ def dx(self, x, u):
             sin_psi,
             cos_psi),
           v_earth)  # Velocity in body-axis
-        fb_drag = self.aero['drag'](v_body)   # drag force in body axis
+        fb_drag = self.parameters['aero']['drag'](v_body)   # drag force in body axis
         fe_drag = np.dot(
           geom.rot_body2earth_fast(
             sin_phi,
@@ -230,9 +230,9 @@ def dx(self, x, u):
         dxdt[7] = ((sin_theta * sin_psi * cos_phi - sin_phi * cos_psi) * T - fe_drag[1]) / m   # Acceleration along y-axis
         dxdt[8] = -self.parameters['gravity'] + (cos_phi * cos_theta * T - fe_drag[2]) / m   # Acceleration along z-axis
 
-        dxdt[9] = ((Iyy - Izz) * q * r + tp * self.geom['arm_length']) / Ixx     # Angular acceleration along body x-axis: roll rate
-        dxdt[10] = ((Izz - Ixx) * p * r + tq * self.geom['arm_length']) / Iyy     # Angular acceleration along body y-axis: pitch rate
-        dxdt[11] = ((Ixx - Iyy) * p * q + tr * 1) / Izz     # Angular acceleration along body z-axis: yaw rate
+        dxdt[9] = ((Iyy - Izz) * q * r + tp * self.parameters['geom']['arm_length']) / Ixx  # Angular acceleration along body x-axis: roll rate
+        dxdt[10] = ((Izz - Ixx) * p * r + tq * self.parameters['geom']['arm_length']) / Iyy  # Angular acceleration along body y-axis: pitch rate
+        dxdt[11] = ((Ixx - Iyy) * p * q + tr * 1) / Izz  # Angular acceleration along body z-axis: yaw rate
         dxdt[12] = 1     # Auxiliary time variable
         dxdt[13] = (u['mission_complete'] - x['mission_complete']) / self.parameters['dt']    # Value to keep track of percentage of mission completed
 
@@ -281,9 +281,9 @@ def linear_model(self, phi, theta, psi, p, q, r, T):
         :param T:         N, scalar, double, thrust
         :return:          Linearized state transition matrix A, n_states x n_states, and linearized input matrix B, n_states x n_inputs
         """
-        m = self.mass['total']
-        Ixx, Iyy, Izz = self.mass['Ixx'], self.mass['Iyy'], self.mass['Izz']
-        length = self.geom['arm_length'] 
+        m = self.parameters['mass']['total']
+        Ixx, Iyy, Izz = self.parameters['mass']['Ixx'], self.parameters['mass']['Iyy'], self.parameters['mass']['Izz']
+        length = self.parameters['geom']['arm_length']
         sin_phi = np.sin(phi)
         cos_phi = np.cos(phi)
         sin_theta = np.sin(theta)
@@ -321,53 +321,59 @@ def linear_model(self, phi, theta, psi, p, q, r, T):
 
         return A, B
 
-    def visualize_traj(self, pred, ref):
+    def visualize_traj(self, pred, ref=None, prefix=''):
         """
         This method provides functionality to visualize a predicted trajectory generated, plotted with the reference trajectory. 
 
         Calling this returns a figure with two subplots: 1) x vs y, and 2) z vs time.
 
-        Parameters
+        Args
         ----------
-        pred : Vehicle model simulation 
-               SimulationResults from simulate_to or simulate_to_threshold for a defined SmallRotorcraft class
+        pred : SimulationResults
+              Results from vehicle model simulation from simulate_to or simulate_to_threshold for a defined SmallRotorcraft class
 
-        ref  : Reference trajectory 
-                dict with keys for each state in the vehicle model and corresponding values as numpy arrays 
+        Keyword Args
+        -------------
+        ref : dict[str, np.ndarray], optional
+              Reference trajectory - dict with keys for each state in the vehicle model and corresponding values as numpy arrays
+        prefix : str, optional
+              Prefix added to keys in predicted values. This is used to plot the trajectory using the results from a composite model
 
         Returns 
         -------
         fig : Visualization of trajectory generation results 
         """
 
-        # Extract reference trajectory information
-        time        = ref['t'].tolist() 
-        ref_x       = ref['x'].tolist() 
-        ref_y       = ref['y'].tolist() 
-        ref_z       = ref['z'].tolist()
-
         # Extract predicted trajectory information
         pred_time = pred.times
-        pred_x = [pred.outputs[iter]['x'] for iter in range(len(pred_time))]
-        pred_y = [pred.outputs[iter]['y'] for iter in range(len(pred_time))]
-        pred_z = [pred.outputs[iter]['z'] for iter in range(len(pred_time))]
+        pred_x = [pred.outputs[iter][prefix+'x'] for iter in range(len(pred_time))]
+        pred_y = [pred.outputs[iter][prefix+'y'] for iter in range(len(pred_time))]
+        pred_z = [pred.outputs[iter][prefix+'z'] for iter in range(len(pred_time))]
 
         # Initialize Figure
         params = dict(figsize=(13, 9), fontsize=14, linewidth=2.0, alpha_preds=0.6)
         fig, (ax1, ax2) = plt.subplots(2)
 
-        # Plot trajectory predictions
-        ax1.plot(ref_x, ref_y, '--', linewidth=params['linewidth'], color='tab:orange', alpha=0.5, label='reference trajectory')
+        # Handle reference information
+        if ref is not None:
+          # Extract reference trajectory information
+          time        = ref['t'].tolist() 
+          ref_x       = ref['x'].tolist() 
+          ref_y       = ref['y'].tolist() 
+          ref_z       = ref['z'].tolist()
+
+          # Plot reference trajectories
+          ax1.plot(ref_x, ref_y, '--', linewidth=params['linewidth'], color='tab:orange', alpha=0.5, label='reference trajectory')
+          ax2.plot(time, ref_z, '-', color='tab:orange', alpha=params['alpha_preds'], linewidth=params['linewidth'], label='reference trajectory')
+
+        # Plot predictions
         ax1.plot(pred_x, pred_y,'-', color='tab:blue', alpha=params['alpha_preds'], linewidth=params['linewidth'], label='prediction')
+        ax2.plot(pred_time, pred_z,'-', color='tab:blue',alpha=params['alpha_preds'], linewidth=params['linewidth'], label='prediction')
 
+        # Add labels
         ax1.set_xlabel('x', fontsize=params['fontsize'])
         ax1.set_ylabel('y', fontsize=params['fontsize'])
         ax1.legend(fontsize=params['fontsize'])
-
-        # Add altitude plot
-        ax2.plot(time, ref_z, '-', color='tab:orange', alpha=params['alpha_preds'], linewidth=params['linewidth'], label='reference trajectory')
-        ax2.plot(pred_time, pred_z,'-', color='tab:blue',alpha=params['alpha_preds'], linewidth=params['linewidth'], label='prediction')
-
         ax2.set_xlabel('time stamp, -', fontsize=params['fontsize'])
         ax2.set_ylabel('z', fontsize=params['fontsize'])