pybricks.hubs: Add update_heading_correction method to imu.

Author: laurensvalk

CHANGELOG.md

@@ -11,6 +11,8 @@
   between reboots.
 - Added `heading_correction` to `hub.imu.settings` to allow for automatic
   correction of the `hub.imu.heading()` value ([support#1678]).
+- Added `update_heading_correction` to interactively set the heading
+  correction value ([support#1678]).
 
 ### Changed
 

bricks/_common/modules/_hub_extra.py

@@ -1,3 +1,4 @@
+from pybricks.parameters import Axis
 from pybricks.tools import wait
 
 
@@ -7,3 +8,108 @@ def light_matrix_text_async(display, text, on, off):
         yield from wait(on)
         display.off()
         yield from wait(off)
+
+
+def imu_update_heading_correction(hub):
+
+    # Number of turns to confirm the result.
+    NUMBER_CONFIRM_TURNS = 5
+
+    # Maximum speed values before we consider the result invalid.
+    MAX_XY_SPEED = 50
+    MAX_Z_SPEED = 800
+
+    # Routine to wait on a button, with some extra time to avoid vibration directly after.
+    def wait_for_click():
+        while hub.buttons.pressed():
+            wait(1)
+        while not hub.buttons.pressed():
+            wait(1)
+        print("Processing...")
+        while hub.buttons.pressed():
+            wait(1)
+        wait(1500)
+
+    # Disable stop buttons
+    hub.system.set_stop_button(None)
+    print(
+        "Put the hub on a flat table. Align against a fixed reference like a wall or heavy book. Press hub button when ready."
+    )
+
+    # Wait for fixed reference and store the initial angle value.
+    wait_for_click()
+    while not hub.imu.ready() or not hub.imu.stationary():
+        wait(1)
+    start_z = hub.imu.rotation(-Axis.Z)
+
+    # Wait for a full rotation and get the result.
+    print(
+        "Keep the hub flat and slide it to make a full turn clockwise. Put it against the same reference. Press hub button when ready."
+    )
+    wait_for_click()
+    one_turn = hub.imu.rotation(-Axis.Z) - start_z
+
+    # Require clockwise...
+    if one_turn < 0:
+        raise ValueError("You turned it the wrong way. Please try again.")
+
+    # Sanity check. Should be close to 360.
+    if not (350 < one_turn < 370):
+        print(one_turn)
+        raise ValueError(
+            "The error was more than 10 degrees, which is unexpected. Please try again."
+        )
+
+    # Instruct to make more turns.
+    print("So far so good! Now make", NUMBER_CONFIRM_TURNS, "full clockwise turns.")
+
+    for i in range(NUMBER_CONFIRM_TURNS):
+
+        # The rotation target is just under a rotation so we can show the next
+        # message to keep going in time, avoiding doubts about what to do.
+        target = one_turn * (i + 2) - 10
+
+        # Wait until the hub reaches the target.
+        while hub.imu.rotation(-Axis.Z) - start_z < target:
+            wait(1)
+
+            if hub.buttons.pressed():
+                raise RuntimeError("Don't press the button until all turns are complete!")
+
+            if abs(hub.imu.angular_velocity(Axis.Z)) > MAX_Z_SPEED:
+                raise RuntimeError("Not so fast! Try again.")
+
+            if (
+                abs(hub.imu.angular_velocity(Axis.X)) + abs(hub.imu.angular_velocity(Axis.Y))
+                > MAX_XY_SPEED
+            ):
+
+                raise RuntimeError("Please keep the hub flat! Try again.")
+
+        # Inform user of status.
+        print("Completed", i + 1, "out of", NUMBER_CONFIRM_TURNS, "turns. ", end="")
+        if i < NUMBER_CONFIRM_TURNS - 1:
+            print("Keep going!")
+        else:
+            print("Put it against the same reference. Press hub button when ready.")
+
+    # Wait for final confirmation.
+    wait_for_click()
+
+    # Verify the result.
+    expected = one_turn * (NUMBER_CONFIRM_TURNS + 1)
+    result = hub.imu.rotation(-Axis.Z) - start_z
+
+    if abs(expected - result) > NUMBER_CONFIRM_TURNS / 2:
+        print(
+            "The",
+            NUMBER_CONFIRM_TURNS,
+            "extra turns where different from the first turn. Try again.",
+        )
+        print("Expected", expected, "but got", result)
+
+    # Get the final result to save.
+    average_turn = result / (NUMBER_CONFIRM_TURNS + 1)
+    print("For every 360-degree turn your gyro Z-axis reports:", average_turn)
+    hub.imu.settings(heading_correction=average_turn)
+    print("Saved! Now the heading() method will report the adjusted value.")

pybricks/common.h

@@ -134,7 +134,7 @@ extern const mp_obj_type_t pb_type_Speaker;
 
 #if PYBRICKS_PY_COMMON_IMU
 
-mp_obj_t pb_type_IMU_obj_new(mp_obj_t top_side_axis, mp_obj_t front_side_axis);
+mp_obj_t pb_type_IMU_obj_new(mp_obj_t hub_in, mp_obj_t top_side_axis, mp_obj_t front_side_axis);
 
 #endif // PYBRICKS_PY_COMMON_IMU
 

pybricks/common/pb_type_imu.c

@@ -14,6 +14,7 @@
 #include <pbio/imu.h>
 
 #include <pbsys/storage_settings.h>
+#include <pbsys/program_stop.h>
 
 #include "py/obj.h"
 
@@ -25,6 +26,7 @@
 
 typedef struct _pb_type_imu_obj_t {
     mp_obj_base_t base;
+    mp_obj_t hub;
 } pb_type_imu_obj_t;
 
 // pybricks._common.IMU.up
@@ -218,6 +220,28 @@ STATIC mp_obj_t pb_type_imu_reset_heading(size_t n_args, const mp_obj_t *pos_arg
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(pb_type_imu_reset_heading_obj, 1, pb_type_imu_reset_heading);
 
+// pybricks._common.IMU.update_heading_correction
+STATIC mp_obj_t pb_type_imu_update_heading_correction(mp_obj_t self_in) {
+    pb_type_imu_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    pb_module_tools_assert_blocking();
+
+    // Disable stop button and cache original setting to restore later.
+    pbio_button_flags_t stop_button = pbsys_program_stop_get_buttons();
+
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+        mp_obj_t func = pb_frozen_function_import(MP_QSTR__hub_extra, MP_QSTR_imu_update_heading_correction);
+        mp_call_function_1(func, self->hub);
+        pbsys_program_stop_set_buttons(stop_button);
+        nlr_pop();
+    } else {
+        pbsys_program_stop_set_buttons(stop_button);
+        nlr_jump(nlr.ret_val);
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(pb_type_imu_update_heading_correction_obj, pb_type_imu_update_heading_correction);
+
 // dir(pybricks.common.IMU)
 STATIC const mp_rom_map_elem_t pb_type_imu_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_acceleration),     MP_ROM_PTR(&pb_type_imu_acceleration_obj)    },
@@ -230,6 +254,7 @@ STATIC const mp_rom_map_elem_t pb_type_imu_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_stationary),       MP_ROM_PTR(&pb_type_imu_stationary_obj)      },
     { MP_ROM_QSTR(MP_QSTR_tilt),             MP_ROM_PTR(&pb_type_imu_tilt_obj)            },
     { MP_ROM_QSTR(MP_QSTR_up),               MP_ROM_PTR(&pb_type_imu_up_obj)              },
+    { MP_ROM_QSTR(MP_QSTR_update_heading_correction), MP_ROM_PTR(&pb_type_imu_update_heading_correction_obj)},
 };
 STATIC MP_DEFINE_CONST_DICT(pb_type_imu_locals_dict, pb_type_imu_locals_dict_table);
 
@@ -244,7 +269,7 @@ STATIC pb_type_imu_obj_t singleton_imu_obj = {
 };
 
 // pybricks._common.IMU.__init__
-mp_obj_t pb_type_IMU_obj_new(mp_obj_t top_side_axis_in, mp_obj_t front_side_axis_in) {
+mp_obj_t pb_type_IMU_obj_new(mp_obj_t hub_in, mp_obj_t top_side_axis_in, mp_obj_t front_side_axis_in) {
 
     // Set user base orientation.
     pbio_geometry_xyz_t front_side_axis;
@@ -256,6 +281,7 @@ mp_obj_t pb_type_IMU_obj_new(mp_obj_t top_side_axis_in, mp_obj_t front_side_axis
     pbio_imu_set_base_orientation(&front_side_axis, &top_side_axis);
 
     // Return singleton instance.
+    singleton_imu_obj.hub = hub_in;
     return MP_OBJ_FROM_PTR(&singleton_imu_obj);
 }
 

pybricks/hubs/pb_type_essentialhub.c

@@ -54,7 +54,7 @@ STATIC mp_obj_t hubs_EssentialHub_make_new(const mp_obj_type_t *type, size_t n_a
     #endif
     self->buttons = pb_type_Keypad_obj_new(pb_type_button_pressed_hub_single_button);
     self->charger = pb_type_Charger_obj_new();
-    self->imu = pb_type_IMU_obj_new(top_side_in, front_side_in);
+    self->imu = pb_type_IMU_obj_new(MP_OBJ_FROM_PTR(self), top_side_in, front_side_in);
     self->light = common_ColorLight_internal_obj_new(pbsys_status_light);
     self->system = MP_OBJ_FROM_PTR(&pb_type_System);
     return MP_OBJ_FROM_PTR(self);

pybricks/hubs/pb_type_primehub.c

@@ -79,7 +79,7 @@ STATIC mp_obj_t hubs_PrimeHub_make_new(const mp_obj_type_t *type, size_t n_args,
     self->buttons = pb_type_Keypad_obj_new(pb_type_primehub_button_pressed);
     self->charger = pb_type_Charger_obj_new();
     self->display = pb_type_LightMatrix_obj_new(pbsys_hub_light_matrix);
-    self->imu = pb_type_IMU_obj_new(top_side_in, front_side_in);
+    self->imu = pb_type_IMU_obj_new(MP_OBJ_FROM_PTR(self), top_side_in, front_side_in);
     self->light = common_ColorLight_internal_obj_new(pbsys_status_light);
     self->speaker = mp_call_function_0(MP_OBJ_FROM_PTR(&pb_type_Speaker));
     self->system = MP_OBJ_FROM_PTR(&pb_type_System);

pybricks/hubs/pb_type_technichub.c

@@ -43,7 +43,7 @@ STATIC mp_obj_t hubs_TechnicHub_make_new(const mp_obj_type_t *type, size_t n_arg
     self->ble = pb_type_BLE_new(broadcast_channel_in, observe_channels_in);
     #endif
     self->button = pb_type_Keypad_obj_new(pb_type_button_pressed_hub_single_button);
-    self->imu = pb_type_IMU_obj_new(top_side_in, front_side_in);
+    self->imu = pb_type_IMU_obj_new(MP_OBJ_FROM_PTR(self), top_side_in, front_side_in);
     self->light = common_ColorLight_internal_obj_new(pbsys_status_light);
     self->system = MP_OBJ_FROM_PTR(&pb_type_System);
     return MP_OBJ_FROM_PTR(self);