decision_stump.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_DECISION_STUMP_DECISION_STUMP_HPP
#define MLPACK_METHODS_DECISION_STUMP_DECISION_STUMP_HPP

#include <mlpack/prereqs.hpp>

namespace mlpack {
namespace decision_stump {

template<typename MatType = arma::mat>
class DecisionStump
{
 public:
  DecisionStump(const MatType& data,
                const arma::Row<size_t>& labels,
                const size_t classes,
                const size_t bucketSize = 10);

  DecisionStump(const DecisionStump<>& other,
                const MatType& data,
                const arma::Row<size_t>& labels,
                const arma::rowvec& weights);

  DecisionStump();

  void Train(const MatType& data,
             const arma::Row<size_t>& labels,
             const size_t classes,
             const size_t bucketSize);

  void Train(const MatType& data,
             const arma::Row<size_t>& labels,
             const arma::rowvec& weights,
             const size_t classes,
             const size_t bucketSize);

  void Classify(const MatType& test, arma::Row<size_t>& predictedLabels);

  size_t SplitDimension() const { return splitDimension; }
  size_t& SplitDimension() { return splitDimension; }

  const arma::vec& Split() const { return split; }
  arma::vec& Split() { return split; }

  const arma::Col<size_t> BinLabels() const { return binLabels; }
  arma::Col<size_t>& BinLabels() { return binLabels; }

  template<typename Archive>
  void Serialize(Archive& ar, const unsigned int /* version */);

 private:
  size_t classes;
  size_t bucketSize;

  size_t splitDimension;
  arma::vec split;
  arma::Col<size_t> binLabels;

  template<bool UseWeights, typename VecType>
  double SetupSplitDimension(const VecType& dimension,
                             const arma::Row<size_t>& labels,
                             const arma::rowvec& weightD);

  template<typename VecType>
  void TrainOnDim(const VecType& dimension,
                  const arma::Row<size_t>& labels);

  void MergeRanges();

  template<typename VecType>
  double CountMostFreq(const VecType& subCols);

  template<typename VecType>
  int IsDistinct(const VecType& featureRow);

  template<bool UseWeights, typename VecType, typename WeightVecType>
  double CalculateEntropy(const VecType& labels,
                          const WeightVecType& weights);

  template<bool UseWeights>
  void Train(const MatType& data,
             const arma::Row<size_t>& labels,
             const arma::rowvec& weights);
};

} // namespace decision_stump
} // namespace mlpack

#include "decision_stump_impl.hpp"

#endif