Commit 30f2da46 authored by Ibrahim's avatar Ibrahim
Browse files

Visualization using FuncAnimation(); 

helpers: bugfix in vehicle_params_factory where Propeller() class was instantiated instead of PropellerParams().
simulation: time attribute for tracking simulation progress.
vehicle: added some placeholder constants for torque/thrust calculations.
visualize: uses FuncAnimation() and optionally draws body axes on vehicle.
parent e4a103f6

README.md

+3 −0
Original line number Diff line number Diff line 
# multirotor



Simulation of multi-rotor unmanned aerial vehicles in python.
 
 No newline at end of file

helpers.py

+3 −4
Original line number Diff line number Diff line
@@ -4,7 +4,6 @@ import numpy as np 
from scipy.optimize import fsolve



from vehicle import PropellerParams, VehicleParams

from simulation import Propeller
@@ -13,7 +12,7 @@ def moment_of_inertia_tensor_from_cooords( 
) -> np.matrix:

    coords = np.asarray(coords)

    masses = np.asarray(point_masses)

    x,y,z = coords[:,0], coords[:,1], coords[:2]

    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))
@@ -53,11 +52,11 @@ def vehicle_params_factory( 
    I = I_body + I_prop



    return VehicleParams(

        propellers = [Propeller(params) for _ in range(n)],

        propellers = [params for _ in range(n)],

        angles = angles,

        distances=np.ones(n) * d_prop,

        mass = n * m_prop + m_body,

        inertia_matrix = I

        inertia_matrix = I,

    )

physics.py

+13 −13
Original line number Diff line number Diff line 
from typing import Iterable, Tuple



import numpy as np

from numba import njit, jit

from numba import njit

from scipy.optimize import fsolve
@@ -38,7 +38,7 @@ def thrust( 
    # 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, 0.1, args=prop_params)

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

    

    # Plug vi back into Thrust equation to solve for T

    Vprime = np.sqrt(u**2 + v**2 + (w - vi)**2)
@@ -60,18 +60,18 @@ def apply_forces_torques( 
    inertia_matrix: np.matrix, inertia_matrix_inverse: np.matrix

) -> np.ndarray:

    # Store state variables in a readable format

    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]

    ub = x[0]       # linear velocity along body-frame-x-axis

    vb = x[1]       # linear velocity along body-frame-y-axis

    wb = x[2]       # linear velocity along body-frame-z-axis (down is positive)

    p = x[3]        # body-frame-x-axis rotation rate

    q = x[4]        # body-frame-y-axis rotation rate

    r = x[5]        # body-frame-z-axis rotation rate

    phi = x[6]      # Roll

    theta = x[7]    # Pitch

    psi = x[8]      # Yaw

    xI = x[9]       # Inertial frame positions

    yI = x[10]

    hI = x[11]

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

    

    # Pre-calculate trig values

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

simulation.py

+3 −0
Original line number Diff line number Diff line
@@ -58,10 +58,12 @@ class Multirotor: 
        self.state: np.ndarray = None

        self.propellers: List[Propeller] = None

        self.propeller_vectors: np.matrix = None

        self.t: float = None

        self.reset()





    def reset(self):

        self.t = 0.

        self.propellers = []

        for params in self.params.propellers:

            self.propellers.append(Propeller(params, self.simulation))
@@ -152,5 +154,6 @@ class Multirotor: 




    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]

        return self.state

visualize.py

+95 −36
Original line number Diff line number Diff line
@@ -5,10 +5,11 @@ import time 


import numpy as np

import matplotlib.pyplot as plt

from matplotlib.animation import FuncAnimation

from mpl_toolkits.mplot3d import Axes3D

from mpl_toolkits.mplot3d.art3d import Line3D



from coords import direction_cosine_matrix, inertial_to_body

from coords import body_to_inertial, direction_cosine_matrix

from vehicle import VehicleParams

from simulation import Multirotor
@@ -20,28 +21,30 @@ class VehicleDrawing: 
    axis: Axes3D = None

    max_frames_per_second: float = 30.

    trace: bool = False

    body_axes: bool = False



    def __post_init__(self):

        self.t = self.vehicle.t

        self.params = self.vehicle.params

        self.lines, self.body_frame_segments, self.trajectory_line = make_drawing(self.params)

        self.interval = 1 / self.max_frames_per_second

        self.is_ipython = is_ipython()

        if self.trace:

        self.arm_lines, self.arm_lines_points, \

        self.trajectory_line, \

        self.axis_lines, self.axis_lines_points = \

            make_drawing(self, self.body_axes)

        self.trajectory = [[], [], []] # [[X,..], [Y,...], [Z,...]]

        else:

            self.ipython_display = None





    def connect(self):

        self.called = 0

        self.ev_cancel = th.Event()

        self.thread = th.Thread(target=self._worker, daemon=False)

        self.thread = th.Thread(target=self._animator, daemon=False)

        self.thread.start()





    def disconnect(self, force=False):

        self.ev_cancel.set()

        self.thread.join(timeout=1.5 * self.interval)

        del self.anim





    def _worker(self):
@@ -50,17 +53,10 @@ class VehicleDrawing: 
        else:

            self.figure = self.axis.figure



        for l in self.lines:

        for l in self.arm_lines:

            self.axis.add_line(l)

        self.axis.add_line(self.trajectory_line)



        # if self.is_ipython:

        #     from IPython.display import display

        #     self.ipython_display = display(self.figure, display_id=True)

        # else:

        self.ipython_display = None

        self.figure.show()



        while not self.ev_cancel.is_set():

            start = time.time()

            position = self.vehicle.position
