Added an implementation to Stratify Data

src/mlpack/core/data/split_data.hpp

@@ -17,6 +17,124 @@
 
 namespace mlpack {
 namespace data {
+
+/**
+ * Given an input dataset and labels, stratify into a training set and test set.
+ * Example usage below. This overload places the stratified dataset into the
+ * four output parameters given (trainData, testData, trainLabel,
+ * and testLabel).
+ *
+ * @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;
+ * math::RandomSeed(100); // Set the seed if you like.
+ *
+ * // Stratify the dataset into a training and test set, with 30% of the data
+ * // being held out for the test set.
+ * StratifiedSplit(input, label, trainData,
+ *                 testData, trainLabel, testLabel, 0.3);
+ * @endcode
+ *
+ * @param input Input dataset to stratify.
+ * @param inputLabel Input labels to stratify.
+ * @param trainData Matrix to store training data into.
+ * @param testData Matrix to store test data into.
+ * @param trainLabel Vector to store training labels into.
+ * @param testLabel Vector to store test labels into.
+ * @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.)
+ */
+template<typename T, typename U>
+void StratifiedSplit(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,
+                     const bool shuffleData = true)
+{
+  /**
+   * Basic idea:
+   * Let us say we have to stratify a dataset based on labels:
+   * 0 0 0 0 0 1 1 1 1 1 1 (5 0s, 6 1s)
+   *
+   * Let our test ratio be 0.2:
+   * We visit each label and keep the count of each label in our unordered map
+   *
+   * Whenever we encounter a label, we calculate
+   * current_count*test_ratio --- labelMap[label]*testRatio and
+   * current_count+1 * test_ratio --- (labelMap[label]+1)*testRatio
+   *
+   * We then static_cast these counts to size_t to remove their decimal points.
+   * If in this case, our integer counts are same then we add to our train set.
+   * If there is a difference in counts, then we add to our test set
+   *
+   * Considering our example
+   * 0 -- train set ( 0*0.2 == 1*0.2 ) (After casting)
+   * 0 -- train set ( 1*0.2 == 2*0.2 ) (After casting)
+   * 0 -- train set ( 2*0.2 == 3*0.2 ) (After casting)
+   * 0 -- train set ( 3*0.2 == 4*0.2 ) (After casting)
+   * 0 -- test set  ( 4*0.2 <  5*0.2 ) (After casting)
+   *
+   * 1 -- train set ( 0*0.2 == 1*0.2 ) (After casting)
+   * 1 -- train set ( 1*0.2 == 2*0.2 ) (After casting)
+   * 1 -- train set ( 2*0.2 == 3*0.2 ) (After casting)
+   * 1 -- train set ( 3*0.2 == 4*0.2 ) (After casting)
+   * 1 -- test set  ( 4*0.2 <  5*0.2 ) (After casting)
+   * 1 -- train set ( 5*0.2 == 6*0.2 ) (After casting)
+   *
+   * Finally
+   * train set,
+   * 0 0 0 0 1 1 1 1 1 (4 0s, 5 1s)
+   *
+   * test set,
+   * 0 1
+   */
+  arma::uvec Indexes;
+  Indexes.set_size(inputLabel.n_cols);
+
+  size_t trainIdx = 0;
+  size_t testIdx = inputLabel.n_cols - 1;
+  std::unordered_map<U, size_t> labelMap;
+
+  arma::uvec order =
+      arma::linspace<arma::uvec>(0, input.n_cols - 1, input.n_cols);
+
+  if (shuffleData)
+  {
+    order = arma::shuffle(order);
+  }
+
+  for (auto i: order)
+  {
+    auto label = inputLabel[i];
+    if (static_cast<size_t>(labelMap[label]*testRatio) <
+        static_cast<size_t>((labelMap[label]+1)*testRatio))
+    {
+      Indexes[testIdx] = i;
+      testIdx -= 1;
+    }
+    else
+    {
+      Indexes[trainIdx] = i;
+      trainIdx += 1;
+    }
+    labelMap[label] += 1;
+  }
+
+  labelMap.clear();
+  testData = input.cols(Indexes.subvec(trainIdx, Indexes.n_rows-1));
+  testLabel = inputLabel.cols(Indexes.subvec(trainIdx, Indexes.n_rows-1));
+  trainData = input.cols(Indexes.subvec(0, trainIdx-1));
+  trainLabel = inputLabel.cols(Indexes.subvec(0, trainIdx-1));
+}
+
 /**
  * Given an input dataset and labels, split into a training set and test set.
  * Example usage below.  This overload places the split dataset into the four
@@ -176,15 +294,24 @@ std::tuple<arma::Mat<T>, arma::Mat<T>, arma::Row<U>, arma::Row<U>>
 Split(const arma::Mat<T>& input,
       const arma::Row<U>& inputLabel,
       const double testRatio,
-      const bool shuffleData = true)
+      const bool shuffleData = true,
+      const bool stratifyData = false)
 {
   arma::Mat<T> trainData;
   arma::Mat<T> testData;
   arma::Row<U> trainLabel;
   arma::Row<U> testLabel;
 
-  Split(input, inputLabel, trainData, testData, trainLabel, testLabel,
-      testRatio, shuffleData);
+  if (stratifyData)
+  {
+    StratifiedSplit(input, inputLabel, trainData, testData, trainLabel,
+        testLabel, testRatio, shuffleData);
+  }
+  else
+  {
+    Split(input, inputLabel, trainData, testData, trainLabel, testLabel,
+        testRatio, shuffleData);
+  }
 
   return std::make_tuple(std::move(trainData),
                          std::move(testData),

src/mlpack/methods/preprocess/preprocess_split_main.cpp

@@ -69,6 +69,21 @@ BINDING_EXAMPLE(
         "test_ratio", 0.3, "training", "X_train", "training_labels", "y_train",
         "test", "X_test", "test_labels", "y_test"));
 
+BINDING_EXAMPLE(
+    "There is an option to stratify the dataset according to the labels." +
+    PRINT_PARAM_STRING("stratify_data") + " option to stratify the data;"
+    "an example to stratify the data is" +
+    "\n\n" +
+    "If we had a dataset " + PRINT_DATASET("X") + " and associated labels " +
+    PRINT_DATASET("y") + ", and we wanted to stratify these into " +
+    PRINT_DATASET("X_train") + ", " + PRINT_DATASET("y_train") + ", " +
+    PRINT_DATASET("X_test") + ", and " + PRINT_DATASET("y_test") + ", with 30% "
+    "of the data in the test set, we could run"
+    "\n\n" +
+    PRINT_CALL("preprocess_split", "input", "X", "input_labels", "y",
+        "test_ratio", 0.3, "training", "X_train", "training_labels", "y_train",
+        "test", "X_test", "test_labels", "y_test", "stratify_data", true));
+
 // See also...
 BINDING_SEE_ALSO("@preprocess_binarize", "#preprocess_binarize");
 BINDING_SEE_ALSO("@preprocess_describe", "#preprocess_describe");
@@ -90,6 +105,7 @@ PARAM_DOUBLE_IN("test_ratio", "Ratio of test set; if not set,"
 
 PARAM_INT_IN("seed", "Random seed (0 for std::time(NULL)).", "s", 0);
 PARAM_FLAG("no_shuffle", "Avoid shuffling and splitting the data.", "S");
+PARAM_FLAG("stratify_data", "Stratify the data according to labels", "z")
 
 using namespace mlpack;
 using namespace mlpack::data;
@@ -102,6 +118,7 @@ static void mlpackMain()
   // Parse command line options.
   const double testRatio = IO::GetParam<double>("test_ratio");
   const bool shuffleData = IO::GetParam<bool>("no_shuffle");
+  const bool stratifyData = IO::GetParam<bool>("stratify_data");
 
   if (IO::GetParam<int>("seed") == 0)
     mlpack::math::RandomSeed(std::time(NULL));
@@ -125,6 +142,7 @@ static void mlpackMain()
   {
     ReportIgnoredParam({{ "input_labels", true }}, "training_labels");
     ReportIgnoredParam({{ "input_labels", true }}, "test_labels");
+    ReportIgnoredParam("stratify_data", "input_labels is not provided");
   }
 
   // Check test_ratio.
@@ -148,11 +166,12 @@ 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);
-    Log::Info << "Training data contains " << get<0>(value).n_cols << " points."
-        << endl;
-    Log::Info << "Test data contains " << get<1>(value).n_cols << " points."
-        << endl;
+    const auto value =
+        data::Split(data, labelsRow, testRatio, !shuffleData, stratifyData);
+    Log::Info << "Training data contains "
+        << get<0>(value).n_cols << " points." << endl;
+    Log::Info << "Test data contains "
+        << get<1>(value).n_cols << " points." << endl;
 
     if (IO::HasParam("training"))
       IO::GetParam<arma::mat>("training") = std::move(get<0>(value));