Add TPOT classifier wrapper, configuration, and tests

- Implement `TPOTClassifierWrapper` in `ml_grid/model_classes/TPOTClassifierWrapper.py` to wrap `TPOTClassifier` for scikit-learn compatibility, handling fitting, prediction, and probability estimation.
- Create `TPOTClassifierClass` in `ml_grid/model_classes/tpot_classifier_class.py` to define hyperparameter spaces for grid search and Bayesian optimization (using `skopt`).
- Add unit tests in `tests/test_tpot_classifier.py` to verify initialization, fitting, prediction, and configuration structure.
- Handle optional TPOT import to avoid hard dependency failures.

Author: SamoraHunter

ml_grid/model_classes/TPOTClassifierWrapper.py

@@ -0,0 +1,114 @@
+"""TPOT Classifier Wrapper.
+
+This module provides a scikit-learn compatible wrapper for TPOTClassifier.
+"""
+
+import logging
+from typing import Optional, Union
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.utils.validation import check_is_fitted
+
+# Attempt to import TPOT
+try:
+    from tpot import TPOTClassifier
+except ImportError:
+    TPOTClassifier = None
+
+logger = logging.getLogger(__name__)
+
+
+class TPOTClassifierWrapper(BaseEstimator, ClassifierMixin):
+    """A scikit-learn compatible wrapper for TPOTClassifier."""
+
+    def __init__(
+        self,
+        generations: int = 5,
+        population_size: int = 20,
+        offspring_size: Optional[int] = None,
+        mutation_rate: float = 0.9,
+        crossover_rate: float = 0.1,
+        scoring: str = "accuracy",
+        cv: int = 5,
+        subsample: float = 1.0,
+        n_jobs: int = -1,
+        max_time_mins: Optional[int] = None,
+        max_eval_time_mins: float = 5,
+        random_state: int = 42,
+        verbosity: int = 2,
+        early_stop: Optional[int] = None,
+    ):
+        self.generations = generations
+        self.population_size = population_size
+        self.offspring_size = offspring_size
+        self.mutation_rate = mutation_rate
+        self.crossover_rate = crossover_rate
+        self.scoring = scoring
+        self.cv = cv
+        self.subsample = subsample
+        self.n_jobs = n_jobs
+        self.max_time_mins = max_time_mins
+        self.max_eval_time_mins = max_eval_time_mins
+        self.random_state = random_state
+        self.verbosity = verbosity
+        self.early_stop = early_stop
+
+        self.model_ = None
+
+    def fit(
+        self,
+        X: Union[np.ndarray, pd.DataFrame],
+        y: Union[np.ndarray, pd.Series],
+        **kwargs,
+    ) -> "TPOTClassifierWrapper":
+        if TPOTClassifier is None:
+            raise ImportError(
+                "TPOT is not installed. Please install it to use TPOTClassifierWrapper."
+            )
+
+        self.model_ = TPOTClassifier(
+            generations=self.generations,
+            population_size=self.population_size,
+            offspring_size=self.offspring_size,
+            mutation_rate=self.mutation_rate,
+            crossover_rate=self.crossover_rate,
+            scoring=self.scoring,
+            cv=self.cv,
+            subsample=self.subsample,
+            n_jobs=self.n_jobs,
+            max_time_mins=self.max_time_mins,
+            max_eval_time_mins=self.max_eval_time_mins,
+            random_state=self.random_state,
+            verbosity=self.verbosity,
+            early_stop=self.early_stop,
+            disable_update_check=True,
+        )
+
+        # TPOT can be slow. For quick checks, it's useful to see it has started.
+        logger.info(
+            f"Starting TPOT fit with generations={self.generations}, population_size={self.population_size}..."
+        )
+
+        self.model_.fit(X, y, **kwargs)
+
+        # After fitting, TPOT stores the best pipeline in the `fitted_pipeline_` attribute.
+        # We must set `classes_` for scikit-learn compatibility (e.g., for check_is_fitted).
+        # While TPOT exposes `self.model_.classes_`, inferring from `y` is a more robust
+        # fallback, consistent with other wrappers in this project.
+        if hasattr(self.model_, "classes_"):
+            self.classes_ = self.model_.classes_
+        else:
+            self.classes_ = np.unique(y)
+
+        logger.info("TPOT fit completed.")
+        return self
+
+    def predict(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
+        check_is_fitted(self, ["model_", "classes_"])
+        return self.model_.predict(X)
+
+    def predict_proba(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
+        check_is_fitted(self, ["model_", "classes_"])
+        return self.model_.predict_proba(X)

ml_grid/model_classes/tpot_classifier_class.py

@@ -0,0 +1,60 @@
+"""TPOT Classifier Configuration.
+
+This module contains the TPOTClassifierClass, which is a configuration
+class for the TPOTClassifierWrapper. It provides parameter spaces for
+grid search and Bayesian optimization, with a focus on providing a fast
+default for unit testing.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+import pandas as pd
+from skopt.space import Categorical, Integer
+
+from ml_grid.model_classes.TPOTClassifierWrapper import TPOTClassifierWrapper
+from ml_grid.util.global_params import global_parameters
+
+logger = logging.getLogger(__name__)
+
+
+class TPOTClassifierClass:
+    """Configuration class for TPOTClassifierWrapper."""
+
+    def __init__(
+        self,
+        X: Optional[pd.DataFrame] = None,
+        y: Optional[pd.Series] = None,
+        parameter_space_size: Optional[str] = None,
+    ):
+        self.X = X
+        self.y = y
+        self.algorithm_implementation = TPOTClassifierWrapper()
+        self.method_name = "TPOTClassifier"
+
+        self.parameter_space: Union[List[Dict[str, Any]], Dict[str, Any]]
+
+        if getattr(global_parameters, "test_mode", False):
+            self.parameter_space = [
+                {"generations": [2], "population_size": [5], "max_time_mins": [1]}
+            ]
+        elif global_parameters.bayessearch:
+            # A slightly larger space for Bayesian search, but still constrained
+            self.parameter_space = {
+                "generations": Integer(5, 100),
+                "population_size": Integer(20, 100),
+                "scoring": Categorical(
+                    ["accuracy", "f1", "roc_auc", "precision", "recall"]
+                ),
+                "max_time_mins": Integer(10, 120),  # Time limit is crucial
+            }
+        else:
+            # Expanded parameter space for grid search
+            self.parameter_space = [
+                {
+                    "generations": [5, 10, 20],
+                    "population_size": [20, 50, 100],
+                    "max_time_mins": [10, 30, 60],
+                    "scoring": ["accuracy", "f1", "roc_auc"],
+                }
+            ]

tests/test_tpot_classifier.py

@@ -0,0 +1,101 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import pandas as pd
+import numpy as np
+
+from ml_grid.model_classes.TPOTClassifierWrapper import TPOTClassifierWrapper
+from ml_grid.model_classes.tpot_classifier_class import TPOTClassifierClass
+
+
+class TestTPOTClassifier(unittest.TestCase):
+    def setUp(self):
+        self.X = pd.DataFrame(
+            {"feature_0": [1.0, 2.0, 3.0, 4.0], "feature_1": [4.0, 3.0, 2.0, 1.0]}
+        )
+        self.y = pd.Series([0, 1, 0, 1], name="target")
+
+    def test_init(self):
+        clf = TPOTClassifierWrapper(generations=10, population_size=50)
+        self.assertEqual(clf.generations, 10)
+        self.assertEqual(clf.population_size, 50)
+        self.assertIsNone(clf.model_)
+
+    @patch("ml_grid.model_classes.TPOTClassifierWrapper.TPOTClassifier")
+    def test_fit(self, mock_tpot_cls):
+        # Setup mocks
+        mock_tpot_instance = MagicMock()
+        mock_tpot_cls.return_value = mock_tpot_instance
+
+        clf = TPOTClassifierWrapper(generations=5, population_size=20)
+
+        # Test fit
+        clf.fit(self.X, self.y)
+
+        # Verify TPOTClassifier init
+        mock_tpot_cls.assert_called_once()
+        _, kwargs = mock_tpot_cls.call_args
+        self.assertEqual(kwargs["generations"], 5)
+        self.assertEqual(kwargs["population_size"], 20)
+        self.assertEqual(kwargs["disable_update_check"], True)
+
+        # Verify fit call
+        mock_tpot_instance.fit.assert_called_once_with(self.X, self.y)
+
+        # Verify attributes set
+        self.assertIsNotNone(clf.model_)
+
+    @patch("ml_grid.model_classes.TPOTClassifierWrapper.TPOTClassifier")
+    def test_predict(self, mock_tpot_cls):
+        # Setup mock
+        mock_tpot_instance = MagicMock()
+        mock_tpot_cls.return_value = mock_tpot_instance
+
+        # Mock predict return
+        mock_tpot_instance.predict.return_value = np.array([0, 1, 0, 1])
+
+        clf = TPOTClassifierWrapper()
+        clf.fit(self.X, self.y)
+
+        preds = clf.predict(self.X)
+
+        self.assertIsInstance(preds, np.ndarray)
+        np.testing.assert_array_equal(preds, np.array([0, 1, 0, 1]))
+
+        # Verify predict called on internal model
+        mock_tpot_instance.predict.assert_called_once_with(self.X)
+
+    def test_missing_tpot(self):
+        # Simulate missing tpot by patching TPOTClassifier to None
+        with patch("ml_grid.model_classes.TPOTClassifierWrapper.TPOTClassifier", None):
+            clf = TPOTClassifierWrapper()
+            with self.assertRaises(ImportError):
+                clf.fit(self.X, self.y)
+
+
+class TestTPOTClassifierClass(unittest.TestCase):
+    def test_structure(self):
+        # Mock global_parameters to control bayessearch flag
+        with patch(
+            "ml_grid.model_classes.tpot_classifier_class.global_parameters"
+        ) as mock_globals:
+            # Case 1: Grid Search (bayessearch = False)
+            mock_globals.bayessearch = False
+
+            config = TPOTClassifierClass()
+            self.assertEqual(config.method_name, "TPOTClassifier")
+            self.assertIsInstance(
+                config.algorithm_implementation, TPOTClassifierWrapper
+            )
+            self.assertIsInstance(config.parameter_space, list)
+
+            # Case 2: Bayes Search (bayessearch = True)
+            mock_globals.bayessearch = True
+
+            config = TPOTClassifierClass()
+            self.assertIsInstance(config.parameter_space, dict)
+            self.assertIn("generations", config.parameter_space)
+            self.assertIn("population_size", config.parameter_space)
+
+
+if __name__ == "__main__":
+    unittest.main()