Making sklearn estimator compliant based on original commit proposed by pizzooid (#50)

Author: maxiimilian

ffx/api.py

@@ -3,6 +3,8 @@
 '''
 
 from sklearn.base import BaseEstimator, RegressorMixin
+from sklearn.utils import check_array, check_X_y
+from sklearn.utils.validation import check_is_fitted
 
 
 def run(train_X, train_y, test_X, test_y, varnames=None, verbose=False):
@@ -15,6 +17,7 @@ class FFXRegressor(BaseEstimator, RegressorMixin):
     '''This class provides a Scikit-learn style estimator.'''
 
     def fit(self, X, y):
+        X, y = check_X_y(X, y, y_numeric=True, multi_output=False)
         # if X is a Pandas DataFrame, we don't have to pass in varnames.
         # otherwise we make up placeholders.
         if hasattr(X, 'columns'):
@@ -25,8 +28,11 @@ def fit(self, X, y):
             X, y, X, y, varnames=varnames
         )
         self.model_ = self.models_[-1]  # pylint: disable=attribute-defined-outside-init
+        return self
 
     def predict(self, X):
+        check_is_fitted(self, "model_")
+        X = check_array(X, accept_sparse=False)
         return self.model_.predict(X)
 
     def complexity(self):

ffx_tests/test_sklearn_api.py

@@ -1,5 +1,6 @@
 import ffx
 import numpy as np
+from sklearn.utils.estimator_checks import check_estimator
 
 EXPECTED_MODELS = [
     (0, 1, '0.298'),
@@ -40,3 +41,9 @@ def test_sklearn_api():
     assert [
         (model.numBases(), model.complexity(), str(model)) for model in FFX.models_
     ] == EXPECTED_MODELS
+
+
+def test_check_estimator():
+    # Pass instance of estimator to run sklearn's built in estimator check
+    check_estimator(ffx.FFXRegressor())
+