ENH Various enhancements to the model_selection module

raghavrv · raghavrv · commit 8f7bea977a25 · 2016-03-10T16:39:09.000+01:00
diff --git a/sklearn/model_selection/_split.py b/sklearn/model_selection/_split.py
@@ -1599,3 +1599,21 @@ def _build_repr(self):
         params[key] = value
 
     return '%s(%s)' % (class_name, _pprint(params, offset=len(class_name)))
+
+
+ALL_CVS = {'KFold': KFold,
+           'LabelKFold': LabelKFold,
+           'LeaveOneLabelOut': LeaveOneLabelOut,
+           'LeaveOneOut': LeaveOneOut,
+           'LeavePLabelOut': LeavePLabelOut,
+           'LeavePOut': LeavePOut,
+           'ShuffleSplit': ShuffleSplit,
+           'LabelShuffleSplit': LabelShuffleSplit,
+           'StratifiedKFold': StratifiedKFold,
+           'StratifiedShuffleSplit': StratifiedShuffleSplit,
+           'PredefinedSplit': PredefinedSplit}
+
+LABEL_CVS = {'LabelKFold': LabelKFold,
+             'LeaveOneLabelOut': LeaveOneLabelOut,
+             'LeavePLabelOut': LeavePLabelOut,
+             'LabelShuffleSplit': LabelShuffleSplit}
diff --git a/sklearn/model_selection/_validation.py b/sklearn/model_selection/_validation.py
@@ -27,39 +27,9 @@
 from ..metrics.scorer import check_scoring
 from ..exceptions import FitFailedWarning
 
-from ._split import KFold
-from ._split import LabelKFold
-from ._split import LeaveOneLabelOut
-from ._split import LeaveOneOut
-from ._split import LeavePLabelOut
-from ._split import LeavePOut
-from ._split import ShuffleSplit
-from ._split import LabelShuffleSplit
-from ._split import StratifiedKFold
-from ._split import StratifiedShuffleSplit
-from ._split import PredefinedSplit
-from ._split import check_cv, _safe_split
-
 __all__ = ['cross_val_score', 'cross_val_predict', 'permutation_test_score',
            'learning_curve', 'validation_curve']
 
-ALL_CVS = {'KFold': KFold,
-           'LabelKFold': LabelKFold,
-           'LeaveOneLabelOut': LeaveOneLabelOut,
-           'LeaveOneOut': LeaveOneOut,
-           'LeavePLabelOut': LeavePLabelOut,
-           'LeavePOut': LeavePOut,
-           'ShuffleSplit': ShuffleSplit,
-           'LabelShuffleSplit': LabelShuffleSplit,
-           'StratifiedKFold': StratifiedKFold,
-           'StratifiedShuffleSplit': StratifiedShuffleSplit,
-           'PredefinedSplit': PredefinedSplit}
-
-LABEL_CVS = {'LabelKFold': LabelKFold,
-             'LeaveOneLabelOut': LeaveOneLabelOut,
-             'LeavePLabelOut': LeavePLabelOut,
-             'LabelShuffleSplit': LabelShuffleSplit}
-
 
 def cross_val_score(estimator, X, y=None, labels=None, scoring=None, cv=None,
                     n_jobs=1, verbose=0, fit_params=None,
diff --git a/sklearn/model_selection/tests/test_search.py b/sklearn/model_selection/tests/test_search.py
@@ -34,11 +34,6 @@
 
 from sklearn.model_selection import KFold
 from sklearn.model_selection import StratifiedKFold
-from sklearn.model_selection import StratifiedShuffleSplit
-from sklearn.model_selection import LeaveOneLabelOut
-from sklearn.model_selection import LeavePLabelOut
-from sklearn.model_selection import LabelKFold
-from sklearn.model_selection import LabelShuffleSplit
 from sklearn.model_selection import GridSearchCV
 from sklearn.model_selection import RandomizedSearchCV
 from sklearn.model_selection import ParameterGrid
@@ -47,6 +42,7 @@
 # TODO Import from sklearn.exceptions once merged.
 from sklearn.base import ChangedBehaviorWarning
 from sklearn.model_selection._validation import FitFailedWarning
+from sklearn.model_selection._split import ALL_CVS, LABEL_CVS
 
 from sklearn.svm import LinearSVC, SVC
 from sklearn.tree import DecisionTreeRegressor
@@ -60,6 +56,14 @@
 from sklearn.pipeline import Pipeline
 
 
+def initialize_cross_validators(CVClass):
+    # set parameters to initialize the cross-validators
+    if CVClass is ALL_CVS['LeavePLabelOut']:
+        return CVClass(n_labels=2)
+    if CVClass is ALL_CVS['LeavePOut']:
+        return CVClass(p=2)
+
+
 # Neither of the following two estimators inherit from BaseEstimator,
 # to test hyperparameter search on user-defined classifiers.
 class MockClassifier(object):
@@ -235,17 +239,16 @@ def test_grid_search_labels():
     clf = LinearSVC(random_state=0)
     grid = {'C': [1]}
 
-    label_cvs = [LeaveOneLabelOut(), LeavePLabelOut(2), LabelKFold(),
-                 LabelShuffleSplit()]
-    for cv in label_cvs:
+    for _, CVClass in LABEL_CVS.iteritems():
+        cv = initialize_cross_validators(CVClass)
         gs = GridSearchCV(clf, grid, cv=cv)
         assert_raise_message(ValueError,
                              "The labels parameter should not be None",
                              gs.fit, X, y)
         gs.fit(X, y, labels)
 
-    non_label_cvs = [StratifiedKFold(), StratifiedShuffleSplit()]
-    for cv in non_label_cvs:
+    for _, CVClass in (set(ALL_CVS.iteritems()) - set(LABEL_CVS.iteritems())):
+        cv = initialize_cross_validators(CVClass)
         gs = GridSearchCV(clf, grid, cv=cv)
         # Should not raise an error
         gs.fit(X, y)
