PhysicsTools/MVAComputer/interface/TreeReader.h - master - githubmirror/cmssw

PhysicsTools/MVAComputer/interface/TreeReader.h
127 строк4 KB
Cms Build
Clang-Format
31 май 2019, 17:22
31 май 2019, 17:2288e58de
#ifndef PhysicsTools_MVAComputer_TreeReader_h
#define PhysicsTools_MVAComputer_TreeReader_h

#include <cstdint>
#include <utility>
#include <string>
#include <vector>
#include <map>

#include <TTree.h>
#include <TBranch.h>

#include "PhysicsTools/MVAComputer/interface/AtomicId.h"
#include "PhysicsTools/MVAComputer/interface/MVAComputer.h"
#include "PhysicsTools/MVAComputer/interface/Variable.h"

namespace PhysicsTools {

  class TreeReader {
  public:
    TreeReader();
    TreeReader(const TreeReader &orig);
    TreeReader(TTree *tree, bool skipTarget = false, bool skipWeight = false);
    virtual ~TreeReader();

    TreeReader &operator=(const TreeReader &orig);

    void setTree(TTree *tree);

    void addBranch(const std::string &expression, AtomicId name = AtomicId(), bool opt = true);
    void addBranch(TBranch *branch, AtomicId name = AtomicId(), bool opt = true);
    template <typename T>
    void addSingle(AtomicId name, const T *value, bool opt = false);
    template <typename T>
    void addMulti(AtomicId name, const std::vector<T> *value);
    void setOptional(AtomicId name, bool opt, double optVal = kOptVal);

    void addTypeSingle(AtomicId name, const void *value, char type, bool opt);
    void addTypeMulti(AtomicId name, const void *value, char type);

    void automaticAdd(bool skipTarget = false, bool skipWeight = false);

    void reset();
    void update();

    uint64_t loop(const MVAComputer *mva);

    double fill(const MVAComputer *mva);

    Variable::ValueList fill();

    std::vector<AtomicId> variables() const;

    static const double kOptVal;

    struct Bool {
      inline Bool() : value(false) {}
      inline operator Bool_t() const { return value; }
      Bool_t value;
    };

    class Value {
    public:
      Value() {}
      Value(int index, bool multiple, bool optional, char type)
          : index(index),
            optional(optional),
            multiple(multiple),
            optVal(TreeReader::kOptVal),
            type(type),
            ptr(nullptr) {}
      ~Value() {}

      void setOpt(bool opt, double optVal) { this->optional = opt, this->optVal = optVal; }
      void setBranchName(const TString &name) { this->name = name; }
      void setPtr(const void *ptr) { this->ptr = ptr; }

      void update(TreeReader *reader) const;
      void fill(AtomicId name, TreeReader *reader) const;

    private:
      TString name;
      int index;
      bool optional;
      bool multiple;
      double optVal;
      char type;
      const void *ptr;
    };

    friend class Value;

  private:
    TTree *tree;

    std::vector<std::pair<void *, std::vector<Double_t> > > multiDouble;
    std::vector<std::pair<void *, std::vector<Float_t> > > multiFloat;
    std::vector<std::pair<void *, std::vector<Int_t> > > multiInt;
    std::vector<std::pair<void *, std::vector<Bool_t> > > multiBool;

    std::vector<Double_t> singleDouble;
    std::vector<Float_t> singleFloat;
    std::vector<Int_t> singleInt;
    std::vector<Bool> singleBool;

    std::map<AtomicId, Value> valueMap;
    Variable::ValueList values;
    bool upToDate;
  };

#define TREEREADER_ADD_IMPL(T)                                            \
  template <>                                                             \
  void TreeReader::addSingle<T>(AtomicId name, const T *value, bool opt); \
                                                                          \
  template <>                                                             \
  void TreeReader::addMulti(AtomicId name, const std::vector<T> *value);

  TREEREADER_ADD_IMPL(Double_t)
  TREEREADER_ADD_IMPL(Float_t)
  TREEREADER_ADD_IMPL(Int_t)
  TREEREADER_ADD_IMPL(Bool_t)

#undef TREEREADER_ADD_IMPL

}  // namespace PhysicsTools

#endif  // PhysicsTools_MVAComputer_TreeReader_h