python: unittest for early stopping of cv

tests/python/test_early_stopping.py

@@ -2,12 +2,13 @@ import xgboost as xgb
 import numpy as np
 from sklearn.datasets import load_digits
 from sklearn.cross_validation import KFold, train_test_split
+from sklearn.metrics import mean_squared_error
 import unittest
 
 rng = np.random.RandomState(1994)
 
-class TestEarlyStopping(unittest.TestCase):
 
+class TestEarlyStopping(unittest.TestCase):
     def test_early_stopping_nonparallel(self):
         digits = load_digits(2)
         X = digits['data']
@@ -30,3 +31,32 @@ class TestEarlyStopping(unittest.TestCase):
 
         # TODO: parallel test for early stopping
         # TODO: comment out for now. Will re-visit later
+
+    def evalerror(self, preds, dtrain):
+        labels = dtrain.get_label()
+        return 'rmse', mean_squared_error(labels, preds)
+
+    def test_cv_early_stopping(self):
+        digits = load_digits(2)
+        X = digits['data']
+        y = digits['target']
+        dm = xgb.DMatrix(X, label=y)
+        params = {'max_depth': 2, 'eta': 1, 'silent': 1, 'objective': 'binary:logistic'}
+
+        import pandas as pd
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, early_stopping_rounds=10)
+        assert cv.shape[0] == 10
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, early_stopping_rounds=5)
+        assert cv.shape[0] == 3
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, early_stopping_rounds=1)
+        assert cv.shape[0] == 1
+
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, feval=self.evalerror,
+                    early_stopping_rounds=10)
+        assert cv.shape[0] == 10
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, feval=self.evalerror,
+                    early_stopping_rounds=1)
+        assert cv.shape[0] == 5
+        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, feval=self.evalerror,
+                    maximize=True, early_stopping_rounds=1)
+        assert cv.shape[0] == 1