Made rotation aware of mag

gaitmap/utils/consts.py

@@ -8,17 +8,21 @@
 SF_ACC = ["acc_x", "acc_y", "acc_z"]
 #: The default names of the Magnetometer columns in the sensor frame
 SF_MAG = ["mag_x", "mag_y", "mag_z"]
-#: The default names of all columns in the sensor frame
-SF_COLS = [*SF_ACC, *SF_GYR, *SF_MAG]
+#: The default names of all columns in the sensor frame excluding the magnetometer
+SF_COLS = [*SF_ACC, *SF_GYR]
+#: The default names of all columns in the sensor frame including the magnetometer
+SF_COLS_WITH_MAG = [*SF_ACC, *SF_GYR, *SF_MAG]
 
 #: The default names of the Gyroscope columns in the body frame
 BF_GYR = ["gyr_pa", "gyr_ml", "gyr_si"]
 #: The default names of the Accelerometer columns in the body frame
 BF_ACC = ["acc_pa", "acc_ml", "acc_si"]
 #: The default names of the Magnetometer columns in the body frame
 BF_MAG = ["mag_pa", "mag_ml", "mag_si"]
-#: The default names of all columns in the body frame
-BF_COLS = [*BF_ACC, *BF_GYR, *BF_MAG]
+#: The default names of all columns in the body frame excluding the magnetometer
+BF_COLS = [*BF_ACC, *BF_GYR]
+#: The default names of all columns in the body frame including the magnetometer
+BF_COLS_WITH_MAG = [*BF_ACC, *BF_GYR, *BF_MAG]
 
 #: The minimal required columns for a stride list
 SL_COLS = ["start", "end"]

gaitmap/utils/rotations.py

@@ -10,7 +10,7 @@
 from numpy.linalg import norm
 from scipy.spatial.transform import Rotation
 
-from gaitmap.utils.consts import GRAV_VEC, SF_ACC, SF_GYR
+from gaitmap.utils.consts import GRAV_VEC, SF_ACC, SF_GYR, SF_MAG
 from gaitmap.utils.datatype_helper import (
     SensorData,
     SingleSensorData,
@@ -71,6 +71,7 @@ def _flip_sensor(data: SingleSensorData, rotation: Optional[Rotation]) -> Single
     """
     if rotation.single is False:
         raise ValueError("Only single rotations are allowed!")
+    has_mag = is_single_sensor_data(data, frame="sensor", check_mag=True)
 
     tol = 10e-9
     rot_matrix = rotation.as_matrix().squeeze()
@@ -89,8 +90,10 @@ def _flip_sensor(data: SingleSensorData, rotation: Optional[Rotation]) -> Single
         return data
 
     orig_col_order = data.columns
-    for sensor in ["acc", "gyr"]:
-        cols = np.array({"acc": SF_ACC, "gyr": SF_GYR}[sensor])
+    sensors = ["acc", "gyr"] + (["mag"] if has_mag else [])
+    rots = {"acc": SF_ACC, "gyr": SF_GYR, "mag": SF_MAG}
+    for sensor in sensors:
+        cols = np.array(rots[sensor])
         rename = {}
         mirror = []
         # We basically iterate over the rotation matrix and find which axis is transformed to which other axis.
@@ -108,12 +111,15 @@ def _flip_sensor(data: SingleSensorData, rotation: Optional[Rotation]) -> Single
 
 
 def _rotate_sensor(data: SingleSensorData, rotation: Optional[Rotation]) -> SingleSensorData:
-    """Rotate the data of a single sensor with acc and gyro."""
+    """Rotate the data of a single sensor with acc, gyro and mag."""
+    has_mag = is_single_sensor_data(data, frame="sensor", check_mag=True)
     data = data.copy()
     if rotation is None:
         return data
     data[SF_GYR] = rotation.apply(data[SF_GYR].to_numpy())
     data[SF_ACC] = rotation.apply(data[SF_ACC].to_numpy())
+    if has_mag:
+        data[SF_MAG] = rotation.apply(data[SF_MAG].to_numpy())
     return data
 
 

tests/test_utils/test_rotations.py

@@ -5,7 +5,7 @@
 from pandas._testing import assert_frame_equal
 from scipy.spatial.transform import Rotation
 
-from gaitmap.utils.consts import SF_ACC, SF_COLS, SF_GYR
+from gaitmap.utils.consts import SF_ACC, SF_COLS, SF_GYR, SF_MAG
 from gaitmap.utils.datatype_helper import MultiSensorData, get_multi_sensor_names
 from gaitmap.utils.exceptions import ValidationError
 from gaitmap.utils.rotations import (
@@ -178,6 +178,15 @@ def test_rotate_sensor(self, cyclic_rotation) -> None:
 
         _compare_cyclic(self.sample_sensor_data, rotated_data)
 
+    def test_rotation_also_works_with_mag(self, cyclic_rotation) -> None:
+        """Test if rotation is correctly applied to gyr and acc of single sensor data."""
+        sample_data = self.sample_sensor_data.copy()
+        sample_data[SF_MAG] = sample_data[SF_ACC] + 6
+        rotated_data = self.single_func(sample_data, cyclic_rotation)
+
+        _compare_cyclic(sample_data, rotated_data)
+        assert set(SF_MAG).issubset(rotated_data.columns)
+
     def test_rotate_dataset_single(self, cyclic_rotation) -> None:
         """Rotate a single dataset with `rotate_dataset`.