TreeReader.cc

Go to the documentation of this file.

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

#include <TString.h>
#include <TTree.h>
#include <TBranch.h>
#include <TLeaf.h>
#include <TList.h>
#include <TKey.h>

#include "FWCore/Utilities/interface/Exception.h"

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

namespace PhysicsTools {

  const double TreeReader::kOptVal = -999.0;

  TreeReader::TreeReader() : tree(nullptr), upToDate(false) {}

  TreeReader::TreeReader(const TreeReader &orig) { this->operator=(orig); }

  TreeReader::TreeReader(TTree *tree, bool skipTarget, bool skipWeight) : tree(tree), upToDate(false) {
    automaticAdd(skipTarget, skipWeight);
  }

  TreeReader::~TreeReader() {}

  TreeReader &TreeReader::operator=(const TreeReader &orig) {
    reset();

    tree = orig.tree;

    multiDouble.resize(orig.multiDouble.size());
    multiFloat.resize(orig.multiFloat.size());
    multiInt.resize(orig.multiInt.size());
    multiBool.resize(orig.multiBool.size());

    singleDouble.resize(orig.singleDouble.size());
    singleFloat.resize(orig.singleFloat.size());
    singleInt.resize(orig.singleInt.size());
    singleBool.resize(orig.singleBool.size());

    valueMap = orig.valueMap;

    return *this;
  }

  void TreeReader::setTree(TTree *tree) {
    this->tree = tree;
    upToDate = false;
  }

  void TreeReader::addBranch(const std::string &expression, AtomicId name, bool opt) {
    if (!tree)
      throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::addBranch." << std::endl;

    TBranch *branch = tree->GetBranch(expression.c_str());
    if (!branch)
      throw cms::Exception("BranchMissing") << "Tree branch \"" << expression << "\" missing." << std::endl;

    addBranch(branch, name, opt);
  }

  void TreeReader::addBranch(TBranch *branch, AtomicId name, bool opt) {
    TString branchName = branch->GetName();
    if (!name)
      name = (const char *)branchName;

    TLeaf *leaf = dynamic_cast<TLeaf *>(branch->GetLeaf(branchName));
    if (!leaf)
      throw cms::Exception("InvalidBranch") << "Tree branch \"" << branchName << "\" has no leaf." << std::endl;

    TString typeName = leaf->GetTypeName();
    char typeId = 0;
    bool multi = false;
    if (typeName == "Double_t" || typeName == "double")
      typeId = 'D';
    else if (typeName == "Float_t" || typeName == "float")
      typeId = 'F';
    else if (typeName == "Int_t" || typeName == "int")
      typeId = 'I';
    else if (typeName == "Bool_t" || typeName == "bool")
      typeId = 'B';
    else {
      multi = true;
      if (typeName == "vector<double>" || typeName == "Vector<Double_t>")
        typeId = 'D';
      else if (typeName == "vector<float>" || typeName == "Vector<Float_t>")
        typeId = 'F';
      else if (typeName == "vector<int>" || typeName == "Vector<Int_t>")
        typeId = 'I';
      else if (typeName == "vector<bool>" || typeName == "Vector<Bool_t>")
        typeId = 'B';
    }

    if (!typeId)
      throw cms::Exception("InvalidBranch") << "Tree branch \"" << branchName
                                            << "\" is of "
                                               "unsupported type \""
                                            << typeName << "\"." << std::endl;

    if (multi)
      addTypeMulti(name, nullptr, typeId);
    else
      addTypeSingle(name, nullptr, typeId, opt);

    valueMap[name].setBranchName(branch->GetName());
  }

  void TreeReader::setOptional(AtomicId name, bool opt, double optVal) {
    std::map<AtomicId, Value>::iterator pos = valueMap.find(name);
    if (pos == valueMap.end())
      throw cms::Exception("UnknownVariable") << "Variable \"" << name
                                              << "\" is not known to the "
                                                 "TreeReader."
                                              << std::endl;

    pos->second.setOpt(opt, optVal);
  }

