pbio/orientation: Move up side getter to pbio.

pybricks · Mar 29, 2023 · acb7d43 · acb7d43
1 parent ea51fea
commit acb7d43
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 43 deletions.
diff --git a/lib/pbio/include/pbio/orientation.h b/lib/pbio/include/pbio/orientation.h
@@ -60,6 +60,8 @@ void pbio_orientation_imu_get_angular_velocity(pbio_geometry_xyz_t *values);
 
 void pbio_orientation_imu_get_acceleration(pbio_geometry_xyz_t *values);
 
+pbio_orientation_side_t pbio_orientation_imu_get_up_side(void);
+
 float pbio_orientation_imu_get_heading(void);
 
 void pbio_orientation_imu_set_heading(float desired_heading);
@@ -79,6 +81,10 @@ static inline void pbio_orientation_imu_get_angular_velocity(pbio_geometry_xyz_t
 static inline void pbio_orientation_imu_get_acceleration(pbio_geometry_xyz_t *values) {
 }
 
+static inline pbio_orientation_side_t pbio_orientation_imu_get_up_side(void) {
+    return PBIO_ORIENTATION_SIDE_TOP;
+}
+
 static inline uint32_t pbio_orientation_imu_get_stationary_count(void) {
     return 0;
 }

diff --git a/lib/pbio/src/orientation.c b/lib/pbio/src/orientation.c
@@ -85,6 +85,34 @@ void pbio_orientation_get_complementary_axis(uint8_t *index, int8_t *sign) {
         );
 }
 
+/**
+ * Gets the side of a unit-sized box through which a given vector passes first.
+ *
+ * @param [in]  vector  The input vector.
+ * @return              The side through which the vector passes first.
+ */
+pbio_orientation_side_t pbio_orientation_side_from_vector(pbio_geometry_xyz_t *vector) {
+
+    // Find index and sign of maximum component
+    float abs_max = 0;
+    uint8_t axis = 0;
+    bool negative = true;
+    for (uint8_t i = 0; i < 3; i++) {
+        if (vector->values[i] > abs_max) {
+            abs_max = vector->values[i];
+            negative = false;
+            axis = i;
+        } else if (-vector->values[i] > abs_max) {
+            abs_max = -vector->values[i];
+            negative = true;
+            axis = i;
+        }
+    }
+
+    // Return as side enum value.
+    return axis | (negative << 2);
+}
+
 #if PBIO_CONFIG_ORIENTATION_IMU
 
 pbdrv_imu_dev_t *imu_dev;
@@ -183,6 +211,20 @@ void pbio_orientation_imu_get_acceleration(pbio_geometry_xyz_t *values) {
     pbio_geometry_vector_map(&pbio_orientation_neutral_orientation, &acceleration, values);
 }
 
+/**
+ * Gets which side of a hub points upwards.
+ *
+ * @param [in] imu_dev      The driver instance.
+ * @param [out] values      An array of 3 32-bit float values to hold the result.
+ */
+pbio_orientation_side_t pbio_orientation_imu_get_up_side(void) {
+    // Up is which side of a unit box intersects the +Z vector first.
+    // So read +Z vector of the inertial frame, in the body frame.
+    // For now, this is the gravity vector. In the future, we can make this
+    // slightly more accurate by using the full IMU orientation.
+    return pbio_orientation_side_from_vector(&acceleration);
+}
+
 static float heading_offset = 0;
 
 /**

diff --git a/pybricks/common/pb_type_imu.c b/pybricks/common/pb_type_imu.c
@@ -65,49 +65,21 @@ STATIC mp_obj_t common_IMU_project_3d_axis(mp_obj_t axis_in, float *values) {
 
 // pybricks._common.IMU.up
 STATIC mp_obj_t common_IMU_up(mp_obj_t self_in) {
-    (void)self_in;
-
-    // Up is which side of a unit box intersects the +Z vector first.
-    // So read +Z vector of the inertial frame, in the body frame.
-    // For now, this is the gravity vector. In the future, we can make this
-    // slightly more accurate by using the full IMU orientation.
-    pbio_geometry_xyz_t acceleration;
-    pbio_orientation_imu_get_acceleration(&acceleration);  // REVISIT: NEED UNROTATED VALUE
-
-    // Find index and sign of maximum component
-    float abs_max = 0;
-    uint8_t axis = 0;
-    bool positive = false;
-    for (uint8_t i = 0; i < 3; i++) {
-        if (acceleration.values[i] > abs_max) {
-            abs_max = acceleration.values[i];
-            positive = true;
-            axis = i;
-        } else if (-acceleration.values[i] > abs_max) {
-            abs_max = -acceleration.values[i];
-            positive = false;
-            axis = i;
-        }
-    }
-
-    // The maximum component dictates which side of a unit box gets intersected
-    // first. So, simply look at axis and sign to give the side.
-    if (axis == 0 && positive) {
-        return MP_OBJ_FROM_PTR(&pb_Side_FRONT_obj);
-    }
-    if (axis == 0 && !positive) {
-        return MP_OBJ_FROM_PTR(&pb_Side_BACK_obj);
-    }
-    if (axis == 1 && positive) {
-        return MP_OBJ_FROM_PTR(&pb_Side_LEFT_obj);
-    }
-    if (axis == 1 && !positive) {
-        return MP_OBJ_FROM_PTR(&pb_Side_RIGHT_obj);
-    }
-    if (axis == 2 && positive) {
-        return MP_OBJ_FROM_PTR(&pb_Side_TOP_obj);
-    } else {
-        return MP_OBJ_FROM_PTR(&pb_Side_BOTTOM_obj);
+    switch (pbio_orientation_imu_get_up_side()) {
+        case PBIO_ORIENTATION_SIDE_FRONT:
+            return MP_OBJ_FROM_PTR(&pb_Side_FRONT_obj);
+        case PBIO_ORIENTATION_SIDE_LEFT:
+            return MP_OBJ_FROM_PTR(&pb_Side_LEFT_obj);
+        case PBIO_ORIENTATION_SIDE_TOP:
+            return MP_OBJ_FROM_PTR(&pb_Side_TOP_obj);
+        case PBIO_ORIENTATION_SIDE_BACK:
+            return MP_OBJ_FROM_PTR(&pb_Side_BACK_obj);
+        case PBIO_ORIENTATION_SIDE_RIGHT:
+            return MP_OBJ_FROM_PTR(&pb_Side_RIGHT_obj);
+        case PBIO_ORIENTATION_SIDE_BOTTOM:
+        // fallthrough
+        default:
+            return MP_OBJ_FROM_PTR(&pb_Side_BOTTOM_obj);
     }
 }
 MP_DEFINE_CONST_FUN_OBJ_1(common_IMU_up_obj, common_IMU_up);