MVAComputer/interface/TreeReader.h

0001 #ifndef PhysicsTools_MVAComputer_TreeReader_h
0002 #define PhysicsTools_MVAComputer_TreeReader_h
0003
0004 #include <cstdint>
0005 #include <utility>
0006 #include <string>
0007 #include <vector>
0008 #include <map>
0009
0010 #include <TTree.h>
0011 #include <TBranch.h>
0012
0013 #include "PhysicsTools/MVAComputer/interface/AtomicId.h"
0014 #include "PhysicsTools/MVAComputer/interface/MVAComputer.h"
0015 #include "PhysicsTools/MVAComputer/interface/Variable.h"
0016
0017 namespace PhysicsTools {
0018
0019   class TreeReader {
0020   public:
0021     TreeReader();
0022     TreeReader(const TreeReader &orig);
0023     TreeReader(TTree *tree, bool skipTarget = false, bool skipWeight = false);
0024     virtual ~TreeReader();
0025
0026     TreeReader &operator=(const TreeReader &orig);
0027
0028     void setTree(TTree *tree);
0029
0030     void addBranch(const std::string &expression, AtomicId name = AtomicId(), bool opt = true);
0031     void addBranch(TBranch *branch, AtomicId name = AtomicId(), bool opt = true);
0032     template <typename T>
0033     void addSingle(AtomicId name, const T *value, bool opt = false);
0034     template <typename T>
0035     void addMulti(AtomicId name, const std::vector<T> *value);
0036     void setOptional(AtomicId name, bool opt, double optVal = kOptVal);
0037
0038     void addTypeSingle(AtomicId name, const void *value, char type, bool opt);
0039     void addTypeMulti(AtomicId name, const void *value, char type);
0040
0041     void automaticAdd(bool skipTarget = false, bool skipWeight = false);
0042
0043     void reset();
0044     void update();
0045
0046     uint64_t loop(const MVAComputer *mva);
0047
0048     double fill(const MVAComputer *mva);
0049
0050     Variable::ValueList fill();
0051
0052     std::vector<AtomicId> variables() const;
0053
0054     static const double kOptVal;
0055
0056     struct Bool {
0057       inline Bool() : value(false) {}
0058       inline operator Bool_t() const { return value; }
0059       Bool_t value;
0060     };
0061
0062     class Value {
0063     public:
0064       Value() {}
0065       Value(int index, bool multiple, bool optional, char type)
0066           : index(index),
0067             optional(optional),
0068             multiple(multiple),
0069             optVal(TreeReader::kOptVal),
0070             type(type),
0071             ptr(nullptr) {}
0072       ~Value() {}
0073
0074       void setOpt(bool opt, double optVal) { this->optional = opt, this->optVal = optVal; }
0075       void setBranchName(const TString &name) { this->name = name; }
0076       void setPtr(const void *ptr) { this->ptr = ptr; }
0077
0078       void update(TreeReader *reader) const;
0079       void fill(AtomicId name, TreeReader *reader) const;
0080
0081     private:
0082       TString name;
0083       int index;
0084       bool optional;
0085       bool multiple;
0086       double optVal;
0087       char type;
0088       const void *ptr;
0089     };
0090
0091     friend class Value;
0092
0093   private:
0094     TTree *tree;
0095
0096     std::vector<std::pair<void *, std::vector<Double_t> > > multiDouble;
0097     std::vector<std::pair<void *, std::vector<Float_t> > > multiFloat;
0098     std::vector<std::pair<void *, std::vector<Int_t> > > multiInt;
0099     std::vector<std::pair<void *, std::vector<Bool_t> > > multiBool;
0100
0101     std::vector<Double_t> singleDouble;
0102     std::vector<Float_t> singleFloat;
0103     std::vector<Int_t> singleInt;
0104     std::vector<Bool> singleBool;
0105
0106     std::map<AtomicId, Value> valueMap;
0107     Variable::ValueList values;
0108     bool upToDate;
0109   };
0110
0111 #define TREEREADER_ADD_IMPL(T)                                            \
0112   template <>                                                             \
0113   void TreeReader::addSingle<T>(AtomicId name, const T *value, bool opt); \
0114                                                                           \
0115   template <>                                                             \
0116   void TreeReader::addMulti(AtomicId name, const std::vector<T> *value);
0117
0118   TREEREADER_ADD_IMPL(Double_t)
0119   TREEREADER_ADD_IMPL(Float_t)
0120   TREEREADER_ADD_IMPL(Int_t)
0121   TREEREADER_ADD_IMPL(Bool_t)
0122
0123 #undef TREEREADER_ADD_IMPL
0124
0125 }  // namespace PhysicsTools
0126
0127 #endif  // PhysicsTools_MVAComputer_TreeReader_h