Feat: Add IMUDataProcessor and use csv logger

Multiple fixes and tiny improvements as well

Scripts/test_imu.py

@@ -4,14 +4,8 @@
 For the pi, you will have to use python3
 """
 
-from airbrakes.imu import IMU
-
-# Should be checked before launch
-UPSIDE_DOWN = True
-# The port that the IMU is connected to
-PORT = "/dev/ttyACM0"
-# The frequency in which the IMU polls data in Hz
-FREQUENCY = 100
+from airbrakes.constants import FREQUENCY, PORT, UPSIDE_DOWN
+from airbrakes.imu.imu import IMU
 
 imu = IMU(PORT, FREQUENCY, UPSIDE_DOWN)
 

Scripts/test_logger.py

@@ -1,15 +1,13 @@
 """Module to test the logger module."""
 
-from airbrakes.imu import IMUDataPacket
+from airbrakes.imu.imu_data_packet import IMUDataPacket
 from airbrakes.logger import Logger
 
-CSV_HEADERS = ["state", "extension", *list(IMUDataPacket(0.0).__slots__)]
-
 
 def main():
-    logger = Logger(CSV_HEADERS)
+    logger = Logger()
 
-    while True:
+    while True:  # TODO: will not work
         logger.log("state", 0.0, IMUDataPacket(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0))
 
 

Scripts/test_servo.py

@@ -4,15 +4,9 @@
 For the pi, you will have to use python3
 """
 
+from airbrakes.constants import MAX_EXTENSION, MIN_EXTENSION, SERVO_PIN
 from airbrakes.servo import Servo
 
-# The pin that the servo's data wire is plugged into, in this case the GPIO 12 pin which is used for PWM
-SERVO_PIN = 12
-
-# The minimum and maximum position of the servo, its range is -1 to 1
-MIN_EXTENSION = -1
-MAX_EXTENSION = 1
-
 servo = Servo(SERVO_PIN, MIN_EXTENSION, MAX_EXTENSION)
 
 print("Type (1) to deploy and (0) to retract the airbrakes.")

airbrakes/airbrakes.py

@@ -2,7 +2,9 @@
 
 from typing import TYPE_CHECKING
 
-from airbrakes.imu import IMU, IMUDataPacket, RollingAverages
+from airbrakes.imu.data_processor import ProcessedIMUData
+from airbrakes.imu.imu import IMU, IMUDataPacket
+from airbrakes.imu.imu_data_packet import EstimatedDataPacket
 from airbrakes.logger import Logger
 from airbrakes.servo import Servo
 from airbrakes.state import StandByState, State
@@ -21,27 +23,28 @@ class AirbrakesContext:
 
     __slots__ = (
         "current_extension",
-        "current_imu_data",
         "imu",
         "logger",
+        "processed_data",
         "servo",
         "shutdown_requested",
         "state",
     )
 
     def __init__(self, logger: Logger, servo: Servo, imu: IMU):
-        self.logger = logger
-        self.servo = servo
-        self.imu = imu
+        self.logger: Logger = logger
+        self.servo: Servo = servo
+        self.imu: IMU = imu
 
         self.state: State = StandByState(self)
         self.shutdown_requested = False
 
-        # Placeholder for the current airbrake extension and IMU data until they are set
-        self.current_extension = 0.0
-        self.current_imu_data = IMUDataPacket(0.0)
+        self.processed_data: ProcessedIMUData = ProcessedIMUData([])
 
