folds.h

Go to the documentation of this file.

#pragma once
 
#include "eigen_compat.h"
#include "utils.h"
#include <Eigen/Core>
#include <cstdint>
#include <vector>
 
namespace slope {
 
using slope::all;

class Folds
{
public:
  Folds()
    : n_folds(0)
    , n_repeats(0)
  {
  }

  Folds(int n_samples, int n_folds, int n_repeats = 1, uint64_t seed = 42)
    : n_folds(n_folds)
    , n_repeats(n_repeats)
  {
    folds.resize(n_repeats);
    for (int rep = 0; rep < n_repeats; ++rep) {
      // Use a different seed for each repetition
      uint64_t rep_seed = seed + rep;
      folds[rep] = createFolds(n_samples, n_folds, rep_seed);
    }
  }

  explicit Folds(const std::vector<std::vector<int>>& user_folds)
    : n_folds(user_folds.size())
    , n_repeats(1)
  {
    folds.emplace_back(user_folds);
  }

  explicit Folds(const std::vector<std::vector<std::vector<int>>>& user_folds)
    : folds(user_folds)
    , n_folds(user_folds[0].size())
    , n_repeats(user_folds.size())
  {
  }

  const std::vector<int>& getTestIndices(size_t fold_idx,
                                         size_t rep_idx = 0) const;

  std::vector<int> getTrainingIndices(size_t fold_idx,
                                      size_t rep_idx = 0) const;

  template<typename T>
  auto split(Eigen::EigenBase<T>& x,
             const Eigen::MatrixXd& y,
             size_t fold_idx,
             size_t rep_idx = 0) const
  {
    auto test_idx = getTestIndices(fold_idx, rep_idx);
    auto train_idx = getTrainingIndices(fold_idx, rep_idx);
 
    auto x_test = subset(x.derived(), test_idx);
    Eigen::MatrixXd y_test = y(test_idx, all);
 
    auto x_train = subset(x.derived(), train_idx);
    Eigen::MatrixXd y_train = y(train_idx, all);
 
    return std::make_tuple(x_train, y_train, x_test, y_test);
  }

  size_t numFolds() const { return n_folds; }

  size_t numRepetitions() const { return n_repeats; }

  size_t numEvals() const { return n_repeats * n_folds; }
 
  std::vector<std::vector<std::vector<int>>>
    folds;               
  std::size_t n_folds;   
  std::size_t n_repeats; 
 
private:
  static std::vector<std::vector<int>> createFolds(int n_samples,
                                                   int n_folds,
                                                   uint64_t seed);
};
 
} // namespace slope