ClassificationTree.hpp

Go to the documentation of this file.

#pragma once
 
#include "models/Tree.hpp"
 
#include <map>
#include <set>
 
namespace ppforest2 {
  class ClassificationTree : public Tree {
  public:
    using Ptr = std::unique_ptr<ClassificationTree>;
    using Tree::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;
 
    ClassificationTree(TreeNode::Ptr root, TrainingSpec::Ptr spec, Groups groups)
        : Tree(std::move(root), std::move(spec))
        , groups(std::move(groups)) {
      invariant(this->training_spec != nullptr, "ClassificationTree requires a non-null TrainingSpec");
      invariant(is_classification(this), "ClassificationTree requires a classification TrainingSpec");
      invariant(!this->groups.empty(), "ClassificationTree requires a non-empty groups set");
    }
 

    static Ptr train(TrainingSpec const& s, FeatureMatrix& x, OutcomeVector& y, GroupPartition const& y_part, RNG& rng);

    types::FeatureVector predict(types::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;
  };
}