Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2024-11-19 23:20:32 

0001 //  .
0002 // ..: P. Chang, philip@physics.ucsd.edu
0003 
0004 #ifndef ttreex_h
0005 #define ttreex_h
0006 
0007 // C/C++
0008 #include <algorithm>
0009 #include <fstream>
0010 #include <iostream>
0011 #include <map>
0012 #include <string>
0013 #include <unordered_map>
0014 #include <vector>
0015 #include <stdarg.h>
0016 #include <functional>
0017 #include <cmath>
0018 #include <utility>
0019 
0020 // ROOT
0021 #include "TBenchmark.h"
0022 #include "TBits.h"
0023 #include "TChain.h"
0024 #include "TFile.h"
0025 #include "TTree.h"
0026 #include "TBranch.h"
0027 #include "TLeaf.h"
0028 #include "TH1.h"
0029 #include "TH1D.h"
0030 #include "TH2D.h"
0031 #include "TChainElement.h"
0032 #include "TTreeCache.h"
0033 #include "TTreePerfStats.h"
0034 #include "TStopwatch.h"
0035 #include "TSystem.h"
0036 #include "TString.h"
0037 #include "TLorentzVector.h"
0038 #include "Math/LorentzVector.h"
0039 #include "Math/GenVector/PtEtaPhiM4D.h"
0040 
0041 #include "printutil.h"
0042 
0043 //#define MAP std::unordered_map
0044 //#define STRING std::string
0045 #define MAP std::map
0046 //#define TTREEXSTRING std::string
0047 #define TTREEXSTRING TString
0048 
0049 ///////////////////////////////////////////////////////////////////////////////////////////////
0050 // LorentzVector typedef that we use very often
0051 ///////////////////////////////////////////////////////////////////////////////////////////////
0052 #ifdef LorentzVectorPtEtaPhiM4D
0053 typedef ROOT::Math::LorentzVector<ROOT::Math::PtEtaPhiM4D<float>> LV;
0054 #else
0055 typedef ROOT::Math::LorentzVector<ROOT::Math::PxPyPzE4D<float>> LV;
0056 #endif
0057 
0058 typedef std::vector<Int_t> VInt;
0059 typedef std::vector<Float_t> VFloat;
0060 
0061 namespace RooUtil {
0062 
0063   ///////////////////////////////////////////////////////////////////////////////////////////////
0064   // TTreeX class
0065   ///////////////////////////////////////////////////////////////////////////////////////////////
0066   // NOTE: This class assumes accessing TTree in the SNT style which uses the following,
0067   // https://github.com/cmstas/Software/blob/master/makeCMS3ClassFiles/makeCMS3ClassFiles.C
0068   // It is assumed that the "template" class passed to this class will have
0069   // 1. "Init(TTree*)"
0070   // 2. "GetEntry(uint)"
0071   // 3. "progress(nevtProc'ed, total)"
0072   class TTreeX {
0073   public:
0074     enum kType {
0075       kInt_t = 1,
0076       kBool_t = 2,
0077       kFloat_t = 3,
0078       kTString = 4,
0079       kLV = 5,
0080       kVecInt_t = 11,
0081       kVecUInt_t = 12,
0082       kVecBool_t = 13,
0083       kVecFloat_t = 14,
0084       kVecTString = 15,
0085       kVecLV = 16
0086     };
0087     typedef std::vector<LV>::const_iterator lviter;
0088 
0089   private:
0090     TTree* ttree;
0091     std::map<TTREEXSTRING, Int_t> mapInt_t;
0092     std::map<TTREEXSTRING, Bool_t> mapBool_t;
0093     std::map<TTREEXSTRING, Float_t> mapFloat_t;
0094     std::map<TTREEXSTRING, TString> mapTString;
0095     std::map<TTREEXSTRING, LV> mapLV;
0096     std::map<TTREEXSTRING, TBits> mapTBits;
0097     std::map<TTREEXSTRING, unsigned long long> mapULL;
0098     std::map<TTREEXSTRING, unsigned int> mapUI;
0099     std::map<TTREEXSTRING, std::vector<Int_t>> mapVecInt_t;
0100     std::map<TTREEXSTRING, std::vector<UInt_t>> mapVecUInt_t;
0101     std::map<TTREEXSTRING, std::vector<Bool_t>> mapVecBool_t;
0102     std::map<TTREEXSTRING, std::vector<Float_t>> mapVecFloat_t;
0103     std::map<TTREEXSTRING, std::vector<TString>> mapVecTString;
0104     std::map<TTREEXSTRING, std::vector<LV>> mapVecLV;
0105     std::map<TTREEXSTRING, std::vector<std::vector<Int_t>>> mapVecVInt;
0106     std::map<TTREEXSTRING, std::vector<std::vector<Float_t>>> mapVecVFloat;
0107 
0108     std::map<TTREEXSTRING, std::function<float()>> mapFloatFunc_t;
0109 
0110     std::map<TTREEXSTRING, Bool_t> mapIsBranchSet;
0111 
0112   public:
0113     TTreeX();
0114     TTreeX(TString treename, TString title);
0115     TTreeX(TTree* tree);
0116     ~TTreeX();
0117     TTree* getTree() { return ttree; }
0118     void setTree(TTree* tree) { ttree = tree; }
0119     void* getValPtr(TString brname);
0120     template <class T>
0121     T* get(TString brname, int entry = -1);
0122     void fill() { ttree->Fill(); }
0123     void write() { ttree->Write(); }
0124 
0125     template <class T>
0126     void createBranch(TString, bool = true);
0127     template <class T>
0128     void setBranch(TString, T, bool = false, bool = false);
0129     template <class T>
0130     const T& getBranch(TString, bool = true);
0131     template <class T>
0132     const T& getBranchLazy(TString);
0133     template <class T>
0134     bool isBranchSet(TString);
0135     template <class T>
0136     T* getBranchAddress(TString);
0137     template <class T>
0138     void createBranch(T&);
0139     template <class T>
0140     bool hasBranch(TString);
0141     template <class T>
0142     void setBranch(T&);
0143     template <class T>
0144     void pushbackToBranch(TString, T);
0145 
0146     void sortVecBranchesByPt(TString, std::vector<TString>, std::vector<TString>, std::vector<TString>);
0147     template <class T>
0148     std::vector<T> sortFromRef(std::vector<T> const& in, std::vector<std::pair<size_t, lviter>> const& reference);
0149     struct ordering {
0150       bool operator()(std::pair<size_t, lviter> const& a, std::pair<size_t, lviter> const& b) {
0151         return (*(a.second)).pt() > (*(b.second)).pt();
0152       }
0153     };
0154     void createFlatBranch(std::vector<TString>, std::vector<TString>, std::vector<TString>, std::vector<TString>, int);
0155     void setFlatBranch(std::vector<TString>, std::vector<TString>, std::vector<TString>, std::vector<TString>, int);
0156 
0157     void clear();
0158     void save(TFile*);
0159   };
0160 
0161   //_________________________________________________________________________________________________
0162   template <>
0163   void TTreeX::setBranch<Int_t>(TString bn, Int_t val, bool force, bool ignore);
0164   template <>
0165   void TTreeX::setBranch<Bool_t>(TString bn, Bool_t val, bool force, bool ignore);
0166   template <>
0167   void TTreeX::setBranch<Float_t>(TString bn, Float_t val, bool force, bool ignore);
0168   template <>
0169   void TTreeX::setBranch<TString>(TString bn, TString val, bool force, bool ignore);
0170   template <>
0171   void TTreeX::setBranch<LV>(TString bn, LV val, bool force, bool ignore);
0172   template <>
0173   void TTreeX::setBranch<TBits>(TString bn, TBits val, bool force, bool ignore);
0174   template <>
0175   void TTreeX::setBranch<unsigned long long>(TString bn, unsigned long long val, bool force, bool ignore);
0176   template <>
0177   void TTreeX::setBranch<unsigned int>(TString bn, unsigned int val, bool force, bool ignore);
0178   template <>
0179   void TTreeX::setBranch<std::vector<Int_t>>(TString bn, std::vector<Int_t> val, bool force, bool ignore);
0180   template <>
0181   void TTreeX::setBranch<std::vector<UInt_t>>(TString bn, std::vector<UInt_t> val, bool force, bool ignore);
0182   template <>
0183   void TTreeX::setBranch<std::vector<Bool_t>>(TString bn, std::vector<Bool_t> val, bool force, bool ignore);
0184   template <>
0185   void TTreeX::setBranch<std::vector<Float_t>>(TString bn, std::vector<Float_t> val, bool force, bool ignore);
0186   template <>
0187   void TTreeX::setBranch<std::vector<TString>>(TString bn, std::vector<TString> val, bool force, bool ignore);
0188   template <>
0189   void TTreeX::setBranch<std::vector<LV>>(TString bn, std::vector<LV> val, bool force, bool ignore);
0190   template <>
0191   void TTreeX::setBranch<std::vector<VInt>>(TString bn, std::vector<VInt> val, bool force, bool ignore);
0192   template <>
0193   void TTreeX::setBranch<std::vector<VFloat>>(TString bn, std::vector<VFloat> val, bool force, bool ignore);
0194   template <>
0195   void TTreeX::pushbackToBranch<Int_t>(TString bn, Int_t val);
0196   template <>
0197   void TTreeX::pushbackToBranch<UInt_t>(TString bn, UInt_t val);
0198   template <>
0199   void TTreeX::pushbackToBranch<Bool_t>(TString bn, Bool_t val);
0200   template <>
0201   void TTreeX::pushbackToBranch<Float_t>(TString bn, Float_t val);
0202   template <>
0203   void TTreeX::pushbackToBranch<TString>(TString bn, TString val);
0204   template <>
0205   void TTreeX::pushbackToBranch<LV>(TString bn, LV val);
0206   template <>
0207   void TTreeX::pushbackToBranch<VInt>(TString bn, VInt val);
0208   template <>
0209   void TTreeX::pushbackToBranch<VFloat>(TString bn, VFloat val);
0210   // functor
0211   template <>
0212   void TTreeX::setBranch<std::function<float()>>(TString bn, std::function<float()> val, bool force, bool ignore);
0213 
0214   //_________________________________________________________________________________________________
0215   template <>
0216   const Int_t& TTreeX::getBranch<Int_t>(TString bn, bool check);
0217   template <>
0218   const Bool_t& TTreeX::getBranch<Bool_t>(TString bn, bool check);
0219   template <>
0220   const Float_t& TTreeX::getBranch<Float_t>(TString bn, bool check);
0221   template <>
0222   const TString& TTreeX::getBranch<TString>(TString bn, bool check);
0223   template <>
0224   const LV& TTreeX::getBranch<LV>(TString bn, bool check);
0225   template <>
0226   const TBits& TTreeX::getBranch<TBits>(TString bn, bool check);
0227   template <>
0228   const unsigned long long& TTreeX::getBranch<unsigned long long>(TString bn, bool check);
0229   template <>
0230   const unsigned int& TTreeX::getBranch<unsigned int>(TString bn, bool check);
0231   template <>
0232   const std::vector<Int_t>& TTreeX::getBranch<std::vector<Int_t>>(TString bn, bool check);
0233   template <>
0234   const std::vector<UInt_t>& TTreeX::getBranch<std::vector<UInt_t>>(TString bn, bool check);
0235   template <>
0236   const std::vector<Bool_t>& TTreeX::getBranch<std::vector<Bool_t>>(TString bn, bool check);
0237   template <>
0238   const std::vector<Float_t>& TTreeX::getBranch<std::vector<Float_t>>(TString bn, bool check);
0239   template <>
0240   const std::vector<TString>& TTreeX::getBranch<std::vector<TString>>(TString bn, bool check);
0241   template <>
0242   const std::vector<LV>& TTreeX::getBranch<std::vector<LV>>(TString bn, bool check);
0243   template <>
0244   const std::vector<VInt>& TTreeX::getBranch<std::vector<VInt>>(TString bn, bool check);
0245   template <>
0246   const std::vector<VFloat>& TTreeX::getBranch<std::vector<VFloat>>(TString bn, bool check);
0247   // functor
0248   template <>
0249   const std::function<float()>& TTreeX::getBranch<std::function<float()>>(TString bn, bool check);
0250 
0251   //_________________________________________________________________________________________________
0252   template <>
0253   const Int_t& TTreeX::getBranchLazy<Int_t>(TString bn);
0254   template <>
0255   const Bool_t& TTreeX::getBranchLazy<Bool_t>(TString bn);
0256   template <>
0257   const Float_t& TTreeX::getBranchLazy<Float_t>(TString bn);
0258   template <>
0259   const TString& TTreeX::getBranchLazy<TString>(TString bn);
0260   template <>
0261   const LV& TTreeX::getBranchLazy<LV>(TString bn);
0262   template <>
0263   const TBits& TTreeX::getBranchLazy<TBits>(TString bn);
0264   template <>
0265   const unsigned long long& TTreeX::getBranchLazy<unsigned long long>(TString bn);
0266   template <>
0267   const unsigned int& TTreeX::getBranchLazy<unsigned int>(TString bn);
0268   template <>
0269   const std::vector<Int_t>& TTreeX::getBranchLazy<std::vector<Int_t>>(TString bn);
0270   template <>
0271   const std::vector<UInt_t>& TTreeX::getBranchLazy<std::vector<UInt_t>>(TString bn);
0272   template <>
0273   const std::vector<Bool_t>& TTreeX::getBranchLazy<std::vector<Bool_t>>(TString bn);
0274   template <>
0275   const std::vector<Float_t>& TTreeX::getBranchLazy<std::vector<Float_t>>(TString bn);
0276   template <>
0277   const std::vector<TString>& TTreeX::getBranchLazy<std::vector<TString>>(TString bn);
0278   template <>
0279   const std::vector<LV>& TTreeX::getBranchLazy<std::vector<LV>>(TString bn);
0280   template <>
0281   const std::vector<VInt>& TTreeX::getBranchLazy<std::vector<VInt>>(TString bn);
0282   template <>
0283   const std::vector<VFloat>& TTreeX::getBranchLazy<std::vector<VFloat>>(TString bn);
0284   // functor
0285   template <>
0286   const std::function<float()>& TTreeX::getBranchLazy<std::function<float()>>(TString bn);
0287 
0288   //_________________________________________________________________________________________________
0289   template <>
0290   bool TTreeX::isBranchSet<Int_t>(TString bn);
0291   template <>
0292   bool TTreeX::isBranchSet<Bool_t>(TString bn);
0293   template <>
0294   bool TTreeX::isBranchSet<Float_t>(TString bn);
0295   template <>
0296   bool TTreeX::isBranchSet<TString>(TString bn);
0297   template <>
0298   bool TTreeX::isBranchSet<LV>(TString bn);
0299   template <>
0300   bool TTreeX::isBranchSet<TBits>(TString bn);
0301   template <>
0302   bool TTreeX::isBranchSet<unsigned long long>(TString bn);
0303   template <>
0304   bool TTreeX::isBranchSet<unsigned int>(TString bn);
0305   template <>
0306   bool TTreeX::isBranchSet<std::vector<Int_t>>(TString bn);
0307   template <>
0308   bool TTreeX::isBranchSet<std::vector<UInt_t>>(TString bn);
0309   template <>
0310   bool TTreeX::isBranchSet<std::vector<Bool_t>>(TString bn);
0311   template <>
0312   bool TTreeX::isBranchSet<std::vector<Float_t>>(TString bn);
0313   template <>
0314   bool TTreeX::isBranchSet<std::vector<TString>>(TString bn);
0315   template <>
0316   bool TTreeX::isBranchSet<std::vector<LV>>(TString bn);
0317   template <>
0318   bool TTreeX::isBranchSet<std::vector<VInt>>(TString bn);
0319   template <>
0320   bool TTreeX::isBranchSet<std::vector<VFloat>>(TString bn);
0321   // functors
0322   template <>
0323   bool TTreeX::isBranchSet<std::function<float()>>(TString bn);
0324 
0325   //_________________________________________________________________________________________________
0326   template <>
0327   Int_t* TTreeX::getBranchAddress<Int_t>(TString bn);
0328   template <>
0329   Bool_t* TTreeX::getBranchAddress<Bool_t>(TString bn);
0330   template <>
0331   Float_t* TTreeX::getBranchAddress<Float_t>(TString bn);
0332 
0333   //_________________________________________________________________________________________________
0334   template <>
0335   void TTreeX::createBranch<Int_t>(TString bn, bool writeToTree);
0336   template <>
0337   void TTreeX::createBranch<Bool_t>(TString bn, bool writeToTree);
0338   template <>
0339   void TTreeX::createBranch<Float_t>(TString bn, bool writeToTree);
0340   template <>
0341   void TTreeX::createBranch<TString>(TString bn, bool writeToTree);
0342   template <>
0343   void TTreeX::createBranch<LV>(TString bn, bool writeToTree);
0344   template <>
0345   void TTreeX::createBranch<TBits>(TString bn, bool writeToTree);
0346   template <>
0347   void TTreeX::createBranch<unsigned long long>(TString bn, bool writeToTree);
0348   template <>
0349   void TTreeX::createBranch<unsigned int>(TString bn, bool writeToTree);
0350   template <>
0351   void TTreeX::createBranch<std::vector<Int_t>>(TString bn, bool writeToTree);
0352   template <>
0353   void TTreeX::createBranch<std::vector<UInt_t>>(TString bn, bool writeToTree);
0354   template <>
0355   void TTreeX::createBranch<std::vector<Bool_t>>(TString bn, bool writeToTree);
0356   template <>
0357   void TTreeX::createBranch<std::vector<Float_t>>(TString bn, bool writeToTree);
0358   template <>
0359   void TTreeX::createBranch<std::vector<TString>>(TString bn, bool writeToTree);
0360   template <>
0361   void TTreeX::createBranch<std::vector<LV>>(TString bn, bool writeToTree);
0362   template <>
0363   void TTreeX::createBranch<std::vector<VInt>>(TString bn, bool writeToTree);
0364   template <>
0365   void TTreeX::createBranch<std::vector<VFloat>>(TString bn, bool writeToTree);
0366   // functors
0367   template <>
0368   void TTreeX::createBranch<std::function<float()>>(
0369       TString bn,
0370       bool writeToTree);  // writeToTree in this specification does not matter it will never write a std::function to TTree
0371 
0372   //_________________________________________________________________________________________________
0373   template <>
0374   bool TTreeX::hasBranch<Int_t>(TString bn);
0375   template <>
0376   bool TTreeX::hasBranch<Bool_t>(TString bn);
0377   template <>
0378   bool TTreeX::hasBranch<Float_t>(TString bn);
0379   template <>
0380   bool TTreeX::hasBranch<TString>(TString bn);
0381   template <>
0382   bool TTreeX::hasBranch<LV>(TString bn);
0383   template <>
0384   bool TTreeX::hasBranch<TBits>(TString bn);
0385   template <>
0386   bool TTreeX::hasBranch<unsigned long long>(TString bn);
0387   template <>
0388   bool TTreeX::hasBranch<unsigned int>(TString bn);
0389   template <>
0390   bool TTreeX::hasBranch<std::vector<Int_t>>(TString bn);
0391   template <>
0392   bool TTreeX::hasBranch<std::vector<UInt_t>>(TString bn);
0393   template <>
0394   bool TTreeX::hasBranch<std::vector<Bool_t>>(TString bn);
0395   template <>
0396   bool TTreeX::hasBranch<std::vector<Float_t>>(TString bn);
0397   template <>
0398   bool TTreeX::hasBranch<std::vector<TString>>(TString bn);
0399   template <>
0400   bool TTreeX::hasBranch<std::vector<LV>>(TString bn);
0401   template <>
0402   bool TTreeX::hasBranch<std::vector<VInt>>(TString bn);
0403   template <>
0404   bool TTreeX::hasBranch<std::vector<VFloat>>(TString bn);
0405   // functors
0406   template <>
0407   bool TTreeX::hasBranch<std::function<float()>>(TString bn);
0408 
0409   //_________________________________________________________________________________________________
0410   template <>
0411   void TTreeX::setBranch<std::map<TTREEXSTRING, std::vector<Int_t>>>(std::map<TTREEXSTRING, std::vector<Int_t>>& objidx);
0412   template <>
0413   void TTreeX::createBranch<std::map<TTREEXSTRING, std::vector<Int_t>>>(
0414       std::map<TTREEXSTRING, std::vector<Int_t>>& objidx);
0415 
0416   template <class T>
0417   std::vector<T> TTreeX::sortFromRef(std::vector<T> const& in,
0418                                      std::vector<std::pair<size_t, TTreeX::lviter>> const& reference) {
0419     std::vector<T> ret(in.size());
0420 
0421     size_t const size = in.size();
0422     for (size_t i = 0; i < size; ++i)
0423       ret[i] = in[reference[i].first];
0424 
0425     return ret;
0426   }
0427 
0428 }  // namespace RooUtil
0429 
0430 //_________________________________________________________________________________________________
0431 template <class T>
0432 T* RooUtil::TTreeX::get(TString brname, int entry) {
0433   if (entry >= 0)
0434     ttree->GetEntry(entry);
0435   return (T*)getValPtr(brname);
0436 }
0437 
0438 #endif
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!