naive_bayes_classifier.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_NAIVE_BAYES_NAIVE_BAYES_CLASSIFIER_HPP
#define MLPACK_METHODS_NAIVE_BAYES_NAIVE_BAYES_CLASSIFIER_HPP

#include <mlpack/prereqs.hpp>

namespace mlpack {
namespace naive_bayes  {

template<typename MatType = arma::mat>
class NaiveBayesClassifier
{
 public:
  NaiveBayesClassifier(const MatType& data,
                       const arma::Row<size_t>& labels,
                       const size_t classes,
                       const bool incrementalVariance = false);

  NaiveBayesClassifier(const size_t dimensionality = 0,
                       const size_t classes = 0);

  void Train(const MatType& data,
             const arma::Row<size_t>& labels,
             const bool incremental = true);

  template<typename VecType>
  void Train(const VecType& point, const size_t label);

  void Classify(const MatType& data, arma::Row<size_t>& results);

  const MatType& Means() const { return means; }
  MatType& Means() { return means; }

  const MatType& Variances() const { return variances; }
  MatType& Variances() { return variances; }

  const arma::vec& Probabilities() const { return probabilities; }
  arma::vec& Probabilities() { return probabilities; }

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

 private:
  MatType means;
  MatType variances;
  arma::vec probabilities;
  size_t trainingPoints;
};

} // namespace naive_bayes
} // namespace mlpack

// Include implementation.
#include "naive_bayes_classifier_impl.hpp"

#endif