Added 'SelectFirstEstimator.eval_rows' attribute

sklearn2pmml/ensemble/__init__.py

@@ -329,15 +329,22 @@ def predict_proba(self, X):
 
 class SelectFirstEstimator(_BaseEnsemble):
 
-	def __init__(self, steps, controller):
+	def __init__(self, steps, controller, eval_rows):
 		super(SelectFirstEstimator, self).__init__(steps, controller)
+		self.eval_rows = eval_rows
+
+	def _eval_step_mask(self, X, predicate):
+		step_func = to_expr_func(predicate)
+		if self.eval_rows:
+			return eval_rows(X, step_func, dtype = bool)
+		else:
+			return step_func(X)
 
 	def fit(self, X, y, **fit_params):
 		X_eval = self._to_evaluation_dataset(X)
 		mask = numpy.zeros(X.shape[0], dtype = bool)
 		for name, estimator, predicate in self.steps:
-			step_mask_func = to_expr_func(predicate)
-			step_mask = eval_rows(X_eval, step_mask_func, dtype = bool)
+			step_mask = self._eval_step_mask(X_eval, predicate)
 			step_mask[mask] = False
 			if numpy.sum(step_mask) < 1:
 				raise ValueError(predicate)
@@ -353,8 +360,7 @@ def _predict(self, X, predict_method):
 		X_eval = self._to_evaluation_dataset(X)
 		mask = numpy.zeros(X.shape[0], dtype = bool)
 		for name, estimator, predicate in self.steps:
-			step_mask_func = to_expr_func(predicate)
-			step_mask = eval_rows(X_eval, step_mask_func, dtype = bool)
+			step_mask = self._eval_step_mask(X_eval, predicate)
 			step_mask[mask] = False
 			if numpy.sum(step_mask) < 1:
 				continue
@@ -376,13 +382,13 @@ def predict(self, X):
 
 class SelectFirstRegressor(SelectFirstEstimator, RegressorMixin):
 
-	def __init__(self, steps, controller = None):
-		super(SelectFirstRegressor, self).__init__(steps, controller)
+	def __init__(self, steps, controller = None, eval_rows = True):
+		super(SelectFirstRegressor, self).__init__(steps, controller, eval_rows)
 
 class SelectFirstClassifier(SelectFirstEstimator, ClassifierMixin):
 
-	def __init__(self, steps, controller = None):
-		super(SelectFirstClassifier, self).__init__(steps, controller)
+	def __init__(self, steps, controller = None, eval_rows = True):
+		super(SelectFirstClassifier, self).__init__(steps, controller, eval_rows)
 
 	def predict_proba(self, X):
 		return self._predict(X, "predict_proba")

sklearn2pmml/ensemble/tests/__init__.py

@@ -116,7 +116,7 @@ def test_fit_predict(self):
 			("negative", DummyClassifier(strategy = "most_frequent"), "X[0] < 0"),
 			("positive", DummyClassifier(strategy = "most_frequent"), "X[0] > 0"),
 			("zero", DummyClassifier(strategy = "constant", constant = 0), str(True))
-		]))
+		], eval_rows = True))
 		params = classifier.get_params(deep = True)
 		self.assertEqual("most_frequent", params["negative__strategy"])
 		self.assertEqual("most_frequent", params["positive__strategy"])
@@ -128,6 +128,15 @@ def test_fit_predict(self):
 		self.assertEqual([-1, 0, -1, 1, -1], pred.tolist())
 		pred_proba = classifier.predict_proba(X)
 		self.assertEqual((5, 2), pred_proba.shape)
+		X = X.values
+		classifier = SelectFirstClassifier([
+			("negative", DummyClassifier(strategy = "most_frequent"), "X[:, 0] < 0"),
+			("positive", DummyClassifier(strategy = "most_frequent"), "X[:, 0] > 0"),
+			("zero", DummyClassifier(strategy = "constant", constant = 0), "X[:, 0] == 0")
+		], eval_rows = False)
+		classifier.fit(X, y)
+		pred = classifier.predict(X)
+		self.assertEqual([-1, 0, -1, 1, -1], pred.tolist())
 
 class SelectFirstRegressorTest(TestCase):
 	pass