Update model classes for test mode and aeon compatibility

SamoraHunter · SamoraHunter · commit 9fe5b4b79b0e · 2026-02-28T19:03:55.000Z
Added test_mode checks to multiple classifier classes (NeuralNetwork, AdaBoost, CatBoost, GradientBoosting, H2O, Keras, LightGBM, RandomForest, XGBoost) to define minimal parameter spaces for faster unit testing.
Updated CNNClassifier_module to use TimeCNNClassifier instead of the deprecated CNNClassifier from aeon, and added a monkeypatch for sklearn.utils.validation.validate_data compatibility.
Updated RocketClassifier import path to match newer aeon structure.
Implemented a dummy classifier fallback for ShapeDTWClassifier since ShapeDTW has been removed in aeon v0.11.0, preventing import errors.
Fixed a type hint syntax error in H2OStackedEnsembleClassifier.
diff --git a/ml_grid/model_classes/H2OStackedEnsembleClassifier.py b/ml_grid/model_classes/H2OStackedEnsembleClassifier.py
@@ -144,7 +144,7 @@ def fit(
             base_models_list = []
             for i, model_wrapper in enumerate(
                 self.base_models
-            ):  # type:H2OBaseClassifier
+            ):  # type: H2OBaseClassifier
                 self.logger.debug(
                     f"Fitting base model {i+1}: {type(model_wrapper).__name__}"
                 )
