DecisionTree.h

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
/// \file    DecisionTree.h
/// \brief   Decision Tree PID
/// \author  Christopher Backhouse - bckhouse@caltech.edu
////////////////////////////////////////////////////////////////////////

#ifndef LEM_DECISIONTREE_H
#define LEM_DECISIONTREE_H

#include <list>
#include <string>
#include <vector>

#include "TMatrixD.h"

class TTree;

namespace lem
{
  /// Decision tree %PID
  namespace dec
  {
    const int kNumPIDVars = 6;

    /// \brief A training or trial event for the decision tree
    ///
    /// For classification you need only fill \a vars
    struct Evt
    {
      double vars[kNumPIDVars];
      int ccnc, pdg;
      bool isSig;
      double weight;
    };

    class Tree;
    class TreeNode;

    /// \brief "Random forest" of decision trees
    ///
    /// All user interaction should be with this class
    class Forest
    {
    public:
      ~Forest();

      /// Calculate the %PID value for \a evt
      double Classify(const Evt& evt) const;

      /// Load %PID from a file
      static Forest* FromFile(const std::string& fname);

      /// \brief Initial training of the %PID
      ///
      /// \param trainEvts The events to be used for training
      /// \param nTrees How many trees to train. The final %PID is an average of all
      /// \param parallel Use multiple cores?
      static Forest Train(std::vector<Evt>& trainEvts,
                          unsigned int nTrees,
                          bool parallel = false);

      /// Write out %PID structure to a file
      void ToFile(const std::string& fname);
    protected:
      Forest(){}

      friend void* thread_func(void*);
      static Tree* TrainSingle(const std::vector<Evt>& trainEvts,
                               const TMatrixD& scaleMat,
                               const TMatrixD& transMat);

      /// Internal helper: train one (sub)tree
      static TreeNode* TrainSingleTransformed(std::vector<std::list<Evt*> >& sorted,
                                              unsigned int depth);

      /// Internal helper. Generate a random set of orthogonal unit vectors
      static TMatrixD RandomOrthoMatrix();

      std::vector<Tree*> fTrees;
    };

    /// Abstract base class for tree cuts and leaves
    class TreeNode
    {
    public:
      virtual ~TreeNode() {}
      /// Estimated figure of merit of this subtree
      virtual double FOM() const = 0;
      /// Calculate %PID value of \a e
      virtual double Classify(const Evt& e) const = 0;
      /// Number of signal events below this point in the tree
      virtual double NSig() const = 0;
      /// Number of background events below this point in the tree
      virtual double NBkg() const = 0;

      virtual TreeNode* Prune(std::vector<Evt>& evts) = 0;

      virtual bool IsLeaf() const = 0;

      /// Serialize to a TTree. Needs assistance from \ref Forest::ToTree
      virtual void ToTree(TTree* tr) const = 0;
    };

    /// A DecisionTree. Forwards most things to the head \ref TreeNode
    class Tree
    {
    public:
      Tree(TreeNode* n, const TMatrixD& mat)
        : fHead(n), fMatrix(mat) {}
      ~Tree() {delete fHead;}

      /// Estimated figure of merit of this tree
      double FOM() const {return fHead->FOM();}
      /// Calculate %PID value of \a e
      double Classify(const Evt& e) const;

      void Prune(std::vector<Evt>& evts) {fHead = fHead->Prune(evts);};

      /// Serialize to a TTree. Needs assistance from \ref Forest::ToTree
      void ToTree(TTree* tr) const {fHead->ToTree(tr);}

      TMatrixD GetMatrix() const {return fMatrix;}
    protected:
      TreeNode* fHead;
      TMatrixD fMatrix;
    };

    /// A cut dividing events into two samples, maximizing FOM
    class Cut: public TreeNode
    {
    public:
      Cut(int d, double v, TreeNode* l, TreeNode* r);
      virtual double FOM() const;
      virtual double Classify(const Evt& e) const;
      virtual double NSig() const;
      virtual double NBkg() const;

      virtual TreeNode* Prune(std::vector<Evt>& evts);

      bool IsLeaf() const {return false;}

      virtual void ToTree(TTree* tr) const;
    protected:
      int fCutDim;              ///< What variable the cut is on
      double fCutVal;           ///< Where the cut is placed
      TreeNode *fLeft, *fRight; ///< Subtrees to descend into based on cut result
    };

    /// %Leaf of a decision tree. No further cuts are made
    class Leaf: public TreeNode
    {
    public:
      Leaf(double s, double b);
      virtual double FOM() const;
      virtual double Classify(const Evt&) const;
      virtual double NSig() const {return fSig;}
      virtual double NBkg() const {return fBkg;}

      virtual TreeNode* Prune(std::vector<Evt>&){return this;}

      bool IsLeaf() const {return true;}

      virtual void ToTree(TTree* tr) const;
    protected:
      double fSig; ///< How many signal events reached this node in training
      double fBkg; ///< How many background events reached this node in training
    };
  } // namespace dec
} // namespace lem

#endif