add test_fit_predict

tests/tpot_tests.py

@@ -603,8 +603,25 @@ def test_fit3():
     assert not (tpot_obj._start_datetime is None)
 
 
+def test_fit_predict():
+    """Assert that the TPOT fit_predict function provides an optimized pipeline and correct output."""
+    tpot_obj = TPOTClassifier(
+        random_state=42,
+        population_size=1,
+        offspring_size=2,
+        generations=1,
+        verbosity=0,
+        config_dict='TPOT light'
+    )
+    result = tpot_obj.fit_predict(training_features, training_target)
+
+    assert isinstance(tpot_obj._optimized_pipeline, creator.Individual)
+    assert not (tpot_obj._start_datetime is None)
+    assert result.shape == (training_features.shape[0],)
+
+
 def test_update_top_pipeline():
-    """Assert that the TPOT _update_top_pipeline updated an optimized pipeline"""
+    """Assert that the TPOT _update_top_pipeline updated an optimized pipeline."""
     tpot_obj = TPOTClassifier(
         random_state=42,
         population_size=1,
@@ -623,7 +640,7 @@ def test_update_top_pipeline():
 
 
 def test_update_top_pipeline_2():
-    """Assert that the TPOT _update_top_pipeline raises RuntimeError when self._pareto_front is empty"""
+    """Assert that the TPOT _update_top_pipeline raises RuntimeError when self._pareto_front is empty."""
     tpot_obj = TPOTClassifier(
         random_state=42,
         population_size=1,
@@ -641,7 +658,7 @@ def pareto_eq(ind1, ind2):
 
 
 def test_update_top_pipeline_3():
-    """Assert that the TPOT _update_top_pipeline raises RuntimeError when self._optimized_pipeline is updated"""
+    """Assert that the TPOT _update_top_pipeline raises RuntimeError when self._optimized_pipeline is not updated."""
     tpot_obj = TPOTClassifier(
         random_state=42,
         population_size=1,

tpot/base.py

@@ -664,7 +664,7 @@ def predict(self, features):
 
         return self.fitted_pipeline_.predict(features)
 
-    def fit_predict(self, features, target):
+    def fit_predict(self, features, target, sample_weight=None, groups=None):
         """Call fit and predict in sequence.
 
         Parameters
@@ -673,14 +673,20 @@ def fit_predict(self, features, target):
             Feature matrix
         target: array-like {n_samples}
             List of class labels for prediction
+        sample_weight: array-like {n_samples}, optional
+            Per-sample weights. Higher weights force TPOT to put more emphasis on those points
+        groups: array-like, with shape {n_samples, }, optional
+            Group labels for the samples used when performing cross-validation.
+            This parameter should only be used in conjunction with sklearn's Group cross-validation
+            functions, such as sklearn.model_selection.GroupKFold
 
         Returns
         ----------
         array-like: {n_samples}
             Predicted target for the provided features
 
         """
-        self.fit(features, target)
+        self.fit(features, target, sample_weight=sample_weight, groups=groups)
         return self.predict(features)
 
     def score(self, testing_features, testing_target):