Changeset 13 for trunk/modules/MadGraphMatchingTreeWriter.cc

Timestamp:

Sep 15, 2008, 10:59:33 AM (16 years ago)

Author:

Pavel Demin

Message:

add configuration for branches

File:

• trunk/modules/MadGraphMatchingTreeWriter.cc (modified) (4 diffs)

trunk/modules/MadGraphMatchingTreeWriter.cc

@@ -40 +40 @@
   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 +93 @@
 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<ExRootCandidate*>(fItInputArrayPartonJets->Next())))
+  // loop over all partons
+  iterator->Reset();
+  while((candidate = static_cast<ExRootCandidate*>(iterator->Next())))
   {
     const TLorentzVector &momentum = candidate->GetP4();
+
+    entry = static_cast<ExRootGenParticle*>(branch->NewEntry());
 
     pt = momentum.Pt();
@@ -91 +124 @@
     rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity());
 
-    if(pt < fJetPTMin) continue;
-    if(TMath::Abs(eta) > fJetEtaMax) continue;
-
-    entryJet = static_cast<ExRootGenJet*>(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<ExRootCandidate*>(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<ExRootMatching*>(iterator->Next())))
+  {
+    entry = static_cast<ExRootMatching*>(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<ExRootCandidate*>(iterator->Next())))
   {
     const TLorentzVector &momentum = candidate->GetP4();
@@ -124 +178 @@
     if(TMath::Abs(eta) > fJetEtaMax) continue;
 
-    entryJet = static_cast<ExRootGenJet*>(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<ExRootMatching*>(fItInputArrayMatching->Next())))
-  {
-    entryMatching = static_cast<ExRootMatching*>(fBranchMatching->NewEntry());
-
-    entryMatching->DMerge = matching->DMerge;
-    entryMatching->YMerge = matching->YMerge;
-  }
-
-  // loop over all partons
-  fItInputArrayPartons->Reset();
-  while((candidate = static_cast<ExRootCandidate*>(fItInputArrayPartons->Next())))
-  {
-    const TLorentzVector &momentum = candidate->GetP4();
-
-    entryParton = static_cast<ExRootGenParticle*>(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<ExRootGenJet*>(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);
+  }
+
+}
+
+//------------------------------------------------------------------------------
+