Merge branch 'master' of github.com:AlexisMignon/xgboost

tests/python/test_cv.py

@@ -0,0 +1,37 @@
+import xgboost as xgb
+import numpy as np
+from sklearn.datasets import load_digits
+from sklearn.cross_validation import KFold, StratifiedKFold, train_test_split
+from sklearn.metrics import mean_squared_error
+import unittest
+
+rng = np.random.RandomState(1994)
+
+
+class TestCrossValidation(unittest.TestCase):
+    def test_cv(self):
+        digits = load_digits(3)
+        X = digits['data']
+        y = digits['target']
+        dm = xgb.DMatrix(X, label=y)
+        
+        params = {
+            'max_depth': 2,
+            'eta': 1,
+            'silent': 1,
+            'objective':
+            'multi:softprob',
+            'num_class': 3
+        }
+
+        seed = 2016
+        nfolds = 5
+        skf = StratifiedKFold(y, n_folds=nfolds, shuffle=True, random_state=seed)
+
+        import pandas as pd
+        cv1 = xgb.cv(params, dm, num_boost_round=10, nfold=nfolds, seed=seed)
+        cv2 = xgb.cv(params, dm, num_boost_round=10, folds=skf, seed=seed)
+        cv3 = xgb.cv(params, dm, num_boost_round=10, nfold=nfolds, stratified=True, seed=seed)
+        assert cv1.shape[0] == cv2.shape[0] and cv2.shape[0] == cv3.shape[0]
+        assert cv2.iloc[-1,0] == cv3.iloc[-1,0]
+