Set .classes_, .coef_ on scikit models correctly to enable prediction (#96)

Author: srowen

python/spark_sklearn/converter.py

@@ -105,14 +105,16 @@ def toSKLearn(self, model):
         :return: scikit-learn model with equivalent predictive behavior.
                  Currently, parameters or arguments for training are not copied.
         """
-        if isinstance(model, LogisticRegressionModel) or isinstance(model, LinearRegressionModel):
-            return self._toSKLGLM(model)
+        if isinstance(model, LogisticRegressionModel) :
+            return self._toSKLGLM(model, True)
+        if isinstance(model, LinearRegressionModel):
+            return self._toSKLGLM(model, False)
         else:
             supported_types = map(lambda t: type(t), self._supported_spark_types)
             raise ValueError("Converter.toSKLearn cannot convert type: %s.  Supported types: %s" %
                              (type(model), ", ".join(supported_types)))
 
-    def _toSKLGLM(self, model):
+    def _toSKLGLM(self, model, is_classifier):
         """ Private method for converting a GLM to a scikit-learn model
         TODO: Add model parameters as well.
         """
@@ -122,7 +124,9 @@ def _toSKLGLM(self, model):
         weights = model.coefficients
         skl = skl_cls()
         skl.intercept_ = np.float64(intercept)
-        skl.coef_ = weights.toArray()
+        skl.coef_ = weights.toArray().reshape(1, -1)
+        if is_classifier:
+            skl.classes_ = np.array([0, 1])
         return skl
 
     def toPandas(self, df):

python/spark_sklearn/converter_test.py

@@ -45,13 +45,17 @@ def test_LogisticRegression_spark2skl(self):
         self.assertTrue(isinstance(skl_lr, SKL_LogisticRegression),
                         "Expected sklearn LogisticRegression but found type %s" % type(skl_lr))
         self._compare_GLMs(skl_lr, lr)
+        # Make sure this doesn't throw an error
+        skl_lr.predict_proba(self.X)
 
     def test_LinearRegression_spark2skl(self):
         lr = LinearRegression().fit(self.df)
         skl_lr = self.converter.toSKLearn(lr)
         self.assertTrue(isinstance(skl_lr, SKL_LinearRegression),
                         "Expected sklearn LinearRegression but found type %s" % type(skl_lr))
         self._compare_GLMs(skl_lr, lr)
+        # Make sure this doesn't throw an error
+        skl_lr.predict(self.X)
 
     def ztest_toPandas(self):
         data = [(Vectors.dense([0.1, 0.2]),),