diff --git a/README.md b/README.md index 262913f..b582ce8 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,25 @@ def main() -> None: if __name__ == "__main__": main() ``` +### Adding a new IMU/Changing the config of an existing IMU +To use an IMU that is not defined in the built-in list of the module (defined IMUs are BNO055, BNO08X, and LSM6DSOX+LIS3MDL), follow these steps to register it to the program: +1. Add the Adafruit library and driver for the IMU to your project dependencies. +2. Define IMU Config / change the config using `replace` (See wiki page for more detail.) Store the new configs in `your_project.imu_configs.py`. +3. Register the model in `pyproject.toml`: +``` +[project.entry-points."imu_module.devices"] +NEW_IMU = "your_project.imu_configs:NEW_IMU" +``` + +where `NEW_IMU` is the IMUConfig name and `your_project.imu_configs` is the file where the IMUConfig is defined. + +or override an existing built-in device: +``` +[project.entry-points."imu_module.device_overrides"] +BNO055 = "your_project.imu_configs:BNO055_CUSTOM" +``` + +4. Save edits. With UV, changes in the `pyproject.toml` are synchronized automatically upon `uv run`. With pip, poetry and conda, a manual module reinstallation may be necessary. ## Program Usage To run the main pipeline for all connected sensors with optional flag `-r` to record data: @@ -98,11 +117,13 @@ make calibrate │ ├── __init__.py │ ├── __main__.py │ ├── base_classes.py +│ ├── builtin_devices.py │ ├── definitions.py │ ├── devices.py │ ├── factory.py │ ├── i2c_bus.py │ ├── orientation_filters.py +│ ├── registry.py │ ├── sensor_manager.py │ ├── utils.py │ └── wrapper.py diff --git a/pyproject.toml b/pyproject.toml index 22abf26..0f04c81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "imu_python" -version = "0.1.0" +version = "0.1.1" description = "IMU sensor codes in Python for the exosuit" readme = "README.md" authors = [ @@ -44,6 +44,12 @@ no_hw = [ # For now this can be empty. ] +[project.entry-points."imu_module.devices"] +# register IMU configs + +[project.entry-points."imu_module.device_overrides"] +# register IMU config overrides + [project.urls] homepage = "https://github.com/TUM-Aries-Lab/imu-module" diff --git a/src/imu_python/base_classes.py b/src/imu_python/base_classes.py index a6605b1..9a89fda 100644 --- a/src/imu_python/base_classes.py +++ b/src/imu_python/base_classes.py @@ -125,6 +125,27 @@ def to_euler(self, seq: str) -> VectorXYZ: euler = rot.as_euler(seq=seq, degrees=False) return VectorXYZ.from_tuple(euler) + def rotate(self, rotation_matrix: NDArray) -> None: + """Rotate the quaternion using a 3x3 rotation matrix. + + :param rotation_matrix: A 3x3 rotation matrix. + :return: None + """ + if rotation_matrix.shape != (3, 3): + msg = f"Expected 3x3 rotation matrix, got {rotation_matrix.shape}" + logger.error(msg) + raise ValueError(msg) + matrix_form = Rot.from_quat( + quat=[self.x, self.y, self.z, self.w], scalar_first=False + ) + rotation = Rot.from_matrix(rotation_matrix) + rotated = rotation * matrix_form + new_quat = Rot.as_quat(rotated, scalar_first=False) + self.x = new_quat[0] + self.y = new_quat[1] + self.z = new_quat[2] + self.w = new_quat[3] + @dataclass(frozen=True) class IMUDeviceData: @@ -168,6 +189,7 @@ class IMUConfig: accel_range_g: Accelerometer range in g. gyro_range_dps: Gyroscope range in degrees per second. filter_config: Configuration for the sensor fusion filter. + scalar_first: Flag for on-board fusion quaternion ordering (True if scalar is the first element, False if last). """ @@ -176,6 +198,7 @@ class IMUConfig: accel_range_g: float gyro_range_dps: float filter_config: FilterConfig = field(default_factory=FilterConfig) + scalar_first: bool = True @dataclass @@ -285,7 +308,8 @@ def magnetic(self) -> tuple[float, float, float] | None: logger.debug("Magnetic data requested") if not self._is_connected: raise OSError(I2C_ERROR, "remote I/O error") - return None + x, y, z = np.random.normal(loc=0, scale=0.1, size=(3,)) + return x, y, z def disconnect(self) -> None: """Simulate a hardware disconnection for testing purposes.""" @@ -301,9 +325,11 @@ class IMUSensorTypes(StrEnum): accel: Accelerometer sensor type. gyro: Gyroscope sensor type. mag: Magnetometer sensor type. + quat: Quaternion (on-board fusion). """ accel = "acceleration" gyro = "gyro" mag = "magnetic" + quat = "quaternion" diff --git a/src/imu_python/builtin_devices.py b/src/imu_python/builtin_devices.py new file mode 100644 index 0000000..7384009 --- /dev/null +++ b/src/imu_python/builtin_devices.py @@ -0,0 +1,205 @@ +"""IMUConfig for known supported IMUs. + +Current configs: BNO055, BNO08x, LSM6DSOX+LIS3MDL +""" + +from imu_python.base_classes import ( + IMUConfig, + IMUParamNames, + IMUSensorTypes, + PreConfigStep, + PreConfigStepType, + SensorConfig, +) +from imu_python.definitions import MOCK_NAME, FilterConfig, IMUDeviceID + +BNO055 = IMUConfig( + devices={ + IMUDeviceID.IMU0: SensorConfig( + name="BNO055", + addresses=[0x28, 0x29], + library="adafruit_bno055", # module import path + module_class="BNO055_I2C", # driver class inside the module + param_names=IMUParamNames(i2c="i2c", address="address"), + pre_config=[ + # Switch to CONFIG mode + PreConfigStep( + name="mode", + args=("CONFIG_MODE",), + step_type=PreConfigStepType.SET, + ), + # Wait for sensor to switch modes + PreConfigStep( + name="time.sleep", + args=(0.025,), + step_type=PreConfigStepType.CALL, + ), + # Set sensor ranges and bandwidths + PreConfigStep( + name="accel_range", + args=("ACCEL_4G",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="gyro_range", + args=("GYRO_2000_DPS",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="accel_bandwidth", + args=("ACCEL_125HZ",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="gyro_bandwidth", + args=("GYRO_116HZ",), + step_type=PreConfigStepType.SET, + ), + # Wait for settings to take effect + PreConfigStep( + name="time.sleep", + args=(0.025,), + step_type=PreConfigStepType.CALL, + ), + # Switch to AMG mode where accel, gyro and mag are on + PreConfigStep( + name="mode", + args=("AMG_MODE",), + step_type=PreConfigStepType.SET, + ), + ], + ), + }, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + IMUSensorTypes.mag: IMUDeviceID.IMU0, + }, + accel_range_g=4.0, + gyro_range_dps=2000.0, + filter_config=FilterConfig(freq_hz=100.0, gain=0.002250), + # Note: Gyro range setting does not actually work on the BNO055 +) + +LSM6DSOX_LIS3MDL = IMUConfig( + devices={ + IMUDeviceID.IMU0: SensorConfig( + name="LSM6DSOX", + addresses=[0x6A, 0x6B], + library="adafruit_lsm6ds.lsm6dsox", + module_class="LSM6DSOX", + param_names=IMUParamNames(i2c="i2c_bus", address="address"), + constants_module="adafruit_lsm6ds", + pre_config=[ + PreConfigStep( + name="accelerometer_range", + args=("AccelRange.RANGE_4G",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="gyro_range", + args=("GyroRange.RANGE_500_DPS",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="accelerometer_data_rate", + args=("Rate.RATE_416_HZ",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="gyro_data_rate", + args=("Rate.RATE_416_HZ",), + step_type=PreConfigStepType.SET, + ), + ], + ), + IMUDeviceID.IMU1: SensorConfig( + name="LIS3MDL", + addresses=[0x1C, 0x1E], + library="adafruit_lis3mdl", + module_class="LIS3MDL", + param_names=IMUParamNames(i2c="i2c_bus", address="address"), + pre_config=[ + PreConfigStep( + name="range", + args=("Range.RANGE_4_GAUSS",), + step_type=PreConfigStepType.SET, + ), + PreConfigStep( + name="data_rate", + args=("Rate.RATE_40_HZ",), + step_type=PreConfigStepType.SET, + ), + ], + ), + }, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + IMUSensorTypes.mag: IMUDeviceID.IMU1, + }, + accel_range_g=4.0, + gyro_range_dps=500.0, + filter_config=FilterConfig(freq_hz=104.0, gain=0.000573), +) + +BNO08X = IMUConfig( + devices={ + IMUDeviceID.IMU0: SensorConfig( + name="BNO08x", + addresses=[0x4A, 0x4B], + library="adafruit_bno08x.i2c", + module_class="BNO08X_I2C", + param_names=IMUParamNames(i2c="i2c_bus", address="address"), + constants_module="adafruit_bno08x", + pre_config=[ + PreConfigStep( + name="enable_feature", + args=("BNO_REPORT_ACCELEROMETER", 10000), + step_type=PreConfigStepType.CALL, + ), + PreConfigStep( + name="enable_feature", + args=("BNO_REPORT_GYROSCOPE", 10000), + step_type=PreConfigStepType.CALL, + ), + PreConfigStep( + name="enable_feature", + args=("BNO_REPORT_MAGNETOMETER",), + step_type=PreConfigStepType.CALL, + ), + PreConfigStep( + name="enable_feature", + args=("BNO_REPORT_ROTATION_VECTOR", 10000), + step_type=PreConfigStepType.CALL, + ), + ], + ), + }, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + IMUSensorTypes.mag: IMUDeviceID.IMU0, + IMUSensorTypes.quat: IMUDeviceID.IMU0, + }, + accel_range_g=8.0, # default 8g, not settable in driver + gyro_range_dps=2000.0, # default 2000dps, not settable in driver +) + +MOCK = IMUConfig( + devices={ + IMUDeviceID.IMU0: SensorConfig( + name=MOCK_NAME, + addresses=[0x00, 0x01], # fake I2C addresses for testing + library="imu_python.base_classes", # module path (corrected) + module_class="AdafruitIMU", # driver class + param_names=IMUParamNames(i2c="i2c", address="address"), + ), + }, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + }, + accel_range_g=8.0, + gyro_range_dps=2000.0, +) diff --git a/src/imu_python/definitions.py b/src/imu_python/definitions.py index 13a6ee9..af3dc92 100644 --- a/src/imu_python/definitions.py +++ b/src/imu_python/definitions.py @@ -86,6 +86,8 @@ def __iter__(self): I2C_ERROR = EREMOTEIO +MOCK_NAME = "MOCK" + @dataclass class IMUUpdateTime: diff --git a/src/imu_python/devices.py b/src/imu_python/devices.py index 09e4013..85176ab 100644 --- a/src/imu_python/devices.py +++ b/src/imu_python/devices.py @@ -1,296 +1,101 @@ -"""Enum registry of IMU device configurations.""" +"""Functions for looking up IMU device configurations from the registry.""" from dataclasses import replace -from enum import Enum from loguru import logger from imu_python.base_classes import ( IMUConfig, - IMUParamNames, - IMUSensorTypes, - PreConfigStep, - PreConfigStepType, - SensorConfig, ) -from imu_python.definitions import FilterConfig, IMUDescriptor, IMUDeviceID +from imu_python.definitions import MOCK_NAME, IMUDescriptor +from imu_python.registry import IMU_DEVICES -class IMUDevices(Enum): - """Enumeration containing configuration for all supported IMU devices.""" +def get_mock() -> tuple[str, IMUConfig]: + """Return a MOCK IMU with name and IMUConfig. - BNO055 = IMUConfig( - devices={ - IMUDeviceID.IMU0: SensorConfig( - name="BNO055", - addresses=[0x28, 0x29], - library="adafruit_bno055", # module import path - module_class="BNO055_I2C", # driver class inside the module - param_names=IMUParamNames(i2c="i2c", address="address"), - pre_config=[ - # Switch to CONFIG mode - PreConfigStep( - name="mode", - args=("CONFIG_MODE",), - step_type=PreConfigStepType.SET, - ), - # Wait for sensor to switch modes - PreConfigStep( - name="time.sleep", - args=(0.025,), - step_type=PreConfigStepType.CALL, - ), - # Set sensor ranges and bandwidths - PreConfigStep( - name="accel_range", - args=("ACCEL_4G",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="gyro_range", - args=("GYRO_2000_DPS",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="accel_bandwidth", - args=("ACCEL_125HZ",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="gyro_bandwidth", - args=("GYRO_116HZ",), - step_type=PreConfigStepType.SET, - ), - # Wait for settings to take effect - PreConfigStep( - name="time.sleep", - args=(0.025,), - step_type=PreConfigStepType.CALL, - ), - # Switch to AMG mode where accel, gyro and mag are on - PreConfigStep( - name="mode", - args=("AMG_MODE",), - step_type=PreConfigStepType.SET, - ), - ], - ), - }, - roles={ - IMUSensorTypes.accel: IMUDeviceID.IMU0, - IMUSensorTypes.gyro: IMUDeviceID.IMU0, - IMUSensorTypes.mag: IMUDeviceID.IMU0, - }, - accel_range_g=4.0, - gyro_range_dps=2000.0, - filter_config=FilterConfig(freq_hz=100.0, gain=0.002250), - # Note: Gyro range setting does not actually work on the BNO055 - ) + :return: tuple with the name and IMUConfig of the MOCK IMU. + """ + mock_name = MOCK_NAME + return mock_name, IMU_DEVICES[mock_name] - LSM6DSOX_LIS3MDL = IMUConfig( - devices={ - IMUDeviceID.IMU0: SensorConfig( - name="LSM6DSOX", - addresses=[0x6A, 0x6B], - library="adafruit_lsm6ds.lsm6dsox", - module_class="LSM6DSOX", - param_names=IMUParamNames(i2c="i2c_bus", address="address"), - constants_module="adafruit_lsm6ds", - pre_config=[ - PreConfigStep( - name="accelerometer_range", - args=("AccelRange.RANGE_4G",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="gyro_range", - args=("GyroRange.RANGE_500_DPS",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="accelerometer_data_rate", - args=("Rate.RATE_416_HZ",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="gyro_data_rate", - args=("Rate.RATE_416_HZ",), - step_type=PreConfigStepType.SET, - ), - ], - ), - IMUDeviceID.IMU1: SensorConfig( - name="LIS3MDL", - addresses=[0x1C, 0x1E], - library="adafruit_lis3mdl", - module_class="LIS3MDL", - param_names=IMUParamNames(i2c="i2c_bus", address="address"), - pre_config=[ - PreConfigStep( - name="range", - args=("Range.RANGE_4_GAUSS",), - step_type=PreConfigStepType.SET, - ), - PreConfigStep( - name="data_rate", - args=("Rate.RATE_40_HZ",), - step_type=PreConfigStepType.SET, - ), - ], - ), - }, - roles={ - IMUSensorTypes.accel: IMUDeviceID.IMU0, - IMUSensorTypes.gyro: IMUDeviceID.IMU0, - IMUSensorTypes.mag: IMUDeviceID.IMU1, - }, - accel_range_g=4.0, - gyro_range_dps=500.0, - filter_config=FilterConfig(freq_hz=104.0, gain=0.000573), - ) - BNO08x = IMUConfig( - devices={ - IMUDeviceID.IMU0: SensorConfig( - name="BNO08x", - addresses=[0x4A, 0x4B], - library="adafruit_bno08x.i2c", - module_class="BNO08X_I2C", - param_names=IMUParamNames(i2c="i2c_bus", address="address"), - constants_module="adafruit_bno08x", - pre_config=[ - PreConfigStep( - name="enable_feature", - args=("BNO_REPORT_ACCELEROMETER", 10000), - step_type=PreConfigStepType.CALL, - ), - PreConfigStep( - name="enable_feature", - args=("BNO_REPORT_GYROSCOPE", 10000), - step_type=PreConfigStepType.CALL, - ), - PreConfigStep( - name="enable_feature", - args=("BNO_REPORT_MAGNETOMETER",), - step_type=PreConfigStepType.CALL, - ), - ], - ), - }, - roles={ - IMUSensorTypes.accel: IMUDeviceID.IMU0, - IMUSensorTypes.gyro: IMUDeviceID.IMU0, - IMUSensorTypes.mag: IMUDeviceID.IMU0, - }, - accel_range_g=8.0, # default 8g, not settable in driver - gyro_range_dps=2000.0, # default 2000dps, not settable in driver - filter_config=FilterConfig( - freq_hz=20.0, gain=0.001538 - ), # 50 ms update interval by default - ) +def _from_address(addr: int) -> tuple[IMUDescriptor, IMUConfig] | None: + """Return a tuple containing IMU anchor information and IMU config based on the given address, or None if unknown. - MOCK = IMUConfig( - devices={ - IMUDeviceID.IMU0: SensorConfig( - name="MOCK", - addresses=[0x00, 0x01], # fake I2C addresses for testing - library="imu_python.base_classes", # module path (corrected) - module_class="AdafruitIMU", # driver class - param_names=IMUParamNames(i2c="i2c", address="address"), - ), - }, - roles={ - IMUSensorTypes.accel: IMUDeviceID.IMU0, - IMUSensorTypes.gyro: IMUDeviceID.IMU0, - }, - accel_range_g=8.0, - gyro_range_dps=2000.0, - ) - - @property - def config(self) -> IMUConfig: - """Return the IMUConfig stored inside the enum member.""" - return self.value - - @staticmethod - def _from_address(addr: int) -> tuple[IMUDescriptor, IMUConfig] | None: - """Return a tuple containing IMU anchor information and IMU config based on the given address, or None if unknown. - - :param addr: I2C address of the device - :return: Information of the IMU (IMUDescriptor, IMUConfig) of the matched device or None - """ - for device in IMUDevices: - base_config: IMUConfig = device.value - - for dev_id, sensor_cfg in base_config.devices.items(): - if addr not in sensor_cfg.addresses: - continue - - # Compute instance index to be addr index - instance_idx = sensor_cfg.addresses.index(addr) - # narrow down to a single address - narrowed_sensor = replace(sensor_cfg, addresses=[addr]) - # narrow down the device dict to contain only this device - devices = {dev_id: narrowed_sensor} - # narrow down the role dict to contain only this device - roles = { - role: target - for role, target in base_config.roles.items() - if target == dev_id - } - # combine the changes - partial_config = replace( - base_config, - devices=devices, - roles=roles, - ) - - # anchor = device name and index of the address in address list - key = IMUDescriptor(name=device.name, index=instance_idx) - - logger.trace( - f"Address 0x{addr:02X} matched to device {device.name} index {instance_idx}" - ) - return key, partial_config - - return None - - @staticmethod - def get_config(addresses: list[int]) -> dict[IMUDescriptor, IMUConfig]: - """Return a dictionary mapping imu name and device index to IMU Configs based on a list of addresses. - - The device index is the same as its address index on the address list, which is used to distinguish between multiple IMU devices of the same model. - It is assumed that a split IMU has either high or low addresses across all of its devices. i.e. - LSM6DSOX+LIS3MDL has the addresses 0x6A and 0x1C or 0x6B and 0x1E. - - :param addresses: list of detected addresses - :return: A dictionary of IMUDescriptor as keys and IMUConfigs as values. - """ - detected: dict[IMUDescriptor, IMUConfig] = {} - - for addr in addresses: - result = IMUDevices._from_address(addr) - if not result: + :param addr: I2C address of the device + :return: Information of the IMU (IMUDescriptor, IMUConfig) of the matched device or None + """ + for device, base_config in IMU_DEVICES.items(): + for dev_id, sensor_cfg in base_config.devices.items(): + if addr not in sensor_cfg.addresses: continue - key, partial = result - - if key not in detected: - detected[key] = partial - else: - # merge partial IMUConfigs if devices belong to the same IMU - base = detected[key] - # add this device to the device list of IMUConfig with the same key - merged_devices = dict(base.devices) - merged_devices.update(partial.devices) - # update roles to include roles of this device - merged_roles = dict(base.roles) - merged_roles.update(partial.roles) - # update IMUConfig - detected[key] = replace( - base, - devices=merged_devices, - roles=merged_roles, - ) - - return detected + # Compute instance index to be addr index + instance_idx = sensor_cfg.addresses.index(addr) + # narrow down to a single address + narrowed_sensor = replace(sensor_cfg, addresses=[addr]) + # narrow down the device dict to contain only this device + devices = {dev_id: narrowed_sensor} + # narrow down the role dict to contain only this device + roles = { + role: target + for role, target in base_config.roles.items() + if target == dev_id + } + # combine the changes + partial_config = replace( + base_config, + devices=devices, + roles=roles, + ) + + # anchor = device name and index of the address in address list + key = IMUDescriptor(name=device, index=instance_idx) + + logger.trace( + f"Address 0x{addr:02X} matched to device {device} index {instance_idx}" + ) + return key, partial_config + + return None + + +def get_config(addresses: list[int]) -> dict[IMUDescriptor, IMUConfig]: + """Return a dictionary mapping imu name and device index to IMU Configs based on a list of addresses. + + The device index is the same as its address index on the address list, which is used to distinguish between multiple IMU devices of the same model. + It is assumed that a split IMU has either high or low addresses across all of its devices. i.e. + LSM6DSOX+LIS3MDL has the addresses 0x6A and 0x1C or 0x6B and 0x1E. + + :param addresses: list of detected addresses + :return: A dictionary of IMUDescriptor as keys and IMUConfigs as values. + """ + detected: dict[IMUDescriptor, IMUConfig] = {} + + for addr in addresses: + result = _from_address(addr) + if not result: + continue + + key, partial = result + + if key not in detected: + detected[key] = partial + else: + # merge partial IMUConfigs if devices belong to the same IMU + base = detected[key] + # add this device to the device list of IMUConfig with the same key + merged_devices = dict(base.devices) + merged_devices.update(partial.devices) + # update roles to include roles of this device + merged_roles = dict(base.roles) + merged_roles.update(partial.roles) + # update IMUConfig + detected[key] = replace( + base, + devices=merged_devices, + roles=merged_roles, + ) + + return detected diff --git a/src/imu_python/factory.py b/src/imu_python/factory.py index 83201c0..42022c6 100644 --- a/src/imu_python/factory.py +++ b/src/imu_python/factory.py @@ -6,7 +6,7 @@ from loguru import logger from imu_python.definitions import CORE_COUNT, GIL_ENABLED, I2CBusID -from imu_python.devices import IMUDevices +from imu_python.devices import get_config, get_mock from imu_python.i2c_bus import I2CBusDescriptor, JetsonBus from imu_python.sensor_manager import IMUManager from imu_python.wrapper import IMUWrapper @@ -20,10 +20,10 @@ def detect_and_create( free_threading: bool = True, log_data: bool = False, calibration_mode: bool = False, - ): + ) -> list[IMUManager]: """Automatically detect addresses on all buses defined in I2CBUSID and create sensor managers. - :param free_threading: + :param free_threading: Flag to enable free threading. :param log_data: Flag to record the IMU data. :param calibration_mode: Flag to use calibration mode. :return: list of IMUManager instances. @@ -77,7 +77,7 @@ def _detect_and_create_per_bus( addresses = IMUFactory.scan_i2c_bus(i2c=i2c_bus) - detected_configs = IMUDevices.get_config(addresses=addresses) + detected_configs = get_config(addresses=addresses) for imu_descriptor, cfg in detected_configs.items(): imu_wrapper = IMUWrapper( @@ -115,9 +115,6 @@ def scan_i2c_bus(i2c: Any) -> list[int]: i2c.unlock() return addresses except Exception as err: - logger.warning( - f"I2C scan failed: {err}. Returning {IMUDevices.MOCK.config} addresses." - ) - return [ - a for d in IMUDevices.MOCK.config.devices.values() for a in d.addresses - ] + name, mock = get_mock() + logger.warning(f"I2C scan failed: {err}. Returning {name} addresses.") + return [a for d in mock.devices.values() for a in d.addresses] diff --git a/src/imu_python/registry.py b/src/imu_python/registry.py new file mode 100644 index 0000000..ef520e6 --- /dev/null +++ b/src/imu_python/registry.py @@ -0,0 +1,61 @@ +"""Registry for device entry points.""" + +from importlib.metadata import entry_points + +from loguru import logger + +from imu_python.base_classes import IMUConfig +from imu_python.builtin_devices import BNO055, BNO08X, LSM6DSOX_LIS3MDL, MOCK +from imu_python.definitions import MOCK_NAME + +IMU_DEVICES: dict[str, IMUConfig] = {} + + +def _load_registry() -> dict[str, IMUConfig]: + """Load all registered IMU devices and overrides from entry points. + + :return: dict of IMU names and IMUConfigs + """ + registry: dict[str, IMUConfig] = {} + + # register built-in IMUs + registry[MOCK_NAME] = MOCK + registry["BNO08X"] = BNO08X + registry["BNO055"] = BNO055 + registry["LSM6DSOX_LIS3MDL"] = LSM6DSOX_LIS3MDL + + for ep in entry_points(group="imu_module.devices"): + config = ep.load() + if not isinstance(config, IMUConfig): + logger.warning( + f"Entry point {ep.name} did not return an IMUConfig, skipping" + ) + continue + registry[ep.name] = config + logger.info(f"loaded IMU config {ep.name}") + + for ep in entry_points(group="imu_module.device_overrides"): + config = ep.load() + if not isinstance(config, IMUConfig): + logger.warning( + f"Entry point {ep.name} did not return an IMUConfig, skipping" + ) + continue + registry[ep.name] = config + logger.info(f"overrode IMU config {ep.name}") + + return registry + + +def reload_registry() -> None: + """Reload all registered IMU devices and overrides. + + IMU config changes are applied only during init time. + :return: None + """ + IMU_DEVICES.clear() + IMU_DEVICES.update(_load_registry()) + + +# Build registry at import time +reload_registry() diff --git a/src/imu_python/sensor_manager.py b/src/imu_python/sensor_manager.py index 29dd2cf..0811dcd 100644 --- a/src/imu_python/sensor_manager.py +++ b/src/imu_python/sensor_manager.py @@ -91,31 +91,37 @@ def _loop(self) -> None: # Attempt to read all sensor data with self.i2c_lock: data = self.imu_wrapper.get_imu_data() + if isinstance(data, IMUData): + with self.data_lock: + self.latest_data = data + if self.log_data: + self.IMUData_log.append(self.latest_data) # Ensure new data - if self._acc_gyro_are_fresh(data): + elif self._acc_gyro_are_fresh(data): logger.debug( f"reading from: {self.imu_descriptor.name} {self.imu_descriptor.index} new data:{data}" ) + + timestamp = time.monotonic() + pose_quat = self.imu_wrapper.filter.update( + timestamp=timestamp, + accel=data.accel.as_array(), + gyro=data.gyro.as_array(), + mag=data.mag.as_array() + if data.mag is not None and self._mag_is_fresh(data) + else None, + clipped=( + data.accel.is_clipped( + sensor_range=self.accel_range_m_s2, + sensor_type="Accel", + ) + or data.gyro.is_clipped( + sensor_range=self.gyro_range_rad_s, + sensor_type="Gyro", + ) + ), + ) with self.data_lock: - timestamp = time.monotonic() - pose_quat = self.imu_wrapper.filter.update( - timestamp=timestamp, - accel=data.accel.as_array(), - gyro=data.gyro.as_array(), - mag=data.mag.as_array() - if data.mag is not None and self._mag_is_fresh(data) - else None, - clipped=( - data.accel.is_clipped( - sensor_range=self.accel_range_m_s2, - sensor_type="Accel", - ) - or data.gyro.is_clipped( - sensor_range=self.gyro_range_rad_s, - sensor_type="Gyro", - ) - ), - ) self.latest_data = IMUData( timestamp=timestamp, quat=pose_quat, diff --git a/src/imu_python/wrapper.py b/src/imu_python/wrapper.py index fdcfc25..73d375e 100644 --- a/src/imu_python/wrapper.py +++ b/src/imu_python/wrapper.py @@ -1,6 +1,7 @@ """Wrapper class for the IMUs.""" import importlib +import time import types from collections.abc import Callable, Iterable from numbers import Number @@ -12,8 +13,10 @@ from imu_python.base_classes import ( AdafruitIMU, IMUConfig, + IMUData, IMUDeviceData, IMUSensorTypes, + Quaternion, SensorConfig, VectorXYZ, ) @@ -57,9 +60,6 @@ def __init__( self.i2c_bus_instance: ExtendedI2C | None = i2c_bus_descriptor.bus_instance self.i2c_bus_id: I2CBusID | None = i2c_bus_descriptor.bus_id self.started: bool = False - self.filter: BaseIMUFilter = MadgwickFilterPyImu( - config=self.config.filter_config - ) self.rotation_matrix: NDArray = DEFAULT_ROTATION_MATRIX self._devices: dict[ IMUDeviceID, AdafruitIMU @@ -90,6 +90,12 @@ def __init__( self.mag_calibration: tuple[NDArray, NDArray] = mag_calibration logger.info(f"Loaded magnetometer calibration for {name}.") + self.scalar_first = config.scalar_first + if IMUSensorTypes.quat not in self.role_to_device_map: + self.filter: BaseIMUFilter = MadgwickFilterPyImu( + config=self.config.filter_config + ) + def reload(self) -> None: """(Re)Initialize the IMU object.""" try: @@ -118,31 +124,45 @@ def _initialize_sensor(self, sensor_config: SensorConfig) -> AdafruitIMU: self._preconfigure_sensor(sensor=sensor, sensor_config=sensor_config) return sensor - def get_imu_data(self) -> IMUDeviceData: + def get_imu_data(self) -> IMUDeviceData | IMUData: """Return acceleration, gyro and magnetic information as an IMUData.""" accel_vector = self.read_sensor(IMUSensorTypes.accel) gyro_vector = self.read_sensor(IMUSensorTypes.gyro) mag_vector = self.read_sensor(IMUSensorTypes.mag) - if accel_vector is None or gyro_vector is None: - raise ValueError("Accel or Gyro reading is invalid.") + quat = self.read_sensor(IMUSensorTypes.quat) + if not isinstance(accel_vector, VectorXYZ): + raise ValueError("Accel reading is invalid.") + if not isinstance(gyro_vector, VectorXYZ): + raise ValueError("Gyro reading is invalid.") accel_vector.rotate(self.rotation_matrix) gyro_vector.rotate(self.rotation_matrix) - if mag_vector is not None: + if isinstance(mag_vector, VectorXYZ): mag_vector = apply_mag_cal( mag_vector=mag_vector, mag_calibration=self.mag_calibration ) mag_vector.rotate(self.rotation_matrix) - return IMUDeviceData( + elif isinstance(mag_vector, Quaternion): + raise ValueError("Mag reading is invalid.") + device_data = IMUDeviceData( accel=accel_vector, gyro=gyro_vector, mag=mag_vector, ) + if isinstance(quat, Quaternion): + quat.rotate(self.rotation_matrix) + return IMUData( + timestamp=time.monotonic(), + device_data=device_data, + quat=quat, + ) + else: + return device_data - def read_sensor(self, attr: IMUSensorTypes) -> VectorXYZ | None: - """Read the IMU attribute and return it as a VectorXYZ. + def read_sensor(self, attr: IMUSensorTypes) -> VectorXYZ | Quaternion | None: + """Read the IMU attribute and return it as a VectorXYZ or Quaternion. :param attr: attribute defined as IMUSensorType - :return: VectorXYZ data or None if not valid or available. + :return: VectorXYZ, Quaternion data or None if not valid or available. """ device_id = self.role_to_device_map.get(attr) if not device_id: @@ -162,12 +182,11 @@ def read_sensor(self, attr: IMUSensorTypes) -> VectorXYZ | None: return None return vector - @staticmethod - def _vectorize(value: Any) -> VectorXYZ | None: - """Vectorize sensor output into a VectorXYZ object. + def _vectorize(self, value: Any) -> VectorXYZ | Quaternion | None: + """Vectorize sensor output into a VectorXYZ Quaternion object. :param value: The raw sensor output value. - :return: VectorXYZ or None if the value is invalid or unavailable. + :return: VectorXYZ, Quaternion, or None if the value is invalid or unavailable. """ if value is None: return None @@ -184,10 +203,18 @@ def _vectorize(value: Any) -> VectorXYZ | None: or not all(isinstance(v, Number) for v in values) ): return None - if len(values) != 3: - return None - - return VectorXYZ.from_tuple(values) + if len(values) == 4: + if self.scalar_first: + quat = Quaternion( + w=values[0], x=values[1], y=values[2], z=values[3] + ) + else: + quat = Quaternion( + w=values[3], x=values[0], y=values[1], z=values[2] + ) + return quat + if len(values) == 3: + return VectorXYZ.from_tuple(values) return None diff --git a/tests/base_classes_test.py b/tests/base_classes_test.py index 84826bc..8f7862e 100644 --- a/tests/base_classes_test.py +++ b/tests/base_classes_test.py @@ -116,3 +116,25 @@ def test_quaternion_to_euler(rot_x_rad: float) -> None: np.testing.assert_almost_equal(euler.x, rot_x_rad) np.testing.assert_almost_equal(euler.y, 0) np.testing.assert_almost_equal(euler.z, 0) + + +def test_quaternion_rotate() -> None: + """Test quaternion rotation.""" + # Arrange + quat = Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + rotation_matrix = np.array( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + + # Act + quat.rotate(rotation_matrix) + + # Assert + np.testing.assert_almost_equal(quat.w, np.sin(np.pi / 4)) + np.testing.assert_almost_equal(quat.x, 0.0) + np.testing.assert_almost_equal(quat.y, 0.0) + np.testing.assert_almost_equal(quat.z, np.cos(np.pi / 4)) diff --git a/tests/devices_test.py b/tests/devices_test.py index ea67f45..b26e000 100644 --- a/tests/devices_test.py +++ b/tests/devices_test.py @@ -5,14 +5,14 @@ import pytest from imu_python.base_classes import IMUConfig, IMUSensorTypes -from imu_python.definitions import IMUDescriptor, IMUDeviceID -from imu_python.devices import IMUDevices +from imu_python.definitions import MOCK_NAME, IMUDescriptor, IMUDeviceID +from imu_python.devices import IMU_DEVICES, _from_address, get_config def test_device_addresses() -> None: """Test the IMU device addresses.""" - for device in IMUDevices: - assert (len(a.addresses) == 2 for a in device.config.devices.values()) + for device in IMU_DEVICES: + assert all(len(a.addresses) == 2 for a in IMU_DEVICES[device].devices.values()) def test_device_from_bad_address() -> None: @@ -21,7 +21,7 @@ def test_device_from_bad_address() -> None: invalid_addr = 0xAA # Act - config = IMUDevices._from_address(addr=invalid_addr) + config = _from_address(addr=invalid_addr) # Assert assert config is None @@ -29,12 +29,12 @@ def test_device_from_bad_address() -> None: @pytest.mark.parametrize( "valid_address", - [addr for dev in IMUDevices.MOCK.config.devices.values() for addr in dev.addresses], + [addr for dev in IMU_DEVICES[MOCK_NAME].devices.values() for addr in dev.addresses], ) def test_device_from_valid_address(valid_address: int) -> None: """Test the IMU device addresses.""" # Act - imu_info = IMUDevices._from_address(addr=valid_address) + imu_info = _from_address(addr=valid_address) # Assert assert imu_info is not None @@ -49,7 +49,7 @@ def test_merge_partial_configs() -> None: # Build a mutated MOCK config with a second device that shares # the same address index (1) for its address list so merging occurs. - config = IMUDevices.MOCK.config + config = IMU_DEVICES[MOCK_NAME] imu0_sensor = config.devices[IMUDeviceID.IMU0] # create a second sensor with a different address pair imu1_sensor = replace(imu0_sensor, name=IMUDeviceID.IMU1, addresses=[0x10, 0x11]) @@ -59,11 +59,11 @@ def test_merge_partial_configs() -> None: mutated_config.roles.update({IMUSensorTypes.mag: IMUDeviceID.IMU1}) # Patch the MOCK member so that there are two devices in MOCK - with patch.object(IMUDevices.MOCK, "_value_", mutated_config): + with patch.dict(IMU_DEVICES, {MOCK_NAME: mutated_config}): # Choose addresses that map to index 1 in both device address lists - detected = IMUDevices.get_config([0x01, 0x11]) + detected = get_config([0x01, 0x11]) - dsc = IMUDescriptor(name=IMUDevices.MOCK.name, index=1) + dsc = IMUDescriptor(MOCK_NAME, index=1) cfg = detected[dsc] assert dsc in detected assert isinstance(cfg, IMUConfig) diff --git a/tests/factory_test.py b/tests/factory_test.py index f7badb9..ed2f0db 100644 --- a/tests/factory_test.py +++ b/tests/factory_test.py @@ -1,13 +1,13 @@ """Test the IMUFactory class.""" -from src.imu_python.devices import IMUDevices +from src.imu_python.devices import IMU_DEVICES, get_mock from src.imu_python.factory import IMUFactory def test_imu_factory() -> None: """Test the IMUFactory class.""" # Arrange - mock_imu_name = IMUDevices.MOCK.name + mock_imu_name, mock_imu_config = get_mock() # Act imu_managers = IMUFactory.detect_and_create() @@ -32,4 +32,5 @@ def test_imu_factory() -> None: getattr(device, attr_name, None) # check that config matches the expected IMU name - assert mock_imu_name in IMUDevices.__members__ + assert mock_imu_name in IMU_DEVICES + assert mock_imu_config == IMU_DEVICES[mock_imu_name] diff --git a/tests/registry_test.py b/tests/registry_test.py new file mode 100644 index 0000000..ec26bf6 --- /dev/null +++ b/tests/registry_test.py @@ -0,0 +1,64 @@ +"""Tests for the device registry loading and overriding mechanism.""" + +from dataclasses import replace +from importlib.metadata import EntryPoint +from unittest.mock import MagicMock, patch + +from imu_python.base_classes import IMUConfig +from imu_python.builtin_devices import MOCK +from imu_python.definitions import MOCK_NAME +from imu_python.registry import _load_registry + + +def test_builtin_devices_are_loaded(): + """Test that the built-in devices are loaded into the registry.""" + registry = _load_registry() + assert len(registry) > 0 + for name, config in registry.items(): + assert isinstance(config, IMUConfig), f"{name} did not load as IMUConfig" + + +def test_known_builtin_present(): + """Test that known built-in devices are present in the registry.""" + registry = _load_registry() + assert MOCK_NAME in registry + + +def test_register_device(): + """Test that a new device can be registered via entry points.""" + new_device_config = MOCK + + mock_ep = MagicMock(spec=EntryPoint) + mock_ep.name = "TestDevice" + mock_ep.load.return_value = new_device_config + + with patch("imu_python.registry.entry_points") as mock_entry_points: + mock_entry_points.side_effect = lambda group: ( + [mock_ep] if group == "imu_module.devices" else [] + ) + registry = _load_registry() + + assert "TestDevice" in registry + assert registry["TestDevice"] == new_device_config + + +def test_override_replaces_builtin(): + """Test that device overrides replace built-in devices.""" + builtin_mock_config = MOCK + modified = replace(builtin_mock_config, accel_range_g=16.0) + + mock_override = MagicMock(spec=EntryPoint) + mock_override.name = MOCK_NAME + mock_override.load.return_value = modified + + with patch("imu_python.registry.entry_points") as mock_ep: + mock_ep.side_effect = lambda group: ( + [] + if group == "imu_module.devices" + else [mock_override] + if group == "imu_module.device_overrides" + else [] + ) + registry = _load_registry() + + assert registry[MOCK_NAME].accel_range_g == 16.0 diff --git a/tests/sensor_manager_test.py b/tests/sensor_manager_test.py index 273bf57..fa14841 100644 --- a/tests/sensor_manager_test.py +++ b/tests/sensor_manager_test.py @@ -2,11 +2,20 @@ import threading import time +from dataclasses import replace +from unittest.mock import MagicMock, patch import pytest +from imu_python.base_classes import ( + IMUData, + IMUDeviceData, + IMUSensorTypes, + Quaternion, + VectorXYZ, +) from imu_python.definitions import IMUDescriptor, IMUDeviceID -from imu_python.devices import IMUDevices +from imu_python.devices import get_mock from imu_python.i2c_bus import I2CBusDescriptor from imu_python.sensor_manager import IMUManager from imu_python.wrapper import IMUWrapper @@ -14,15 +23,16 @@ @pytest.fixture def imu_setup() -> IMUManager: - """Fixture providing sensor_manager for tests.""" + """Fixture providing manager for tests.""" + name, config = get_mock() wrapper = IMUWrapper( - config=IMUDevices.MOCK.config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + config=config, + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) lock = threading.Lock() - sensor_manager = IMUManager(imu_wrapper=wrapper, i2c_lock=lock) - return sensor_manager + manager = IMUManager(imu_wrapper=wrapper, i2c_lock=lock) + return manager def test_manager_get_data(imu_setup: IMUManager) -> None: @@ -43,6 +53,67 @@ def test_manager_get_data(imu_setup: IMUManager) -> None: assert data.device_data.mag is None # mock IMU has no mag device +def test_manager_get_data_mag() -> None: + """Test if manager can get data with mag.""" + # Arrange + name, config = get_mock() + config_mag = replace( + config, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + IMUSensorTypes.mag: IMUDeviceID.IMU0, + }, + ) + wrapper = IMUWrapper( + config=config_mag, + imu_descriptor=IMUDescriptor(name=name, index=0), + i2c_bus_descriptor=I2CBusDescriptor(None, None), + calibration_mode=True, # ignore calibration requirement for the purpose of the test + ) + lock = threading.Lock() + manager = IMUManager(imu_wrapper=wrapper, i2c_lock=lock) + + # Act + manager.start() + time.sleep(0.01) # wait for data to be read + data = manager.get_data() + manager.stop() + + # Assert + assert data is not None + assert data.device_data.mag is not None + + +def test_manager_get_on_board_quat(imu_setup: IMUManager) -> None: + """Test if manager can get on-board fusion quaternion.""" + # Arrange + sensor_manager = imu_setup + + expected_quat = Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) + expected_imu_data = IMUData( + timestamp=0.0, + quat=expected_quat, + device_data=IMUDeviceData( + accel=VectorXYZ(x=0.0, y=0.0, z=0.0), + gyro=VectorXYZ(x=0.0, y=0.0, z=0.0), + ), + ) + + with patch.object( + sensor_manager.imu_wrapper, "get_imu_data", return_value=expected_imu_data + ): + # Act + sensor_manager.start() + time.sleep(0.01) # wait for data to be read + data = sensor_manager.get_data() + sensor_manager.stop() + + # Assert + assert data is not None and data == expected_imu_data + assert data.quat == expected_quat + + def test_manager_apply_remapping(imu_setup: IMUManager) -> None: """Test if manager applies a rotation_matrix to wrapper.""" import numpy as np @@ -93,11 +164,12 @@ def test_manager_pauses_during_disconnect(imu_setup: IMUManager) -> None: def test_manager_records_data() -> None: """Test if manager records data when logging is enabled.""" # Arrange - from unittest.mock import MagicMock, patch + from unittest.mock import patch + name, config = get_mock() wrapper = IMUWrapper( - config=IMUDevices.MOCK.config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + config=config, + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) mock_writer = MagicMock() @@ -118,3 +190,16 @@ def test_manager_records_data() -> None: assert isinstance(appended_arg, list) assert len(appended_arg) > 0 mock_writer.save_dataframe.assert_called_once() + + +def test_manager_set_core_affinity(imu_setup: IMUManager) -> None: + """Test if manager sets core affinity correctly.""" + # Arrange + sensor_manager = imu_setup + core_id = 0 + + # Act + sensor_manager.set_core_affinity(core_id) + + # Assert + assert sensor_manager.core_id == core_id diff --git a/tests/wrapper_test.py b/tests/wrapper_test.py index f32c3b9..4082dfa 100644 --- a/tests/wrapper_test.py +++ b/tests/wrapper_test.py @@ -1,6 +1,5 @@ """Test the factory and manager for the imu sensor objects.""" -import copy from dataclasses import replace from unittest.mock import MagicMock, patch @@ -9,14 +8,16 @@ from imu_python.base_classes import ( IMUConfig, + IMUData, IMUSensorTypes, PreConfigStep, PreConfigStepType, + Quaternion, VectorXYZ, ) from imu_python.calibration.mag_calibration import apply_mag_cal from imu_python.definitions import IMUDescriptor, IMUDeviceID -from imu_python.devices import IMUDevices +from imu_python.devices import get_mock from imu_python.i2c_bus import I2CBusDescriptor from imu_python.wrapper import IMUWrapper @@ -36,34 +37,37 @@ def mutate_sensor( return replace(cfg, devices=new_devices) -def test_imu_wrapper() -> None: - """Test the imu wrapper class.""" - # Arrange - config = IMUDevices.MOCK.config +@pytest.fixture +def wrapper_setup() -> IMUWrapper: + """Fixture providing wrapper for tests.""" + name, config = get_mock() - # Act wrapper = IMUWrapper( config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) + return wrapper + + +def test_imu_wrapper(wrapper_setup: IMUWrapper) -> None: + """Test the imu wrapper class.""" + # Arrange + wrapper = wrapper_setup + + # Act wrapper.reload() # Assert assert wrapper.started -def test_imu_wrapper_attr_with_no_role() -> None: +def test_imu_wrapper_attr_with_no_role(wrapper_setup: IMUWrapper) -> None: """Test the imu wrapper class read attribute with no role.""" # Arrange - config = IMUDevices.MOCK.config + wrapper = wrapper_setup # Act - wrapper = IMUWrapper( - config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), - i2c_bus_descriptor=I2CBusDescriptor(None, None), - ) wrapper.reload() assert wrapper.read_sensor(IMUSensorTypes.mag) is None @@ -72,13 +76,15 @@ def test_imu_wrapper_attr_with_no_role() -> None: def test_imu_wrapper_attr_with_no_device() -> None: """Test the imu wrapper class read attribute with no device.""" # Arrange - config = IMUDevices.MOCK.config - config.roles.update({IMUSensorTypes.mag: IMUDeviceID.IMU1}) + name, config = get_mock() + config_with_mag = replace( + config, roles={**config.roles, IMUSensorTypes.mag: IMUDeviceID.IMU1} + ) # Act wrapper = IMUWrapper( - config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + config=config_with_mag, + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) wrapper.reload() @@ -212,12 +218,12 @@ def test_imu_wrapper_attr_with_no_device() -> None: ) def test_imu_wrapper_reload_fails(reason, mutate_config): """Test if wrapper raises runtime error with bad IMU Configs.""" - config = copy.deepcopy(IMUDevices.MOCK.config) + name, config = get_mock() config = mutate_config(config) wrapper = IMUWrapper( config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) @@ -228,7 +234,7 @@ def test_imu_wrapper_reload_fails(reason, mutate_config): def test_pre_config_with_mock() -> None: """Test if the IMU is pre-configured properly with mock.""" # Arrange - config = IMUDevices.MOCK.config + name, config = get_mock() config = mutate_sensor( cfg=config, device_id=IMUDeviceID.IMU0, @@ -259,7 +265,7 @@ def test_pre_config_with_mock() -> None: wrapper = IMUWrapper( config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) imu = MagicMock() @@ -296,7 +302,7 @@ def test_pre_config_with_mock() -> None: def test_pre_config_string(): """Test if the IMU is pre-configured properly with a string argument.""" # Arrange - config = IMUDevices.MOCK.config + name, config = get_mock() config = mutate_sensor( cfg=config, device_id=IMUDeviceID.IMU0, @@ -309,7 +315,7 @@ def test_pre_config_string(): wrapper = IMUWrapper( config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) @@ -323,7 +329,7 @@ def test_pre_config_string(): def test_pre_config_time_sleep(): """Test if time.sleep can be called in pre-configuration.""" # Arrange - config = IMUDevices.MOCK.config + name, config = get_mock() config = mutate_sensor( cfg=config, device_id=IMUDeviceID.IMU0, @@ -336,7 +342,7 @@ def test_pre_config_time_sleep(): wrapper = IMUWrapper( config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) @@ -378,13 +384,12 @@ def __iter__(self): ((1.0, None, 3.0), None), ((1.0, "a", 3.0), None), ((1.0, 2.0), None), - ((1.0, 2.0, 3.0, 4.0), None), ((1.0, 2.0, 3.0), (1.0, 2.0, 3.0)), ], ) -def test_vectorize_parametrized(value, expected) -> None: +def test_vectorize_parametrized(value, expected, wrapper_setup: IMUWrapper) -> None: """Parametrized tests covering all None-return cases and a positive case.""" - v = IMUWrapper._vectorize(value) + v = wrapper_setup._vectorize(value) if expected is None: assert v is None else: @@ -395,15 +400,15 @@ def test_vectorize_parametrized(value, expected) -> None: def test_apply_mag_calibration() -> None: """Test if magnetometer calibration is applied correctly.""" # Arrange - config = IMUDevices.MOCK.config - config.roles.update({IMUSensorTypes.mag: IMUDeviceID.IMU0}) + name, config = get_mock() + config_mag = replace(config, roles={IMUSensorTypes.mag: IMUDeviceID.IMU0}) mag_calibration = ( np.array([1.0, 2.0, 3.0]), np.array([[0.4, 0.0, 0.0], [0.0, 0.5, 0.0], [0.0, 0.0, 0.6]]), ) wrapper = IMUWrapper( - config=config, - imu_descriptor=IMUDescriptor(name="MOCK", index=0), + config=config_mag, + imu_descriptor=IMUDescriptor(name=name, index=0), i2c_bus_descriptor=I2CBusDescriptor(None, None), ) wrapper.mag_calibration = mag_calibration @@ -417,3 +422,44 @@ def test_apply_mag_calibration() -> None: # Assert assert np.allclose(calibrated_mag.as_array(), expected) + + +def test_on_board_fusion() -> None: + """Test if on-board fusion is applied correctly.""" + # Arrange + name, config = get_mock() + config_quat = replace( + config, + roles={ + IMUSensorTypes.accel: IMUDeviceID.IMU0, + IMUSensorTypes.gyro: IMUDeviceID.IMU0, + IMUSensorTypes.mag: IMUDeviceID.IMU0, + IMUSensorTypes.quat: IMUDeviceID.IMU0, + }, + ) + wrapper = IMUWrapper( + config=config_quat, + imu_descriptor=IMUDescriptor(name=name, index=0), + i2c_bus_descriptor=I2CBusDescriptor(None, None), + calibration_mode=True, # ignore calibration requirement + ) + + expected_quat = Quaternion(w=0.7071, x=0.7071, y=0.0, z=0.0) + + def read_sensor_side_effect(sensor_type): + if sensor_type == IMUSensorTypes.quat: + return expected_quat + else: + return VectorXYZ(x=0.0, y=0.0, z=0.0) + + # Act + with patch.object(wrapper, "read_sensor", side_effect=read_sensor_side_effect): + data = wrapper.get_imu_data() + + # Assert + assert isinstance(data, IMUData) + assert isinstance(data.quat, Quaternion) + assert np.allclose( + (data.quat.w, data.quat.x, data.quat.y, data.quat.z), + (expected_quat.w, expected_quat.x, expected_quat.y, expected_quat.z), + ) diff --git a/uv.lock b/uv.lock index 1220ec3..18608b7 100644 --- a/uv.lock +++ b/uv.lock @@ -644,7 +644,7 @@ wheels = [ [[package]] name = "imu-python" -version = "0.1.0" +version = "0.1.1" source = { editable = "." } dependencies = [ { name = "ahrs" },