Fork me on GitHub

source: svn/trunk/modules/TreeWriter.cc@ 1323

Last change on this file since 1323 was 1323, checked in by Pavel Demin, 11 years ago

add vertex block

  • Property svn:keywords set to Id Revision Date
File size: 16.0 KB
Line 
1
2/** \class TreeWriter
3 *
4 * Fills ROOT tree branches.
5 *
6 * $Date: 2013-11-08 16:09:58 +0000 (Fri, 08 Nov 2013) $
7 * $Revision: 1323 $
8 *
9 *
10 * \author P. Demin - UCL, Louvain-la-Neuve
11 *
12 */
13
14#include "modules/TreeWriter.h"
15
16#include "classes/DelphesClasses.h"
17#include "classes/DelphesFactory.h"
18#include "classes/DelphesFormula.h"
19
20#include "ExRootAnalysis/ExRootResult.h"
21#include "ExRootAnalysis/ExRootFilter.h"
22#include "ExRootAnalysis/ExRootClassifier.h"
23#include "ExRootAnalysis/ExRootTreeBranch.h"
24
25#include "TROOT.h"
26#include "TMath.h"
27#include "TString.h"
28#include "TFormula.h"
29#include "TRandom3.h"
30#include "TObjArray.h"
31#include "TDatabasePDG.h"
32#include "TLorentzVector.h"
33
34#include <algorithm>
35#include <stdexcept>
36#include <iostream>
37#include <sstream>
38
39using namespace std;
40
41//------------------------------------------------------------------------------
42
43TreeWriter::TreeWriter()
44{
45}
46
47//------------------------------------------------------------------------------
48
49TreeWriter::~TreeWriter()
50{
51}
52
53//------------------------------------------------------------------------------
54
55void TreeWriter::Init()
56{
57 fClassMap[GenParticle::Class()] = &TreeWriter::ProcessParticles;
58 fClassMap[Track::Class()] = &TreeWriter::ProcessTracks;
59 fClassMap[Tower::Class()] = &TreeWriter::ProcessTowers;
60 fClassMap[Photon::Class()] = &TreeWriter::ProcessPhotons;
61 fClassMap[Electron::Class()] = &TreeWriter::ProcessElectrons;
62 fClassMap[Muon::Class()] = &TreeWriter::ProcessMuons;
63 fClassMap[Jet::Class()] = &TreeWriter::ProcessJets;
64 fClassMap[MissingET::Class()] = &TreeWriter::ProcessMissingET;
65 fClassMap[ScalarHT::Class()] = &TreeWriter::ProcessScalarHT;
66 fClassMap[Rho::Class()] = &TreeWriter::ProcessRho;
67 fClassMap[Weight::Class()] = &TreeWriter::ProcessWeight;
68
69 TBranchMap::iterator itBranchMap;
70 map< TClass *, TProcessMethod >::iterator itClassMap;
71
72 // read branch configuration and
73 // import array with output from filter/classifier/jetfinder modules
74
75 ExRootConfParam param = GetParam("Branch");
76 Long_t i, size;
77 TString branchName, branchClassName, branchInputArray;
78 TClass *branchClass;
79 TObjArray *array;
80 ExRootTreeBranch *branch;
81
82 size = param.GetSize();
83 for(i = 0; i < size/3; ++i)
84 {
85 branchInputArray = param[i*3].GetString();
86 branchName = param[i*3 + 1].GetString();
87 branchClassName = param[i*3 + 2].GetString();
88
89 branchClass = gROOT->GetClass(branchClassName);
90
91 if(!branchClass)
92 {
93 cout << "** ERROR: cannot find class '" << branchClassName << "'" << endl;
94 continue;
95 }
96
97 itClassMap = fClassMap.find(branchClass);
98 if(itClassMap == fClassMap.end())
99 {
100 cout << "** ERROR: cannot create branch for class '" << branchClassName << "'" << endl;
101 continue;
102 }
103
104 array = ImportArray(branchInputArray);
105 branch = NewBranch(branchName, branchClass);
106
107 fBranchMap.insert(make_pair(branch, make_pair(itClassMap->second, array)));
108 }
109
110}
111
112//------------------------------------------------------------------------------
113
114void TreeWriter::Finish()
115{
116
117}
118
119//------------------------------------------------------------------------------
120
121void TreeWriter::FillParticles(Candidate *candidate, TRefArray *array)
122{
123 TIter it1(candidate->GetCandidates());
124 it1.Reset();
125 array->Clear();
126 while((candidate = static_cast<Candidate*>(it1.Next())))
127 {
128 TIter it2(candidate->GetCandidates());
129
130 // particle
131 if(candidate->GetCandidates()->GetEntriesFast() == 0)
132 {
133 array->Add(candidate);
134 continue;
135 }
136
137 // track
138 candidate = static_cast<Candidate*>(candidate->GetCandidates()->At(0));
139 if(candidate->GetCandidates()->GetEntriesFast() == 0)
140 {
141 array->Add(candidate);
142 continue;
143 }
144
145 // tower
146 it2.Reset();
147 while((candidate = static_cast<Candidate*>(it2.Next())))
148 {
149 array->Add(candidate->GetCandidates()->At(0));
150 }
151 }
152}
153
154//------------------------------------------------------------------------------
155
156void TreeWriter::ProcessParticles(ExRootTreeBranch *branch, TObjArray *array)
157{
158 TIter iterator(array);
159 Candidate *candidate = 0;
160 GenParticle *entry = 0;
161 Double_t pt, signPz, cosTheta, eta, rapidity;
162
163 // loop over all particles
164 iterator.Reset();
165 while((candidate = static_cast<Candidate*>(iterator.Next())))
166 {
167 const TLorentzVector &momentum = candidate->Momentum;
168 const TLorentzVector &position = candidate->Position;
169
170 entry = static_cast<GenParticle*>(branch->NewEntry());
171
172 entry->SetBit(kIsReferenced);
173 entry->SetUniqueID(candidate->GetUniqueID());
174
175 pt = momentum.Pt();
176 cosTheta = TMath::Abs(momentum.CosTheta());
177 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
178 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
179 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
180
181 entry->PID = candidate->PID;
182
183 entry->Status = candidate->Status;
184 entry->IsPU = candidate->IsPU;
185
186 entry->M1 = candidate->M1;
187 entry->M2 = candidate->M2;
188
189 entry->D1 = candidate->D1;
190 entry->D2 = candidate->D2;
191
192 entry->Charge = candidate->Charge;
193 entry->Mass = candidate->Mass;
194
195 entry->E = momentum.E();
196 entry->Px = momentum.Px();
197 entry->Py = momentum.Py();
198 entry->Pz = momentum.Pz();
199
200 entry->Eta = eta;
201 entry->Phi = momentum.Phi();
202 entry->PT = pt;
203
204 entry->Rapidity = rapidity;
205
206 entry->X = position.X();
207 entry->Y = position.Y();
208 entry->Z = position.Z();
209 entry->T = position.T();
210 }
211}
212
213//------------------------------------------------------------------------------
214
215void TreeWriter::ProcessVertices(ExRootTreeBranch *branch, TObjArray *array)
216{
217 TIter iterator(array);
218 Candidate *candidate = 0;
219 Vertex *entry = 0;
220
221 // loop over all vertices
222 iterator.Reset();
223 while((candidate = static_cast<Candidate*>(iterator.Next())))
224 {
225 const TLorentzVector &position = candidate->Position;
226
227 entry = static_cast<Vertex*>(branch->NewEntry());
228
229 entry->X = position.X();
230 entry->Y = position.Y();
231 entry->Z = position.Z();
232 entry->T = position.T();
233 }
234}
235
236//------------------------------------------------------------------------------
237
238void TreeWriter::ProcessTracks(ExRootTreeBranch *branch, TObjArray *array)
239{
240 TIter iterator(array);
241 Candidate *candidate = 0;
242 Candidate *particle = 0;
243 Track *entry = 0;
244 Double_t pt, signz, cosTheta, eta, rapidity;
245
246 // loop over all tracks
247 iterator.Reset();
248 while((candidate = static_cast<Candidate*>(iterator.Next())))
249 {
250 const TLorentzVector &position = candidate->Position;
251
252 cosTheta = TMath::Abs(position.CosTheta());
253 signz = (position.Pz() >= 0.0) ? 1.0 : -1.0;
254 eta = (cosTheta == 1.0 ? signz*999.9 : position.Eta());
255 rapidity = (cosTheta == 1.0 ? signz*999.9 : position.Rapidity());
256
257 entry = static_cast<Track*>(branch->NewEntry());
258
259 entry->SetBit(kIsReferenced);
260 entry->SetUniqueID(candidate->GetUniqueID());
261
262 entry->PID = candidate->PID;
263
264 entry->Charge = candidate->Charge;
265
266 entry->EtaOuter = eta;
267 entry->PhiOuter = position.Phi();
268
269 entry->XOuter = position.X();
270 entry->YOuter = position.Y();
271 entry->ZOuter = position.Z();
272
273 const TLorentzVector &momentum = candidate->Momentum;
274
275 pt = momentum.Pt();
276 cosTheta = TMath::Abs(momentum.CosTheta());
277 signz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
278 eta = (cosTheta == 1.0 ? signz*999.9 : momentum.Eta());
279 rapidity = (cosTheta == 1.0 ? signz*999.9 : momentum.Rapidity());
280
281 entry->Eta = eta;
282 entry->Phi = momentum.Phi();
283 entry->PT = pt;
284
285 particle = static_cast<Candidate*>(candidate->GetCandidates()->At(0));
286 const TLorentzVector &initialPosition = particle->Position;
287
288 entry->X = initialPosition.X();
289 entry->Y = initialPosition.Y();
290 entry->Z = initialPosition.Z();
291
292 entry->Particle = particle;
293 }
294}
295
296//------------------------------------------------------------------------------
297
298void TreeWriter::ProcessTowers(ExRootTreeBranch *branch, TObjArray *array)
299{
300 TIter iterator(array);
301 Candidate *candidate = 0;
302 Tower *entry = 0;
303 Double_t pt, signPz, cosTheta, eta, rapidity;
304
305 // loop over all towers
306 iterator.Reset();
307 while((candidate = static_cast<Candidate*>(iterator.Next())))
308 {
309 const TLorentzVector &momentum = candidate->Momentum;
310
311 pt = momentum.Pt();
312 cosTheta = TMath::Abs(momentum.CosTheta());
313 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
314 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
315 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
316
317 entry = static_cast<Tower*>(branch->NewEntry());
318
319 entry->SetBit(kIsReferenced);
320 entry->SetUniqueID(candidate->GetUniqueID());
321
322 entry->Eta = eta;
323 entry->Phi = momentum.Phi();
324 entry->ET = pt;
325 entry->E = momentum.E();
326 entry->Eem = candidate->Eem;
327 entry->Ehad = candidate->Ehad;
328 entry->Edges[0] = candidate->Edges[0];
329 entry->Edges[1] = candidate->Edges[1];
330 entry->Edges[2] = candidate->Edges[2];
331 entry->Edges[3] = candidate->Edges[3];
332
333 FillParticles(candidate, &entry->Particles);
334 }
335}
336
337//------------------------------------------------------------------------------
338
339void TreeWriter::ProcessPhotons(ExRootTreeBranch *branch, TObjArray *array)
340{
341 TIter iterator(array);
342 Candidate *candidate = 0;
343 Photon *entry = 0;
344 Double_t pt, signPz, cosTheta, eta, rapidity;
345
346 array->Sort();
347
348 // loop over all photons
349 iterator.Reset();
350 while((candidate = static_cast<Candidate*>(iterator.Next())))
351 {
352 TIter it1(candidate->GetCandidates());
353 const TLorentzVector &momentum = candidate->Momentum;
354
355 pt = momentum.Pt();
356 cosTheta = TMath::Abs(momentum.CosTheta());
357 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
358 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
359 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
360
361 entry = static_cast<Photon*>(branch->NewEntry());
362
363 entry->Eta = eta;
364 entry->Phi = momentum.Phi();
365 entry->PT = pt;
366 entry->E = momentum.E();
367
368 entry->EhadOverEem = candidate->Eem > 0.0 ? candidate->Ehad/candidate->Eem : 999.9;
369
370 FillParticles(candidate, &entry->Particles);
371 }
372}
373
374//------------------------------------------------------------------------------
375
376void TreeWriter::ProcessElectrons(ExRootTreeBranch *branch, TObjArray *array)
377{
378 TIter iterator(array);
379 Candidate *candidate = 0;
380 Electron *entry = 0;
381 Double_t pt, signPz, cosTheta, eta, rapidity;
382
383 array->Sort();
384
385 // loop over all electrons
386 iterator.Reset();
387 while((candidate = static_cast<Candidate*>(iterator.Next())))
388 {
389 const TLorentzVector &momentum = candidate->Momentum;
390
391 pt = momentum.Pt();
392 cosTheta = TMath::Abs(momentum.CosTheta());
393 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
394 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
395 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
396
397 entry = static_cast<Electron*>(branch->NewEntry());
398
399 entry->Eta = eta;
400 entry->Phi = momentum.Phi();
401 entry->PT = pt;
402
403 entry->Charge = candidate->Charge;
404
405 entry->EhadOverEem = 0.0;
406
407 entry->Particle = candidate->GetCandidates()->At(0);
408 }
409}
410
411//------------------------------------------------------------------------------
412
413void TreeWriter::ProcessMuons(ExRootTreeBranch *branch, TObjArray *array)
414{
415 TIter iterator(array);
416 Candidate *candidate = 0;
417 Muon *entry = 0;
418 Double_t pt, signPz, cosTheta, eta, rapidity;
419
420 array->Sort();
421
422 // loop over all muons
423 iterator.Reset();
424 while((candidate = static_cast<Candidate*>(iterator.Next())))
425 {
426 const TLorentzVector &momentum = candidate->Momentum;
427
428 pt = momentum.Pt();
429 cosTheta = TMath::Abs(momentum.CosTheta());
430 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
431 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
432 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
433
434 entry = static_cast<Muon*>(branch->NewEntry());
435
436 entry->SetBit(kIsReferenced);
437 entry->SetUniqueID(candidate->GetUniqueID());
438
439 entry->Eta = eta;
440 entry->Phi = momentum.Phi();
441 entry->PT = pt;
442
443 entry->Charge = candidate->Charge;
444
445 entry->Particle = candidate->GetCandidates()->At(0);
446 }
447}
448
449//------------------------------------------------------------------------------
450
451void TreeWriter::ProcessJets(ExRootTreeBranch *branch, TObjArray *array)
452{
453 TIter iterator(array);
454 Candidate *candidate = 0, *constituent = 0;
455 Jet *entry = 0;
456 Double_t pt, signPz, cosTheta, eta, rapidity;
457 Double_t ecalEnergy, hcalEnergy;
458
459 array->Sort();
460
461 // loop over all jets
462 iterator.Reset();
463 while((candidate = static_cast<Candidate*>(iterator.Next())))
464 {
465 TIter itConstituents(candidate->GetCandidates());
466 const TLorentzVector &momentum = candidate->Momentum;
467
468 pt = momentum.Pt();
469 cosTheta = TMath::Abs(momentum.CosTheta());
470 signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
471 eta = (cosTheta == 1.0 ? signPz*999.9 : momentum.Eta());
472 rapidity = (cosTheta == 1.0 ? signPz*999.9 : momentum.Rapidity());
473
474 entry = static_cast<Jet*>(branch->NewEntry());
475
476 entry->Eta = eta;
477 entry->Phi = momentum.Phi();
478 entry->PT = pt;
479
480 entry->Mass = momentum.M();
481
482 entry->DeltaEta = candidate->DeltaEta;
483 entry->DeltaPhi = candidate->DeltaPhi;
484
485 entry->BTag = candidate->BTag;
486 entry->TauTag = candidate->TauTag;
487
488 entry->Charge = candidate->Charge;
489
490 itConstituents.Reset();
491 entry->Constituents.Clear();
492 ecalEnergy = 0.0;
493 hcalEnergy = 0.0;
494 while((constituent = static_cast<Candidate*>(itConstituents.Next())))
495 {
496 entry->Constituents.Add(constituent);
497 ecalEnergy += constituent->Eem;
498 hcalEnergy += constituent->Ehad;
499 }
500
501 entry->EhadOverEem = ecalEnergy > 0.0 ? hcalEnergy/ecalEnergy : 999.9;
502
503 FillParticles(candidate, &entry->Particles);
504 }
505}
506
507//------------------------------------------------------------------------------
508
509void TreeWriter::ProcessMissingET(ExRootTreeBranch *branch, TObjArray *array)
510{
511 Candidate *candidate = 0;
512 MissingET *entry = 0;
513
514 // get the first entry
515 if((candidate = static_cast<Candidate*>(array->At(0))))
516 {
517 const TLorentzVector &momentum = candidate->Momentum;
518
519 entry = static_cast<MissingET*>(branch->NewEntry());
520
521 entry->Eta = (-momentum).Eta();
522 entry->Phi = (-momentum).Phi();
523 entry->MET = momentum.Pt();
524 }
525}
526
527//------------------------------------------------------------------------------
528
529void TreeWriter::ProcessScalarHT(ExRootTreeBranch *branch, TObjArray *array)
530{
531 Candidate *candidate = 0;
532 ScalarHT *entry = 0;
533
534 // get the first entry
535 if((candidate = static_cast<Candidate*>(array->At(0))))
536 {
537 const TLorentzVector &momentum = candidate->Momentum;
538
539 entry = static_cast<ScalarHT*>(branch->NewEntry());
540
541 entry->HT = momentum.Pt();
542 }
543}
544
545//------------------------------------------------------------------------------
546
547void TreeWriter::ProcessRho(ExRootTreeBranch *branch, TObjArray *array)
548{
549 TIter iterator(array);
550 Candidate *candidate = 0;
551 Rho *entry = 0;
552
553 // loop over all rho
554 iterator.Reset();
555 while((candidate = static_cast<Candidate*>(iterator.Next())))
556 {
557 const TLorentzVector &momentum = candidate->Momentum;
558
559 entry = static_cast<Rho*>(branch->NewEntry());
560
561 entry->Rho = momentum.E();
562 entry->Edges[0] = candidate->Edges[0];
563 entry->Edges[1] = candidate->Edges[1];
564 }
565}
566
567//------------------------------------------------------------------------------
568
569void TreeWriter::ProcessWeight(ExRootTreeBranch *branch, TObjArray *array)
570{
571 Candidate *candidate = 0;
572 Weight *entry = 0;
573
574 // get the first entry
575 if((candidate = static_cast<Candidate*>(array->At(0))))
576 {
577 const TLorentzVector &momentum = candidate->Momentum;
578
579 entry = static_cast<Weight*>(branch->NewEntry());
580
581 entry->Weight = momentum.E();
582 }
583}
584
585//------------------------------------------------------------------------------
586
587void TreeWriter::Process()
588{
589 TBranchMap::iterator itBranchMap;
590 ExRootTreeBranch *branch;
591 TProcessMethod method;
592 TObjArray *array;
593
594 for(itBranchMap = fBranchMap.begin(); itBranchMap != fBranchMap.end(); ++itBranchMap)
595 {
596 branch = itBranchMap->first;
597 method = itBranchMap->second.first;
598 array = itBranchMap->second.second;
599
600 (this->*method)(branch, array);
601 }
602}
603
604//------------------------------------------------------------------------------
Note: See TracBrowser for help on using the repository browser.