Significant speed up by vectorizing wrench computation.

rotorpy/vehicles/multirotor.py

@@ -4,6 +4,8 @@ import scipy.integrate
 from scipy.spatial.transform import Rotation
 from rotorpy.vehicles.hummingbird_params import quad_params
 
+import time
+
 """
 Multirotor models
 """
@@ -46,6 +48,7 @@ class Multirotor(object):
                                 'cmd_ctbm': the controller commands a collective thrust and moments on the x/y/z body axes
                                 'cmd_ctatt': the controller commands a collective thrust and attitude (as a quaternion).
                                 'cmd_vel': the controller commands a velocity vector in the body frame. 
+        aero: boolean, determines whether or not aerodynamic drag forces are computed. 
     """
     def __init__(self, quad_params, initial_state = {'x': np.array([0,0,0]),
                                             'v': np.zeros(3,),
@@ -54,6 +57,7 @@ class Multirotor(object):
                                             'wind': np.array([0,0,0]),  # Since wind is handled elsewhere, this value is overwritten
                                             'rotor_speeds': np.array([1788.53, 1788.53, 1788.53, 1788.53])},
                        control_abstraction='cmd_motor_speeds',
+                       aero = True,  
                 ):
         """
         Initialize quadrotor physical parameters.
@@ -127,6 +131,8 @@ class Multirotor(object):
         self.kp_att = 544           # The attitude P gain (for cmd_vel and cmd_ctatt)
         self.kd_att = 46.64         # The attitude D gain (for cmd_vel and cmd_ctatt)
 
+        self.aero = aero
+
     def extract_geometry(self):
         """
         Extracts the geometry in self.rotors for efficient use later on in the computation of 
@@ -210,9 +216,6 @@ class Multirotor(object):
 
         R = Rotation.from_quat(state['q']).as_matrix()
 
-        # Compute airspeed vector in the body frame
-        body_airspeed_vector = R.T@(inertial_velocity - wind_velocity)
-
         # Rotor speed derivative
         rotor_accel = (1/self.tau_m)*(cmd_rotor_speeds - rotor_speeds)
 
@@ -222,8 +225,11 @@ class Multirotor(object):
         # Orientation derivative.
         q_dot = quat_dot(state['q'], state['w'])
 
+        # Compute airspeed vector in the body frame
+        body_airspeed_vector = R.T@(inertial_velocity - wind_velocity)
+
         # Compute total wrench in the body frame based on the current rotor speeds and their location w.r.t. CoM
-        (FtotB, Mtot) = self.compute_body_wrench(state['w'], rotor_speeds, body_airspeed_vector)
+        (FtotB, MtotB) = self.compute_body_wrench(state['w'], rotor_speeds, body_airspeed_vector)
 
         # Rotate the force from the body frame to the inertial frame
         Ftot = R@FtotB
@@ -234,7 +240,7 @@ class Multirotor(object):
         # Angular velocity derivative.
         w = state['w']
         w_hat = Multirotor.hat_map(w)
-        w_dot = self.inv_inertia @ (Mtot - w_hat @ (self.inertia @ w))
+        w_dot = self.inv_inertia @ (MtotB - w_hat @ (self.inertia @ w))
 
         # NOTE: the wind dynamics are currently handled in the wind_profile object. 
         # The line below doesn't do anything, as the wind state is assigned elsewhere. 
@@ -260,32 +266,32 @@ class Multirotor(object):
         The net moment Mtot is represented in the body frame. 
         """
 
-        FtotB = np.zeros((3,))
+        # Get the local airspeeds for each rotor
-        MtotB = np.zeros((3,))
+        local_airspeeds = body_airspeed_vector[:, np.newaxis] + Multirotor.hat_map(body_rates)@(self.rotor_geometry.T) 
 
-        for i in range(self.num_rotors):
+        # Compute the thrust of each rotor, assuming that the rotors all point in the body z direction!
-            # Loop through each rotor, compute the forces 
+        T = np.array([0, 0, self.k_eta])[:, np.newaxis]*rotor_speeds**2
         
-            r = self.rotor_geometry[i,:] # the position of rotor i relative to the CoM, in body coordinates
+        # Add in aero wrenches (if applicable)
-            
+        if self.aero:
-            # Compute the local airspeed by adding on the rotational component to the airspeed.
+            # Parasitic drag force acting at the CoM
-            local_airspeed_vector = body_airspeed_vector + Multirotor.hat_map(body_rates)@r
-
-            T = np.array([0, 0, self.k_eta*rotor_speeds[i]**2])             # thrust vector in body frame
-            H = -rotor_speeds[i]*self.rotor_drag_matrix@local_airspeed_vector     # rotor drag force
-
-            # Compute the moments
-            M_force = Multirotor.hat_map(r)@(T+H)
-            M_yaw = self.rotor_dir[i]*np.array([0, 0, self.k_m*rotor_speeds[i]**2]) 
-            M_flap = -rotor_speeds[i]*self.k_flap*Multirotor.hat_map(local_airspeed_vector)@np.array([0,0,1])
-
-            FtotB += (T+H)
-            MtotB += (M_force + M_yaw + M_flap)
-
-        # Compute the drag force acting on the frame
             D = -Multirotor._norm(body_airspeed_vector)*self.drag_matrix@body_airspeed_vector
+            # Rotor drag (aka H force) acting at each propeller hub.
+            H = -rotor_speeds*(self.rotor_drag_matrix@local_airspeeds)
+            # Pitching flapping moment acting at each propeller hub.
+            M_flap = -self.k_flap*rotor_speeds*((Multirotor.hat_map(local_airspeeds.T).transpose(2, 0, 1))@np.array([0,0,1])).T
+        else:
+            D = np.zeros(3,)
+            H = np.zeros((3,self.num_rotors))
+            M_flap = np.zeros((3,self.num_rotors))
 
-        FtotB += D
+        # Compute the moments due to the rotor thrusts, rotor drag (if applicable), and rotor drag torques
+        M_force = -np.einsum('ijk, ik->j', Multirotor.hat_map(self.rotor_geometry), T+H)
+        M_yaw = self.rotor_dir*(np.array([0, 0, self.k_m])[:, np.newaxis]*rotor_speeds**2)
+
+        # Sum all elements to compute the total body wrench
+        FtotB = np.sum(T + H, axis=1) + D
+        MtotB = M_force + np.sum(M_yaw + M_flap, axis=1)
 
         return (FtotB, MtotB)
 
@@ -396,7 +402,13 @@ class Multirotor(object):
     def hat_map(cls, s):
         """
         Given vector s in R^3, return associate skew symmetric matrix S in R^3x3
+        In the vectorized implementation, we assume that s is in the shape (N arrays, 3)
         """
+        if len(s.shape) > 1:  # Vectorized implementation
+            return np.array([[ np.zeros(s.shape[0]), -s[:,2],  s[:,1]],
+                             [ s[:,2],     np.zeros(s.shape[0]), -s[:,0]],
+                             [-s[:,1],  s[:,0],     np.zeros(s.shape[0])]])
+        else:
             return np.array([[    0, -s[2],  s[1]],
                              [ s[2],     0, -s[0]],
                              [-s[1],  s[0],     0]])