Added Stratified split implementation

src/mlpack/core/data/CMakeLists.txt

@@ -27,6 +27,7 @@ set(SOURCES
   save_image.cpp
   serialization_template_version.hpp
   split_data.hpp
+  stratified_split_data.hpp
   imputer.hpp
   binarize.hpp
   string_encoding.hpp

src/mlpack/core/data/stratified_split_data.hpp

@@ -57,6 +57,8 @@ void StratifiedSplit(const arma::Mat<T>& input,
                      const double testRatio,
                      const bool shuffleData = true)
 {
+  arma::uvec trainIndexes;
+  arma::uvec testIndexes;
   if (shuffleData)
   {
     arma::uvec order = arma::shuffle(arma::linspace<arma::uvec>(
@@ -64,6 +66,70 @@ void StratifiedSplit(const arma::Mat<T>& input,
     input = input.cols(order);
     inputLabel = inputLabel.cols(order);
   }
+
   arma::Row<U> uniqueLabel = arma::unique(inputLabel);
-  //for (U )
+
+  for (typename U label : uniqueLabel)
+  {
+    arma::uvec uniqueIndexes = arma::find(inputLabel == label);
+
+    const size_t testStrataSize =
+        static_cast<size_t>(uniqueIndexes.n_rows*testRatio);
+    const size_t trainStrataSize = uniqueIndexes.n_rows - testSize;
+
+    arma::uvec testStrataIndexes =
+        uniqueIndexes.subvec(0, testStrataSize - 1);
+    arma::uvec trainStrataIndexes =
+        uniqueIndexes.subvec(testStrataSize - 1, uniqueIndexes.n_rows - 1);
+
+    testIndexes = join_cols(testIndexes, testStrataIndexes);
+    trainIndexes = join_cols(trainIndexes, trainStrataIndexes);
+  }
+
+  testData = input.cols(testIndexes);
+  testLabel = inputLabel.cols(testIndexes);
+  trainData = input.cols(trainIndexes);
+  trainLabel = inputLabel.cols(trainIndexes);
+}
+
+/**
+ * Given an input dataset and labels, split into a training set and test set.
+ * Example usage below.  This overload returns the split dataset as a std::tuple
+ * with four elements: an arma::Mat<T> containing the training data, an
+ * arma::Mat<T> containing the test data, an arma::Row<U> containing the
+ * training labels, and an arma::Row<U> containing the test labels.
+ *
+ * @code
+ * arma::mat input = loadData();
+ * arma::Row<size_t> label = loadLabel();
+ * auto splitResult = Split(input, label, 0.2);
+ * @endcode
+ *
+ * @param input Input dataset to split.
+ * @param inputLabel Input labels to split.
+ * @param testRatio Percentage of dataset to use for test set (between 0 and 1).
+ * @param shuffleData If true, the sample order is shuffled; otherwise, each
+ *       sample is visited in linear order. (Default true).
+ * @return std::tuple containing trainData (arma::Mat<T>), testData
+ *      (arma::Mat<T>), trainLabel (arma::Row<U>), and testLabel (arma::Row<U>).
+ */
+template<typename T, typename U>
+std::tuple<arma::Mat<T>, arma::Mat<T>, arma::Row<U>, arma::Row<U>>
+StratifiedSplit(const arma::Mat<T>& input,
+                const arma::Row<U>& inputLabel,
+                const double testRatio,
+                const bool shuffleData = true)
+{
+  arma::Mat<T> trainData;
+  arma::Mat<T> testData;
+  arma::Row<U> trainLabel;
+  arma::Row<U> testLabel;
+
+  StratifiedSplit(input, inputLabel, trainData, testData, trainLabel, testLabel,
+      testRatio, shuffleData);
+
+  return std::make_tuple(std::move(trainData),
+                         std::move(testData),
+                         std::move(trainLabel),
+                         std::move(testLabel));
 }

src/mlpack/methods/preprocess/preprocess_split_main.cpp

@@ -148,7 +148,14 @@ static void mlpackMain()
         IO::GetParam<arma::Mat<size_t>>("input_labels");
     arma::Row<size_t> labelsRow = labels.row(0);
 
-    const auto value = data::Split(data, labelsRow, testRatio, !shuffleData);
+    if(IO::HasParam("stratify"))
+    {
+      const auto value =
+          data::StratifiedSplit(data, labelsRow, testRatio, !shuffleData);
+    }
+    else {
+      const auto value = data::Split(data, labelsRow, testRatio, !shuffleData);
+    }
     Log::Info << "Training data contains " << get<0>(value).n_cols << " points."
         << endl;
     Log::Info << "Test data contains " << get<1>(value).n_cols << " points."