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 @@
                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 @@
         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 @@
         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 @@
             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 @@
             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
+
 
 class Vector:
     """Pure python basic R3 vector class designed for simple operations.
@@ -152,7 +154,10 @@
     @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 __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 @@
         """
         return self.__matmul__(other)
 
+    def round(self, decimals=0):
+        """Round the 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,72 @@
             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_norm = (q_w**2 + q_x**2 + q_y**2 + q_z**2) ** 0.5
+        if q_norm == 0:
+            return Matrix.identity()
+        q_w /= q_norm
+        q_x /= q_norm
+        q_y /= q_norm
+        q_z /= q_norm
+        # precompute 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 @@
 
         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 @@
             }
             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 @@
                 "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 @@
         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 @@
             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.