  void TreeReader::addTypeSingle(AtomicId name, const void *value, char type, bool opt) {
    std::map<AtomicId, Value>::const_iterator pos = valueMap.find(name);
    if (pos != valueMap.end())
      throw cms::Exception("DuplicateVariable") << "Duplicate Variable \"" << name << "\"." << std::endl;

    if (type != 'D' && type != 'F' && type != 'I' && type != 'B')
      throw cms::Exception("InvalidType") << "Unsupported type '" << type
                                          << "' in call to"
                                             "TreeReader::addTypeSingle."
                                          << std::endl;

    int index = -1;
    if (!value) {
      switch (type) {
        case 'D':
          index = (int)singleDouble.size();
          singleDouble.push_back(Double_t());
          break;
        case 'F':
          index = (int)singleFloat.size();
          singleFloat.push_back(Float_t());
          break;
        case 'I':
          index = (int)singleInt.size();
          singleInt.push_back(Int_t());
          break;
        case 'B':
          index = (int)singleBool.size();
          singleBool.push_back(Bool());
          break;
      }
    }

    valueMap[name] = Value(index, false, opt, type);
    if (value)
      valueMap[name].setPtr(value);

    upToDate = false;
  }

  template <typename T>
  static std::pair<void *, std::vector<T> > makeMulti() {
    return std::pair<void *, std::vector<T> >(nullptr, std::vector<T>());
  }

  void TreeReader::addTypeMulti(AtomicId name, const void *value, char type) {
    std::map<AtomicId, Value>::const_iterator pos = valueMap.find(name);
    if (pos != valueMap.end())
      throw cms::Exception("DuplicateVariable") << "Duplicate Variable \"" << name << "\"." << std::endl;

    if (type != 'D' && type != 'F' && type != 'I' && type != 'B')
      throw cms::Exception("InvalidType") << "Unsupported type '" << type
                                          << "' in call to"
                                             "TreeReader::addTypeMulti."
                                          << std::endl;

    int index = -1;
    if (!value) {
      switch (type) {
        case 'D':
          index = (int)multiDouble.size();
          multiDouble.push_back(makeMulti<Double_t>());
          break;
        case 'F':
          index = (int)multiFloat.size();
          multiFloat.push_back(makeMulti<Float_t>());
          break;
        case 'I':
          index = (int)multiInt.size();
          multiInt.push_back(makeMulti<Int_t>());
          break;
        case 'B':
          index = (int)multiBool.size();
          multiBool.push_back(makeMulti<Bool_t>());
          break;
      }
    }

    valueMap[name] = Value(index, true, false, type);
    if (value)
      valueMap[name].setPtr(value);

    upToDate = false;
  }

  void TreeReader::automaticAdd(bool skipTarget, bool skipWeight) {
    if (!tree)
      throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;

    TIter iter(tree->GetListOfBranches());
    TObject *obj;
    while ((obj = iter())) {
      TBranch *branch = dynamic_cast<TBranch *>(obj);
      if (!branch)
        continue;

      if (skipTarget && !std::strcmp(branch->GetName(), "__TARGET__"))
        continue;

      if (skipWeight && !std::strcmp(branch->GetName(), "__WEIGHT__"))
        continue;

      addBranch(branch);
    }
  }

  void TreeReader::reset() {
    multiDouble.clear();
    multiFloat.clear();
    multiInt.clear();
    multiBool.clear();

    singleDouble.clear();
    singleFloat.clear();
    singleInt.clear();
    singleBool.clear();

    valueMap.clear();

    upToDate = false;
  }

