Index: /trunk/modules/MadGraphMatchingTreeWriter.cc =================================================================== --- /trunk/modules/MadGraphMatchingTreeWriter.cc (revision 12) +++ /trunk/modules/MadGraphMatchingTreeWriter.cc (revision 13) @@ -40,22 +40,51 @@ fJetEtaMax = GetDouble("JetEtaMax", 4.5); + fClassMap[ExRootGenParticle::Class()] = &MadGraphMatchingTreeWriter::ProcessPartons; + + fClassMap[ExRootMatching::Class()] = &MadGraphMatchingTreeWriter::ProcessMatching; + + fClassMap[ExRootGenJet::Class()] = &MadGraphMatchingTreeWriter::ProcessJets; + + TBranchMap::iterator itBranchMap; + map< TClass *, TProcessMethod >::iterator itClassMap; + + // read branch configuration and // import array with output from filter/classifier/jetfinder modules - fInputArrayPartonJets = ImportArray(GetString("InputArrayPartonJets", "partonjetfinder/candidates")); - fItInputArrayPartonJets = fInputArrayPartonJets->MakeIterator(); - - fInputArrayHadronJets = ImportArray(GetString("InputArrayHadronJets", "hadronjetfinder/candidates")); - fItInputArrayHadronJets = fInputArrayHadronJets->MakeIterator(); - - fInputArrayMatching = ImportArray(GetString("InputArrayMatching", "partonjetfinder/matching")); - fItInputArrayMatching = fInputArrayMatching->MakeIterator(); - - fInputArrayPartons = ImportArray(GetString("InputArrayPartons", "initstateselection/candidates")); - fItInputArrayPartons = fInputArrayPartons->MakeIterator(); - - fBranchPartonJets = NewBranch("PartonJet", ExRootGenJet::Class()); - fBranchHadronJets = NewBranch("HadronJet", ExRootGenJet::Class()); - fBranchMatching = NewBranch("Match", ExRootMatching::Class()); - fBranchPartons = NewBranch("Parton", ExRootGenParticle::Class()); + ExRootConfParam param = GetParam("Branch"); + Long_t i, size; + TString branchName, branchClassName, branchInputArray; + TClass *branchClass; + const TObjArray *array; + ExRootTreeBranch *branch; + + size = param.GetSize(); + for(i = 0; i < size; ++i) + { + branchName = param[i][0].GetString(); + branchClassName = param[i][1].GetString(); + branchInputArray = param[i][2].GetString(); + + branchClass = gROOT->GetClass(branchClassName); + + if(!branchClass) + { + cout << "** ERROR: cannot find class '" << branchClassName << "'" << endl; + continue; + } + + itClassMap = fClassMap.find(branchClass); + if(itClassMap == fClassMap.end()) + { + cout << "** ERROR: cannot create branch for class '" << branchClassName << "'" << endl; + continue; + } + + array = ImportArray(branchInputArray); + branch = NewBranch(branchName, branchClass); + + fBranchMap.insert(make_pair(branch, make_pair(itClassMap->second, array->MakeIterator()))); + } + } @@ -64,25 +93,29 @@ void MadGraphMatchingTreeWriter::Finish() { - if(fItInputArrayPartonJets) delete fItInputArrayPartonJets; - if(fItInputArrayHadronJets) delete fItInputArrayHadronJets; - if(fItInputArrayMatching) delete fItInputArrayMatching; - if(fItInputArrayPartons) delete fItInputArrayPartons; -} - -//------------------------------------------------------------------------------ - -void MadGraphMatchingTreeWriter::Process() + TBranchMap::iterator itBranchMap; + TIterator *iterator; + + for(itBranchMap = fBranchMap.begin(); itBranchMap != fBranchMap.end(); ++itBranchMap) + { + iterator = itBranchMap->second.second; + if(iterator) delete iterator; + } +} + +//------------------------------------------------------------------------------ + +void MadGraphMatchingTreeWriter::ProcessPartons(ExRootTreeBranch *branch, TIterator *iterator) { ExRootCandidate *candidate = 0; - ExRootGenParticle *entryParton = 0; - ExRootMatching *matching = 0, *entryMatching = 0; - ExRootGenJet *entryJet = 0; + ExRootGenParticle *entry = 0; Double_t pt, signPz, eta, rapidity; - // loop over all parton jets - fItInputArrayPartonJets->Reset(); - while((candidate = static_cast(fItInputArrayPartonJets->Next()))) + // loop over all partons + iterator->Reset(); + while((candidate = static_cast(iterator->Next()))) { const TLorentzVector &momentum = candidate->GetP4(); + + entry = static_cast(branch->NewEntry()); pt = momentum.Pt(); @@ -91,26 +124,47 @@ rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity()); - if(pt < fJetPTMin) continue; - if(TMath::Abs(eta) > fJetEtaMax) continue; - - entryJet = static_cast(fBranchPartonJets->NewEntry()); - - entryJet->E = momentum.E(); - entryJet->Px = momentum.Px(); - entryJet->Py = momentum.Py(); - entryJet->Pz = momentum.Pz(); - - entryJet->Eta = eta; - entryJet->Phi = momentum.Phi(); - entryJet->PT = pt; - - entryJet->Rapidity = rapidity; - - entryJet->Mass = momentum.M(); - } - - // loop over all hadron jets - fItInputArrayHadronJets->Reset(); - while((candidate = static_cast(fItInputArrayHadronJets->Next()))) + entry->PID = candidate->GetType()->PdgCode(); + + entry->E = momentum.E(); + entry->Px = momentum.Px(); + entry->Py = momentum.Py(); + entry->Pz = momentum.Pz(); + + entry->Eta = eta; + entry->Phi = momentum.Phi(); + entry->PT = pt; + + entry->Rapidity = rapidity; + } +} + +//------------------------------------------------------------------------------ + +void MadGraphMatchingTreeWriter::ProcessMatching(ExRootTreeBranch *branch, TIterator *iterator) +{ + ExRootMatching *matching = 0, *entry = 0; + + // loop over all matching + iterator->Reset(); + while((matching = static_cast(iterator->Next()))) + { + entry = static_cast(branch->NewEntry()); + + entry->DMerge = matching->DMerge; + entry->YMerge = matching->YMerge; + } +} + +//------------------------------------------------------------------------------ + +void MadGraphMatchingTreeWriter::ProcessJets(ExRootTreeBranch *branch, TIterator *iterator) +{ + ExRootCandidate *candidate = 0; + ExRootGenJet *entry = 0; + Double_t pt, signPz, eta, rapidity; + + // loop over all jets + iterator->Reset(); + while((candidate = static_cast(iterator->Next()))) { const TLorentzVector &momentum = candidate->GetP4(); @@ -124,58 +178,42 @@ if(TMath::Abs(eta) > fJetEtaMax) continue; - entryJet = static_cast(fBranchHadronJets->NewEntry()); - - entryJet->E = momentum.E(); - entryJet->Px = momentum.Px(); - entryJet->Py = momentum.Py(); - entryJet->Pz = momentum.Pz(); - - entryJet->Eta = eta; - entryJet->Phi = momentum.Phi(); - entryJet->PT = pt; - - entryJet->Rapidity = rapidity; - - entryJet->Mass = momentum.M(); - } - - // loop over all matching - fItInputArrayMatching->Reset(); - while((matching = static_cast(fItInputArrayMatching->Next()))) - { - entryMatching = static_cast(fBranchMatching->NewEntry()); - - entryMatching->DMerge = matching->DMerge; - entryMatching->YMerge = matching->YMerge; - } - - // loop over all partons - fItInputArrayPartons->Reset(); - while((candidate = static_cast(fItInputArrayPartons->Next()))) - { - const TLorentzVector &momentum = candidate->GetP4(); - - entryParton = static_cast(fBranchPartons->NewEntry()); - - pt = momentum.Pt(); - signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0; - eta = (pt == 0.0 ? signPz*999.9 : momentum.Eta()); - rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity()); - - entryParton->PID = candidate->GetType()->PdgCode(); - - entryParton->E = momentum.E(); - entryParton->Px = momentum.Px(); - entryParton->Py = momentum.Py(); - entryParton->Pz = momentum.Pz(); - - entryParton->Eta = eta; - entryParton->Phi = momentum.Phi(); - entryParton->PT = pt; - - entryParton->Rapidity = rapidity; - } -} - -//------------------------------------------------------------------------------ - + entry = static_cast(branch->NewEntry()); + + entry->E = momentum.E(); + entry->Px = momentum.Px(); + entry->Py = momentum.Py(); + entry->Pz = momentum.Pz(); + + entry->Eta = eta; + entry->Phi = momentum.Phi(); + entry->PT = pt; + + entry->Rapidity = rapidity; + + entry->Mass = momentum.M(); + } +} + +//------------------------------------------------------------------------------ + +void MadGraphMatchingTreeWriter::Process() +{ + + TBranchMap::iterator itBranchMap; + ExRootTreeBranch *branch; + TProcessMethod method; + TIterator *iterator; + + for(itBranchMap = fBranchMap.begin(); itBranchMap != fBranchMap.end(); ++itBranchMap) + { + branch = itBranchMap->first; + method = itBranchMap->second.first; + iterator = itBranchMap->second.second; + + (this->*method)(branch, iterator); + } + +} + +//------------------------------------------------------------------------------ + Index: /trunk/modules/MadGraphMatchingTreeWriter.h =================================================================== --- /trunk/modules/MadGraphMatchingTreeWriter.h (revision 12) +++ /trunk/modules/MadGraphMatchingTreeWriter.h (revision 13) @@ -4,12 +4,9 @@ #include "ExRootAnalysis/ExRootModule.h" -#include "TString.h" - #include -class TH1; +class TClass; class TIterator; -class TObjArray; -class TClonesArray; +class ExRootTreeBranch; class MadGraphMatchingTreeWriter: public ExRootModule @@ -26,21 +23,18 @@ private: - Double_t fJetPTMin, fJetEtaMax; + void ProcessPartons(ExRootTreeBranch *branch, TIterator *iterator); + void ProcessMatching(ExRootTreeBranch *branch, TIterator *iterator); + void ProcessJets(ExRootTreeBranch *branch, TIterator *iterator); - const TObjArray *fInputArrayPartonJets; //! - const TObjArray *fInputArrayHadronJets; //! - const TObjArray *fInputArrayMatching; //! - const TObjArray *fInputArrayPartons; //! + Double_t fJetPTMin, fJetEtaMax; //! +#ifndef __CINT__ + typedef void (MadGraphMatchingTreeWriter::*TProcessMethod)(ExRootTreeBranch*, TIterator*); //! - TIterator *fItInputArrayPartonJets; //! - TIterator *fItInputArrayHadronJets; //! - TIterator *fItInputArrayMatching; //! - TIterator *fItInputArrayPartons; //! - - ExRootTreeBranch *fBranchPartonJets; //! - ExRootTreeBranch *fBranchHadronJets; //! - ExRootTreeBranch *fBranchMatching; //! - ExRootTreeBranch *fBranchPartons; //! + typedef std::map > TBranchMap; //! + TBranchMap fBranchMap; //! + + std::map< TClass *, TProcessMethod > fClassMap; //! +#endif ClassDef(MadGraphMatchingTreeWriter, 1) };