Significant speed up by vectorizing wrench computation.

This commit is contained in:
spencerfolk
2024-01-03 15:30:03 -05:00
parent 128d500953
commit bd0f7182cc

View File

@@ -4,6 +4,8 @@ import scipy.integrate
from scipy.spatial.transform import Rotation from scipy.spatial.transform import Rotation
from rotorpy.vehicles.hummingbird_params import quad_params from rotorpy.vehicles.hummingbird_params import quad_params
import time
""" """
Multirotor models 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_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_ctatt': the controller commands a collective thrust and attitude (as a quaternion).
'cmd_vel': the controller commands a velocity vector in the body frame. '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]), def __init__(self, quad_params, initial_state = {'x': np.array([0,0,0]),
'v': np.zeros(3,), '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 '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])}, 'rotor_speeds': np.array([1788.53, 1788.53, 1788.53, 1788.53])},
control_abstraction='cmd_motor_speeds', control_abstraction='cmd_motor_speeds',
aero = True,
): ):
""" """
Initialize quadrotor physical parameters. 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.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.kd_att = 46.64 # The attitude D gain (for cmd_vel and cmd_ctatt)
self.aero = aero
def extract_geometry(self): def extract_geometry(self):
""" """
Extracts the geometry in self.rotors for efficient use later on in the computation of 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() 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 speed derivative
rotor_accel = (1/self.tau_m)*(cmd_rotor_speeds - rotor_speeds) rotor_accel = (1/self.tau_m)*(cmd_rotor_speeds - rotor_speeds)
@@ -222,8 +225,11 @@ class Multirotor(object):
# Orientation derivative. # Orientation derivative.
q_dot = quat_dot(state['q'], state['w']) 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 # 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 # Rotate the force from the body frame to the inertial frame
Ftot = R@FtotB Ftot = R@FtotB
@@ -234,7 +240,7 @@ class Multirotor(object):
# Angular velocity derivative. # Angular velocity derivative.
w = state['w'] w = state['w']
w_hat = Multirotor.hat_map(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. # 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. # 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. 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
# Add in aero wrenches (if applicable)
if self.aero:
# Parasitic drag force acting at the CoM
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))
r = self.rotor_geometry[i,:] # the position of rotor i relative to the CoM, in body coordinates # 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)
# Compute the local airspeed by adding on the rotational component to the airspeed. M_yaw = self.rotor_dir*(np.array([0, 0, self.k_m])[:, np.newaxis]*rotor_speeds**2)
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 # Sum all elements to compute the total body wrench
H = -rotor_speeds[i]*self.rotor_drag_matrix@local_airspeed_vector # rotor drag force FtotB = np.sum(T + H, axis=1) + D
MtotB = M_force + np.sum(M_yaw + M_flap, axis=1)
# 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
FtotB += D
return (FtotB, MtotB) return (FtotB, MtotB)
@@ -396,10 +402,16 @@ class Multirotor(object):
def hat_map(cls, s): def hat_map(cls, s):
""" """
Given vector s in R^3, return associate skew symmetric matrix S in R^3x3 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)
""" """
return np.array([[ 0, -s[2], s[1]], if len(s.shape) > 1: # Vectorized implementation
[ s[2], 0, -s[0]], return np.array([[ np.zeros(s.shape[0]), -s[:,2], s[:,1]],
[-s[1], s[0], 0]]) [ 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]])
@classmethod @classmethod
def _pack_state(cls, state): def _pack_state(cls, state):