diff --git a/Airbrakes/airbrakes.py b/Airbrakes/airbrakes.py index 9b29a294..f48acce0 100644 --- a/Airbrakes/airbrakes.py +++ b/Airbrakes/airbrakes.py @@ -1,10 +1,15 @@ """Module which provides a high level interface to the air brakes system on the rocket.""" +from typing import TYPE_CHECKING + from Airbrakes.imu import IMU, IMUDataPacket from Airbrakes.logger import Logger from Airbrakes.servo import Servo from Airbrakes.state import StandByState, State +if TYPE_CHECKING: + import collections + class AirbrakesContext: """ @@ -44,20 +49,19 @@ def update(self): """ # Gets the current extension and IMU data, the states will use these values self.current_extension = self.servo.current_extension - data_packets = [] # 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. - while len(data_packets) < 50: - # get_imu_data() gets the "first" item in the queue, i.e, it *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 behind on processing - data_packets.append(self.imu.get_imu_data()) + # 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 + # behind on processing + data_packets: collections.deque[IMUDataPacket] = self.imu.get_imu_data_packets(50) # Logs the current state, extension, and IMU data # TODO: Compute state(s) for given IMU data - self.logger.log(self.state.get_name(), self.current_extension, data_packets) + self.logger.log(self.state.get_name(), self.current_extension, data_packets.copy()) self.state.update() diff --git a/Airbrakes/imu.py b/Airbrakes/imu.py index 7ec4f118..4040ac9c 100644 --- a/Airbrakes/imu.py +++ b/Airbrakes/imu.py @@ -1,5 +1,6 @@ """Module for interacting with the IMU (Inertial measurement unit) on the rocket.""" +import collections import multiprocessing import mscl @@ -92,7 +93,7 @@ def _fetch_data_loop(self, port: str, frequency: int, _: bool): # 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 + # 2ms on average, so we use a timeout of 1000 / frequency = 10ms which should be more # than enough. # (help needed: what happens if the timeout is hit? An empty list? Exception?) packets: mscl.MipDataPackets = node.getDataPackets(timeout) @@ -130,17 +131,48 @@ def _fetch_data_loop(self, port: str, frequency: int, _: bool): # Put the latest data into the shared queue self.data_queue.put(imu_data_packet) - def get_imu_data(self) -> IMUDataPacket: + def get_imu_data_packet(self, block: bool = True) -> IMUDataPacket: """ - Gets the latest data from the IMU. + Gets the last available data packet from the IMU. - Note: If `get_imu_data` is called slower than the frequency set, the data will not be the - latest, but the first in the queue. + Note: If `get_imu_data_packet` is called slower than the frequency set, the data will not + be the latest, but the first in the queue. + + :param block: If True, the function will block until data is available. If False, it may + raise a `multiprocessing.queue.Empty` exception if no data is available. :return: an IMUDataPacket object containing the latest data from the IMU. If a value is - not available, it will be None. + not available, it will be None. + + :raises multiprocessing.queue.Empty: If block is False and there is no data available. """ - return self.data_queue.get() + return self.data_queue.get() if block else self.data_queue.get_nowait() + + def get_imu_data_packets(self, num_packets: int, block: bool = True) -> collections.deque[IMUDataPacket]: + """Returns a specified amount of data packets from the IMU. + + :param num_packets: The number of packets to return. + :param block: If True, the function will block until the specified number of packets are + available. Otherwise, it will return the available packets, which may be less than + `num_packets`. + + :return: A deque containing the specified number of data packets, or less, depending on + `block`. + """ + + data_packets = collections.deque() + + if block: + while len(data_packets) < num_packets: + data_packets.append(self.get_imu_data_packet()) + else: + try: + for _ in range(num_packets): + data_packets.append(self.get_imu_data_packet(block=False)) + except multiprocessing.queues.Empty: + pass + + return data_packets def stop(self): """ diff --git a/Airbrakes/logger.py b/Airbrakes/logger.py index 6a334531..4ed1734c 100644 --- a/Airbrakes/logger.py +++ b/Airbrakes/logger.py @@ -1,5 +1,6 @@ """Module for logging data to a CSV file in real time.""" +import collections import logging import multiprocessing from pathlib import Path @@ -62,7 +63,7 @@ def _logging_loop(self): message = self.log_queue.get() logger.info(message) - def log(self, state: str, extension: float, imu_data_list: list[IMUDataPacket]): + def log(self, state: str, extension: float, imu_data_list: collections.deque[IMUDataPacket]): """ Logs the current state, extension, and IMU data to the CSV file. :param state: the current state of the airbrakes state machine diff --git a/Scripts/test_imu.py b/Scripts/test_imu.py index e12490a2..e5cc862b 100644 --- a/Scripts/test_imu.py +++ b/Scripts/test_imu.py @@ -16,4 +16,4 @@ imu = IMU(PORT, FREQUENCY, UPSIDE_DOWN) while True: - print(imu.get_imu_data()) + print(imu.get_imu_data_packet())