diff --git a/ml_grid/model_classes/NeuralNetworkClassifier_class.py b/ml_grid/model_classes/NeuralNetworkClassifier_class.py
@@ -56,7 +56,18 @@ def __init__(
 
         from ml_grid.util.global_params import global_parameters
 
-        if global_parameters.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "hidden_layer_sizes": ["(8,)"],
+                    "dropout_rate": [0.0],
+                    "learning_rate": [0.01],
+                    "activation_func": ["relu"],
+                    "epochs": [1],
+                    "batch_size": [32],
+                }
+            ]
+        elif global_parameters.bayessearch:
             from skopt.space import Categorical, Integer, Real
 
             self.parameter_space = [
diff --git a/ml_grid/model_classes/adaboost_classifier_class.py b/ml_grid/model_classes/adaboost_classifier_class.py
@@ -49,7 +49,13 @@ def __init__(
         self.parameter_space: List[Dict[str, Any]]
         is_bayes_search = global_parameters.bayessearch
 
-        if is_bayes_search:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "n_estimators": [2],
+                }
+            ]
+        elif is_bayes_search:
             # For BayesSearchCV, define the search space using skopt.space objects.
             # We define separate spaces for each estimator type.
             self.parameter_space = [
diff --git a/ml_grid/model_classes/catboost_classifier_class.py b/ml_grid/model_classes/catboost_classifier_class.py
@@ -50,7 +50,14 @@ def __init__(
 
         self.parameter_space: Union[List[Dict[str, Any]], Dict[str, Any]]
         # Define parameter space for Bayesian search or traditional grid search
-        if global_parameters.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "iterations": [2],
+                    "depth": [4],
+                }
+            ]
+        elif global_parameters.bayessearch:
             self.parameter_space = {
                 "iterations": Integer(100, 1000),
                 "learning_rate": Real(0.01, 0.3, prior="uniform"),
diff --git a/ml_grid/model_classes/gradientboosting_classifier_class.py b/ml_grid/model_classes/gradientboosting_classifier_class.py
@@ -55,7 +55,14 @@ def __init__(
         )
         self.parameter_space: Dict[str, Any]
 
-        if global_params.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "n_estimators": [2],
+                    "max_depth": [2],
+                }
+            ]
+        elif global_parameters.bayessearch:
             # Define the parameter space for Bayesian optimization
             self.parameter_space = {
                 "ccp_alpha": Real(0.0, 1.0, prior="uniform"),
diff --git a/ml_grid/model_classes/h2o_classifier_class.py b/ml_grid/model_classes/h2o_classifier_class.py
@@ -53,7 +53,15 @@ def __init__(
         self.method_name: str = "H2OAutoMLClassifier"
         self.parameter_space: List[Dict[str, Any]]
 
-        if global_params.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "max_runtime_secs": [5],
+                    "max_models": [1],
+                    "nfolds": [2],
+                }
+            ]
+        elif global_parameters.bayessearch:
             # Define the parameter space for Bayesian optimization
             self.parameter_space = [
                 {
diff --git a/ml_grid/model_classes/keras_classifier_class.py b/ml_grid/model_classes/keras_classifier_class.py
@@ -17,6 +17,8 @@ class for a Keras Sequential model wrapped by KerasClassifier. It provides
 from keras.optimizers import Adam
 from scikeras.wrappers import KerasClassifier
 
+from ml_grid.util.global_params import global_parameters
+
 
 def create_model(
     layers: int = 1,
@@ -155,5 +157,15 @@ def floorer(t):
             # dropout_val = np.logspace(-1, -3, 2)
         }
 
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = {
+                "layers": [1],
+                "epochs": [1],
+                "batch_size": [32],
+                "l1_reg": [0.0],
+                "l2_reg": [0.0],
+                "width": [4],
+            }
+
     # The duplicate create_model method has been removed. The module-level
     # function will be used instead.
diff --git a/ml_grid/model_classes/knn_wrapper_class.py b/ml_grid/model_classes/knn_wrapper_class.py
@@ -4,6 +4,7 @@
 This module provides a scikit-learn compatible wrapper for the
 simbsig.neighbors.KNeighborsClassifier, with GPU support.
 """
+
 import logging
 from typing import Any, Dict, Optional, Union
 
diff --git a/ml_grid/model_classes/light_gbm_class.py b/ml_grid/model_classes/light_gbm_class.py
@@ -47,7 +47,14 @@ def __init__(
 
         self.parameter_space: Union[Dict[str, Any], List[Dict[str, Any]]]
 
-        if global_params.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "n_estimators": [2],
+                    "num_leaves": [2],
+                }
+            ]
+        elif global_parameters.bayessearch:
             self.parameter_space = {
                 "boosting_type": Categorical(("gbdt", "dart", "goss")),
                 "num_leaves": Integer(2, 100),
diff --git a/ml_grid/model_classes/randomforest_classifier_class.py b/ml_grid/model_classes/randomforest_classifier_class.py
@@ -54,7 +54,14 @@ def __init__(
         self.parameter_space: Dict[str, Any]
 
         # Define the parameter space for Bayesian and traditional search
-        if global_parameters.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "n_estimators": [2],
+                    "max_depth": [2],
+                }
+            ]
+        elif global_parameters.bayessearch:
             # Bayesian Optimization: Adjust parameters if traditional doesn't use param_dict
             self.parameter_space = {
                 "bootstrap": Categorical([True, False]),
diff --git a/ml_grid/model_classes/xgb_classifier_class.py b/ml_grid/model_classes/xgb_classifier_class.py
@@ -85,7 +85,14 @@ def patch_max_bin(param_value: Any) -> Union[int, Real, Integer, Any]:
                 return param_value
 
         # Set up the parameter space based on the selected optimization method
-        if global_parameters.bayessearch:
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {
+                    "n_estimators": [2],
+                    "max_depth": [2],
+                }
+            ]
+        elif global_parameters.bayessearch:
             # Bayesian Optimization: Define parameter space using Real and Categorical
             self.parameter_space = {
                 "objective": Categorical(
diff --git a/ml_grid/model_classes_time_series/CNNClassifier_module.py b/ml_grid/model_classes_time_series/CNNClassifier_module.py
@@ -1,6 +1,12 @@
 from typing import Any, Dict, List
 
-from aeon.classification.deep_learning import CNNClassifier
+# Monkeypatch sklearn.utils.validation for aeon compatibility
+import sklearn.utils.validation
+
+if not hasattr(sklearn.utils.validation, "validate_data"):
+    sklearn.utils.validation.validate_data = sklearn.utils.validation.check_X_y
+
+from aeon.classification.deep_learning import TimeCNNClassifier
 
 from ml_grid.pipeline.data import pipe
 from ml_grid.util.param_space import ParamSpace
@@ -19,7 +25,7 @@ class CNNClassifier_class:
             for the classifier.
     """
 
-    algorithm_implementation: CNNClassifier
+    algorithm_implementation: TimeCNNClassifier
     method_name: str
     parameter_space: Dict[str, List[Any]]
 
@@ -40,7 +46,7 @@ def __init__(self, ml_grid_object: pipe):
 
         log_epoch = param_space.param_dict.get("log_epoch")
 
-        self.algorithm_implementation = CNNClassifier()
+        self.algorithm_implementation = TimeCNNClassifier()
         self.method_name = "CNNClassifier"
 
         self.parameter_space = {
diff --git a/ml_grid/model_classes_time_series/rocketClassifier_module.py b/ml_grid/model_classes_time_series/rocketClassifier_module.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict, List
 
-from aeon.classification.convolution_based._rocket_classifier import RocketClassifier
+from aeon.classification.convolution_based import RocketClassifier
 
 from ml_grid.pipeline.data import pipe
 
diff --git a/ml_grid/model_classes_time_series/shapeDTWClassifier_module.py b/ml_grid/model_classes_time_series/shapeDTWClassifier_module.py
@@ -1,7 +1,30 @@
 from typing import Any, Dict, List
+import numpy as np
+from sklearn.base import BaseEstimator, ClassifierMixin
 
-from aeon.classification.distance_based import ShapeDTW
 
+class _DummyClassifier(BaseEstimator, ClassifierMixin):
+    """A dummy classifier to act as a placeholder for a missing model."""
+
+    def fit(self, X, y):
+        self.classes_ = np.unique(y)
+        return self
+
+    def predict(self, X):
+        return np.zeros(len(X), dtype=int)
+
+    def predict_proba(self, X):
+        proba = np.zeros((len(X), len(self.classes_)))
+        if proba.shape[1] > 0:
+            proba[:, 0] = 1.0
+        return proba
+
+
+try:
+    from aeon.classification.distance_based import ShapeDTW
+except ImportError:
+    # ShapeDTW was removed in aeon v0.11.0. Use a dummy placeholder.
+    ShapeDTW = _DummyClassifier
 from ml_grid.pipeline.data import pipe
 
 
@@ -19,27 +42,29 @@ class ShapeDTW_class:
             for the classifier.
     """
 
-    algorithm_implementation: ShapeDTW
+    algorithm_implementation: Any
     method_name: str
     parameter_space: Dict[str, List[Any]]
 
     def __init__(self, ml_grid_object: pipe):
         """Initializes the ShapeDTW_class.
-
         Args:
             ml_grid_object (pipe): An instance of the main data pipeline object.
         """
+        if ShapeDTW == _DummyClassifier:
+            ml_grid_object.logger.warning(
+                "ShapeDTW not found in aeon (removed in v0.11.0). "
+                "Using a dummy classifier to prevent a crash. This model will be skipped with a low score."
+            )
+            self.algorithm_implementation = _DummyClassifier()
+            self.method_name = "ShapeDTW"
+            self.parameter_space = {}
+        else:
+            n_jobs_model_val = ml_grid_object.global_params.n_jobs_model_val
 
-        n_jobs_model_val = ml_grid_object.global_params.n_jobs_model_val
-
-        self.algorithm_implementation = ShapeDTW()
-        self.method_name = "ShapeDTW"
-        self.parameter_space = {
-            "n_neighbours": [-1],
-            "subsequence_length": ["sqrt(n_timepoints)"],
-            "shape_descriptor_function": ["raw"],
-            "params": [None],
-            "shape_descriptor_functions": [["raw", "derivative"]],
-            "metric_params": [None],
-            "n_jobs": [n_jobs_model_val],
-        }
+            self.algorithm_implementation = ShapeDTW()
+            self.method_name = "ShapeDTW"
+            self.parameter_space = {
+                "shape_descriptor_function": ["raw", "derivative", "paa", "dwt"],
+                "n_jobs": [n_jobs_model_val],
+            }

Original file line number	Diff line number	Diff line change
`@@ -144,7 +144,7 @@ def fit(`
`144`	`144`	`base_models_list = []`
`145`	`145`	`for i, model_wrapper in enumerate(`
`146`	`146`	`self.base_models`
`147`		`- ): # type:H2OBaseClassifier`
	`147`	`+ ): # type: H2OBaseClassifier`
`148`	`148`	`self.logger.debug(`
`149`	`149`	`f"Fitting base model {i+1}: {type(model_wrapper).__name__}"`
`150`	`150`	`)`