Commit 56aaaae2 authored by Ibrahim's avatar Ibrahim
Browse files

Used numba to improve Multirotor simulation step() ~3x from ~2.78ms to ~0.89ms.

parent 2f041c45
__pycache__/
.ipynb_checkpoints/
.vscode/
import numpy as np
from numpy import sin, cos, tan, sec
from numpy import sin, cos, tan
def body_to_inertial(coords: np.ndarray, *, rotations=None, dcm=None, dcm_inverse=None) -> np.ndarray:
......@@ -58,7 +58,7 @@ def angular_to_euler_rate(angular_velocity: np.ndarray, orientation: np.ndarray)
p, q, r = angular_velocity
roll_rate = p + q * tan(pitch) * (q * sin(roll) + r * cos(roll))
pitch_rate = q * cos(roll) - r * sin(roll)
yaw_rate = sec(pitch) * (q * sin(roll) + r * cos(roll))
yaw_rate = (q * sin(roll) + r * cos(roll)) / cos(pitch)
return np.asarray([roll_rate, pitch_rate, yaw_rate])
......
from typing import Iterable
import numpy as np
from vehicle import PropellerParams, VehicleParams
from simulation import Propeller
def moment_of_inertia_tensor_from_cooords(point_masses: Iterable[float], coords: Iterable[np.ndarray]) -> np.matrix:
coords = np.asarray(coords)
masses = np.asarray(point_masses)
x,y,z = coords[:,0], coords[:,1], coords[:2]
Ixx = np.sum(masses * (y**2 + z**2))
Iyy = np.sum(masses * (x**2 + z**2))
Izz = np.sum(masses * (x**2 + y**2))
Ixy = Iyx = -np.sum(x * y * masses)
Iyz = Izy = -np.sum(y * z * masses)
Ixz = Izx = -np.sum(x * z * masses)
return np.asmatrix([
[Ixx, Ixy, Ixz],
[Iyx, Iyy, Iyz],
[Izx, Izy, Izz]
])
def vehicle_params_factory(n: int, m_prop: float, d_prop: float, params: PropellerParams,
m_body: float, body_shape: str='point', body_size: float=0.1):
angles = np.linspace(0, 2 * np.pi, num=n, endpoint=False)
masses = [m_prop] * n
x = np.cos(angles) * d_prop
y = np.sin(angles) * d_prop
z = np.zeros_like(y)
coords = np.vstack((x, y, z)).T
I_prop = moment_of_inertia_tensor_from_cooords(masses, coords)
if body_shape == 'point':
I_body = np.asmatrix(np.zeros((3,3)))
elif body_shape == 'sphere_solid':
I_body = np.asmatrix(np.eye(3) * 0.4 * m_body * body_size**2)
elif body_shape == 'sphere_shell':
I_body = np.asmatrix(np.eye(3) * 2 * m_body * body_size**2 / 3)
elif body_shape == 'cube':
I_body = np.asmatrix(np.eye(3) * 2 * m_body * body_size**2 / 12)
I = I_body + I_prop
return VehicleParams(
propellers = [Propeller(params) for _ in range(n)],
angles = angles,
distances=np.ones(n) * d_prop,
mass = n * m_prop + m_body,
inertia_matrix = I
)
from typing import List
from typing import Iterable
import numpy as np
from numpy import (cos, sin)
from numba import njit, jit
from scipy.optimize import fsolve
from vehicle import PropellerParams, SimulationParams, VehicleParams
from coords import rotating_frame_derivative
@njit
def thrustEqn(vi, *prop_params):
# Unpack parameters
......@@ -28,196 +26,36 @@ def thrustEqn(vi, *prop_params):
class Propeller:
def __init__(self, params: PropellerParams) -> None:
self.params: PropellerParams = params
self.speed: float = None
"Revolutions per minute"
self._last_induced_velocity = 0.
"Guess for the induced velocity about the propeller"
self.theta0 = 2*np.arctan2(self.params.p_pitch, (2 * np.pi * 3/4 * self.params.p_diameter/2))
self.theta1 = -4 / 3 * np.arctan2(self.params.p_pitch, 2 * np.pi * 3/4 * self.params.p_diameter/2)
def step(self, u: float, dt: float) -> float:
return self.apply_speed(u)
def apply_speed(self, speed: float) -> float:
self.speed = speed
return self.speed
def thrust(self, airstream_velocity: np.ndarray=(0,0,0)) -> float:
u, v, w = airstream_velocity
# Inputs: Current state x[k], Commanded Propeller RPM inputs u[k],
# Propeller location distances dx, dy (m)
# Returns: Thrust vector for 4 propellers (Newtons)
# Convert commanded RPM to rad/s
Omega = 2 * np.pi / 60 * self.speed
#Collect propeller config, state, and input parameters
prop_params = (self.params.R,self.params.A,self.params.rho,self.params.a,self.params.b,self.params.c,self.params.eta,self.theta0,self.theta1,u,v,w,Omega)
# Numerically solve for propeller induced velocity, vi
# using nonlinear root finder, fsolve, and prop_params
vi = fsolve(thrustEqn, self._last_induced_velocity, args=prop_params)
self._last_induced_velocity = vi
# Plug vi back into Thrust equation to solve for T
Vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)
Thrust = self.params.eta * 2 * vi * self.params.rho * self.params.A * Vprime
return Thrust
def torque(self, position_vector: np.ndarray, thrust: float) -> float:
thrust = np.asarray([0, 0, -thrust])
return np.cross(position_vector, thrust)
@property
def state(self) -> float:
return self.speed
class Multirotor:
def __init__(self, vehicle: VehicleParams, simulation: SimulationParams) -> None:
self.vehicle: VehicleParams = vehicle
self.simulation: SimulationParams = simulation
self.state: np.ndarray = None
self.propellers: List[Propeller] = None
self.propeller_vectors: np.matrix = None
def reset(self):
self.propellers = []
for params in self.vehicle.propellers:
self.propellers.append(Propeller(params))
x = cos(self.vehicle.angles) * self.vehicle.distances
y = sin(self.vehicle.angles) * self.vehicle.distances
z = np.zeros_like(y)
self.propeller_vectors = np.asmatrix(np.vstack((x, y, z)))
self.inertial_matrix_inverse = np.asmatrix(np.linalg.inv(self.vehicle.inertia_matrix))
self.state = np.zeros(12)
@property
def position(self) -> np.ndarray:
"""Navigation coordinates (height = - z coordinate)"""
return np.asfarray(self.state[9], self.state[10], -self.state[11])
@property
def velocity(self) -> np.ndarray:
"""Body-frame velocity"""
return self.state[:3]
@property
def orientation(self) -> np.ndarray:
"""Euler rotations (roll, pitch, yaw)"""
return self.state[6:9]
@property
def angular_velocity(self) -> np.ndarray:
"""Angular rate of body frame axes (not same as rate of roll, pitch, yaw)"""
return self.state[3:6]
def apply_forces_torques(self, forces: np.ndarray, torques: np.ndarray):
# Store state variables in a readable format
x = self.state
ub = x[0]
vb = x[1]
wb = x[2]
p = x[3]
q = x[4]
r = x[5]
phi = x[6]
theta = x[7]
psi = x[8]
xI = x[9]
yI = x[10]
hI = x[11]
# Pre-calculate trig values
cphi = np.cos(phi); sphi = np.sin(phi)
cthe = np.cos(theta); sthe = np.sin(theta)
cpsi = np.cos(psi); spsi = np.sin(psi)
fx, fy, fz = forces
tx, ty, tz = torques
I = self.vehicle.inertia_matrix
I_inv = self.inertial_matrix_inverse
# Calculate the derivative of the state matrix using EOM
xdot = np.zeros_like(self.state)
xdot[0] = -self.simulation.g * sthe + r * vb - q * wb # = udot
xdot[1] = self.simulation.g * sphi * cthe - r * ub + p * wb # = vdot
xdot[2] = 1/self.vehicle.mass * (fz) + self.simulation.g * cphi * cthe + q * ub - p * vb # = wdot
# xdot[3] = 1/Ixx * (tx + (Iyy - Izz) * q * r) # = pdot
# xdot[4] = 1/Iyy * (ty + (Izz - Ixx) * p * r) # = qdot
# xdot[5] = 1/Izz * (tz + (Ixx - Iyy) * p * q) # = rdot
xdot[3:6] = I_inv @ (torques - np.cross(x[3:6], I @ x[3:6]))
xdot[6] = p + (q*sphi + r*cphi) * sthe / cthe # = phidot
xdot[7] = q * cphi - r * sphi # = thetadot
xdot[8] = (q * sphi + r * cphi) / cthe # = psidot
xdot[9] = cthe*cpsi*ub + (-cphi * spsi + sphi*sthe*cpsi) * vb + \
(sphi*spsi+cphi*sthe*cpsi) * wb # = xIdot
xdot[10] = cthe*spsi * ub + (cphi*cpsi+sphi*sthe*spsi) * vb + \
(-sphi*cpsi+cphi*sthe*spsi) * wb # = yIdot
xdot[11] = (-sthe * ub + sphi*cthe * vb + cphi*cthe * wb) # = zIdot
return xdot
def thrust(speed, airstream_velocity, R, A, rho, a, b, c, eta, theta0, theta1) -> float:
u, v, w = airstream_velocity
# Convert commanded RPM to rad/s
Omega = 2 * np.pi / 60 * speed
#Collect propeller config, state, and input parameters
prop_params = (R,A,rho,a,b,c,eta,theta0,theta1,u,v,w,Omega)
# Numerically solve for propeller induced velocity, vi
# using nonlinear root finder, fsolve, and prop_params
# TODO: numba jit gives error for this function ('Untyped global name fsolve')
vi = fsolve(thrustEqn, 0.1, args=prop_params)
# Plug vi back into Thrust equation to solve for T
Vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)
Thrust = eta * 2 * vi * rho * A * Vprime
return Thrust
def apply_propeller_speeds(self, speeds: np.ndarray):
linear_vel_body = self.state[:3]
angular_vel_body = self.state[3:6]
airstream_velocity_inertial = rotating_frame_derivative(
self.propeller_vectors,
linear_vel_body,
angular_vel_body)
thrust = np.zeros(3, len(self.propellers))
torque = np.zeros_like(thrust)
for i, speed, prop in enumerate(zip(
speeds,
self.propellers)):
speed = prop.apply_speed(speed)
thrust[2, i] = -prop.thrust(airstream_velocity_inertial[:, i])
torque[:, i] = prop.torque(self.propeller_vectors[:,i], -thrust[2:,i])
forces = thrust.sum(axis=1)
torques = np.cross(self.propeller_vectors, thrust)
return self.apply_forces_torques(forces, torques)
def step(self):
pass
@njit
def torque(position_vector: np.ndarray, thrust: float) -> np.ndarray:
thrust = np.asarray([0, 0, -thrust])
return np.cross(position_vector, thrust)
def stateDerivative(x,u,physics=None,geometry=None):
# Inputs: state vector (x), input vector (u)
# Returns: time derivative of state vector (xdot)
# State Vector Reference:
#idx 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11
#x = [u, v, w, p, q, r, phi, the, psi, xE, yE, hE]
@njit
def apply_forces_torques(forces: np.ndarray, torques: np.ndarray, x: np.ndarray,
g: float, mass: float, inertia_matrix: np.matrix, inertia_matrix_inverse: np.matrix):
# Store state variables in a readable format
ub = x[0]
vb = x[1]
......@@ -228,44 +66,43 @@ def stateDerivative(x,u,physics=None,geometry=None):
phi = x[6]
theta = x[7]
psi = x[8]
xE = x[9]
yE = x[10]
hE = x[11]
# Calculate forces and moments from propeller inputs (u)
F1 = fthrust(x, u[0], dx, dy)
F2 = fthrust(x, u[1], -dx, -dy)
F3 = fthrust(x, u[2], dx, -dy)
F4 = fthrust(x, u[3], -dx, dy)
Fz = F1 + F2 + F3 + F4
L = (F2 + F3) * dy - (F1 + F4) * dy
M = (F1 + F3) * dx - (F2 + F4) * dx
N = -T(F1,dx,dy) - T(F2,dx,dy) + T(F3,dx,dy) + T(F4,dx,dy)
xI = x[9]
yI = x[10]
hI = x[11]
# Pre-calculate trig values
cphi = np.cos(phi); sphi = np.sin(phi)
cthe = np.cos(theta); sthe = np.sin(theta)
cpsi = np.cos(psi); spsi = np.sin(psi)
fx, fy, fz = forces
tx, ty, tz = torques
I = inertia_matrix
I_inv = inertia_matrix_inverse
# Calculate the derivative of the state matrix using EOM
xdot = np.zeros(12)
xdot = np.zeros_like(x)
xdot[0] = -g * sthe + r * vb - q * wb # = udot
xdot[1] = g * sphi*cthe - r * ub + p * wb # = vdot
xdot[2] = 1/m * (-Fz) + g*cphi*cthe + q * ub - p * vb # = wdot
xdot[3] = 1/Ixx * (L + (Iyy - Izz) * q * r) # = pdot
xdot[4] = 1/Iyy * (M + (Izz - Ixx) * p * r) # = qdot
xdot[5] = 1/Izz * (N + (Ixx - Iyy) * p * q) # = rdot
xdot[1] = g * sphi * cthe - r * ub + p * wb # = vdot
xdot[2] = 1/mass * (fz) + g * cphi * cthe + q * ub - p * vb # = wdot
# xdot[3] = 1/Ixx * (tx + (Iyy - Izz) * q * r) # = pdot
# xdot[4] = 1/Iyy * (ty + (Izz - Ixx) * p * r) # = qdot
# xdot[5] = 1/Izz * (tz + (Ixx - Iyy) * p * q) # = rdot
xdot[3:6] = I_inv @ (torques - np.cross(x[3:6], I @ x[3:6]))
xdot[6] = p + (q*sphi + r*cphi) * sthe / cthe # = phidot
xdot[7] = q * cphi - r * sphi # = thetadot
xdot[8] = (q * sphi + r * cphi) / cthe # = psidot
xdot[9] = cthe*cpsi*ub + (-cphi*spsi + sphi*sthe*cpsi) * vb + \
(sphi*spsi+cphi*sthe*cpsi) * wb # = xEdot
xdot[9] = cthe*cpsi*ub + (-cphi * spsi + sphi*sthe*cpsi) * vb + \
(sphi*spsi+cphi*sthe*cpsi) * wb # = xIdot
xdot[10] = cthe*spsi * ub + (cphi*cpsi+sphi*sthe*spsi) * vb + \
(-sphi*cpsi+cphi*sthe*spsi) * wb # = yEdot
(-sphi*cpsi+cphi*sthe*spsi) * wb # = yIdot
xdot[11] = -1*(-sthe * ub + sphi*cthe * vb + cphi*cthe * wb) # = hEdot
xdot[11] = (-sthe * ub + sphi*cthe * vb + cphi*cthe * wb) # = zIdot
return xdot
from typing import List, Tuple
import numpy as np
from numpy import (cos, sin)
from scipy.optimize import fsolve
from scipy.integrate import odeint
from vehicle import PropellerParams, SimulationParams, VehicleParams
from coords import rotating_frame_derivative, angular_to_euler_rate
from physics import thrust, torque, apply_forces_torques
class Propeller:
def __init__(self, params: PropellerParams, simulation: SimulationParams) -> None:
self.params: PropellerParams = params
self.simulation: SimulationParams = simulation
self.speed: float = None
"Revolutions per minute"
self._last_induced_velocity = 0.1
"Guess for the induced velocity about the propeller"
def reset(self):
self.speed = 0.
self._last_induced_velocity = 0.1
def step(self, u: float, dt: float) -> float:
return self.apply_speed(u)
def apply_speed(self, speed: float) -> float:
self.speed = speed
return self.speed
@property
def state(self) -> float:
return self.speed
class Multirotor:
def __init__(self, params: VehicleParams, simulation: SimulationParams) -> None:
self.params: VehicleParams = params
self.simulation: SimulationParams = simulation
self.state: np.ndarray = None
self.propellers: List[Propeller] = None
self.propeller_vectors: np.matrix = None
def reset(self):
self.propellers = []
for params in self.params.propellers:
self.propellers.append(Propeller(params, self.simulation))
x = cos(self.params.angles) * self.params.distances
y = sin(self.params.angles) * self.params.distances
z = np.zeros_like(y)
self.propeller_vectors = np.vstack((x, y, z))
self.inertial_matrix_inverse = np.asmatrix(np.linalg.inv(self.params.inertia_matrix))
self.state = np.zeros(12)
@property
def position(self) -> np.ndarray:
"""Navigation coordinates (height = - z coordinate)"""
return np.asfarray(self.state[9], self.state[10], -self.state[11])
@property
def velocity(self) -> np.ndarray:
"""Body-frame velocity"""
return self.state[:3]
@property
def orientation(self) -> np.ndarray:
"""Euler rotations (roll, pitch, yaw)"""
return self.state[6:9]
@property
def angular_velocity(self) -> np.ndarray:
"""Angular rate of body frame axes (not same as rate of roll, pitch, yaw)"""
return self.state[3:6]
@property
def euler_rate(self) -> np.ndarray:
"""Euler rate of vehicle d(roll, pitch, yaw)/dt"""
return angular_to_euler_rate(self.angular_velocity, self.orientation)
def get_forces_torques(self, speeds: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
linear_vel_body = self.state[:3]
angular_vel_body = self.state[3:6]
airstream_velocity_inertial = rotating_frame_derivative(
self.propeller_vectors,
linear_vel_body,
angular_vel_body)
thrust_vec = np.zeros((3, len(self.propellers)))
torque_vec = np.zeros_like(thrust_vec)
for i, (speed, prop) in enumerate(zip(
speeds,
self.propellers)):
speed = prop.apply_speed(speed)
thrust_vec[2, i] = -thrust(speed,airstream_velocity_inertial[:, i],prop.params.R, prop.params.A, self.simulation.rho, prop.params.a, prop.params.b, prop.params.c, prop.params.eta, prop.params.theta0, prop.params.theta1)
torque_vec[:, i] = torque(self.propeller_vectors[:,i], -thrust_vec[2,i])
forces = thrust_vec.sum(axis=1)
torques = torque_vec.sum(axis=1)
return forces, torques
def dxdt(self, x: np.ndarray, t: float, u: np.ndarray):
forces, torques = self.get_forces_torques(u)
xdot = apply_forces_torques(forces, torques, self.state, self.simulation.g, self.params.mass, self.params.inertia_matrix, self.params.inertia_matrix_inverse)
return xdot
def step(self, u: np.ndarray, jit=False):
self.state = odeint(self.dxdt, self.state, (0, self.simulation.dt), args=(u,))[-1]
return self.state
from typing import List
from typing import Iterable, List, Union
from dataclasses import dataclass
import numpy as np
@dataclass
class PropellerParams:
R: float = 0.0762 # propeller length/ disk radius (m)
A: float = np.pi * R ** 2
rho: float = 1.225 #kg/m^3 at MSL
a: float = 5.7 # Lift curve slope used in example in Stevens & Lewis
b: float = 2 # number of blades
c: float = 0.0274 # mean chord length (m)
eta: float = 1 # propeller efficiency
# Manufacturer propeller length x pitch specification:
p_diameter: float = 6 #inches
p_pitch: float = 3 #inches
diameter: float = 6 #inches
pitch: float = 3 #inches
a: float = 5.7
"Lift curve slope used in example in Stevens & Lewis"
b: float = 2
"Number of blades"
c: float = 0.0274
"Mean chord length (m)"
eta: float = 1.
"Propeller efficiency"
def __post_init__(self):
self.R = self.diameter * 0.0254
self.A = np.pi * self.R**2
self.theta0 = 2*np.arctan2(self.pitch, (2 * np.pi * 3/4 * self.diameter/2))
self.theta1 = -4 / 3 * np.arctan2(self.pitch, 2 * np.pi * 3/4 * self.diameter/2)
......@@ -31,6 +38,10 @@ class VehicleParams:
inertia_matrix: np.matrix = np.eye(3)
def __post_init__(self):
self.inertia_matrix_inverse = np.linalg.inv(self.inertia_matrix)
@dataclass
class SimulationParams:
......@@ -39,3 +50,5 @@ class SimulationParams:
"""Timestep of simulation"""
g: float = 9.81
"""Gravitational acceleration"""
rho: float = 1.225
"Air density kg/m^3 at MSL"
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment