Implement Apogee Prediction

README.md

@@ -11,7 +11,7 @@ _A video of our air brakes extending and retracting_
 As per the finite state machine design pattern, we have a context class which links everything together. Every loop, the context:
 
 1. **Gets data from the IMU**
-2. **Processes the data** in the Data Processor (calculates speed, averages, maximums, etc.)
+2. **Processes the data** in the Data Processor (calculates velocity, averages, maximums, etc.)
 3. **Updates the current state** with the processed data
 4. **Controls the servo extension** based on the current state's instructions (e.g., extends air brakes to slow down the rocket)
 5. **Logs all data** from the IMU, Data Processor, Servo, and States
@@ -37,6 +37,7 @@ flowchart TD
     Airbrakes --> Update[update]:::bubble
     %% IMU Data Packet Flow
     IMUDataPacket --> Update
+    Apogee_Predictor --> Update
 
     %% States as individual nodes
     State((State)):::circular
@@ -59,18 +60,20 @@ flowchart TD
     %% Connections with Labels
     Airbrakes ---|Child Process| Logger((Logger)):::circular
     Airbrakes ---|Child Process| IMU((IMU)):::circular
+    Airbrakes ---|Child Process| Apogee_Predictor((Apogee Predictor)):::circular
     IMU((IMU)):::circular ---|Fetch Packets| IMUDataPacket[(IMU Data Packet)]:::outputSquare
 
     %% IMU Data Processing
     IMUDataPacket --> IMUDataProcessor[IMU Data Processor]:::circular
-    %%IMUDataProcessor --> ProcessedData[(Processed Data Packet)]:::outputSquare
-    IMUDataProcessor --> Speed[(Speed)]:::outputSquare
+    IMUDataProcessor --> Velocity[(Velocity)]:::outputSquare
     IMUDataProcessor --> Altitude[(Altitude)]:::outputSquare
+    IMUDataProcessor --> Rotated_Accel[(Rotated Acceleration)]:::outputSquare
 
-    Speed -->  ProcessedData[(Processed Data Packet)]:::outputSquare
+    Velocity -->  ProcessedData[(Processed Data Packet)]:::outputSquare
     Altitude -->  ProcessedData[(Processed Data Packet)]:::outputSquare
+    Rotated_Accel -->  ProcessedData[(Processed Data Packet)]:::outputSquare
 
-     ProcessedData[(Processed Data Packet)]:::outputSquare --> Update
+    ProcessedData[(Processed Data Packet)]:::outputSquare --> Update
 
     %% Logging Mechanism
     Logger --> LogFunction[log]:::bubble

airbrakes/airbrakes.py

@@ -1,19 +1,20 @@
 """Module which provides a high level interface to the air brakes system on the rocket."""
 
+from collections import deque
 from typing import TYPE_CHECKING
 
+from airbrakes.data_handling.apogee_predictor import ApogeePredictor
 from airbrakes.data_handling.data_processor import IMUDataProcessor
 from airbrakes.data_handling.imu_data_packet import EstimatedDataPacket
 from airbrakes.data_handling.logger import Logger
-from airbrakes.hardware.imu import IMU, IMUDataPacket
+from airbrakes.hardware.imu import IMU
 from airbrakes.hardware.servo import Servo
-from airbrakes.state import StandByState, State
+from airbrakes.state import StandbyState, State
 from constants import ServoExtension
 
 if TYPE_CHECKING:
-    from collections import deque
-
     from airbrakes.data_handling.processed_data_packet import ProcessedDataPacket
+    from airbrakes.hardware.imu import IMUDataPacket
 
 
 class AirbrakesContext:
@@ -25,42 +26,59 @@ class AirbrakesContext:
     """
 
     __slots__ = (
+        "apogee_predictor",
         "current_extension",
         "data_processor",
+        "est_data_packets",
         "imu",
+        "imu_data_packets",
         "logger",
+        "processed_data_packets",
         "servo",
         "shutdown_requested",
         "state",
     )
 
-    def __init__(self, servo: Servo, imu: IMU, logger: Logger, data_processor: IMUDataProcessor) -> None:
+    def __init__(
+        self,
+        servo: Servo,
+        imu: IMU,
+        logger: Logger,
+        data_processor: IMUDataProcessor,
+        apogee_predictor: ApogeePredictor,
+    ) -> None:
         """
         Initializes the airbrakes context with the specified hardware objects, logger, and data processor. The state
-        machine starts in the StandByState, which is the initial state of the airbrakes system.
+        machine starts in the StandbyState, which is the initial state of the airbrakes system.
         :param servo: The servo object that controls the extension of the airbrakes. This can be a real servo or a mock
         servo.
         :param imu: The IMU object that reads data from the rocket's IMU. This can be a real IMU or a mock IMU.
         :param logger: The logger object that logs data to a CSV file.
         :param data_processor: The data processor object that processes IMU data on a higher level.
+        :param apogee_predictor: The apogee predictor object that predicts the apogee of the rocket.
         """
         self.servo = servo
         self.imu = imu
         self.logger = logger
         self.data_processor = data_processor
+        self.apogee_predictor = apogee_predictor
 
         # Placeholder for the current airbrake extension until they are set
         self.current_extension: ServoExtension = ServoExtension.MIN_EXTENSION
-
-        self.state: State = StandByState(self)
+        # The rocket starts in the StandbyState
+        self.state: State = StandbyState(self)
         self.shutdown_requested = False
+        self.imu_data_packets: deque[IMUDataPacket] = deque()
+        self.processed_data_packets: deque[ProcessedDataPacket] = deque()
+        self.est_data_packets: list[EstimatedDataPacket] = []
 
     def start(self) -> None:
         """
         Starts the IMU and logger processes. This is called before the main while loop starts.
         """
         self.imu.start()
         self.logger.start()
+        self.apogee_predictor.start()
 
     def stop(self) -> None:
         """
@@ -72,6 +90,7 @@ def stop(self) -> None:
         self.retract_airbrakes()
         self.imu.stop()
         self.logger.stop()
+        self.apogee_predictor.stop()
         self.shutdown_requested = True
 
     def update(self) -> None:
@@ -84,32 +103,42 @@ def update(self) -> None:
         # *may* not be the most recent data. But we want continuous 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
-        imu_data_packets: deque[IMUDataPacket] = self.imu.get_imu_data_packets()
+        self.imu_data_packets = self.imu.get_imu_data_packets()
 
         # This should never happen, but if it does, we want to not error out and wait for packets
-        if not imu_data_packets:
+        if not self.imu_data_packets:
             return
 
         # Split the data packets into estimated and raw data packets for use in processing and logging
-        est_data_packets = [
+        self.est_data_packets = [
             data_packet
-            for data_packet in imu_data_packets.copy()
+            for data_packet in self.imu_data_packets.copy()
             if isinstance(data_packet, EstimatedDataPacket)
             # The copy() above is critical to ensure the data here is not modified by the data processor
         ]
 
         # Update the processed data with the new data packets. We only care about EstimatedDataPackets
-        self.data_processor.update(est_data_packets)
+        self.data_processor.update(self.est_data_packets)
 
         # Get the processed data packets from the data processor, this will have the same length as the number of
         # EstimatedDataPackets in data_packets
-        processed_data_packets: deque[ProcessedDataPacket] = self.data_processor.get_processed_data_packets()
+        self.processed_data_packets = self.data_processor.get_processed_data_packets()
 
         # Update the state machine based on the latest processed data
         self.state.update()
 
         # Logs the current state, extension, IMU data, and processed data
-        self.logger.log(self.state.name[0], self.current_extension.value, imu_data_packets, processed_data_packets)
+        self.logger.log(
+            self.state.name,
+            self.current_extension.value,
+            self.imu_data_packets,
+            self.processed_data_packets,
+            self.apogee_predictor.apogee,
+        )
+        # CAUTION: You would need to copy() self.processed_data_packets in the log line above
+        # if you want to reference it anywhere else after update(), because it would be modified
+        # by the logger. e.g. in tests. For tests, the workaround would be to monkeypatch
+        # the log() method to not modify the processed_data_packets.
 
     def extend_airbrakes(self) -> None:
         """
@@ -124,3 +153,14 @@ def retract_airbrakes(self) -> None:
         """
         self.servo.set_retracted()
         self.current_extension = ServoExtension.MIN_EXTENSION
+
+    def predict_apogee(self) -> None:
+        """
+        Predicts the apogee of the rocket based on the current processed data. This
+        should only be called in the coast state, before we start controlling the airbrakes.
+        """
+        # We have to only only run this for estimated data packets, otherwise we send duplicate
+        # data to the predictor (because for a raw data packet, we still have the 'old' processed_data_packets)
+        # This would result in a very slow convergence and inaccurate predictions.
+        if self.est_data_packets:
+            self.apogee_predictor.update(self.processed_data_packets)