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 ModelMatType = arma::mat>
class NaiveBayesClassifier
{
 public:
  // Convenience typedef.
  typedef typename ModelMatType::elem_type ElemType;

  template<typename MatType>
  NaiveBayesClassifier(const MatType& data,
                       const arma::Row<size_t>& labels,
                       const size_t numClasses,
                       const bool incrementalVariance = false);

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

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

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

  template<typename VecType>
  size_t Classify(const VecType& point) const;

  template<typename VecType, typename ProbabilitiesVecType>
  void Classify(const VecType& point,
                size_t& prediction,
                ProbabilitiesVecType& probabilities) const;

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

  template<typename MatType, typename ProbabilitiesMatType>
  void Classify(const MatType& data,
                arma::Row<size_t>& predictions,
                ProbabilitiesMatType& probabilities) const;

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

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

  const ModelMatType& Probabilities() const { return probabilities; }
  ModelMatType& Probabilities() { return probabilities; }

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

 private:
  ModelMatType means;
  ModelMatType variances;
  ModelMatType probabilities;
  size_t trainingPoints;

  template<typename MatType>
  void LogLikelihood(const MatType& data,
                     ModelMatType& logLikelihoods) const;
};

} // namespace naive_bayes
} // namespace mlpack

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

#endif