Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-02-14 13:32:42

0001 #include <cstdint>
0002 #include <utility>
0003 #include <cstring>
0004 #include <string>
0005 #include <vector>
0006 #include <map>
0007 
0008 #include <TString.h>
0009 #include <TTree.h>
0010 #include <TBranch.h>
0011 #include <TLeaf.h>
0012 #include <TList.h>
0013 #include <TKey.h>
0014 
0015 #include "FWCore/Utilities/interface/Exception.h"
0016 
0017 #include "PhysicsTools/MVAComputer/interface/AtomicId.h"
0018 #include "PhysicsTools/MVAComputer/interface/MVAComputer.h"
0019 #include "PhysicsTools/MVAComputer/interface/TreeReader.h"
0020 
0021 namespace PhysicsTools {
0022 
0023   const double TreeReader::kOptVal = -999.0;
0024 
0025   TreeReader::TreeReader() : tree(nullptr), upToDate(false) {}
0026 
0027   TreeReader::TreeReader(const TreeReader &orig) { this->operator=(orig); }
0028 
0029   TreeReader::TreeReader(TTree *tree, bool skipTarget, bool skipWeight) : tree(tree), upToDate(false) {
0030     automaticAdd(skipTarget, skipWeight);
0031   }
0032 
0033   TreeReader::~TreeReader() {}
0034 
0035   TreeReader &TreeReader::operator=(const TreeReader &orig) {
0036     reset();
0037 
0038     tree = orig.tree;
0039 
0040     multiDouble.resize(orig.multiDouble.size());
0041     multiFloat.resize(orig.multiFloat.size());
0042     multiInt.resize(orig.multiInt.size());
0043     multiBool.resize(orig.multiBool.size());
0044 
0045     singleDouble.resize(orig.singleDouble.size());
0046     singleFloat.resize(orig.singleFloat.size());
0047     singleInt.resize(orig.singleInt.size());
0048     singleBool.resize(orig.singleBool.size());
0049 
0050     valueMap = orig.valueMap;
0051 
0052     return *this;
0053   }
0054 
0055   void TreeReader::setTree(TTree *tree) {
0056     this->tree = tree;
0057     upToDate = false;
0058   }
0059 
0060   void TreeReader::addBranch(const std::string &expression, AtomicId name, bool opt) {
0061     if (!tree)
0062       throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::addBranch." << std::endl;
0063 
0064     TBranch *branch = tree->GetBranch(expression.c_str());
0065     if (!branch)
0066       throw cms::Exception("BranchMissing") << "Tree branch \"" << expression << "\" missing." << std::endl;
0067 
0068     addBranch(branch, name, opt);
0069   }
0070 
0071   void TreeReader::addBranch(TBranch *branch, AtomicId name, bool opt) {
0072     TString branchName = branch->GetName();
0073     if (!name)
0074       name = (const char *)branchName;
0075 
0076     TLeaf *leaf = dynamic_cast<TLeaf *>(branch->GetLeaf(branchName));
0077     if (!leaf)
0078       throw cms::Exception("InvalidBranch") << "Tree branch \"" << branchName << "\" has no leaf." << std::endl;
0079 
0080     TString typeName = leaf->GetTypeName();
0081     char typeId = 0;
0082     bool multi = false;
0083     if (typeName == "Double_t" || typeName == "double")
0084       typeId = 'D';
0085     else if (typeName == "Float_t" || typeName == "float")
0086       typeId = 'F';
0087     else if (typeName == "Int_t" || typeName == "int")
0088       typeId = 'I';
0089     else if (typeName == "Bool_t" || typeName == "bool")
0090       typeId = 'B';
0091     else {
0092       multi = true;
0093       if (typeName == "vector<double>" || typeName == "Vector<Double_t>")
0094         typeId = 'D';
0095       else if (typeName == "vector<float>" || typeName == "Vector<Float_t>")
0096         typeId = 'F';
0097       else if (typeName == "vector<int>" || typeName == "Vector<Int_t>")
0098         typeId = 'I';
0099       else if (typeName == "vector<bool>" || typeName == "Vector<Bool_t>")
0100         typeId = 'B';
0101     }
0102 
0103     if (!typeId)
0104       throw cms::Exception("InvalidBranch") << "Tree branch \"" << branchName
0105                                             << "\" is of "
0106                                                "unsupported type \""
0107                                             << typeName << "\"." << std::endl;
0108 
0109     if (multi)
0110       addTypeMulti(name, nullptr, typeId);
0111     else
0112       addTypeSingle(name, nullptr, typeId, opt);
0113 
0114     valueMap[name].setBranchName(branch->GetName());
0115   }
0116 
0117   void TreeReader::setOptional(AtomicId name, bool opt, double optVal) {
0118     std::map<AtomicId, Value>::iterator pos = valueMap.find(name);
0119     if (pos == valueMap.end())
0120       throw cms::Exception("UnknownVariable") << "Variable \"" << name
0121                                               << "\" is not known to the "
0122                                                  "TreeReader."
0123                                               << std::endl;
0124 
0125     pos->second.setOpt(opt, optVal);
0126   }
0127 
0128   void TreeReader::addTypeSingle(AtomicId name, const void *value, char type, bool opt) {
0129     std::map<AtomicId, Value>::const_iterator pos = valueMap.find(name);
0130     if (pos != valueMap.end())
0131       throw cms::Exception("DuplicateVariable") << "Duplicate Variable \"" << name << "\"." << std::endl;
0132 
0133     if (type != 'D' && type != 'F' && type != 'I' && type != 'B')
0134       throw cms::Exception("InvalidType") << "Unsupported type '" << type
0135                                           << "' in call to"
0136                                              "TreeReader::addTypeSingle."
0137                                           << std::endl;
0138 
0139     int index = -1;
0140     if (!value) {
0141       switch (type) {
0142         case 'D':
0143           index = (int)singleDouble.size();
0144           singleDouble.push_back(Double_t());
0145           break;
0146         case 'F':
0147           index = (int)singleFloat.size();
0148           singleFloat.push_back(Float_t());
0149           break;
0150         case 'I':
0151           index = (int)singleInt.size();
0152           singleInt.push_back(Int_t());
0153           break;
0154         case 'B':
0155           index = (int)singleBool.size();
0156           singleBool.push_back(Bool());
0157           break;
0158       }
0159     }
0160 
0161     valueMap[name] = Value(index, false, opt, type);
0162     if (value)
0163       valueMap[name].setPtr(value);
0164 
0165     upToDate = false;
0166   }
0167 
0168   template <typename T>
0169   static std::pair<void *, std::vector<T> > makeMulti() {
0170     return std::pair<void *, std::vector<T> >(nullptr, std::vector<T>());
0171   }
0172 
0173   void TreeReader::addTypeMulti(AtomicId name, const void *value, char type) {
0174     std::map<AtomicId, Value>::const_iterator pos = valueMap.find(name);
0175     if (pos != valueMap.end())
0176       throw cms::Exception("DuplicateVariable") << "Duplicate Variable \"" << name << "\"." << std::endl;
0177 
0178     if (type != 'D' && type != 'F' && type != 'I' && type != 'B')
0179       throw cms::Exception("InvalidType") << "Unsupported type '" << type
0180                                           << "' in call to"
0181                                              "TreeReader::addTypeMulti."
0182                                           << std::endl;
0183 
0184     int index = -1;
0185     if (!value) {
0186       switch (type) {
0187         case 'D':
0188           index = (int)multiDouble.size();
0189           multiDouble.push_back(makeMulti<Double_t>());
0190           break;
0191         case 'F':
0192           index = (int)multiFloat.size();
0193           multiFloat.push_back(makeMulti<Float_t>());
0194           break;
0195         case 'I':
0196           index = (int)multiInt.size();
0197           multiInt.push_back(makeMulti<Int_t>());
0198           break;
0199         case 'B':
0200           index = (int)multiBool.size();
0201           multiBool.push_back(makeMulti<Bool_t>());
0202           break;
0203       }
0204     }
0205 
0206     valueMap[name] = Value(index, true, false, type);
0207     if (value)
0208       valueMap[name].setPtr(value);
0209 
0210     upToDate = false;
0211   }
0212 
0213   void TreeReader::automaticAdd(bool skipTarget, bool skipWeight) {
0214     if (!tree)
0215       throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;
0216 
0217     TIter iter(tree->GetListOfBranches());
0218     TObject *obj;
0219     while ((obj = iter())) {
0220       TBranch *branch = dynamic_cast<TBranch *>(obj);
0221       if (!branch)
0222         continue;
0223 
0224       if (skipTarget && !std::strcmp(branch->GetName(), "__TARGET__"))
0225         continue;
0226 
0227       if (skipWeight && !std::strcmp(branch->GetName(), "__WEIGHT__"))
0228         continue;
0229 
0230       addBranch(branch);
0231     }
0232   }
0233 
0234   void TreeReader::reset() {
0235     multiDouble.clear();
0236     multiFloat.clear();
0237     multiInt.clear();
0238     multiBool.clear();
0239 
0240     singleDouble.clear();
0241     singleFloat.clear();
0242     singleInt.clear();
0243     singleBool.clear();
0244 
0245     valueMap.clear();
0246 
0247     upToDate = false;
0248   }
0249 
0250   void TreeReader::update() {
0251     if (!tree)
0252       throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;
0253 
0254     for (std::map<AtomicId, Value>::iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
0255       iter->second.update(this);
0256 
0257     upToDate = true;
0258   }
0259 
0260   uint64_t TreeReader::loop(const MVAComputer *mva) {
0261     if (!tree)
0262       throw cms::Exception("NoTreeAvailable") << "No TTree set in TreeReader::automaticAdd." << std::endl;
0263 
0264     if (!upToDate)
0265       update();
0266 
0267     Long64_t entries = tree->GetEntries();
0268     for (Long64_t entry = 0; entry < entries; entry++) {
0269       tree->GetEntry(entry);
0270       fill(mva);
0271     }
0272 
0273     return entries;
0274   }
0275 
0276   double TreeReader::fill(const MVAComputer *mva) {
0277     for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
0278       iter->second.fill(iter->first, this);
0279 
0280     double result = mva->eval(values);
0281     values.clear();
0282 
0283     return result;
0284   }
0285 
0286   Variable::ValueList TreeReader::fill() {
0287     for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
0288       iter->second.fill(iter->first, this);
0289 
0290     Variable::ValueList result = values;
0291     values.clear();
0292 
0293     return result;
0294   }
0295 
0296   std::vector<AtomicId> TreeReader::variables() const {
0297     std::vector<AtomicId> result;
0298     for (std::map<AtomicId, Value>::const_iterator iter = valueMap.begin(); iter != valueMap.end(); iter++)
0299       result.push_back(iter->first);
0300 
0301     return result;
0302   }
0303 
0304   void TreeReader::Value::update(TreeReader *reader) const {
0305     if (ptr)
0306       return;
0307 
0308     void *value = nullptr;
0309     if (multiple) {
0310       switch (type) {
0311         case 'D':
0312           reader->multiDouble[index].first = &reader->multiDouble[index].second;
0313           value = &reader->multiDouble[index].first;
0314           break;
0315         case 'F':
0316           reader->multiFloat[index].first = &reader->multiFloat[index].second;
0317           value = &reader->multiFloat[index].first;
0318           break;
0319         case 'I':
0320           reader->multiInt[index].first = &reader->multiInt[index].second;
0321           value = &reader->multiInt[index].first;
0322           break;
0323         case 'B':
0324           reader->multiBool[index].first = value;
0325           value = &reader->multiBool[index].first;
0326           break;
0327       }
0328     } else {
0329       switch (type) {
0330         case 'D':
0331           value = &reader->singleDouble[index];
0332           break;
0333         case 'F':
0334           value = &reader->singleFloat[index];
0335           break;
0336         case 'I':
0337           value = &reader->singleInt[index];
0338           break;
0339         case 'B':
0340           value = &reader->singleBool[index];
0341           break;
0342       }
0343     }
0344 
0345     reader->tree->SetBranchAddress(name, value);
0346   }
0347 
0348   void TreeReader::Value::fill(AtomicId name, TreeReader *reader) const {
0349     if (multiple) {
0350       switch (type) {
0351         case 'D': {
0352           const std::vector<Double_t> *values = static_cast<const std::vector<Double_t> *>(ptr);
0353           if (!values)
0354             values = &reader->multiDouble[index].second;
0355           for (std::vector<Double_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
0356             reader->values.add(name, *iter);
0357           break;
0358         }
0359         case 'F': {
0360           const std::vector<Float_t> *values = static_cast<const std::vector<Float_t> *>(ptr);
0361           if (!values)
0362             values = &reader->multiFloat[index].second;
0363           for (std::vector<Float_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
0364             reader->values.add(name, *iter);
0365           break;
0366         }
0367         case 'I': {
0368           const std::vector<Int_t> *values = static_cast<const std::vector<Int_t> *>(ptr);
0369           if (!values)
0370             values = &reader->multiInt[index].second;
0371           for (std::vector<Int_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
0372             reader->values.add(name, *iter);
0373           break;
0374         }
0375         case 'B': {
0376           const std::vector<Bool_t> *values = static_cast<const std::vector<Bool_t> *>(ptr);
0377           if (!values)
0378             values = &reader->multiBool[index].second;
0379           for (std::vector<Bool_t>::const_iterator iter = values->begin(); iter != values->end(); iter++)
0380             reader->values.add(name, *iter);
0381           break;
0382         }
0383       }
0384     } else {
0385       double value = 0.0;
0386 
0387       switch (type) {
0388         case 'D':
0389           value = ptr ? *(const Double_t *)ptr : reader->singleDouble[index];
0390           break;
0391         case 'F':
0392           value = ptr ? *(const Float_t *)ptr : reader->singleFloat[index];
0393           break;
0394         case 'I':
0395           value = ptr ? *(const Int_t *)ptr : reader->singleInt[index];
0396           break;
0397         case 'B':
0398           value = ptr ? *(const Bool_t *)ptr : reader->singleBool[index];
0399           break;
0400       }
0401 
0402       if (!optional || value != optVal)
0403         reader->values.add(name, value);
0404     }
0405   }
0406 
0407 #define TREEREADER_ADD_IMPL(T, C)                                          \
0408   template <>                                                              \
0409   void TreeReader::addSingle<T>(AtomicId name, const T *value, bool opt) { \
0410     addTypeSingle(name, value, C, opt);                                    \
0411   }                                                                        \
0412                                                                            \
0413   template <>                                                              \
0414   void TreeReader::addMulti(AtomicId name, const std::vector<T> *value) {  \
0415     addTypeMulti(name, value, C);                                          \
0416   }
0417 
0418   TREEREADER_ADD_IMPL(Double_t, 'D')
0419   TREEREADER_ADD_IMPL(Float_t, 'F')
0420   TREEREADER_ADD_IMPL(Int_t, 'I')
0421   TREEREADER_ADD_IMPL(Bool_t, 'B')
0422 
0423 #undef TREEREADER_ADD_IMPL
0424 
0425 }  // namespace PhysicsTools