simulation/physics: making dxdt() method pure for integration and adding...

def thrustEqn(vi, *prop_params):

    R,A,rho,a,b,c,eta,theta0,theta1,u,v,w,Omega = prop_params

    # Calculate local airflow velocity at propeller with vi, V'

    Vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)

    # Calculate local airflow velocity at propeller with vi, v'

    vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)

    # Calculate Thrust averaged over one revolution of propeller using vi

    Thrust = 1/4 * rho * a * b * c * R * \

          (u**2 + v**2) * (theta0 + 1/2 * theta1) )

    # Calculate residual for equation: Thrust = mass flow rate * delta Velocity

    residual = eta * 2 * vi * rho * A * Vprime - Thrust

    residual = eta * 2 * vi * rho * A * vprime - Thrust

    return residual

def thrust(

    speed, airstream_velocity, R, A, rho, a, b, c, eta, theta0, theta1,

    vi_guess=0.1

    prop_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

    # omega = 2 * np.pi / 60 * prop_speed # Commented out - already in rad/s

    omega = prop_speed

    vi_guess = airstream_velocity[2]

    #Collect propeller config, state, and input parameters

    prop_params = (R,A,rho,a,b,c,eta,theta0,theta1,u,v,w,Omega)

    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 fsolve ('Untyped global name fsolve')

    vi = fsolve(thrustEqn, vi_guess, args=prop_params)

    vi = fsolve(thrustEqn, vi_guess, args=prop_params, xtol=1e-4)[0]

    # 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

    vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)

    thrust = eta * 2 * vi * rho * A * vprime

    return thrust

@njit

def torque(position_vector: np.ndarray, thrust: float) -> np.ndarray:

    thrust = np.asarray([0, 0, -thrust])

    return np.cross(position_vector, thrust)

def torque(

    position_vector: np.ndarray, force: np.ndarray,

    moment_of_inertia: float, prop_angular_acceleration: float

) -> np.ndarray:

    # Total moments in the body frame

    # yaw moments

    # tau = I . d omega/dt

    tau_rot = moment_of_inertia * prop_angular_acceleration

    # tau = r x F

    tau = np.cross(position_vector, force)

    # print(moment_of_inertia, prop_angular_acceleration)

    tau[2] = tau[2] + tau_rot

    return tau

    zI = x[11]      # In inertial frame, down is positive z

    # 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)

    cphi = np.cos(phi);   sphi = np.sin(phi)    # roll

    cthe = np.cos(theta); sthe = np.sin(theta)  # pitch

    cpsi = np.cos(psi);   spsi = np.sin(psi)    # yaw

    fx, fy, fz = forces

    tx, ty, tz = torques

    # Calculate the derivative of the state matrix using EOM

    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[0] = 1/mass * (fx) - g * sthe + r * vb - q * wb  # = udot

    xdot[1] = 1/mass * (fy) + 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]))

    # print(I.shape, x[3:6].shape, (I @ x[3:6]).shape)

    gyro = np.cross(x[3:6], I @ x[3:6])

    # gyro = np.cross(x[3:6], I @ x[3:6]).squeeze()

    xdot[3:6] = I_inv @ (torques - gyro) # TODO: verify whether it is + or - gyro

    xdot[6] = p + (q*sphi + r*cphi) * sthe / cthe  # = phidot

    xdot[7] = q * cphi - r * sphi  # = thetadot

from numpy import (cos, sin)

from scipy.integrate import odeint

from vehicle import PropellerParams, SimulationParams, VehicleParams

from coords import body_to_inertial, rotating_frame_derivative, angular_to_euler_rate

from vehicle import MotorParams, PropellerParams, SimulationParams, VehicleParams

from coords import body_to_inertial, direction_cosine_matrix, rotating_frame_derivative, angular_to_euler_rate

from physics import thrust, torque, apply_forces_torques

