extratrees: Return probabilities in predict

Author: jonnor

src/emlearn_extratrees/extratrees.c

@@ -269,7 +269,7 @@ static mp_obj_t extratrees_model_train_step(mp_obj_t self_obj) {
 static MP_DEFINE_CONST_FUN_OBJ_1(extratrees_model_train_step_obj, extratrees_model_train_step);
 
 // Predict using the model (returns class probabilities)
-static mp_obj_t extratrees_model_predict_proba(size_t n_args, const mp_obj_t *args) {
+static mp_obj_t extratrees_model_predict(size_t n_args, const mp_obj_t *args) {
     if (n_args != 3) {
         mp_raise_ValueError(MP_ERROR_TEXT("Expected 3 arguments: self, features, probabilities"));
     }
@@ -309,37 +309,9 @@ static mp_obj_t extratrees_model_predict_proba(size_t n_args, const mp_obj_t *ar
 
     return mp_obj_new_int(predicted_class);
 }
-static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(extratrees_model_predict_proba_obj, 3, 3, extratrees_model_predict_proba);
-
-// Predict using the model (returns only class label)
-static mp_obj_t extratrees_model_predict(size_t n_args, const mp_obj_t *args) {
-    if (n_args != 2) {
-        mp_raise_ValueError(MP_ERROR_TEXT("Expected 2 arguments: self, features"));
-    }
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(extratrees_model_predict_obj, 3, 3, extratrees_model_predict);
 
-    mp_obj_extratrees_model_t *o = MP_OBJ_TO_PTR(args[0]);
-    EmlExtraTreesModel *model = &o->model;
-    EmlExtraTreesWorkspace *workspace = &o->workspace;
-
-    // Extract features buffer pointer and verify typecode
-    mp_buffer_info_t features_bufinfo;
-    mp_get_buffer_raise(args[1], &features_bufinfo, MP_BUFFER_READ);
-    if (features_bufinfo.typecode != 'h') {  // int16_t
-        mp_raise_ValueError(MP_ERROR_TEXT("features expecting int16 array"));
-    }
-    const int16_t *features = features_bufinfo.buf;
-    const int n_features = features_bufinfo.len / sizeof(int16_t);
-
-    if (n_features != model->n_features) {
-        mp_raise_ValueError(MP_ERROR_TEXT("Feature count mismatch"));
-    }
-
-    // Make prediction using pre-allocated workspace arrays
-    int16_t predicted_class = eml_extratrees_predict_proba(model, features, workspace->probabilities, workspace->votes);
-
-    return mp_obj_new_int(predicted_class);
-}
-static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(extratrees_model_predict_obj, 2, 2, extratrees_model_predict);
+// (old predict_proba removed - predict now handles both)
 
 // Get number of features
 static mp_obj_t extratrees_model_get_n_features(mp_obj_t self_obj) {
@@ -396,7 +368,7 @@ static MP_DEFINE_CONST_FUN_OBJ_1(extratrees_model_get_n_trees_trained_obj, extra
 #if MICROPY_ENABLE_DYNRUNTIME
 
 // Module setup
-mp_map_elem_t extratrees_model_locals_dict_table[12];
+mp_map_elem_t extratrees_model_locals_dict_table[11];
 static MP_DEFINE_CONST_DICT(extratrees_model_locals_dict, extratrees_model_locals_dict_table);
 
 // Module setup entrypoint
@@ -413,16 +385,15 @@ mp_obj_t mpy_init(mp_obj_fun_bc_t *self, size_t n_args, size_t n_kw, mp_obj_t *a
     // methods
     extratrees_model_locals_dict_table[0] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_make_new), MP_OBJ_FROM_PTR(&extratrees_model_make_new_obj) };
     extratrees_model_locals_dict_table[1] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_predict), MP_OBJ_FROM_PTR(&extratrees_model_predict_obj) };
-    extratrees_model_locals_dict_table[2] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_predict_proba), MP_OBJ_FROM_PTR(&extratrees_model_predict_proba_obj) };
-    extratrees_model_locals_dict_table[3] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train), MP_OBJ_FROM_PTR(&extratrees_model_train_obj) };
-    extratrees_model_locals_dict_table[4] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train_init), MP_OBJ_FROM_PTR(&extratrees_model_train_init_obj) };
-    extratrees_model_locals_dict_table[5] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train_step), MP_OBJ_FROM_PTR(&extratrees_model_train_step_obj) };
-    extratrees_model_locals_dict_table[6] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR___del__), MP_OBJ_FROM_PTR(&extratrees_model_del_obj) };
-    extratrees_model_locals_dict_table[7] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_features), MP_OBJ_FROM_PTR(&extratrees_model_get_n_features_obj) };
-    extratrees_model_locals_dict_table[8] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_classes), MP_OBJ_FROM_PTR(&extratrees_model_get_n_classes_obj) };
-    extratrees_model_locals_dict_table[9] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_trees), MP_OBJ_FROM_PTR(&extratrees_model_get_n_trees_obj) };
-    extratrees_model_locals_dict_table[10] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_nodes_used), MP_OBJ_FROM_PTR(&extratrees_model_get_n_nodes_used_obj) };
-    extratrees_model_locals_dict_table[11] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_trees_trained), MP_OBJ_FROM_PTR(&extratrees_model_get_n_trees_trained_obj) };
+    extratrees_model_locals_dict_table[2] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train), MP_OBJ_FROM_PTR(&extratrees_model_train_obj) };
+    extratrees_model_locals_dict_table[3] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train_init), MP_OBJ_FROM_PTR(&extratrees_model_train_init_obj) };
+    extratrees_model_locals_dict_table[4] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_train_step), MP_OBJ_FROM_PTR(&extratrees_model_train_step_obj) };
+    extratrees_model_locals_dict_table[5] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR___del__), MP_OBJ_FROM_PTR(&extratrees_model_del_obj) };
+    extratrees_model_locals_dict_table[6] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_features), MP_OBJ_FROM_PTR(&extratrees_model_get_n_features_obj) };
+    extratrees_model_locals_dict_table[7] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_classes), MP_OBJ_FROM_PTR(&extratrees_model_get_n_classes_obj) };
+    extratrees_model_locals_dict_table[8] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_trees), MP_OBJ_FROM_PTR(&extratrees_model_get_n_trees_obj) };
+    extratrees_model_locals_dict_table[9] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_nodes_used), MP_OBJ_FROM_PTR(&extratrees_model_get_n_nodes_used_obj) };
+    extratrees_model_locals_dict_table[10] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_get_n_trees_trained), MP_OBJ_FROM_PTR(&extratrees_model_get_n_trees_trained_obj) };
 
     MP_OBJ_TYPE_SET_SLOT(&extratrees_model_type, locals_dict, (void*)&extratrees_model_locals_dict, 10);
 
