fix issue #50 and #51

AxeldeRomblay · Oct 11, 2017 · bbb77f8 · bbb77f8
1 parent 925c227
commit bbb77f8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 11 deletions.
diff --git a/python-package/mlbox/preprocessing/drift/drift_estimator.py b/python-package/mlbox/preprocessing/drift/drift_estimator.py
@@ -107,7 +107,7 @@ def fit(self, df_train, df_test):
                                         X=X_tmp,
                                         y=self.__target,
                                         cv=self.__cv,
-                                        method="predict_proba")
+                                        method="predict_proba")[:,1]
 
         del df_train["target"]
         del df_test["target"]

diff --git a/python-package/mlbox/preprocessing/drift/drift_threshold.py b/python-package/mlbox/preprocessing/drift/drift_threshold.py
@@ -4,7 +4,6 @@
 # License: BSD 3 clause
 
 from joblib import Parallel, delayed
-import numpy as np
 from sklearn.tree import DecisionTreeClassifier
 
 from .drift_estimator import DriftEstimator
@@ -148,16 +147,9 @@ def fit(self, df_train, df_test):
 
         self.__Ddrifts = dict()
 
-        np.random.seed(self.random_state)
-
-        train_range = np.random.permutation(np.arange(self.subsample *
-                                                          len(df_train)))
-        test_range = np.random.permutation(np.arange(self.subsample *
-                                                         len(df_test)))
-
         Ldrifts = Parallel(n_jobs=self.n_jobs)(delayed(sync_fit)
-                                               (df_train.iloc[train_range][[col]],
-                                                df_test.iloc[test_range][[col]],
+                                               (df_train.sample(frac=self.subsample)[[col]],
+                                                df_test.sample(frac=self.subsample)[[col]],
                                                 self.estimator,
                                                 self.n_folds,
                                                 self.stratify,