Commit 7f31ddbc authored by Ibrahim's avatar Ibrahim
Browse files

Utility function changes; 

simulation.py: allocation matrices are calculated only once at __init__, and not during reset(),
visualize.py: convenience function for plotting wind quivers,
trajectories.py: figure eight curve function,
parent 09aaadc8

multirotor/helpers.py

+9 −1
Changes for multirotor/helpers.py: 9 added lines, 1 removed line.
Original line number Diff line number Diff line
@@ -338,6 +338,13 @@ class DataLog: 
        self.controller = controller





    def __getitem__(self, item):

        if isinstance(item, slice):

            d = DataLog(other_vars=self._args)

            pass

        # TODO: Finish. Return a DataLog with the requested slice



        

    def log(self, **kwargs):

        """

        Add the state and action variables from the Multirotor and Controller.
@@ -417,13 +424,14 @@ class DataLog: 
            self.target = SimpleNamespace(**d)



        if relative:

            last_time = self.times[old_len - 1]

            # last_pos = self.position[old_len - 1]

            if not self._arrayed:

                last_time = self._times[old_len - 1]

                for i in range(old_len, len(self)):

                        # self._states[i][:3] += last_pos

                        self._times[i] += last_time

            elif self._arrayed:

                last_time = self.times[old_len - 1]

                # self.states[old_len:,:3] += last_pos

                self.times[old_len:] += last_time

multirotor/physics.py

+5 −2
Changes for multirotor/physics.py: 5 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -102,8 +102,8 @@ def torque( 






@njit

def apply_forces_torques(

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

) -> np.ndarray:
@@ -185,3 +185,6 @@ def apply_forces_torques( 
    xdot[9:12] = I_inv @ (torques - gyro)

    

    return xdot





apply_forces_torques = njit(_apply_forces_torques)

multirotor/simulation.py

+11 −2
Changes for multirotor/simulation.py: 11 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -275,6 +275,7 @@ class Multirotor: 
            self.battery = Battery(self.params.battery, self.simulation)

        else:

            self.battery = Battery(BatteryParams(max_voltage=np.inf), self.simulation)

        self.alloc, self.alloc_inverse = control_allocation_matrix(self.params)

        self.reset()
@@ -296,10 +297,9 @@ class Multirotor: 
        if self.battery is not None:

            self.battery.reset()



        self.alloc, self.alloc_inverse = control_allocation_matrix(self.params)

        self.alloc = self.alloc.astype(self.dtype)

        self.params.propeller_vectors = self.params.propeller_vectors.astype(self.dtype)

        self.alloc_inverse = self.alloc_inverse.astype(self.dtype)

        self.params.propeller_vectors = self.params.propeller_vectors.astype(self.dtype)

        self.params.inertia_matrix = self.params.inertia_matrix.astype(self.dtype)

        self.params.inertia_matrix_inverse = self.params.inertia_matrix_inverse.astype(self.dtype)

        self.state = np.zeros(12, dtype=self.dtype)
@@ -375,6 +375,15 @@ class Multirotor: 
        return np.sum(duty_cycle * currents)





    @property

    def speeds(self) -> np.ndarray:

        return np.asarray([p.speed for p in self.propellers], self.dtype)

    @speeds.setter

    def speeds(self, speeds: np.ndarray):

        for s, p in zip(speeds, self.propellers):

            p.step(s, max_voltage=self.battery.voltage)





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

        """

        Calculate the forces and torques acting on the vehicle's center of gravity

multirotor/trajectories.py

+26 −1
Changes for multirotor/trajectories.py: 26 added lines, 1 removed line.
Original line number Diff line number Diff line
@@ -101,7 +101,7 @@ class Trajectory: 
            for i, (p1, p2) in enumerate(zip(points[:-1], points[1:])):

                pos_vec = p2 - p1

                dist = np.linalg.norm(pos_vec)

                unit_vec = pos_vec / dist

                unit_vec = pos_vec / (dist + 1e-6)

                # num = int(dist / self.resolution) + 1

                number = dist // self.resolution

                remainder  = dist % self.resolution
@@ -227,3 +227,28 @@ class GuidedTrajectory: 


    def reached(self, wp: np.ndarray) -> bool:

        return np.linalg.norm(self.vehicle.position - wp) <= self.proximity







def eight_curve(a: float=10, N:int=20) -> np.ndarray:

    """

    Generate a list of points following the Eight curve pattern.



    Parameters

    ----------

    a : int, optional

        The scale of the curve, by default 50

    N : int, optional

        Number of points to generate, by default 20



    Returns

    -------

    np.ndarray

        A Nx3 array of points.

    """

    wp = np.zeros((N, 3), np.float32)

    t = np.linspace(0, 2 * np.pi, N)

    wp[:,0] = a * np.sin(t)

    wp[:,1] = a * np.sin(t) * np.cos(t)

    wp[:,2] = np.zeros(N)

    return wp

multirotor/visualize.py

+70 −17
Changes for multirotor/visualize.py: 70 added lines, 17 removed lines.
Original line number Diff line number Diff line 
from typing import Tuple, Union

from typing import Tuple, Union, Dict

import threading as th

import multiprocessing as mp

import queue
@@ -7,22 +7,21 @@ import time 


import numpy as np

import matplotlib.pyplot as plt