@@ -70,32 +66,86 @@ class VehicleDrawing: 
            self.ev_cancel.wait(self.interval - (end - start))





    def _animator(self):

        if self.axis is None:

            self.figure, self.axis = make_fig((-10,10), (-10,10), (-10,10))

        else:

            self.figure = self.axis.figure

        

        def init_func():

            for l in self.arm_lines:

                self.axis.add_line(l)

            self.axis.add_line(self.trajectory_line)

            for l in self.axis_lines:

                self.axis.add_line(l)

            return (*self.arm_lines, self.trajectory_line, self.axis_lines)

        

        def func(f):

            if self.ev_cancel.is_set():

                raise th.ThreadError('Canceling animation update.')

            if self.vehicle.t != self.t:

                if self.vehicle.t < self.t:

                    self.trajectory = [[], [], []] # likely vehicle is reset, so reset trajectory

                self.t = self.vehicle.t

                return update_drawing(self, self.vehicle.position, self.vehicle.orientation)

            else:

                return []



        self.anim = FuncAnimation(

            self.figure,

            func=func,

            init_func=init_func,

            blit=True,

            repeat=False,

            interval=self.interval * 1000)

        return self.anim

        



def make_drawing(params: VehicleParams):

    arms = np.zeros((len(params.propellers) * 2, 3)) # [2 points/ propeller, axis]



def make_drawing(drawing: VehicleDrawing, body_axes: bool=False):

    params = drawing.params

    arm_lines_points = np.zeros((len(params.propellers) * 2, 3)) # [2 points/ propeller, axis]

    x = params.distances * np.cos(params.angles)

    y = params.distances * np.sin(params.angles)

    arms[1::2,0] = x

    arms[1::2,1] = y

    lines = []

    arm_lines_points[1::2,0] = x

    arm_lines_points[1::2,1] = y

    arm_lines = []

    for i in range(len(params.propellers)):

        lines.append(

        arm_lines.append(

            Line3D(

                arms[2*i:2*i+2,0],

                arms[2*i:2*i+2,1],

                arms[2*i:2*i+2,2],

                arm_lines_points[2*i:2*i+2,0],

                arm_lines_points[2*i:2*i+2,1],

                arm_lines_points[2*i:2*i+2,2],

                antialiased=False))

    trajectory = Line3D([], [], [])

    return lines, arms, trajectory



    trajectory_line = Line3D([], [], [], linewidth=0.5, color='black', linestyle=':')



    axis_lines_points = np.zeros((6, 3)) # [2 points/ axis, axis]

    axis_lines_points[1::2] = np.eye(3)

    axis_lines = []

    for i, c in enumerate(['r', 'g', 'b']):

        if body_axes:

            axis_lines.append(Line3D(

                axis_lines_points[2*i:2*i+2,0],

                axis_lines_points[2*i:2*i+2,1],

                axis_lines_points[2*i:2*i+2,2],

                antialiased=False,

                linewidth=0.5,

                color=c

            )) 

        else:

            axis_lines.append(Line3D([], [], []))



    return arm_lines, arm_lines_points, trajectory_line, axis_lines, axis_lines_points







def update_drawing(drawing: VehicleDrawing, position: np.ndarray, orientation: np.ndarray):

    dcm = direction_cosine_matrix(orientation[0], orientation[1], orientation[2])

    arms = np.copy(drawing.body_frame_segments)

    arms = inertial_to_body(arms.T, dcm).T

    arms = np.copy(drawing.arm_lines_points)

    arms = body_to_inertial(arms.T, dcm).T

    arms += position

    for i, l in enumerate(drawing.lines):

    for i, l in enumerate(drawing.arm_lines):

        j = 2*i

        l.set_data(arms[j:j+2, 0], arms[j:j+2, 1])

        l.set_3d_properties(arms[j:j+2, 2])
@@ -107,18 +157,27 @@ def update_drawing(drawing: VehicleDrawing, position: np.ndarray, orientation: n 
        drawing.trajectory_line.set_data(drawing.trajectory[0], drawing.trajectory[1])

        drawing.trajectory_line.set_3d_properties(drawing.trajectory[2])

    

    drawing.figure.canvas.draw_idle()

    drawing.figure.canvas.flush_events()

    plt.pause(1e-7)

    if drawing.body_axes:

        axes = np.copy(drawing.axis_lines_points)

        axes[1::2] = position

        axes[0::2] = dcm.T + position

        for i, l in enumerate(drawing.axis_lines):

            l.set_data(axes[i*2:i*2+2,0], axes[i*2:i*2+2,1])

            l.set_3d_properties(axes[i*2:i*2+2,2])

    # drawing.figure.canvas.draw_idle()

    # drawing.figure.canvas.flush_events()

    # plt.pause(1e-7)



    if drawing.ipython_display is not None:

        drawing.ipython_display.update(drawing.figure)

    return (*drawing.arm_lines, drawing.trajectory_line, drawing.axis_lines)







def make_fig(xlim, ylim, zlim) -> Tuple[plt.Figure, Axes3D]:

    fig = plt.figure()

    ax = fig.add_subplot(projection='3d', xlim=xlim, ylim=ylim, zlim=zlim)

    ax.set_xlabel('x')

    ax.set_ylabel('y')

    ax.set_zlabel('z')

    return fig, ax