  void TreeReader::update() {
    if (!tree)
      throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;

    for (std::map<AtomicId, Value>::iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
      iter->second.update(this);

    upToDate = true;
  }

  uint64_t TreeReader::loop(const MVAComputer *mva) {
    if (!tree)
      throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;

    if (!upToDate)
      update();

    Long64_t entries = tree->GetEntries();
    for (Long64_t entry = 0; entry < entries; entry++) {
      tree->GetEntry(entry);
      fill(mva);
    }

    return entries;
  }

  double TreeReader::fill(const MVAComputer *mva) {
    for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
      iter->second.fill(iter->first, this);

    double result = mva->eval(values);
    values.clear();

    return result;
  }

  Variable::ValueList TreeReader::fill() {
    for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
      iter->second.fill(iter->first, this);

    Variable::ValueList result = values;
    values.clear();

    return result;
  }

  std::vector<AtomicId> TreeReader::variables() const {
    std::vector<AtomicId> result;
    for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
      result.push_back(iter->first);

    return result;
  }

  void TreeReader::Value::update(TreeReader *reader) const {
    if (ptr)
      return;

    void *value = nullptr;
    if (multiple) {
      switch (type) {
        case 'D':
          reader->multiDouble[index].first = &reader->multiDouble[index].second;
          value = &reader->multiDouble[index].first;
          break;
        case 'F':
          reader->multiFloat[index].first = &reader->multiFloat[index].second;
          value = &reader->multiFloat[index].first;
          break;
        case 'I':
          reader->multiInt[index].first = &reader->multiInt[index].second;
          value = &reader->multiInt[index].first;
          break;
        case 'B':
          reader->multiBool[index].first = value;
          value = &reader->multiBool[index].first;
          break;
      }
    } else {
      switch (type) {
        case 'D':
          value = &reader->singleDouble[index];
          break;
        case 'F':
          value = &reader->singleFloat[index];
          break;
        case 'I':
          value = &reader->singleInt[index];
          break;
        case 'B':
          value = &reader->singleBool[index];
          break;
      }
    }

    reader->tree->SetBranchAddress(name, value);
  }

  void TreeReader::Value::fill(AtomicId name, TreeReader *reader) const {
    if (multiple) {
      switch (type) {
        case 'D': {
          const std::vector<Double_t> *values = static_cast<const std::vector<Double_t> *>(ptr);
          if (!values)
            values = &reader->multiDouble[index].second;
          for (std::vector<Double_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
            reader->values.add(name, *iter);
          break;
        }
        case 'F': {
          const std::vector<Float_t> *values = static_cast<const std::vector<Float_t> *>(ptr);
          if (!values)
            values = &reader->multiFloat[index].second;
          for (std::vector<Float_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
            reader->values.add(name, *iter);
          break;
        }
        case 'I': {
          const std::vector<Int_t> *values = static_cast<const std::vector<Int_t> *>(ptr);
          if (!values)
            values = &reader->multiInt[index].second;
          for (std::vector<Int_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
            reader->values.add(name, *iter);
          break;
        }
        case 'B': {
          const std::vector<Bool_t> *values = static_cast<const std::vector<Bool_t> *>(ptr);
          if (!values)
            values = &reader->multiBool[index].second;
          for (std::vector<Bool_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
            reader->values.add(name, *iter);
          break;
        }
      }
    } else {
      double value = 0.0;

      switch (type) {
        case 'D':
          value = ptr ? *(const Double_t *)ptr : reader->singleDouble[index];
          break;
        case 'F':
          value = ptr ? *(const Float_t *)ptr : reader->singleFloat[index];
          break;
        case 'I':
          value = ptr ? *(const Int_t *)ptr : reader->singleInt[index];
          break;
        case 'B':
          value = ptr ? *(const Bool_t *)ptr : reader->singleBool[index];
          break;
      }

      if (!optional || value != optVal)
        reader->values.add(name, value);
    }
  }

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

  TREEREADER_ADD_IMPL(Double_t, 'D')
  TREEREADER_ADD_IMPL(Float_t, 'F')
  TREEREADER_ADD_IMPL(Int_t, 'I')
  TREEREADER_ADD_IMPL(Bool_t, 'B')

#undef TREEREADER_ADD_IMPL

}  // namespace PhysicsTools