MomentumScaleCalibration/interface/RootTreeHandler.h

0001 #include <iostream>
0002
0003 #include <TFile.h>
0004 #include <TTree.h>
0005
0006 #include <MuonAnalysis/MomentumScaleCalibration/interface/MuonPair.h>
0007 #include <MuonAnalysis/MomentumScaleCalibration/interface/GenMuonPair.h>
0008 #include <MuonAnalysis/MomentumScaleCalibration/interface/MuScleFitProvenance.h>
0009 #include <TH1F.h>
0010 #include <cstdlib>
0011 #include <vector>
0012
0013 typedef std::vector<std::pair<lorentzVector, lorentzVector> > MuonPairVector;
0014 typedef std::vector<std::pair<MuScleFitMuon, MuScleFitMuon> > MuonPairExtendedVector;
0015
0016 /**
0017  * This class can be used to save all the muon pairs (and gen muon pairs if any) to a root tree. <br>
0018  * The writeTree method gets the name of the file to store the tree and the savedPair (and possibly genPair)
0019  * vector of muon pairs. <br>
0020  * Likewise, the readTree method takes the same arguments. It reads back from the file with the given name the
0021  * pairs and stores them in the given savedPair (and genPair) vector.
0022  */
0023
0024 class RootTreeHandler {
0025 public:
0026   // void writeTree( const TString & fileName, const MuonPairVector * savedPair, const int muonType = 0,
0027   //                 const MuonPairVector * genPair = 0, const bool saveAll = false )
0028   void writeTree(const TString& fileName,
0029                  const std::vector<MuonPair>* savedPair,
0030                  const int muonType = 0,
0031                  const std::vector<GenMuonPair>* genPair = nullptr,
0032                  const bool saveAll = false) {
0033     lorentzVector emptyLorentzVector(0, 0, 0, 0);
0034     TFile* f1 = new TFile(fileName, "RECREATE");
0035     TTree* tree = new TTree("T", "Muon pairs");
0036     MuonPair* muonPair = new MuonPair;
0037     GenMuonPair* genMuonPair = new GenMuonPair;
0038     // MuonPair * genMuonPair = new MuonPair;
0039     tree->Branch("event", "MuonPair", &muonPair);
0040     if (genPair != nullptr) {
0041       tree->Branch("genEvent", "GenMuonPair", &genMuonPair);
0042       // tree->Branch("genEvent", "MuonPair", &genMuonPair);
0043
0044       if (savedPair->size() != genPair->size()) {
0045         std::cout << "Error: savedPair size (" << savedPair->size() << ") and genPair size (" << genPair->size()
0046                   << ") are different. This is severe and I will not write the tree." << std::endl;
0047         exit(1);
0048       }
0049     }
0050     std::cout << "savedPair->size() is " << savedPair->size() << std::endl;
0051     std::vector<MuonPair>::const_iterator muonPairIt = savedPair->begin();
0052     unsigned int iev = 0;
0053     for (; muonPairIt != savedPair->end(); ++muonPairIt, ++iev) {
0054       if (saveAll || ((muonPairIt->mu1.p4() != emptyLorentzVector) && (muonPairIt->mu2.p4() != emptyLorentzVector))) {
0055         // muonPair->setPair(muonType, std::make_pair(muonPairIt->first, muonPairIt->second));
0056         muonPair->copy(*muonPairIt);
0057
0058         // if( genPair != 0 && genPair->size() != 0 ) {
0059         //   genMuonPair->setPair(muonId, std::make_pair((*genPair)[iev].first, (*genPair)[iev].second));
0060         //   genMuonPair->mu1 = ((*genPair)[iev].first);
0061         //   genMuonPair->mu2 = ((*genPair)[iev].second);
0062         // }
0063         if (genPair != nullptr) {
0064           genMuonPair->copy((*genPair)[iev]);
0065         }
0066
0067         tree->Fill();
0068       }
0069       // // Tree filled. Clear the map for the next event.
0070       // muonPair->muonPairs.clear();
0071     }
0072
0073     // Save provenance information in the TFile
0074     TH1F muonTypeHisto("MuonType", "MuonType", 40, -20, 20);
0075     muonTypeHisto.Fill(muonType);
0076     muonTypeHisto.Write();
0077     MuScleFitProvenance provenance(muonType);
0078     provenance.Write();
0079
0080     f1->Write();
0081     f1->Close();
0082   }
0083
0084   // void readTree( const int maxEvents, const TString & fileName, MuonPairVector * savedPair,
0085   //                const int muonType, MuonPairVector * genPair = 0 )
0086   void readTree(const int maxEvents,
0087                 const TString& fileName,
0088                 MuonPairVector* savedPair,
0089                 const int muonType,
0090                 std::vector<std::pair<unsigned int, unsigned long long> >* evtRun,
0091                 MuonPairVector* genPair = nullptr) {
0092     TFile* file = TFile::Open(fileName, "READ");
0093     if (file->IsOpen()) {
0094       TTree* tree = (TTree*)file->Get("T");
0095       MuonPair* muonPair = nullptr;
0096       GenMuonPair* genMuonPair = nullptr;
0097       // MuonPair * genMuonPair = 0;
0098       tree->SetBranchAddress("event", &muonPair);
0099       if (genPair != nullptr) {
0100         tree->SetBranchAddress("genEvent", &genMuonPair);
0101       }
0102
0103       Long64_t nentries = tree->GetEntries();
0104       if ((maxEvents != -1) && (nentries > maxEvents))
0105         nentries = maxEvents;
0106       for (Long64_t i = 0; i < nentries; ++i) {
0107         tree->GetEntry(i);
0108         //std::cout << "Reco muon1, pt = " << muonPair->mu1 << "; Reco muon2, pt = " << muonPair->mu2 << std::endl;
0109         savedPair->push_back(std::make_pair(muonPair->mu1.p4(), muonPair->mu2.p4()));
0110         evtRun->push_back(std::make_pair(muonPair->event.event(), muonPair->event.run()));
0111         // savedPair->push_back(muonPair->getPair(muonType));
0112         if (genPair != nullptr) {
0113           genPair->push_back(std::make_pair(genMuonPair->mu1.p4(), genMuonPair->mu2.p4()));
0114           //std::cout << "Gen muon1, pt = " << genMuonPair->mu1 << "; Gen muon2, pt = " << genMuonPair->mu2 << std::endl;
0115           // genPair->push_back(genMuonPair->getPair(muonId));
0116         }
0117       }
0118     } else {
0119       std::cout << "ERROR: no file " << fileName
0120                 << " found. Please, correct the file name or specify an empty field in the InputRootTreeFileName "
0121                    "parameter to read events from the edm source."
0122                 << std::endl;
0123       exit(1);
0124     }
0125     file->Close();
0126   }
0127
0128   void readTree(const int maxEvents,
0129                 const TString& fileName,
0130                 MuonPairExtendedVector* savedPair,
0131                 const int muonType,
0132                 std::vector<std::pair<unsigned int, unsigned long long> >* evtRun,
0133                 MuonPairExtendedVector* genPair = nullptr) {
0134     TFile* file = TFile::Open(fileName, "READ");
0135     if (file->IsOpen()) {
0136       TTree* tree = (TTree*)file->Get("T");
0137       MuonPair* muonPair = nullptr;
0138       GenMuonPair* genMuonPair = nullptr;
0139       tree->SetBranchAddress("event", &muonPair);
0140       if (genPair != nullptr) {
0141         tree->SetBranchAddress("genEvent", &genMuonPair);
0142       }
0143       Long64_t nentries = tree->GetEntries();
0144       if ((maxEvents != -1) && (nentries > maxEvents))
0145         nentries = maxEvents;
0146       for (Long64_t i = 0; i < nentries; ++i) {
0147         tree->GetEntry(i);
0148         //std::cout << "Reco muon1, pt = " << muonPair->mu1 << "; Reco muon2, pt = " << muonPair->mu2 << std::endl;
0149         savedPair->push_back(std::make_pair(muonPair->mu1, muonPair->mu2));
0150         evtRun->push_back(std::make_pair(muonPair->event.event(), muonPair->event.run()));
0151         // savedPair->push_back(muonPair->getPair(muonType));
0152         if (genPair != nullptr) {
0153           genPair->push_back(std::make_pair(genMuonPair->mu1, genMuonPair->mu2));
0154           //std::cout << "Gen muon1, pt = " << genMuonPair->mu1 << "; Gen muon2, pt = " << genMuonPair->mu2 << std::endl;
0155           // genPair->push_back(genMuonPair->getPair(muonId));
0156         }
0157       }
0158     } else {
0159       std::cout << "ERROR: no file " << fileName
0160                 << " found. Please, correct the file name or specify an empty field in the InputRootTreeFileName "
0161                    "parameter to read events from the edm source."
0162                 << std::endl;
0163       exit(1);
0164     }
0165     file->Close();
0166   }
0167
0168   /// Used to read the external trees
0169   void readTree(const int maxEvents,
0170                 const TString& fileName,
0171                 std::vector<MuonPair>* savedPair,
0172                 const int muonType,
0173                 std::vector<GenMuonPair>* genPair = nullptr) {
0174     TFile* file = TFile::Open(fileName, "READ");
0175     if (file->IsOpen()) {
0176       TTree* tree = (TTree*)file->Get("T");
0177       MuonPair* muonPair = nullptr;
0178       GenMuonPair* genMuonPair = nullptr;
0179       tree->SetBranchAddress("event", &muonPair);
0180       if (genPair != nullptr) {
0181         tree->SetBranchAddress("genEvent", &genMuonPair);
0182       }
0183
0184       Long64_t nentries = tree->GetEntries();
0185       if ((maxEvents != -1) && (nentries > maxEvents))
0186         nentries = maxEvents;
0187       for (Long64_t i = 0; i < nentries; ++i) {
0188         tree->GetEntry(i);
0189         savedPair->push_back(*muonPair);
0190         if (genPair != nullptr) {
0191           genPair->push_back(*genMuonPair);
0192         }
0193       }
0194     } else {
0195       std::cout << "ERROR: no file " << fileName
0196                 << " found. Please, correct the file name or specify an empty field in the InputRootTreeFileName "
0197                    "parameter to read events from the edm source."
0198                 << std::endl;
0199       exit(1);
0200     }
0201     file->Close();
0202   }
0203 };