code/rooutil/looper.h

0001 //  .
0002 // ..: P. Chang, philip@physics.ucsd.edu
0003
0004 #ifndef looper_cc
0005 #define looper_cc
0006
0007 // C/C++
0008 #include <unistd.h>
0009 #include <algorithm>
0010 #include <fstream>
0011 #include <iostream>
0012 #include <map>
0013 #include <string>
0014 #include <unordered_map>
0015 #include <vector>
0016 #include <stdarg.h>
0017 #include <functional>
0018 #include <cmath>
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
0040 #include "printutil.h"
0041 #include "eventindexmap.h"
0042 #include "treeutil.h"
0043
0044 //#include "cpptqdm/tqdm.h"
0045
0046 namespace RooUtil {
0047
0048   ///////////////////////////////////////////////////////////////////////////////////////////////
0049   // Looper class
0050   ///////////////////////////////////////////////////////////////////////////////////////////////
0051   class Looper {
0052     // Members
0053     TChain* tchain;
0054     TBenchmark* bmark;
0055     TObjArray* listOfFiles;
0056     TObjArrayIter* fileIter;
0057     TFile* tfile;
0058     TTree* ttree;
0059     TTreePerfStats* ps;
0060     unsigned int nEventsTotalInChain;
0061     unsigned int nEventsTotalInTree;
0062     int nEventsToProcess;
0063     unsigned int nEventsProcessed;
0064     unsigned int indexOfEventInTTree;
0065     bool fastmode;
0066     TreeUtil* treeclass;
0067     TStopwatch my_timer;
0068     int bar_id;
0069     int print_rate;
0070     bool doskim;
0071     TString skimfilename;
0072     TFile* skimfile;
0073     TTree* skimtree;
0074     unsigned int nEventsSkimmed;
0075     std::vector<TString> skimbrfiltpttn;
0076     bool silent;
0077     bool isinit;
0078     bool use_treeclass_progress;
0079     bool isnewfileopened;
0080     //        bool use_tqdm_progress_bar;
0081     unsigned int nskipped_batch;
0082     unsigned int nskipped;
0083     unsigned int nbatch_skip_threshold;
0084     unsigned int nbatch_to_skip;
0085     unsigned int nskipped_threshold;
0086     unsigned int ncounter;
0087     //        tqdm bar;
0088     EventIndexMap eventindexmap;
0089     TEventList* teventlist;
0090
0091   public:
0092     // Functions
0093     Looper();
0094     Looper(TChain* chain, TreeUtil* treeclass, int nEventsToProcess = -1);
0095     ~Looper();
0096     void init(TChain* chain, TreeUtil* treeclass, int nEventsToProcess);
0097     void setTChain(TChain* c);
0098     void setTreeClass(TreeUtil* t);
0099     void printCurrentEventIndex();
0100     void setSilent(bool s = true) { silent = s; }
0101     void setEventIndexMap(TString file) { eventindexmap.load(file); }
0102     bool allEventsInTreeProcessed();
0103     bool allEventsInChainProcessed();
0104     bool nextEvent();
0105     bool isNewFileInChain();
0106     TTree* getTree() { return ttree; }
0107     TChain* getTChain() { return tchain; }
0108     unsigned int getNEventsProcessed() { return nEventsProcessed; }
0109     void setSkim(TString ofilename);
0110     void setSkimBranchFilterPattern(std::vector<TString> x) { skimbrfiltpttn = x; }
0111     void fillSkim();
0112     void saveSkim();
0113     TTree* getSkimTree() { return skimtree; }
0114     void setSkimMaxSize(Long64_t maxsize) { skimtree->SetMaxTreeSize(maxsize); }
0115     TTreePerfStats* getTTreePerfStats() { return ps; }
0116     unsigned int getCurrentEventIndex() { return indexOfEventInTTree - 1; }
0117     TFile* getCurrentFile() { return tfile; }
0118     TString getCurrentFileBaseName() { return gSystem->BaseName(tfile->GetName()); }
0119     TString getCurrentFileName() { return TString(tfile->GetName()); }
0120     TString getListOfFileNames();
0121     TString getCurrentFileTitle() { return TString(tfile->GetTitle()); }
0122     unsigned int getNEventsTotalInChain() { return nEventsTotalInChain; }
0123     void setNbatchToSkip(unsigned int n) { nbatch_to_skip = n; }
0124     void setNbadEventThreshold(unsigned int n) { nskipped_threshold = n; }
0125     void setNbadEventThresholdToTriggerBatchSkip(unsigned int n) { nbatch_skip_threshold = n; }
0126     bool handleBadEvent();
0127     void printStatus();
0128     void printSkippedBadEventStatus();
0129     void setFastMode(bool f = true) { fastmode = f; }
0130     void addCount() { ncounter++; }
0131     void resetCounter() { ncounter = 0; }
0132     bool doesBranchExist(TString bname);
0133     TString getSkimFileName() { return skimfilename; }
0134     TFile* getSkimFile() { return skimfile; }
0135
0136   private:
0137     void setFileList();
0138     void setNEventsToProcess();
0139     bool nextTree();
0140     bool nextEventInTree();
0141     void initProgressBar();
0142     void printProgressBar(bool force = false);
0143     void createSkimTree();
0144     void copyAddressesToSkimTree();
0145   };
0146
0147 }  // namespace RooUtil
0148
0149 #endif