@@ -436,7 +407,6 @@ mp_obj_t mpy_init(mp_obj_fun_bc_t *self, size_t n_args, size_t n_kw, mp_obj_t *a
 static const mp_rom_map_elem_t extratrees_model_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_make_new), MP_ROM_PTR(&extratrees_model_make_new_obj) },
     { MP_ROM_QSTR(MP_QSTR_predict), MP_ROM_PTR(&extratrees_model_predict_obj) },
-    { MP_ROM_QSTR(MP_QSTR_predict_proba), MP_ROM_PTR(&extratrees_model_predict_proba_obj) },
     { MP_ROM_QSTR(MP_QSTR_train), MP_ROM_PTR(&extratrees_model_train_obj) },
     { MP_ROM_QSTR(MP_QSTR_train_init), MP_ROM_PTR(&extratrees_model_train_init_obj) },
     { MP_ROM_QSTR(MP_QSTR_train_step), MP_ROM_PTR(&extratrees_model_train_step_obj) },

tests/test_extratrees.py

@@ -2,6 +2,16 @@
 import array
 import emlearn_extratrees
 
+
+def argmax(arr, n):
+    best_idx = 0
+    best_val = arr[0]
+    for i in range(1, n):
+        if arr[i] > best_val:
+            best_val = arr[i]
+            best_idx = i
+    return best_idx
+
 def test_single_tree_prediction():
     """Test with just one tree to isolate prediction issues"""
     print("=== Single Tree Prediction Debug ===")
@@ -33,7 +43,7 @@ def test_single_tree_prediction():
             expected = y[i]
 
             test_features = array.array('h', [x1, x2])