from tqdm.autonotebook import tqdm

from matplotlib.animation import FuncAnimation

from mpl_toolkits.mplot3d import Axes3D

from mpl_toolkits.mplot3d.art3d import Line3D

from matplotlib.lines import Line2D



from .coords import body_to_inertial, direction_cosine_matrix

from .simulation import Multirotor

from .trajectories import Trajectory, GuidedTrajectory

from .controller import Controller

from .vehicle import VehicleParams

from .helpers import DataLog







def plot_datalog(log: DataLog, figsize=(21,10.5),

    plots=('pos', 'vel', 'ctrl', 'traj'),

    nrows=2, ncols=None):

    nrows=2, ncols=None) -> Dict[str, plt.Axes]:

    """

    Plot recorded values from a Multirotor's flight. Including:
@@ -39,6 +38,11 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
        The datalog, where `datalog.done_logging()` has been called.

    figsize : tuple, optional

        The x/y dimensions of the figure, by default (21,10.5)

    

    Returns

    -------

    Dict : Dict[str, plt.Axes]

        A dictionary of plot names mapping to Axes

    """

    nplots = len(plots)

    if nrows is None and ncols is not None:
@@ -52,6 +56,7 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
    n = len(log)

    hasctrl = log.controller is not None

    plot_number = 1

    axes = {}



    # Positions

    if 'pos' in plots:
@@ -76,6 +81,7 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
        plt.legend(handles=plt.gca().lines[:3] + lines, ncol=2)

        plt.title('Position and Orientation')

        plot_number += 1

        axes['pos'] = plt.gca()



    if 'vel' in plots:

        plt.subplot(*plot_grid, plot_number)
@@ -94,6 +100,7 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
        plt.legend()

        plt.title('Velocities')

        plot_number += 1

        axes['vel'] = plt.gca()



    if 'ctrl' in plots:

        plt.subplot(*plot_grid, plot_number)
@@ -106,6 +113,7 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
        plt.ylabel('Torque /Nm')

        plt.legend(handles=plt.gca().lines + l, ncol=2)

        plot_number += 1

        axes['ctrl'] = plt.gca()



    if 'traj' in plots:

        plt.subplot(*plot_grid, plot_number)
@@ -117,8 +125,50 @@ def plot_datalog(log: DataLog, figsize=(21,10.5), 
        plt.ylabel('Y /m')

        plt.legend()

        plot_number += 1

        axes['traj'] = plt.gca()



    plt.tight_layout()

    return axes







def get_wind_quiver(heading: str, ax: plt.Axes, n=5, dim=2):

    """

    Create arrays of x,y,z coordinates for a quiver plot of wind.



    Parameters

    ----------

    heading : str

        The heading of the wind, e.g. '5@45' for 5N wind from 45 degrees.

    ax : plt.Axes

        The axes on which to plot the quiver.

    n : int, optional

        Size of arrays, by default 5

    dim : int, optional

        Dimension of quiver (2 or 3), by default 2



    Returns

    -------

    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

        x,y,[z],dx,dy,[dz] coordinates for quiver plot.

    """

    magnitude, angle = heading.split('@')

    magnitude = float(magnitude)

    if magnitude==0:

        return (0,0,0,0) if dim==2 else (0,0,0,0,0,0)

    angle = float(angle) * np.pi / 180

    dx, dy = -np.cos(angle), -np.sin(angle)

    xlim, ylim = ax.get_xlim(), ax.get_ylim()

    if dim==3:

        zlim = ax.get_zlim()

        dz = 0

        x,y,z = np.meshgrid(np.linspace(*xlim, num=n), np.linspace(*ylim, num=n), np.linspace(*zlim, num=n),

                      indexing='xy')

        return x,y,z,dx,dy,dz

    else:

        x,y = np.meshgrid(np.linspace(*xlim, num=n), np.linspace(*ylim, num=n),

                      indexing='xy')

        return x,y,dx,dy
@@ -138,7 +188,7 @@ class VehicleDrawing: 
        self.arm_lines, self.arm_lines_points, \

        self.trajectory_line, \

        self.axis_lines, self.axis_lines_points = \

            make_drawing(self, self.body_axes)

            make_drawing(self.params, self.body_axes)

        self.trajectory = [[], [], []] # [[X,..], [Y,...], [Z,...]]
@@ -260,39 +310,42 @@ class VehicleDrawing: 
        





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]

def make_drawing(params: VehicleParams, body_axes: bool=False, make_2d: bool=False, scale_arms=1.):

    Line = Line3D if not make_2d else Line2D

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

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

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

    arm_lines_points[1::2,0] = x

    arm_lines_points[1::2,1] = y

    arm_lines_points *=  scale_arms

    arm_lines = []

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

        arm_lines.append(

            Line3D(

            Line(

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

                antialiased=False),

                **({'zs':arm_lines_points[2*i:2*i+2,2]} if not make_2d else {}),

            )



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

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

                           **({'zs':[]} if not make_2d else {}))



    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.append(Line(

                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

                color=c,

                **({'zs':axis_lines_points[2*i:2*i+2,2]} if not make_2d else {}),

            )) 

        else:

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

            axis_lines.append(Line([], [], **({'zs':[]} if not make_2d else {})))



    return arm_lines, arm_lines_points, trajectory_line, axis_lines, axis_lines_points