Merge pull request #523 from stereomatchingkiss/split_data

add train test split

src/mlpack/core/util/split_data.hpp

@@ -0,0 +1,103 @@
+#ifndef __MLPACK_CORE_UTIL_SPLIT_DATA_HPP
+#define __MLPACK_CORE_UTIL_SPLIT_DATA_HPP
+
+#include <mlpack/core.hpp>
+
+#include <tuple>
+
+namespace mlpack {
+namespace util {
+
+/**
+ *Split training data and test data
+ *@param input input data want to split
+ *@param label input label want to split
+ *@param trainData training data split by input
+ *@param testData test data split by input
+ *@param trainLabel train label split by input
+ *@param testLabel test label split by input
+ *@param testRatio the ratio of test data
+ *@code
+ *arma::mat input = loadData();
+ *arma::Row<size_t> label = loadLabel();
+ *arma::mat trainData;
+ *arma::mat testData;
+ *arma::Row<size_t> trainLabel;
+ *arma::Row<size_t> testLabel;
+ *arma::arma_rng::set_seed(100); //set the seed if you like
+ *TrainTestSplit(input, label, trainData,
+ *               testData, trainLabel, testLabel);
+ *@endcode
+ */
+template<typename T, typename U>
+void TrainTestSplit(const arma::Mat<T> &input,
+                    const arma::Row<U> &inputLabel,
+                    arma::Mat<T> &trainData,
+                    arma::Mat<T> &testData,
+                    arma::Row<U> &trainLabel,
+                    arma::Row<U> &testLabel,
+                    const double testRatio)
+{
+  size_t const testSize =
+      static_cast<size_t>(input.n_cols * testRatio);
+  const size_t trainSize = input.n_cols - testSize;
+  trainData.set_size(input.n_rows, trainSize);
+  testData.set_size(input.n_rows, testSize);
+  trainLabel.set_size(trainSize);
+  testLabel.set_size(testSize);
+
+  const arma::Col<size_t> order =
+      arma::shuffle(arma::linspace<arma::Col<size_t>>(0, input.n_cols - 1,
+                                                      input.n_cols));
+
+  for(size_t i = 0; i != trainSize; ++i)
+  {
+    trainData.col(i) = input.col(order[i]);
+    trainLabel(i) = inputLabel(order[i]);
+  }
+
+  for(size_t i = 0; i != testSize; ++i)
+  {
+    testData.col(i) = input.col(order[i + trainSize]);
+    testLabel(i) = inputLabel(order[i + trainSize]);
+  }
+}
+
+/**
+ *Overload of Split, if you do not like to pass in
+ *so many param, you could call this api instead
+ *@param input input data want to split
+ *@param label input label want to split
+ *@return They are trainData, testData, trainLabel and
+ *testLabel
+ *@code
+ *arma::mat input = loadData();
+ *arma::Row<size_t> label = loadLabel();
+ *auto splitResult = TrainTestSplit(input, label, 0.2);
+ *@endcode
+ */
+template<typename T,typename U>
+std::tuple<arma::Mat<T>, arma::Mat<T>,
+arma::Row<U>, arma::Row<U>>
+TrainTestSplit(const arma::Mat<T> &input,
+               const arma::Row<U> &inputLabel,
+               const double testRatio)
+{
+  arma::Mat<T> trainData;
+  arma::Mat<T> testData;
+  arma::Row<U> trainLabel;
+  arma::Row<U> testLabel;
+
+  TrainTestSplit(input, inputLabel,
+                 trainData, testData,
+                 trainLabel, testLabel,
+                 testRatio);
+
+  return std::make_tuple(trainData, testData,
+                         trainLabel, testLabel);
+}
+
+} // namespace util
+} // namespace mlpack
+
+#endif

src/mlpack/prereqs.hpp

@@ -77,6 +77,7 @@
 // it's part of the C++11 standard.
 #ifdef _MSC_VER
   #pragma warning(disable : 4519)
+  #define ARMA_USE_CXX11
 #endif
 
 #endif

src/mlpack/tests/CMakeLists.txt

@@ -67,6 +67,7 @@ add_executable(mlpack_test
   sort_policy_test.cpp
   sparse_autoencoder_test.cpp
   sparse_coding_test.cpp
+  split_data_test.cpp
   termination_policy_test.cpp
   tree_test.cpp
   tree_traits_test.cpp

src/mlpack/tests/split_data_test.cpp

@@ -0,0 +1,59 @@
+/**
+ * @file sparse_autoencoder_test.cpp
+ * @author Siddharth Agrawal
+ *
+ * Test the SparseAutoencoder class.
+ */
+
+#include <mlpack/core.hpp>
+#include <mlpack/core/util/split_data.hpp>
+
+#include <boost/test/unit_test.hpp>
+#include "old_boost_test_definitions.hpp"
+
+using namespace mlpack;
+using namespace arma;
+
+BOOST_AUTO_TEST_SUITE(SplitDataTest);
+
+/**
+ * compare the data after train test split
+ * @param inputData The original data set before split
+ * @param compareData The data want to compare with the inputData,
+ * it could be train data or test data
+ * @param inputLabel The label of the compareData
+ */
+void CompareData(arma::mat const &inputData, arma::mat const &compareData,
+                 arma::Row<size_t> const &inputLabel)
+{
+  for(size_t i = 0; i != compareData.n_cols; ++i){
+    arma::mat const &lhsCol = inputData.col(inputLabel(i));
+    arma::mat const &rhsCol = compareData.col(i);
+    for(size_t j = 0; j != lhsCol.n_rows; ++j){
+      BOOST_REQUIRE_CLOSE(lhsCol(j), rhsCol(j), 1e-5);
+    }
+  }
+}
+
+BOOST_AUTO_TEST_CASE(SplitDataSplitResultMat)
+{    
+  arma::mat input(2,10);
+  input.randu();
+  using Labels = arma::Row<size_t>;
+  //set the labels range same as the col, so the CompareData
+  //can compare the data after TrainTestSplit are valid or not
+  Labels const labels =
+          arma::linspace<Labels>(0, input.n_cols-1,
+                                 input.n_cols);
+
+  auto const value = util::TrainTestSplit(input, labels, 0.2);
+  BOOST_REQUIRE(std::get<0>(value).n_cols == 8);
+  BOOST_REQUIRE(std::get<1>(value).n_cols == 2);
+  BOOST_REQUIRE(std::get<2>(value).n_cols == 8);
+  BOOST_REQUIRE(std::get<3>(value).n_cols == 2);
+
+  CompareData(input, std::get<0>(value), std::get<2>(value));
+  CompareData(input, std::get<1>(value), std::get<3>(value));
+}
+
+BOOST_AUTO_TEST_SUITE_END();