ENH: Sensors

rocketpy/__init__.py

@@ -37,4 +37,5 @@
     Tail,
     TrapezoidalFins,
 )
+from .sensors import Accelerometer, Gyroscope, Sensors
 from .simulation import Flight

rocketpy/control/controller.py

@@ -1,3 +1,5 @@
+from inspect import signature
+
 from ..prints.controller_prints import _ControllerPrints
 
 
@@ -52,6 +54,10 @@ def __init__(
                the controller function can interact with. The objects are
                listed in the same order as they are provided in the
                `interactive_objects`.
+            7. `sensors` (list): A list of sensors that are attached to the
+                rocket. The most recent measurements of the sensors are provided
+                with the ``sensor.measurement`` attribute. The sensors are
+                listed in the same order as they are added to the rocket
 
             This function will be called during the simulation at the specified
             sampling rate. The function should evaluate and change the interactive
@@ -78,7 +84,7 @@ def __init__(
         None
         """
         self.interactive_objects = interactive_objects
-        self.controller_function = controller_function
+        self.controller_function = self.__init_controller_function(controller_function)
         self.sampling_rate = sampling_rate
         self.name = name
         self.prints = _ControllerPrints(self)
@@ -88,7 +94,44 @@ def __init__(
         else:
             self.observed_variables = []
 
-    def __call__(self, time, state_vector, state_history):
+    def __init_controller_function(self, controller_function):
+        """Checks number of arguments of the controller function and initializes
+        it with the correct number of arguments. This is a workaround to allow
+        the controller function to receive sensors without breaking changes"""
+        sig = signature(controller_function)
+        if len(sig.parameters) == 6:
+
+            def new_controller_function(
+                time,
+                sampling_rate,
+                state_vector,
+                state_history,
+                observed_variables,
+                interactive_objects,
+                sensors,
+            ):
+                return controller_function(
+                    time,
+                    sampling_rate,
+                    state_vector,
+                    state_history,
+                    observed_variables,
+                    interactive_objects,
+                )
+
+        elif len(sig.parameters) == 7:
+            new_controller_function = controller_function
+        else:
+            raise ValueError(
+                "The controller function must have 6 or 7 arguments. "
+                "The arguments must be in the following order: "
+                "(time, sampling_rate, state_vector, state_history, "
+                "observed_variables, interactive_objects, sensors)."
+                "Sensors argument is optional."
+            )
+        return new_controller_function
+
+    def __call__(self, time, state_vector, state_history, sensors):
         """Call the controller function. This is used by the simulation class.
 
         Parameters
@@ -104,6 +147,11 @@ def __call__(self, time, state_vector, state_history):
             history is a list of every state vector of every step of the
             simulation. The state history is a list of lists, where each
             sublist is a state vector and is ordered from oldest to newest.
+        sensors : list
+            A list of sensors that are attached to the rocket. The most recent
+            measurements of the sensors are provided with the
+            ``sensor.measurement`` attribute. The sensors are listed in the same
+            order as they are added to the rocket.
 
         Returns
         -------
@@ -116,6 +164,7 @@ def __call__(self, time, state_vector, state_history):
             state_history,
             self.observed_variables,
             self.interactive_objects,
+            sensors,
         )
         if observed_variables is not None:
             self.observed_variables.append(observed_variables)

rocketpy/mathutils/vector_matrix.py

@@ -2,6 +2,8 @@
 from functools import cached_property
 from itertools import product
 
+from rocketpy.tools import euler_to_quaternions, normalize_quaternions
+
 
 class Vector:
     """Pure python basic R3 vector class designed for simple operations.
@@ -152,7 +154,10 @@ def __len__(self):
     @cached_property
     def unit_vector(self):
         """R3 vector with the same direction of self, but normalized."""
-        return self / abs(self)
+        try:
+            return self / abs(self)
+        except ZeroDivisionError:
+            return self
 
     @cached_property
     def cross_matrix(self):
@@ -238,6 +243,30 @@ def __xor__(self, other):
             ]
         )
 
+    def __and__(self, other):
+        """Element wise multiplication between two R3 vectors. Also known as
+        Hadamard product.
+
+        Parameters
+        ----------
+        other : Vector
+            R3 vector to be multiplied with self.
+
+        Returns
+        -------
+        Vector
+            R3 vector resulting from the element wise multiplication between
+            self and other.
+
+        Examples
+        --------
+        >>> v = Vector([1, 7, 3])
+        >>> u = Vector([2, 5, 6])
+        >>> (v & u)
+        Vector(2, 35, 18)
+        """
+        return Vector([self.x * other[0], self.y * other[1], self.z * other[2]])
+
     def __matmul__(self, other):
         """Dot product between two R3 vectors."""
         return self.x * other.x + self.y * other.y + self.z * other.z
@@ -920,6 +949,47 @@ def dot(self, other):
         """
         return self.__matmul__(other)
 
+    def round(self, decimals=0):
+        """Round all the values matrix to a given number of decimals.
+
+        Parameters
+        ----------
+        decimals : int, optional
+            Number of decimal places to round to. Defaults to 0.
+
+        Returns
+        -------
+        Matrix
+            The rounded matrix.
+
+        Examples
+        --------
+        >>> M = Matrix([[1.1234, 2.3456, 3.4567], [4.5678, 5.6789, 6.7890], [7.8901, 8.9012, 9.0123]])
+        >>> M.round(2)
+        Matrix([1.12, 2.35, 3.46],
+               [4.57, 5.68, 6.79],
+               [7.89, 8.9, 9.01])
+        """
+        return Matrix(
+            [
+                [
+                    round(self.xx, decimals),
+                    round(self.xy, decimals),
+                    round(self.xz, decimals),
+                ],
+                [
+                    round(self.yx, decimals),
+                    round(self.yy, decimals),
+                    round(self.yz, decimals),
+                ],
+                [
+                    round(self.zx, decimals),
+                    round(self.zy, decimals),
+                    round(self.zz, decimals),
+                ],
+            ]
+        )
+
     def __str__(self):
         return (
             f"[{self.xx}, {self.xy}, {self.xz}]\n"
@@ -955,34 +1025,66 @@ def transformation(quaternion):
             The quaternion representing the rotation from frame A to frame B.
             Example: (cos(phi/2), 0, 0, sin(phi/2)) represents a rotation of
             phi around the z-axis.
-            Note: the quaternion must be normalized.
 
         Returns
         -------
         Matrix
             The transformation matrix from frame B to frame A.
         """
-        q_w, q_x, q_y, q_z = quaternion
+        # normalize quaternion
+        q_w, q_x, q_y, q_z = normalize_quaternions(quaternion)
+
+        # pre-compute common terms
+        q_x2 = q_x**2
+        q_y2 = q_y**2
+        q_z2 = q_z**2
+        q_wx = q_w * q_x
+        q_wy = q_w * q_y
+        q_wz = q_w * q_z
+        q_xy = q_x * q_y
+        q_xz = q_x * q_z
+        q_yz = q_y * q_z
         return Matrix(
             [
                 [
-                    1 - 2 * (q_y**2 + q_z**2),
-                    2 * (q_x * q_y - q_w * q_z),
-                    2 * (q_x * q_z + q_w * q_y),
+                    1 - 2 * (q_y2 + q_z2),
+                    2 * (q_xy - q_wz),
+                    2 * (q_xz + q_wy),
                 ],
                 [
-                    2 * (q_x * q_y + q_w * q_z),
-                    1 - 2 * (q_x**2 + q_z**2),
-                    2 * (q_y * q_z - q_w * q_x),
+                    2 * (q_xy + q_wz),
+                    1 - 2 * (q_x2 + q_z2),
+                    2 * (q_yz - q_wx),
                 ],
                 [
-                    2 * (q_x * q_z - q_w * q_y),
-                    2 * (q_y * q_z + q_w * q_x),
-                    1 - 2 * (q_x**2 + q_y**2),
+                    2 * (q_xz - q_wy),
+                    2 * (q_yz + q_wx),
+                    1 - 2 * (q_x2 + q_y2),
                 ],
             ]
         )
 
+    @staticmethod
+    def transformation_euler_angles(roll, pitch, yaw):
+        """Returns the transformation Matrix from frame B to frame A, where B
+        is rotated by the Euler angles roll, pitch and yaw with respect to A.
+
+        Parameters
+        ----------
+        roll : float
+            The roll angle in degrees.
+        pitch : float
+            The pitch angle in degrees.
+        yaw : float
+            The yaw angle in degrees.
+
+        Returns
+        -------
+        Matrix
+            The transformation matrix from frame B to frame A.
+        """
+        return Matrix.transformation(euler_to_quaternions(roll, pitch, yaw))
+
 
 if __name__ == "__main__":
     import doctest

rocketpy/plots/rocket_plots.py

@@ -3,6 +3,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+from rocketpy.mathutils.vector_matrix import Vector
 from rocketpy.motors import EmptyMotor, HybridMotor, LiquidMotor, SolidMotor
 from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
 
@@ -168,7 +169,7 @@ def thrust_to_weight(self):
 
         return None
 
-    def draw(self, vis_args=None):
+    def draw(self, vis_args=None, plane="xz"):
         """Draws the rocket in a matplotlib figure.
 
         Parameters
@@ -188,6 +189,9 @@ def draw(self, vis_args=None):
             }
             A full list of color names can be found at:
             https://matplotlib.org/stable/gallery/color/named_colors
+        plane : str, optional
+            Plane in which the rocket will be drawn. Default is 'xz'. Other
+            options is 'yz'. Used only for sensors representation.
         """
         if vis_args is None:
             vis_args = {
@@ -201,10 +205,8 @@ def draw(self, vis_args=None):
                 "line_width": 1.0,
             }
 
-        # Create the figure and axis
-        _, ax = plt.subplots(figsize=(8, 6), facecolor="#EEEEEE")
+        fig, ax = plt.subplots(figsize=(8, 6), facecolor=vis_args["background"])
         ax.set_aspect("equal")
-        ax.set_facecolor(vis_args["background"])
         ax.grid(True, linestyle="--", linewidth=0.5)
 
         csys = self.rocket._csys
@@ -216,6 +218,7 @@ def draw(self, vis_args=None):
         self._draw_motor(last_radius, last_x, ax, vis_args)
         self._draw_rail_buttons(ax, vis_args)
         self._draw_center_of_mass_and_pressure(ax)
+        self._draw_sensor(ax, self.rocket.sensors, plane, vis_args)
 
         plt.title("Rocket Representation")
         plt.xlim()
@@ -552,6 +555,56 @@ def _draw_center_of_mass_and_pressure(self, ax):
             cp, 0, label="Static Center of Pressure", color="red", s=10, zorder=10
         )
 
+    def _draw_sensor(self, ax, sensors, plane, vis_args):
+        """Draw the sensor as a small thick line at the position of the sensor,
+        with a vector pointing in the direction normal of the sensor. Get the
+        normal vector from the sensor orientation matrix."""
+        colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+        for i, sensor_pos in enumerate(sensors):
+            sensor = sensor_pos[0]
+            pos = sensor_pos[1]
+            if plane == "xz":
+                # z position of the sensor is the x position in the plot
+                x_pos = pos[2]
+                normal_x = sensor.normal_vector.z
+                # x position of the sensor is the y position in the plot
+                y_pos = pos[0]
+                normal_y = sensor.normal_vector.x
+            elif plane == "yz":
+                # z position of the sensor is the x position in the plot
+                x_pos = pos[2]
+                normal_x = sensor.normal_vector.z
+                # y position of the sensor is the y position in the plot
+                y_pos = pos[1]
+                normal_y = sensor.normal_vector.y
+            else:
+                raise ValueError("Plane must be 'xz' or 'yz'.")
+
+            # line length is 2/5 of the rocket radius
+            line_length = self.rocket.radius / 2.5
+
+            ax.plot(
+                [x_pos, x_pos],
+                [y_pos + line_length, y_pos - line_length],
+                linewidth=2,
+                color=colors[(i + 1) % len(colors)],
+                zorder=10,
+                label=sensor.name,
+            )
+            ax.quiver(
+                x_pos,
+                y_pos,
+                normal_x,
+                normal_y,
+                color=colors[(i + 1) % len(colors)],
+                scale_units="xy",
+                angles="xy",
+                minshaft=2,
+                headwidth=2,
+                headlength=4,
+                zorder=10,
+            )
+
     def all(self):
         """Prints out all graphs available about the Rocket. It simply calls
         all the other plotter methods in this class.