-            predicted = model.predict_proba(test_features, probabilities)
+            predicted = model.predict(test_features, probabilities)
 
             print("  ({}, {}) -> pred={}, exp={}, probs=[{:.3f}, {:.3f}] {}".format(
                 x1, x2, predicted, expected,
@@ -68,12 +78,11 @@ def test_class_bias():
 
     # Test on a clear class 1 example
     test_features = array.array('h', [300, 300])
-    predicted = model.predict(test_features)
+    probabilities = array.array('f', [0.0, 0.0])
+    model.predict(test_features, probabilities)
+    predicted = argmax(probabilities, 2)
 
     print("Prediction for (300,300): {} (should strongly favor class 1)".format(predicted))
-    
-    probabilities = array.array('f', [0.0, 0.0])
-    model.predict_proba(test_features, probabilities)
     print("Probabilities: [{:.3f}, {:.3f}]".format(probabilities[0], probabilities[1]))
 
 def test_manual_verification():
@@ -122,7 +131,7 @@ def test_manual_verification():
     all_correct = True
     for x1_val, expected in test_cases:
         test_features = array.array('h', [x1_val, 500])  # x2 irrelevant
-        predicted = model.predict_proba(test_features, probabilities)
+        predicted = model.predict(test_features, probabilities)
 
         is_correct = predicted == expected
         if not is_correct:
@@ -180,12 +189,12 @@ def test_train_step_by_step():
     # Verify predictions still work
     probabilities = array.array('f', [0.0, 0.0])
     test_features = array.array('h', [0, 0])
-    predicted = model.predict_proba(test_features, probabilities)
+    predicted = model.predict(test_features, probabilities)
     print("Predict (0,0): {} probs=[{:.3f}, {:.3f}]".format(predicted, probabilities[0], probabilities[1]))
     assert predicted == 0, "Expected class 0"
 
     test_features = array.array('h', [300, 300])
-    predicted = model.predict_proba(test_features, probabilities)
+    predicted = model.predict(test_features, probabilities)
     print("Predict (300,300): {} probs=[{:.3f}, {:.3f}]".format(predicted, probabilities[0], probabilities[1]))
     assert predicted == 1, "Expected class 1"
 
@@ -216,12 +225,15 @@ def test_train_generator():
     assert model.get_n_trees_trained() == 5
 
     # Verify predictions
+    probabilities = array.array('f', [0.0, 0.0])
     test_features = array.array('h', [0, 0])
-    predicted = model.predict(test_features)
+    model.predict(test_features, probabilities)
+    predicted = argmax(probabilities, 2)
     assert predicted == 0, "Expected class 0"
 
     test_features = array.array('h', [300, 300])
-    predicted = model.predict(test_features)
+    model.predict(test_features, probabilities)
+    predicted = argmax(probabilities, 2)
     assert predicted == 1, "Expected class 1"
 
     print("✓ Generator training works")

tests/test_extratrees_cancer.py

@@ -66,7 +66,7 @@ def test_real_dataset():
         end_idx = start_idx + n_features
         features = array.array('h', X_test_flat[start_idx:end_idx])
 
-        predicted = model.predict_proba(features, probabilities)
+        predicted = model.predict(features, probabilities)
         actual = y_test[i]
 
         if predicted == actual:

tests/test_extratrees_wine.py

@@ -82,7 +82,7 @@ def test_wine():
         end_idx = start_idx + n_features
         features = array.array('h', X_test_flat[start_idx:end_idx])
 
-        predicted = model.predict_proba(features, probabilities)
+        predicted = model.predict(features, probabilities)
         actual = y_test[i]
 
         # Track per-class stats

tests/test_extratrees_xor.py

@@ -52,7 +52,7 @@ def test_xor_comprehensive():
 
     for features, expected in test_cases:
         test_features = array.array('h', features)
-        predicted = model.predict_proba(test_features, probabilities)
+        predicted = model.predict(test_features, probabilities)
         is_correct = predicted == expected
         if is_correct:
             correct += 1
@@ -76,7 +76,7 @@ def test_xor_comprehensive():
 
     for features, expected in interpolation_cases:
         test_features = array.array('h', features)
-        predicted = model.predict_proba(test_features, probabilities)
+        predicted = model.predict(test_features, probabilities)
         confidence = max(probabilities[0], probabilities[1])
 
         if expected == "?":
@@ -120,7 +120,7 @@ def test_xor_robustness():
 
         for features, expected in test_cases:
             test_features = array.array('h', features)
-            predicted = model.predict_proba(test_features, probabilities)
+            predicted = model.predict(test_features, probabilities)
             if predicted == expected:
                 correct += 1
 
@@ -178,7 +178,7 @@ def test_xor_different_values():
 
         for features, expected in test_cases:
             test_features = array.array('h', features)
-            predicted = model.predict_proba(test_features, probabilities)
+            predicted = model.predict(test_features, probabilities)
             if predicted == expected:
                 correct += 1