split_data.hpp

Go to the documentation of this file.

#ifndef MLPACK_CORE_DATA_SPLIT_DATA_HPP
#define MLPACK_CORE_DATA_SPLIT_DATA_HPP

#include <mlpack/prereqs.hpp>

namespace mlpack {
namespace data {
template<typename T, typename U>
void Split(const arma::Mat& input,
           const arma::Row& inputLabel,
           arma::Mat& trainData,
           arma::Mat& testData,
           arma::Row& trainLabel,
           arma::Row& testLabel,
           const double testRatio,
           const bool shuffleData = true)
{
  const size_t 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);

  if (shuffleData)
  {
    arma::uvec order = arma::shuffle(arma::linspace(
        0, input.n_cols - 1, input.n_cols));
    if (trainSize > 0)
    {
      trainData = input.cols(order.subvec(0, trainSize - 1));
      trainLabel = inputLabel.cols(order.subvec(0, trainSize - 1));
    }
    if (trainSize < input.n_cols)
    {
    testData = input.cols(order.subvec(trainSize, input.n_cols - 1));
    testLabel = inputLabel.cols(order.subvec(trainSize, input.n_cols - 1));
    }
  }
  else
  {
    if (trainSize > 0)
    {
      trainData = input.cols(0, trainSize - 1);
      trainLabel = inputLabel.subvec(0, trainSize - 1);
    }
    if (trainSize < input.n_cols)
    {
      testData = input.cols(trainSize , input.n_cols - 1);
      testLabel = inputLabel.subvec(trainSize , input.n_cols - 1);
    }
  }
}

template<typename T>
void Split(const arma::Mat& input,
           arma::Mat& trainData,
           arma::Mat& testData,
           const double testRatio,
           const bool shuffleData = true)
{
  const size_t 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);

  if (shuffleData)
  {
    arma::uvec order = arma::shuffle(arma::linspace(
        0, input.n_cols - 1, input.n_cols));

    if (trainSize > 0)
      trainData = input.cols(order.subvec(0, trainSize - 1));

    if (trainSize < input.n_cols)
      testData = input.cols(order.subvec(trainSize, input.n_cols - 1));
  }
  else
  {
    if (trainSize > 0)
      trainData = input.cols(0, trainSize - 1);

    if (trainSize < input.n_cols)
      testData = input.cols(trainSize , input.n_cols - 1);
  }
}

template<typename T, typename U>
std::tuple, arma::Mat, arma::Row, arma::Row>
Split(const arma::Mat& input,
      const arma::Row& inputLabel,
      const double testRatio,
      const bool shuffleData = true)
{
  arma::Mat trainData;
  arma::Mat testData;
  arma::Row trainLabel;
  arma::Row testLabel;

  Split(input, inputLabel, trainData, testData, trainLabel, testLabel,
      testRatio, shuffleData);

  return std::make_tuple(std::move(trainData),
                         std::move(testData),
                         std::move(trainLabel),
                         std::move(testLabel));
}

template<typename T>
std::tuple, arma::Mat>
Split(const arma::Mat& input,
      const double testRatio,
      const bool shuffleData = true)
{
  arma::Mat trainData;
  arma::Mat testData;
  Split(input, trainData, testData, testRatio, shuffleData);

  return std::make_tuple(std::move(trainData),
                         std::move(testData));
}

} // namespace data
} // namespace mlpack

#endif