extratrees: Use keyword arguments for new

Author: jonnor

src/emlearn_extratrees/emlearn_extratrees.py

@@ -7,6 +7,36 @@
     pass
 
 
+def new(n_features, n_classes, *,
+        n_trees=10, max_depth=10, min_samples_leaf=1,
+        n_thresholds=10, subsample_ratio=1.0, feature_subsample_ratio=1.0,
+        max_nodes=1000, max_samples=1000,
+        rng_seed=42, use_global_feature_range=False):
+    """Create a new ExtraTrees model.
+
+    Args:
+        n_features: Number of input features (required)
+        n_classes: Number of output classes (required)
+        n_trees: Number of trees in the ensemble
+        max_depth: Maximum tree depth
+        min_samples_leaf: Minimum samples at a leaf node
+        n_thresholds: Random thresholds drawn per feature split
+        subsample_ratio: Fraction of samples used per tree (0.0-1.0)
+        feature_subsample_ratio: Fraction of features considered per split (0.0-1.0)
+        max_nodes: Max pre-allocated nodes (memory)
+        max_samples: Max pre-allocated samples (memory)
+        rng_seed: Random number generator seed
+        use_global_feature_range: Use global feature min/max instead of per-node
+    """
+    return make_new(
+        n_features, n_classes,
+        n_trees, max_depth, min_samples_leaf, n_thresholds,
+        subsample_ratio, feature_subsample_ratio,
+        max_nodes, max_samples,
+        rng_seed, int(use_global_feature_range)
+    )
+
+
 def train_steps(model, X, y):
     """Generator that trains one node at a time, yielding after each step.
 

src/emlearn_extratrees/extratrees.c

@@ -37,25 +37,24 @@ static const mp_obj_type_t extratrees_model_type;
 #endif
 
 // Create a new instance
-static mp_obj_t extratrees_model_new(size_t n_args, const mp_obj_t *args) {
-    // Args: n_features, n_classes, [n_trees], [max_depth], [min_samples_leaf], [n_thresholds], 
-    //       [subsample_ratio], [feature_subsample_ratio], [max_nodes], [max_samples], [rng_seed], [use_global_feature_range]
-    if (n_args < 2 || n_args > 12) {
-        mp_raise_ValueError(MP_ERROR_TEXT("Expected 2-12 arguments: n_features, n_classes, [n_trees=10], [max_depth=10], [min_samples_leaf=1], [n_thresholds=10], [subsample_ratio=1.0], [feature_subsample_ratio=1.0], [max_nodes=1000], [max_samples=1000], [rng_seed=42], [use_global_feature_range=0]"));
+static mp_obj_t extratrees_model_make_new(size_t n_args, const mp_obj_t *args) {
+    // All 12 args provided by Python wrapper (defaults handled there)
+    if (n_args != 12) {
+        mp_raise_ValueError(MP_ERROR_TEXT("Expected 12 arguments"));
     }
 
     mp_int_t n_features = mp_obj_get_int(args[0]);
     mp_int_t n_classes = mp_obj_get_int(args[1]);
-    mp_int_t n_trees = (n_args > 2) ? mp_obj_get_int(args[2]) : 10;
-    mp_int_t max_depth = (n_args > 3) ? mp_obj_get_int(args[3]) : 10;
-    mp_int_t min_samples_leaf = (n_args > 4) ? mp_obj_get_int(args[4]) : 1;
-    mp_int_t n_thresholds = (n_args > 5) ? mp_obj_get_int(args[5]) : 10;
-    float subsample_ratio = (n_args > 6) ? mp_obj_get_float_to_f(args[6]) : 1.0f;
-    float feature_subsample_ratio = (n_args > 7) ? mp_obj_get_float_to_f(args[7]) : 1.0f;
-    mp_int_t max_nodes = (n_args > 8) ? mp_obj_get_int(args[8]) : 1000;
-    mp_int_t max_samples = (n_args > 9) ? mp_obj_get_int(args[9]) : 1000;
-    mp_int_t rng_seed = (n_args > 10) ? mp_obj_get_int(args[10]) : 42;
-    mp_int_t use_global_feature_range = (n_args > 11) ? mp_obj_get_int(args[11]) : 0;
+    mp_int_t n_trees = mp_obj_get_int(args[2]);
+    mp_int_t max_depth = mp_obj_get_int(args[3]);
+    mp_int_t min_samples_leaf = mp_obj_get_int(args[4]);
+    mp_int_t n_thresholds = mp_obj_get_int(args[5]);
+    float subsample_ratio = mp_obj_get_float_to_f(args[6]);
+    float feature_subsample_ratio = mp_obj_get_float_to_f(args[7]);
+    mp_int_t max_nodes = mp_obj_get_int(args[8]);
+    mp_int_t max_samples = mp_obj_get_int(args[9]);
+    mp_int_t rng_seed = mp_obj_get_int(args[10]);
+    mp_int_t use_global_feature_range = mp_obj_get_int(args[11]);
 
     // Allocate space
     mp_obj_extratrees_model_t *o = \
@@ -116,7 +115,7 @@ static mp_obj_t extratrees_model_new(size_t n_args, const mp_obj_t *args) {
     return MP_OBJ_FROM_PTR(o);
 }
 // Define a Python reference to the function above
-static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(extratrees_model_new_obj, 2, 12, extratrees_model_new);
+static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(extratrees_model_make_new_obj, 12, 12, extratrees_model_make_new);
 
 // Delete an instance
 static mp_obj_t extratrees_model_del(mp_obj_t self_obj) {
@@ -397,32 +396,33 @@ 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[11];
+mp_map_elem_t extratrees_model_locals_dict_table[12];
 static MP_DEFINE_CONST_DICT(extratrees_model_locals_dict, extratrees_model_locals_dict_table);
 
 // Module setup entrypoint
 mp_obj_t mpy_init(mp_obj_fun_bc_t *self, size_t n_args, size_t n_kw, mp_obj_t *args) {
     // This must be first, it sets up the globals dict and other things
     MP_DYNRUNTIME_INIT_ENTRY
 
-    mp_store_global(MP_QSTR_new, MP_OBJ_FROM_PTR(&extratrees_model_new_obj));
+    mp_store_global(MP_QSTR_make_new, MP_OBJ_FROM_PTR(&extratrees_model_make_new_obj));
 
     extratrees_model_type.base.type = (void*)&mp_fun_table.type_type;
     extratrees_model_type.flags = MP_TYPE_FLAG_ITER_IS_CUSTOM;
     extratrees_model_type.name = MP_QSTR_extratrees;
 
     // methods
-    extratrees_model_locals_dict_table[0] = (mp_map_elem_t){ MP_OBJ_NEW_QSTR(MP_QSTR_predict), MP_OBJ_FROM_PTR(&extratrees_model_predict_obj) };
-    extratrees_model_locals_dict_table[1] = (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[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) };
+    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) };
 
     MP_OBJ_TYPE_SET_SLOT(&extratrees_model_type, locals_dict, (void*)&extratrees_model_locals_dict, 10);
 
@@ -434,6 +434,7 @@ mp_obj_t mpy_init(mp_obj_fun_bc_t *self, size_t n_args, size_t n_kw, mp_obj_t *a
 
 // User C module mode
 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) },
@@ -457,7 +458,7 @@ static MP_DEFINE_CONST_OBJ_TYPE(
 
 // Define module object
 static const mp_rom_map_elem_t emlearn_extratrees_globals_table[] = {
-    { MP_ROM_QSTR(MP_QSTR_new), MP_ROM_PTR(&extratrees_model_new_obj) },
+    { MP_ROM_QSTR(MP_QSTR_make_new), MP_ROM_PTR(&extratrees_model_make_new_obj) },
 };
 static MP_DEFINE_CONST_DICT(emlearn_extratrees_globals, emlearn_extratrees_globals_table);
 

tests/test_extratrees.py

@@ -16,7 +16,7 @@ def test_single_tree_prediction():
     print("Data: (0,0)->0, (1000,1000)->1")
 
     # Single tree with enough depth and thresholds
-    model = emlearn_extratrees.new(2, 2, 1, 5, 1, 100)
+    model = emlearn_extratrees.new(2, 2, n_trees=1, max_depth=5, n_thresholds=100)
     model.train(X, y)
 
     nodes_used = model.get_n_nodes_used()
@@ -63,7 +63,7 @@ def test_class_bias():
 
     print("Data: 1 sample class 0, 5 samples class 1")
 
-    model = emlearn_extratrees.new(2, 2, 3, 5, 1, 50)
+    model = emlearn_extratrees.new(2, 2, n_trees=3, max_depth=5, n_thresholds=50)
     model.train(X, y)
 
     # Test on a clear class 1 example
@@ -103,7 +103,7 @@ def test_manual_verification():
         print("  ({}, {}) -> {}".format(x1, x2, label))
 
     # Train with parameters that should definitely work
-    model = emlearn_extratrees.new(2, 2, 10, 10, 1, 100)
+    model = emlearn_extratrees.new(2, 2, n_trees=10, max_depth=10, n_thresholds=100)
     model.train(X, y)
 
     print("\nNodes used: {}".format(model.get_n_nodes_used()))
@@ -160,7 +160,7 @@ def test_train_step_by_step():
     y = array.array('h', [0, 0, 1, 1])
 
     # 3 trees
-    model = emlearn_extratrees.new(2, 2, 3, 5, 1, 50)
+    model = emlearn_extratrees.new(2, 2, n_trees=3, max_depth=5, n_thresholds=50)
     model.train_init(X, y)
 
     steps = 0
@@ -203,7 +203,7 @@ def test_train_generator():
     ])
     y = array.array('h', [0, 0, 1, 1])
 
-    model = emlearn_extratrees.new(2, 2, 5, 5, 1, 50)
+    model = emlearn_extratrees.new(2, 2, n_trees=5, max_depth=5, n_thresholds=50)
 
     trees_progress = []
     for trees_done in emlearn_extratrees.train_steps(model, X, y):
@@ -239,11 +239,11 @@ def test_train_step_same_as_train():
     y = array.array('h', [0, 0, 1, 1])
 
     # Bulk train
-    model_bulk = emlearn_extratrees.new(2, 2, 3, 5, 1, 50, 1.0, 1.0, 1000, 100, 42)
+    model_bulk = emlearn_extratrees.new(2, 2, n_trees=3, max_depth=5, n_thresholds=50, max_samples=100, rng_seed=42)
     model_bulk.train(X, y)
 
     # Step-by-step train
-    model_step = emlearn_extratrees.new(2, 2, 3, 5, 1, 50, 1.0, 1.0, 1000, 100, 42)
+    model_step = emlearn_extratrees.new(2, 2, n_trees=3, max_depth=5, n_thresholds=50, max_samples=100, rng_seed=42)
     model_step.train_init(X, y)
     while not model_step.train_step():
         pass

tests/test_extratrees_cancer.py

@@ -47,17 +47,10 @@ def test_real_dataset():
 
     # Create model
     model = emlearn_extratrees.new(
-        30,    # n_features
-        2,     # n_classes  
-        20,    # n_trees
-        12,    # max_depth
-        2,     # min_samples_leaf
-        15,    # n_thresholds
-        0.8,   # subsample_ratio
-        0.7,   # feature_subsample_ratio  
-        3000,  # max_nodes
-        500,   # max_samples
-        42     # rng_seed
+        30, 2,
+        n_trees=20, max_depth=12, min_samples_leaf=2,
+        n_thresholds=15, subsample_ratio=0.8, feature_subsample_ratio=0.7,
+        max_nodes=3000, max_samples=500, rng_seed=42
     )
 
     print("Training...")

tests/test_extratrees_wine.py

@@ -56,17 +56,10 @@ def test_wine():
 
     # Create model
     model = emlearn_extratrees.new(
-        n_features,  # n_features
-        n_classes,   # n_classes
-        20,          # n_trees 
-        12,          # max_depth
-        2,           # min_samples_leaf
-        15,          # n_thresholds
-        0.8,         # subsample_ratio
-        1.0,         # feature_subsample_ratio
-        3000,        # max_nodes
-        500,         # max_samples
-        42           # rng_seed
+        n_features, n_classes,
+        n_trees=20, max_depth=12, min_samples_leaf=2,
+        n_thresholds=15, subsample_ratio=0.8, feature_subsample_ratio=1.0,
+        max_nodes=3000, max_samples=500, rng_seed=42
     )
 
     train_start = time.ticks_ms()

tests/test_extratrees_xor.py

@@ -29,18 +29,10 @@ def test_xor_comprehensive():
 
     # Test with ensemble of trees (now that individual trees work)
     model = emlearn_extratrees.new(
-        2,      # n_features
-        2,      # n_classes  
-        10,     # n_trees (ensemble)
-        8,      # max_depth
-        1,      # min_samples_leaf
-        10,     # n_thresholds
-        0.8,    # subsample_ratio (80% for diversity)
-        1.0,    # feature_subsample_ratio (use both features)
-        500,    # max_nodes
-        100,    # max_samples
-        42      # rng_seed
-    )
+        2, 2,
+        n_trees=10, max_depth=8, n_thresholds=10,
+        subsample_ratio=0.8, feature_subsample_ratio=1.0,
+        max_nodes=500, max_samples=100, rng_seed=42)
 
     model.train(X, y)
 
@@ -118,7 +110,7 @@ def test_xor_robustness():
     for n_trees, max_depth, desc in configs:
         print(f"\nTesting {desc}:")
 
-        model = emlearn_extratrees.new(2, 2, n_trees, max_depth, 1, 8, 0.9, 1.0, 1000, 100, 123)
+        model = emlearn_extratrees.new(2, 2, n_trees=n_trees, max_depth=max_depth, n_thresholds=8, subsample_ratio=0.9, max_samples=100, rng_seed=123)
         model.train(X, y)
 
         # Test all XOR cases
@@ -170,7 +162,7 @@ def test_xor_different_values():
         X = array.array('h', X_data)
         y = array.array('h', y_data)
 
-        model = emlearn_extratrees.new(2, 2, 12, 10, 1, 10, 0.8, 1.0, 800, 100, 456)
+        model = emlearn_extratrees.new(2, 2, n_trees=12, max_depth=10, n_thresholds=10, subsample_ratio=0.8, max_nodes=800, max_samples=100, rng_seed=456)
         model.train(X, y)
 
         # Test