optimize spike

Author: techcatgato

pybricks/experimental/odometry.c

@@ -1,44 +1,91 @@
 // SPDX-License-Identifier: MIT
 #include "py/mpconfig.h"
+
+#if PYBRICKS_PY_EXPERIMENTAL
+
 #include "py/mphal.h"
 #include "py/runtime.h"
 #include <math.h>
-#include "pybricks/experimental/platform_math.h"
+
+// Core Pybricks headers for bare-metal hardware access
+#include <pybricks/common.h>
+#include <pbio/servo.h>
 #include "pybricks/experimental/odometry.h"
 
-mp_obj_t calculate_odometry(int num_iters, float wheel_circ, float axle_track, mp_obj_t right_angle_func, mp_obj_t left_angle_func) {
-    float deg_to_mm = wheel_circ * 0.0027777778f;
+// Hardware acceleration for the Spike Prime (FPU + RAM execution)
+#if defined(__ARM_ARCH_7M__) || defined(__ARM_ARCH_7EM__)
+#define ACCEL_RAM __attribute__((section(".data"), noinline))
+#else
+#define ACCEL_RAM
+#endif
+
+ACCEL_RAM static float fast_sin_internal(float theta) {
+    float x = theta * 0.159154943f;
+    x = theta - (float)((int)(x + (x > 0.00f ? 0.5f : -0.5f))) * 6.2831853f;
+
+    if (x > 1.5707963f) {
+        x = 3.1415926f - x;
+    } else if (x < -1.5707963f) {
+        x = -3.1415926f - x;
+    }
+
+    float x2 = x * x;
+    return x * (0.99999906f + x2 * (-0.16665554f + x2 * (0.00831190f + x2 * -0.00018488f)));
+}
+
+mp_obj_t calculate_odometry(int num_iters, float wheel_circ, float axle_track, mp_obj_t right_motor_obj, mp_obj_t left_motor_obj) {
+
+    float deg_to_mm = wheel_circ * 0.0027777778f; 
     float inv_axle_track = 1.0f / axle_track;
-    float rx = 0.0f, ry = 0.0f, rh = 0.0f;
 
-    int32_t last_r = mp_obj_get_int(mp_call_function_0(right_angle_func));
-    int32_t last_l = mp_obj_get_int(mp_call_function_0(left_angle_func));
+    float rx = 0.0f;
+    float ry = 0.0f;
+    float rh = 0.0f;
+
+    // THE MAGIC: Extract the raw Pybricks hardware struct from the Python object.
+    // This entirely removes the Python interpreter from the loop.
+    pbio_servo_t *srv_r = ((pb_type_Motor_obj_t *)MP_OBJ_TO_PTR(right_motor_obj))->srv;
+    pbio_servo_t *srv_l = ((pb_type_Motor_obj_t *)MP_OBJ_TO_PTR(left_motor_obj))->srv;
+
+    int32_t last_r, last_l, unused_rate;
+    // Direct memory read to initialize
+    pbio_servo_get_state_user(srv_r, &last_r, &unused_rate);
+    pbio_servo_get_state_user(srv_l, &last_l, &unused_rate);
+
     uint32_t start_time = mp_hal_ticks_ms();
 
     for (int i = 0; i < num_iters; i++) {
-        int32_t cur_r = mp_obj_get_int(mp_call_function_0(right_angle_func));
-        int32_t cur_l = mp_obj_get_int(mp_call_function_0(left_angle_func));
+        int32_t cur_r, cur_l;
+        
+        // 1. Direct memory read (Instantaneous, 0 Python overhead)
+        pbio_servo_get_state_user(srv_r, &cur_r, &unused_rate);
+        pbio_servo_get_state_user(srv_l, &cur_l, &unused_rate);
 
+        // 2. FPU Delta Math
         float dR = (float)(cur_r - last_r) * deg_to_mm;
         float dL = (float)(cur_l - last_l) * deg_to_mm;
         float dD = (dR + dL) * 0.5f;
         float dH = (dR - dL) * inv_axle_track;
 
+        // 3. RK2 Integration
         if (dD != 0.0f || dH != 0.0f) {
             float avg_h = rh + (dH * 0.5f);
-            rx += dD * pb_fast_cos(avg_h);
-            ry += dD * pb_fast_sin(avg_h);
+            rx += dD * fast_sin_internal(avg_h + 1.5707963f); // cos
+            ry += dD * fast_sin_internal(avg_h);              // sin
             rh += dH;
         }
+
         last_r = cur_r;
         last_l = cur_l;
 
+        // 4. Polling (bitwise check is faster than modulo)
         if ((i & 0x3FF) == 0) {
             mp_handle_pending(true);
         }
     }
 
     uint32_t dur = mp_hal_ticks_ms() - start_time;
+
     mp_obj_t tuple[5] = {
         mp_obj_new_float_from_f((float)dur * 0.001f),
         mp_obj_new_int(num_iters),
@@ -48,3 +95,35 @@ mp_obj_t calculate_odometry(int num_iters, float wheel_circ, float axle_track, m
     };
     return mp_obj_new_tuple(5, tuple);
 }
+
+// Optimized Pure Pursuit Logic using Dot Products
+mp_obj_t get_pure_pursuit_multipliers(float tx, float ty, float rx_pos, float ry_pos, float rh_ang, float track) {
+    float x_dif = tx - rx_pos;
+    float y_dif = ty - ry_pos;
+
+    float cos_h = fast_sin_internal(rh_ang + 1.5707963f);
+    float sin_h = fast_sin_internal(rh_ang);
+
+    float relative_y = (y_dif * cos_h) - (x_dif * sin_h);
+
+    float m_left = 1.0f;
+    float m_right = 1.0f;
+    
+    float abs_rel_y = relative_y < 0.0f ? -relative_y : relative_y;
+
+    if (abs_rel_y > 0.001f) {
+        float dist_sq = (x_dif * x_dif) + (y_dif * y_dif);
+        float radius = -(dist_sq / (2.0f * relative_y));
+        float two_r = 2.0f * radius;
+        m_right = two_r / (two_r + track);
+        m_left = two_r / (two_r - track);
+    }
+
+    mp_obj_t tuple[2] = {
+        mp_obj_new_float_from_f(m_left),
+        mp_obj_new_float_from_f(m_right)
+    };
+    return mp_obj_new_tuple(2, tuple);
+}
+
+#endif

pybricks/experimental/pb_module_experimental.c

@@ -1,26 +1,46 @@
 #include "py/mpconfig.h"
 #include "py/obj.h"
 #include "py/runtime.h"
-#include "pybricks/experimental/odometry.h"
 
-static mp_obj_t experimental_odometry_benchmark(size_t n_args, const mp_obj_t *args) {
-    int num_iters = mp_obj_get_int(args[0]);
-    float wheel_circ = mp_obj_get_float(args[1]);
-    float axle_track = mp_obj_get_float(args[2]);
-    mp_obj_t right_func = args[3];
-    mp_obj_t left_func = args[4];
+#if PYBRICKS_PY_EXPERIMENTAL
+
+// Forward declare the C functions
+extern mp_obj_t calculate_odometry(int num_iters, float wheel_circ, float axle_track, mp_obj_t right_motor, mp_obj_t left_motor);
+extern mp_obj_t get_pure_pursuit_multipliers(float tx, float ty, float rx_pos, float ry_pos, float rh_ang, float track);
 
-    return calculate_odometry(num_iters, wheel_circ, axle_track, right_func, left_func);
+static mp_obj_t experimental_odometry_benchmark(size_t n_args, const mp_obj_t *args) {
+    return calculate_odometry(
+        mp_obj_get_int(args[0]),
+        mp_obj_get_float(args[1]),
+        mp_obj_get_float(args[2]),
+        args[3], // Raw Python Motor Object (Right)
+        args[4]  // Raw Python Motor Object (Left)
+    );
 }
 static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(experimental_odometry_benchmark_obj, 5, 5, experimental_odometry_benchmark);
 
+static mp_obj_t experimental_pure_pursuit_logic(size_t n_args, const mp_obj_t *args) {
+    return get_pure_pursuit_multipliers(
+        mp_obj_get_float(args[0]), 
+        mp_obj_get_float(args[1]), 
+        mp_obj_get_float(args[2]), 
+        mp_obj_get_float(args[3]), 
+        mp_obj_get_float(args[4]), 
+        mp_obj_get_float(args[5])
+    );
+}
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(experimental_pure_pursuit_logic_obj, 6, 6, experimental_pure_pursuit_logic);
+
 const mp_rom_map_elem_t pb_module_experimental_globals_table[] = {
-    { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_experimental) },
-    { MP_ROM_QSTR(MP_QSTR_odometry_benchmark), MP_ROM_PTR(&experimental_odometry_benchmark_obj) },
+    { MP_ROM_QSTR(MP_QSTR___name__),            MP_ROM_QSTR(MP_QSTR_experimental) },
+    { MP_ROM_QSTR(MP_QSTR_odometry_benchmark),  MP_ROM_PTR(&experimental_odometry_benchmark_obj) },
+    { MP_ROM_QSTR(MP_QSTR_pure_pursuit_logic),  MP_ROM_PTR(&experimental_pure_pursuit_logic_obj) },
 };
 MP_DEFINE_CONST_DICT(pb_module_experimental_globals, pb_module_experimental_globals_table);
 
 const mp_obj_module_t pb_module_experimental = {
     .base = { &mp_type_module },
     .globals = (mp_obj_dict_t *)&pb_module_experimental_globals,
-};
+};
+
+#endif // PYBRICKS_PY_EXPERIMENTAL