-    def update(self):
+        # Placeholder for the current airbrake extension until they are set
+        self.current_extension: float = 0.0
+
+    def update(self) -> None:
         """
         Called every loop iteration from the main process. Depending on the current state, it will
         do different things. It is what controls the airbrakes and chooses when to move to the next
@@ -50,30 +53,30 @@ def update(self):
         # Gets the current extension and IMU data, the states will use these values
         self.current_extension = self.servo.current_extension
 
-        # Let's get 50 data packets to ensure we have enough data to work with.
-        # 50 is an arbitrary number for now - if the time resolution between each data packet is
-        # 2ms, then we have 2*50 = 100ms of data to work with at once.
         # get_imu_data_packets() gets from the "first" item in the queue, i.e, the set of data
-        # *may* not be the most recent data. But we want continous data for proper state and
-        # apogee calculation, so we don't need to worry about that, as long as we're not too
+        # *may* not be the most recent data. But we want continous data for state, apogee,
+        # and logging purposes, so we don't need to worry about that, as long as we're not too
         # behind on processing
         data_packets: collections.deque[IMUDataPacket] = self.imu.get_imu_data_packets()
 
+        # Update the processed data with the new data packets. We only care about EstimatedDataPackets
+        self.processed_data.update(
+            data_packet for data_packet in data_packets if isinstance(data_packet, EstimatedDataPacket)
+        )
         # Logs the current state, extension, and IMU data
-        RollingAverages(data_packets.copy())
         # TODO: Compute state(s) for given IMU data
         self.logger.log(self.state.get_name(), self.current_extension, data_packets.copy())
 
         self.state.update()
 
-    def set_airbrake_extension(self, extension: float):
+    def set_airbrake_extension(self, extension: float) -> None:
         """
         Sets the airbrake extension via the servo. It will be called by the states.
         :param extension: the extension of the airbrakes, between 0 and 1
         """
         self.servo.set_extension(extension)
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         """
         Handles shutting down the airbrakes. This will cause the main loop to break.
         """

airbrakes/constants.py

@@ -0,0 +1,27 @@
+"""Contains the constants used in the airbrakes module"""
+
+from airbrakes.imu.imu_data_packet import EstimatedDataPacket, IMUDataPacket, RawDataPacket
+
+# The pin that the servo's data wire is plugged into, in this case the GPIO 12 pin which is used for PWM
+SERVO_PIN = 12
+
+# The minimum and maximum position of the servo, its range is -1 to 1
+MIN_EXTENSION = -1
+MAX_EXTENSION = 1
+
+# Should be checked before launch
+UPSIDE_DOWN = True  # TODO: Currently not factored in the implementation
+# The port that the IMU is connected to
+PORT = "/dev/ttyACM0"
+
+# The frequency in Hz that the IMU will be polled at
+FREQUENCY = 100  # TODO: Remove this since we don't/can't control the frequency from the code.
+
+# The headers for the CSV file
+CSV_HEADERS = [
+    "State",
+    "Extension",
+    *list(IMUDataPacket.__slots__),
+    *list(RawDataPacket.__slots__),
+    *list(EstimatedDataPacket.__slots__),
+]

airbrakes/imu/data_processor.py

@@ -0,0 +1,75 @@
+"""Module for processing IMU data on a higher level."""
+
+import statistics as stats
+from collections.abc import Sequence
+
+from airbrakes.imu.imu_data_packet import EstimatedDataPacket
+
+
+class ProcessedIMUData:
+    """Performs high level calculations on the data packets received from the IMU. Includes
+    calculation the rolling averages of acceleration, maximum altitude so far, etc, from the set of
+    data points.
+
+    Args:
+        data_points (Sequence[EstimatedDataPacket]): A list of EstimatedDataPacket objects
+            to process.
+    """
+
+    __slots__ = ("_avg_accel", "_avg_accel_mag", "_max_altitude", "data_points")
+
+    def __init__(self, data_points: Sequence[EstimatedDataPacket]):
+        self.data_points: Sequence[EstimatedDataPacket] = data_points
+        self._avg_accel: tuple[float, float, float] = (0.0, 0.0, 0.0)
+        self._avg_accel_mag: float = 0.0
+        self._max_altitude: float = 0.0
+
+    def update_data(self, data_points: Sequence[EstimatedDataPacket]) -> None:
+        """Updates the data points to process. This will recalculate the averages and maximum
+        altitude."""
+        self.data_points = data_points
+        a_x, a_y, a_z = self.compute_averages()
+        self._avg_accel = (a_x, a_y, a_z)
+        self._avg_accel_mag = (self._avg_accel**2 + self._avg_accel**2 + self._avg_accel**2) ** 0.5
+        self._max_altitude = max(*(data_point.altitude for data_point in self.data_points), self._max_altitude)
+
+    def compute_averages(self) -> tuple[float, float, float]:
+        """Calculates the average acceleration and acceleration magnitude of the data points."""
+        # calculate the average acceleration in the x, y, and z directions
+        # TODO: Test what these accel values actually look like
+        x_accel = stats.fmean(data_point.estCompensatedAccelX for data_point in self.data_points)
+        y_accel = stats.fmean(data_point.estCompensatedAccelY for data_point in self.data_points)
+        z_accel = stats.fmean(data_point.estCompensatedAccelZ for data_point in self.data_points)
+        # TODO: Calculate avg velocity if that's also available
+        return x_accel, y_accel, z_accel
+
+    @property
+    def avg_acceleration_z(self) -> float:
+        """Returns the average acceleration in the z direction of the data points, in m/s^2."""
+        return self._avg_accel[-1]
+
+    @property
+    def avg_acceleration(self) -> tuple[float, float, float]:
+        """Returns the averaged acceleration as a vector of the data points, in m/s^2."""
+        return self._avg_accel
+
+    @property
+    def avg_acceleration_mag(self) -> float:
+        """Returns the magnitude of the acceleration vector of the data points, in m/s^2."""
+        return self._avg_accel_mag
+
+    @property
+    def max_altitude(self) -> float:
+        """Returns the highest altitude attained by the rocket for the entire flight so far,
+        in meters.
+        """
+        return self._max_altitude
+
+    def __str__(self) -> str:
+        """Returns a string representation of the ProcessedIMUData object."""
+        return (
+            f"{self.__class__.__name__}("
+            f"avg_acceleration={self.avg_acceleration}, "
+            f"avg_acceleration_mag={self.avg_acceleration_mag}, "
+            f"max_altitude={self.max_altitude})"
+        )

airbrakes/imu/imu.py

@@ -2,33 +2,12 @@
 
 import collections
 import multiprocessing
-from typing import Literal
 
 import mscl
 
 from airbrakes.imu.imu_data_packet import EstimatedDataPacket, IMUDataPacket, RawDataPacket
 
 
-class RollingAverages:
-    """Calculates the rolling averages of acceleration, (and?) from the set of data points"""
-
-    def __init__(self, data_points: list[IMUDataPacket]):
-        self.data_points = data_points
-
-    def add_estimated_data_packet(self):
-        pass
-
-    def calculate_average(self, field: Literal["acceleration"]) -> None:
-        if field == "acceleration":
-            self.averaged_acceleration = sum(data_point.acceleration for data_point in self.data_points) / len(
-                self.data_points
-            )
-
-    @property
-    def averaged_acceleration(self):
-        return self.averaged_acceleration
-
-
 class IMU:
     """
     Represents the IMU on the rocket. It's used to get the current acceleration of the rocket. This is used to interact
