ClassificationForest.hpp

Go to the documentation of this file.

#pragma once
 
#include "models/Forest.hpp"
#include "stats/GroupPartition.hpp"
 
#include <map>
#include <set>
 
namespace ppforest2 {
  class ClassificationForest : public Forest {
  public:
    using Ptr = std::unique_ptr<ClassificationForest>;
    using Forest::predict;
 
    using FeatureMatrix = types::FeatureMatrix;
    using FeatureVector = types::FeatureVector;
    using OutcomeVector = types::OutcomeVector;
    using Outcome       = types::Outcome;
    using RNG           = stats::RNG;
 
    using Groups         = std::set<types::GroupId>;
    using GroupIndices   = std::map<types::GroupId, int>;
    using GroupPartition = stats::GroupPartition;
 

    Groups groups;
 
    ClassificationForest(TrainingSpec::Ptr spec, Groups groups)
        : Forest(std::move(spec))
        , groups(std::move(groups)) {
      invariant(this->training_spec != nullptr, "ClassificationForest requires a non-null TrainingSpec");
      invariant(is_classification(this), "ClassificationForest requires a classification TrainingSpec");
      invariant(!this->groups.empty(), "ClassificationForest requires a non-empty groups set");
    }
 
    static Ptr train(TrainingSpec const& spec, FeatureMatrix const& x, OutcomeVector const& y);
 
    types::Outcome predict(FeatureVector const& x) const override;

    FeatureVector predict(FeatureVector const& x, Proportions) const;

    FeatureVector predict(FeatureVector const& x, Proportions, GroupIndices const& indices) const;

    FeatureMatrix predict(FeatureMatrix const& x, Proportions) const;

    FeatureMatrix predict(FeatureMatrix const& x, Proportions, GroupIndices const& indices) const;
 
    void accept(Model::Visitor& visitor) const override;
 
  protected:
    BaggedTree::Ptr train_tree(FeatureMatrix const& x, OutcomeVector const& y, RNG& rng) const override;
 
  private:
    GroupPartition const* y_part = nullptr;
  };
}