Now compatible with AliESDEvent. (A. Dainese)
[u/mrichter/AliRoot.git] / STEER / AliMCEventHandler.cxx
CommitLineData
5fe09262 1/**************************************************************************
2 * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3 * *
4 * Author: The ALICE Off-line Project. *
5 * Contributors are mentioned in the code where appropriate. *
6 * *
7 * Permission to use, copy, modify and distribute this software and its *
8 * documentation strictly for non-commercial purposes is hereby granted *
9 * without fee, provided that the above copyright notice appears in all *
10 * copies and that both the copyright notice and this permission notice *
11 * appear in the supporting documentation. The authors make no claims *
12 * about the suitability of this software for any purpose. It is *
13 * provided "as is" without express or implied warranty. *
14 **************************************************************************/
15
16/* $Id$ */
17//---------------------------------------------------------------------------------
18// Class AliMCEventHandler
19// This class gives access to MC truth during the analysis.
20// Monte Carlo truth is containe in the kinematics tree (produced particles) and
21// the tree of reference hits.
22//
23// Origin: Andreas Morsch, CERN, andreas.morsch@cern.ch
24//---------------------------------------------------------------------------------
25
26
27
28#include "AliMCEventHandler.h"
29#include "AliTrackReference.h"
30#include "AliHeader.h"
31#include "AliStack.h"
2c36081e 32#include "AliLog.h"
5fe09262 33
34#include <TTree.h>
35#include <TFile.h>
36#include <TParticle.h>
37#include <TClonesArray.h>
38#include <TDirectoryFile.h>
39#include <TArrow.h>
40#include <TMarker.h>
41#include <TH2F.h>
42
43
44ClassImp(AliMCEventHandler)
45
46AliMCEventHandler::AliMCEventHandler() :
d2f1d9ef 47 AliVEventHandler(),
5fe09262 48 fFileE(0),
49 fFileK(0),
50 fFileTR(0),
51 fTmpFileTR(0),
52 fTreeE(0),
53 fTreeK(0),
54 fTreeTR(0),
55 fTmpTreeTR(0),
56 fStack(0),
57 fHeader(0),
58 fTrackReferences(0),
59 fNEvent(-1),
60 fEvent(-1),
61 fNprimaries(-1),
62 fNparticles(-1),
9aea8469 63 fPathName("./"),
64 fExtension(""),
65 fFileNumber(0),
66 fEventsPerFile(0)
5fe09262 67{
68 // Default constructor
69}
70
71AliMCEventHandler::AliMCEventHandler(const char* name, const char* title) :
d2f1d9ef 72 AliVEventHandler(name, title),
5fe09262 73 fFileE(0),
74 fFileK(0),
75 fFileTR(0),
76 fTmpFileTR(0),
77 fTreeE(0),
78 fTreeK(0),
79 fTreeTR(0),
80 fTmpTreeTR(0),
81 fStack(0),
82 fHeader(new AliHeader()),
83 fTrackReferences(new TClonesArray("AliTrackReference", 200)),
84 fNEvent(-1),
85 fEvent(-1),
86 fNprimaries(-1),
87 fNparticles(-1),
9aea8469 88 fPathName("./"),
89 fExtension(""),
90 fFileNumber(0),
91 fEventsPerFile(0)
5fe09262 92{
93 // Constructor
94}
95AliMCEventHandler::~AliMCEventHandler()
96{
97 // Destructor
98 delete fFileE;
99 delete fFileK;
100 delete fFileTR;
101}
102
103Bool_t AliMCEventHandler::InitIO(Option_t* /*opt*/)
104{
105 // Initialize input
106 //
0c476ce8 107 fFileE = TFile::Open(Form("%sgalice.root", fPathName));
2c36081e 108 if (!fFileE) AliFatal(Form("AliMCEventHandler:galice.root not found in directory %s ! \n", fPathName));
5fe09262 109
110 fFileE->GetObject("TE", fTreeE);
111 fTreeE->SetBranchAddress("Header", &fHeader);
112 fNEvent = fTreeE->GetEntries();
113 //
114 // Tree K
0c476ce8 115 fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName, fExtension));
2c36081e 116 if (!fFileK) AliFatal(Form("AliMCEventHandler:Kinematics.root not found in directory %s ! \n", fPathName));
9aea8469 117 fEventsPerFile = fFileK->GetNkeys() - fFileK->GetNProcessIDs();
5fe09262 118 //
119 // Tree TR
0c476ce8 120 fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName, fExtension));
2c36081e 121 if (!fFileTR) AliWarning(Form("AliMCEventHandler:TrackRefs.root not found in directory %s ! \n", fPathName));
5fe09262 122 //
123 // Reset the event number
2c36081e 124 fEvent = -1;
125 fFileNumber = 0;
9aea8469 126
2c36081e 127 AliInfo(Form("AliMCEventHandler:Number of events in this directory %5d \n", fNEvent));
5fe09262 128 return kTRUE;
129
130}
131
9aea8469 132Bool_t AliMCEventHandler::GetEvent(Int_t iev)
133{
134 // Load the event number iev
2c36081e 135 //
136 // Calculate the file number
9aea8469 137 Int_t inew = iev/fEventsPerFile;
138 if (inew != fFileNumber) {
139 fFileNumber = inew;
140 if (!OpenFile(fFileNumber)){
141 return kFALSE;
142 }
143 }
2c36081e 144 // Folder name
5fe09262 145 char folder[20];
9aea8469 146 sprintf(folder, "Event%d", iev);
2c36081e 147
5fe09262 148 // TreeE
9aea8469 149 fTreeE->GetEntry(iev);
5fe09262 150 fStack = fHeader->Stack();
151 // Tree K
152 TDirectoryFile* dirK = 0;
153 fFileK->GetObject(folder, dirK);
9aea8469 154 if (!dirK) {
2c36081e 155 AliWarning(Form("AliMCEventHandler: Event #%5d not found\n", iev));
9aea8469 156 return kFALSE;
157 }
5fe09262 158 dirK->GetObject("TreeK", fTreeK);
159 fStack->ConnectTree(fTreeK);
160 fStack->GetEvent();
2d8f26f6 161
5fe09262 162 //Tree TR
2d8f26f6 163
07bd4750 164 if (fFileTR) {
2c36081e 165 TDirectoryFile* dirTR = 0;
166 fFileTR->GetObject(folder, dirTR);
167 dirTR->GetObject("TreeTR", fTreeTR);
168 if (fTreeTR->GetBranch("AliRun")) {
169 // This is an old format with one branch per detector not in synch with TreeK
170 ReorderAndExpandTreeTR();
2d8f26f6 171 delete dirTR;
2c36081e 172 } else {
173 // New format
174 fTreeTR->SetBranchAddress("TrackReferences", &fTrackReferences);
175 }
5fe09262 176 }
2d8f26f6 177
5fe09262 178 //
179 fNparticles = fStack->GetNtrack();
180 fNprimaries = fStack->GetNprimary();
2d8f26f6 181 AliInfo(Form("AliMCEventHandler: Event#: %5d Number of particles: %5d (all) %5d (primaries)\n",
182 fEvent, fNparticles, fNprimaries));
5fe09262 183
184 return kTRUE;
9aea8469 185}
186
187Bool_t AliMCEventHandler::OpenFile(Int_t i)
188{
189 // Open file i
190 Bool_t ok = kTRUE;
191 if (i > 0) {
192 fExtension = Form("%d", i);
193 } else {
194 fExtension = "";
195 }
196
197
198 delete fFileK;
3b7d72ed 199 fFileK = TFile::Open(Form("%sKinematics%s.root", fPathName, fExtension));
9aea8469 200 if (!fFileK) {
2c36081e 201 AliFatal(Form("AliMCEventHandler:Kinematics%s.root not found in directory %s ! \n", fExtension, fPathName));
9aea8469 202 ok = kFALSE;
203 }
204
205 delete fFileTR;
3b7d72ed 206 fFileTR = TFile::Open(Form("%sTrackRefs%s.root", fPathName, fExtension));
9aea8469 207 if (!fFileTR) {
2c36081e 208 AliWarning(Form("AliMCEventHandler:TrackRefs%s.root not found in directory %s ! \n", fExtension, fPathName));
9aea8469 209 ok = kFALSE;
210 }
211
212 return ok;
213}
214
215Bool_t AliMCEventHandler::BeginEvent()
216{
217 // Read the next event
218 fEvent++;
219 if (fEvent >= fNEvent) {
2c36081e 220 AliWarning(Form("AliMCEventHandler: Event number out of range %5d\n", fEvent));
9aea8469 221 return kFALSE;
222 }
223 return GetEvent(fEvent);
5fe09262 224}
225
226Int_t AliMCEventHandler::GetParticleAndTR(Int_t i, TParticle*& particle, TClonesArray*& trefs)
227{
228 // Retrieve entry i
2d8f26f6 229 if (i < 0 || i >= fNparticles) {
2c36081e 230 AliWarning(Form("AliMCEventHandler::GetEntry: Index out of range"));
231 particle = 0;
232 trefs = 0;
233 return (-1);
234 }
235 particle = fStack->Particle(i);
236 if (fFileTR) {
5fe09262 237 fTreeTR->GetEntry(fStack->TreeKEntry(i));
2c36081e 238 trefs = fTrackReferences;
239 return trefs->GetEntries();
5fe09262 240 } else {
2c36081e 241 trefs = 0;
242 return -1;
5fe09262 243 }
5fe09262 244}
245
2d8f26f6 246void AliMCEventHandler::DrawCheck(Int_t i, Int_t search)
5fe09262 247{
248 // Retrieve entry i and draw momentum vector and hits
2c36081e 249 if (!fFileTR) {
250 AliWarning("No Track Reference information available");
251 return;
252 }
253
5fe09262 254 if (i > -1 && i < fNparticles) {
255 fTreeTR->GetEntry(fStack->TreeKEntry(i));
256 } else {
2c36081e 257 AliWarning("AliMCEventHandler::GetEntry: Index out of range");
5fe09262 258 }
2c36081e 259
36e4cfbb 260 Int_t nh = fTrackReferences->GetEntries();
261
5fe09262 262
36e4cfbb 263 if (search) {
2d8f26f6 264 while(nh <= search && i < fNparticles - 1) {
36e4cfbb 265 i++;
266 fTreeTR->GetEntry(fStack->TreeKEntry(i));
267 nh = fTrackReferences->GetEntries();
268 }
269 printf("Found Hits at %5d\n", i);
270 }
271 TParticle* particle = fStack->Particle(i);
272
5fe09262 273 TH2F* h = new TH2F("", "", 100, -500, 500, 100, -500, 500);
274 Float_t x0 = particle->Vx();
275 Float_t y0 = particle->Vy();
276
277 Float_t x1 = particle->Vx() + particle->Px() * 50.;
278 Float_t y1 = particle->Vy() + particle->Py() * 50.;
279
280 TArrow* a = new TArrow(x0, y0, x1, y1, 0.01);
281 h->Draw();
282 a->SetLineColor(2);
283
284 a->Draw();
285
286 for (Int_t ih = 0; ih < nh; ih++) {
287 AliTrackReference* ref = (AliTrackReference*) fTrackReferences->At(ih);
288 TMarker* m = new TMarker(ref->X(), ref->Y(), 20);
289 m->Draw();
290 m->SetMarkerSize(0.4);
291
292 }
293}
294
890126ab 295Bool_t AliMCEventHandler::Notify(const char *path)
5fe09262 296{
297 // Notify about directory change
890126ab 298 // The directory is taken from the 'path' argument
36e4cfbb 299 // Reconnect trees
890126ab 300
301 printf("AliMCEventHandler::Notify() file: %s\n", path);
302 fPathName = Form("%s", path);
5fe09262 303 ResetIO();
304 InitIO("");
305 return kTRUE;
306}
307
308void AliMCEventHandler::ResetIO()
309{
36e4cfbb 310 // Reset files
5fe09262 311 if (fFileE) delete fFileE;
312 if (fFileK) delete fFileK;
313 if (fFileTR) delete fFileTR;
314}
315
316
317Bool_t AliMCEventHandler::FinishEvent()
318{
2d8f26f6 319 // Reset the stack
2d8f26f6 320 Stack()->Reset();
321
5fe09262 322 return kTRUE;
323}
324
325Bool_t AliMCEventHandler::Terminate()
326{
327 // Dummy
328 return kTRUE;
329}
330
331Bool_t AliMCEventHandler::TerminateIO()
332{
333 // Dummy
334 return kTRUE;
335}
336
337void AliMCEventHandler::ReorderAndExpandTreeTR()
338{
339//
340// Reorder and expand the track reference tree in order to match the kinematics tree.
341// Copy the information from different branches into one
342//
343// TreeTR
2d8f26f6 344 if (fTmpTreeTR) {
345 fTmpTreeTR->Delete("all");
346 }
347
5fe09262 348 if (fTmpFileTR) {
349 fTmpFileTR->Close();
350 delete fTmpFileTR;
351 }
352
353 fTmpFileTR = new TFile("TrackRefsTmp.root", "recreate");
354 fTmpTreeTR = new TTree("TreeTR", "Track References");
355 if (!fTrackReferences) fTrackReferences = new TClonesArray("AliTrackReference", 100);
356 fTmpTreeTR->Branch("TrackReferences", "TClonesArray", &fTrackReferences, 4000);
2d8f26f6 357
5fe09262 358//
808ba8bf 359 fTreeTR->SetBranchStatus("*", 0);
360 fTreeTR->SetBranchStatus("AliRun", 1);
361 fTreeTR->SetBranchStatus("ITS", 1);
362 fTreeTR->SetBranchStatus("TPC", 1);
363 fTreeTR->SetBranchStatus("TRD", 1);
364 fTreeTR->SetBranchStatus("TOF", 1);
365 fTreeTR->SetBranchStatus("FRAME", 1);
366 fTreeTR->SetBranchStatus("MUON", 1);
367
5fe09262 368 TClonesArray* trefs[7];
369 for (Int_t i = 0; i < 7; i++) trefs[i] = 0;
370 if (fTreeTR){
371 // make branch for central track references
372 if (fTreeTR->GetBranch("AliRun")) fTreeTR->SetBranchAddress("AliRun", &trefs[0]);
373 if (fTreeTR->GetBranch("ITS")) fTreeTR->SetBranchAddress("ITS", &trefs[1]);
374 if (fTreeTR->GetBranch("TPC")) fTreeTR->SetBranchAddress("TPC", &trefs[2]);
375 if (fTreeTR->GetBranch("TRD")) fTreeTR->SetBranchAddress("TRD", &trefs[3]);
376 if (fTreeTR->GetBranch("TOF")) fTreeTR->SetBranchAddress("TOF", &trefs[4]);
377 if (fTreeTR->GetBranch("FRAME")) fTreeTR->SetBranchAddress("FRAME", &trefs[5]);
378 if (fTreeTR->GetBranch("MUON")) fTreeTR->SetBranchAddress("MUON", &trefs[6]);
379 }
380
381 Int_t np = fStack->GetNprimary();
382 Int_t nt = fTreeTR->GetEntries();
383 //
384 // Loop over tracks and find the secondaries with the help of the kine tree
385 Int_t ifills = 0;
386 Int_t it = 0;
387 Int_t itlast = 0;
388
389 for (Int_t ip = np - 1; ip > -1; ip--) {
390 TParticle *part = fStack->Particle(ip);
391// printf("Particle %5d %5d %5d %5d %5d %5d \n",
392// ip, part->GetPdgCode(), part->GetFirstMother(), part->GetFirstDaughter(),
393// part->GetLastDaughter(), part->TestBit(kTransportBit));
36e4cfbb 394
395 // Determine range of secondaries produced by this primary during transport
5fe09262 396 Int_t dau1 = part->GetFirstDaughter();
36e4cfbb 397 if (dau1 < np) continue; // This particle has no secondaries produced during transport
5fe09262 398 Int_t dau2 = -1;
5fe09262 399 if (dau1 > -1) {
400 Int_t inext = ip - 1;
401 while (dau2 < 0) {
402 if (inext >= 0) {
403 part = fStack->Particle(inext);
404 dau2 = part->GetFirstDaughter();
405 if (dau2 == -1 || dau2 < np) {
406 dau2 = -1;
407 } else {
408 dau2--;
409 }
410 } else {
411 dau2 = fStack->GetNtrack() - 1;
412 }
413 inext--;
414 } // find upper bound
415 } // dau2 < 0
36e4cfbb 416
5fe09262 417
418// printf("Check (1) %5d %5d %5d %5d %5d \n", ip, np, it, dau1, dau2);
36e4cfbb 419//
420// Loop over reference hits and find secondary label
421// First the tricky part: find the entry in treeTR than contains the hits or
422// make sure that no hits exist.
423//
5fe09262 424 Bool_t hasHits = kFALSE;
425 Bool_t isOutside = kFALSE;
36e4cfbb 426
5fe09262 427 it = itlast;
5fe09262 428 while (!hasHits && !isOutside && it < nt) {
429 fTreeTR->GetEntry(it++);
430 for (Int_t ib = 0; ib < 7; ib++) {
431 if (!trefs[ib]) continue;
432 Int_t nh = trefs[ib]->GetEntries();
433 for (Int_t ih = 0; ih < nh; ih++) {
434 AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
435 Int_t label = tr->Label();
436 if (label >= dau1 && label <= dau2) {
437 hasHits = kTRUE;
438 itlast = it - 1;
36e4cfbb 439 break;
5fe09262 440 }
441 if (label > dau2 || label < ip) {
442 isOutside = kTRUE;
443 itlast = it - 1;
36e4cfbb 444 break;
5fe09262 445 }
446 } // hits
36e4cfbb 447 if (hasHits || isOutside) break;
5fe09262 448 } // branches
449 } // entries
450
451 if (!hasHits) {
36e4cfbb 452 // Write empty entries
5fe09262 453 for (Int_t id = dau1; (id <= dau2); id++) {
454 fTmpTreeTR->Fill();
455 ifills++;
456 }
457 } else {
36e4cfbb 458 // Collect all hits
5fe09262 459 fTreeTR->GetEntry(itlast);
460 for (Int_t id = dau1; (id <= dau2) && (dau1 > -1); id++) {
461 for (Int_t ib = 0; ib < 7; ib++) {
462 if (!trefs[ib]) continue;
463 Int_t nh = trefs[ib]->GetEntries();
464 for (Int_t ih = 0; ih < nh; ih++) {
465 AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
466 Int_t label = tr->Label();
467 // Skip primaries
468 if (label == ip) continue;
469 if (label > dau2 || label < dau1)
36e4cfbb 470 printf("AliMCEventHandler::Track Reference Label out of range !: %5d %5d %5d %5d \n",
471 itlast, label, dau1, dau2);
5fe09262 472 if (label == id) {
473 // secondary found
474 tr->SetDetectorId(ib-1);
475 Int_t nref = fTrackReferences->GetEntriesFast();
476 TClonesArray &lref = *fTrackReferences;
477 new(lref[nref]) AliTrackReference(*tr);
478 }
479 } // hits
480 } // branches
481 fTmpTreeTR->Fill();
482 fTrackReferences->Clear();
483 ifills++;
484 } // daughters
485 } // has hits
486 } // tracks
5fe09262 487 //
488 // Now loop again and write the primaries
36e4cfbb 489 //
5fe09262 490 it = nt - 1;
491 for (Int_t ip = 0; ip < np; ip++) {
492 Int_t labmax = -1;
493 while (labmax < ip && it > -1) {
494 fTreeTR->GetEntry(it--);
495 for (Int_t ib = 0; ib < 7; ib++) {
496 if (!trefs[ib]) continue;
497 Int_t nh = trefs[ib]->GetEntries();
498 //
499 // Loop over reference hits and find primary labels
500 for (Int_t ih = 0; ih < nh; ih++) {
501 AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
502 Int_t label = tr->Label();
503 if (label < np && label > labmax) {
504 labmax = label;
505 }
506
507 if (label == ip) {
508 tr->SetDetectorId(ib-1);
509 Int_t nref = fTrackReferences->GetEntriesFast();
510 TClonesArray &lref = *fTrackReferences;
511 new(lref[nref]) AliTrackReference(*tr);
512 }
513 } // hits
514 } // branches
515 } // entries
516 it++;
517 fTmpTreeTR->Fill();
518 fTrackReferences->Clear();
519 ifills++;
520 } // tracks
521 // Check
808ba8bf 522 delete fTreeTR;
523 for (Int_t ib = 0; ib < 7; ib++) {
524 if (trefs[ib]) delete trefs[ib];
525 }
526
5fe09262 527 if (ifills != fStack->GetNtrack())
36e4cfbb 528 printf("AliMCEventHandler:Number of entries in TreeTR (%5d) unequal to TreeK (%5d) \n",
529 ifills, fStack->GetNtrack());
5fe09262 530 fTreeTR = fTmpTreeTR;
5fe09262 531}