@@ -44,7 +23,6 @@ class IMU:
     RAW_DESCRIPTOR_SET = 128
 
     __slots__ = (
-        "connection",
         "data_fetch_process",
         "data_queue",
         "running",
@@ -73,8 +51,9 @@ def _fetch_data_loop(self, port: str, frequency: int, _: bool):
         while self.running.value:
             # Get the latest data packets from the IMU, with the help of `getDataPackets`.
             # `getDataPackets` accepts a timeout in milliseconds.
-            # Testing has shown that the maximum rate at which we can fetch data is roughly every
-            # 2ms on average, so we use a timeout of 1000 / frequency = 10ms which should be more
+            # During IMU configuration (outside of this code), we set the sampling rate of the IMU
+            # as 1ms for RawDataPackets, and 2ms for EstimatedDataPackets.
+            # So we use a timeout of 1000 / frequency = 10ms which should be more
             # than enough. If the timeout is hit, the function will return an empty list.
 
             packets: mscl.MipDataPackets = node.getDataPackets(timeout)
@@ -108,7 +87,8 @@ def _fetch_data_loop(self, port: str, frequency: int, _: bool):
                                 case "estAttitudeUncertQuaternion" | "estOrientQuaternion":
                                     matrix = data_point.as_Matrix()
                                     # Converts the [4x1] matrix to a tuple
-                                    # TODO: maybe we should just make these be stored in the data packet with individual attributes
+                                    # TODO: maybe we should just make these be stored in the data
+                                    # packet with individual attributes
                                     quaternion_tuple = tuple(matrix[i, 0] for i in range(matrix.rows()))
                                     setattr(imu_data_packet, channel, quaternion_tuple)
                                 case _:
@@ -120,7 +100,7 @@ def _fetch_data_loop(self, port: str, frequency: int, _: bool):
                 self.data_queue.put(imu_data_packet)
                 # TODO: this is where we should calculate the rolling averages
 
-    def get_imu_data_packet(self) -> IMUDataPacket:
+    def get_imu_data_packet(self) -> IMUDataPacket | EstimatedDataPacket:
         """
         Gets the last available data packet from the IMU.
 
@@ -132,8 +112,8 @@ def get_imu_data_packet(self) -> IMUDataPacket:
         """
         return self.data_queue.get()
 
-    def get_imu_data_packets(self) -> collections.deque[IMUDataPacket]:
-        """Returns a specified amount of data packets from the IMU.
+    def get_imu_data_packets(self) -> collections.deque[IMUDataPacket | EstimatedDataPacket]:
+        """Returns all available data packets from the IMU.
 
         :return: A deque containing the specified number of data packets
         """

airbrakes/imu/imu_data_packet.py

@@ -1,20 +1,30 @@
 """Module for describing the datapackets from the IMU"""
 
 
+def _mro_slots(obj) -> list[str]:
+    """Helper function to get all __slots__ from the MRO (Method Resolution Order) of an object"""
+    return [
+        attr
+        for cls in obj.__class__.__mro__[:-1]  # :-1 doesn't include the object class
+        for attr in cls.__slots__
+    ]
+
+
 class IMUDataPacket:
     """
-    Represents a collection of data packets from the IMU. It contains the acceleration, velocity, altitude, yaw, pitch,
-    roll of the rocket and the timestamp of the data. The attributes should be named the same as they are when sent from
-    the IMU--this just means they're going to be in camelCase.
+    Base class representing a collection of data packets from the IMU.
+    The attributes should be named the same as they are when sent from the IMU -- this just means
+    they're going to be in camelCase.
     """
 
     __slots__ = ("timestamp",)
 
     def __init__(self, timestamp: float):
+        # TODO: This likely doesn't work as there is no "timestamp" field in the data packets
         self.timestamp = timestamp
 
     def __str__(self):
-        attributes = ", ".join([f"{attr}={getattr(self, attr)}" for attr in self.__slots__])
+        attributes = ", ".join(f"{attr}={getattr(self, attr)}" for attr in _mro_slots(self))
         return f"{self.__class__.__name__}({attributes})"
 
 
@@ -26,8 +36,8 @@ class RawDataPacket(IMUDataPacket):
 
     __slots__ = (
         "gpsCorrelTimestampFlags",
-        "gpsCorrelTimestampTow",
-        "gpsCorrelTimestampWeekNum",
+        "gpsCorrelTimestampTow",  # Time of week
+        "gpsCorrelTimestampWeekNum",  # Week number
         "scaledAccelX",
         "scaledAccelY",
         "scaledAccelZ",
@@ -64,8 +74,9 @@ def __init__(
 
 class EstimatedDataPacket(IMUDataPacket):
     """
-    Represents an estimated data packet from the IMU. These values are the processed values of the raw data
-    that are supposed to be more accurate/smoothed. It contains a timestamp and the estimated values of the relevant data points.
+    Represents an estimated data packet from the IMU. These values are the processed values of the
+    raw data that are supposed to be more accurate/smoothed. It contains a timestamp and the
+    estimated values of the relevant data points.
     """
 
     __slots__ = (
@@ -77,8 +88,8 @@ class EstimatedDataPacket(IMUDataPacket):
         "estCompensatedAccelY",
         "estCompensatedAccelZ",
         "estFilterDynamicsMode",
-        "estFilterGpsTimeTow",
-        "estFilterGpsTimeWeekNum",
+        "estFilterGpsTimeTow",  # Time of week
+        "estFilterGpsTimeWeekNum",  # Week number
         "estFilterState",
         "estFilterStatusFlags",
         "estOrientQuaternion",