class Propeller:

    """

    Models the thrust and aerodynamics of the propeller blades spinning at a 

    certain rate.

    """

    def __init__(self, params: PropellerParams, simulation: SimulationParams) -> None:

        self.params: PropellerParams = params

        self.simulation: SimulationParams = simulation

        self.speed: float = None

        "Revolutions per minute"

        self._induced_velocity_guess = 0.1

        "Guess for the induced velocity about the propeller"

        self.speed: float = 0.

        "Radians per second"

    def reset(self):

        self.speed = 0.

        self._induced_velocity_guess = 0.1

    def step(self, u: float) -> float:

    def apply_speed(self, speed: float) -> float:

        return speed

    def step(self, speed: float) -> float:

        self.speed = speed

        return self.speed

    def thrust(self, speed=None, airstream_velocity: np.ndarray=np.zeros(3)) -> float:

    def thrust(self, speed, airstream_velocity: np.ndarray=np.zeros(3)) -> float:

        p = self.params

        speed = speed or self.speed

        return thrust(

            speed, airstream_velocity,

            p.R, p.A, self.simulation.rho, p.a, p.b, p.c, p.eta, p.theta0, p.theta1,

            self._induced_velocity_guess)

            p.R, p.A, self.simulation.rho, p.a, p.b, p.c, p.eta, p.theta0, p.theta1

        )

    @property

        return self.speed

    @property

    def rpm(self) -> float:

        return self.speed * 60. / (2. * np.pi)

class Motor:

    """

    Models the electronic and mechanical characterists of the rotation of the 

    propeller shaft.

    """

    def __init__(self, params: MotorParams, simulation: SimulationParams) -> None:

        self.params = params

        self.simulation = simulation

        self.speed: float = 0.

        "Radians per second"

    def reset(self) -> float:

        self.speed = 0.

    def step(self, u: float) -> float:

        pass

class Multirotor:

        self.state: np.ndarray = None

        self.propellers: List[Propeller] = None

        self.propeller_vectors: np.matrix = None

        self.t: float = None

        self.t: float = 0.

        self.reset()

    @property

    def navigation_velocity(self) -> np.ndarray:

        v_inertial = body_to_inertial(self.velocity)

        dcm = direction_cosine_matrix(*self.orientation)

        v_inertial = body_to_inertial(self.velocity, dcm)

        v_inertial[2] *= -1 # convert inertial to navigation frame (h = -z)

        return v_inertial

    @property

    def angular_velocity(self) -> np.ndarray:

    def angular_rate(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)

        return angular_to_euler_rate(self.angular_rate, self.orientation)

    @property

    def weight(self) -> float:

        return self.simulation.g * self.params.mass

    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]

    def get_forces_torques(self, speeds: np.ndarray, state: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:

        linear_vel_body = state[:3]

        angular_vel_body = state[3:6]

        airstream_velocity_inertial = rotating_frame_derivative(

            self.propeller_vectors,

            linear_vel_body,

        thrust_vec = np.zeros((3, len(self.propellers)))

        torque_vec = np.zeros_like(thrust_vec)

        for i, (speed, prop) in enumerate(zip(

        for i, (speed, prop, clockwise) in enumerate(zip(

                speeds,

                self.propellers)):

                self.propellers,

                self.params.clockwise)

        ):

            last_speed = prop.speed

            speed = prop.apply_speed(speed)

            angular_acc = (speed - last_speed) / self.simulation.dt

            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,

                        prop._induced_velocity_guess)

            # Store solution as initial guess for next step's solution

            torque_vec[:, i] = torque(self.propeller_vectors[:,i], -thrust_vec[2,i])

                prop.params.theta0, prop.params.theta1

            )

            torque_vec[:, i] = torque(

                self.propeller_vectors[:,i], thrust_vec[:,i],

                prop.params.moment_of_inertia, clockwise * angular_acc

            )

            # print(clockwise * angular_acc)

        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)

        # This method must not have any side-effects. It should not change the

        # state of the vehicle. This method is called multiple times from the 

        # same state by the odeint() function, and the results should be consistent.

        forces, torques = self.get_forces_torques(u, x)

        xdot = apply_forces_torques(

            forces, torques, self.state, self.simulation.g,

            forces, torques, x, self.simulation.g,

            self.params.mass, self.params.inertia_matrix, self.params.inertia_matrix_inverse)

        return xdot

    def step(self, u: np.ndarray):

        self.t += self.simulation.dt

        self.state = odeint(self.dxdt, self.state, (0, self.simulation.dt), args=(u,))[-1]

        for u_, prop in zip(u, self.propellers):

            prop.step(u_)

        return self.state

    def allocate_control(self, thrust: float, torques: np.ndarray) -> np.ndarray:

        # TODO: njit it? np.linalg.lstsq can be compiled

        vec = np.asarray([-thrust, *torques])

        # return np.sqrt(np.linalg.lstsq(self.params.alloc, vec, rcond=None)[0])

        return np.sqrt(

            np.clip(self.params.alloc_inverse @ vec, a_min=0., a_max=None)

        )