d22440c421b969d6265385541782891c97b365fe
[u/mrichter/AliRoot.git] / STEER / AliMCEvent.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-2007, 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 //-------------------------------------------------------------------------
19 //     Class for Kinematic Events
20 //     Author: Andreas Morsch, CERN
21 //-------------------------------------------------------------------------
22 #include <TArrow.h>
23 #include <TMarker.h>
24 #include <TH2F.h>
25 #include <TTree.h>
26 #include <TFile.h>
27 #include <TParticle.h>
28 #include <TClonesArray.h>
29 #include <TRefArray.h>
30
31 #include "AliLog.h"
32 #include "AliMCEvent.h"
33 #include "AliStack.h"
34 #include "AliTrackReference.h"
35 #include "AliHeader.h"
36 #include "AliGenEventHeader.h"
37
38
39 AliMCEvent::AliMCEvent():
40     AliVEvent(),
41     fStack(0),
42     fMCParticles(new TClonesArray("AliMCParticle",1000)),
43     fMCParticleMap(0),
44     fHeader(0),
45     fTRBuffer(0),
46     fTrackReferences(new TClonesArray("AliTrackReference", 1000)),
47     fTreeTR(0),
48     fTmpTreeTR(0),
49     fTmpFileTR(0),
50     fNprimaries(-1),
51     fNparticles(-1)
52 {
53     // Default constructor
54 }
55
56 AliMCEvent::AliMCEvent(const AliMCEvent& mcEvnt) :
57     AliVEvent(mcEvnt),
58     fStack(0),
59     fMCParticles(0),
60     fMCParticleMap(0),
61     fHeader(0),
62     fTRBuffer(0),
63     fTrackReferences(0),
64     fTreeTR(0),
65     fTmpTreeTR(0),
66     fTmpFileTR(0),
67     fNprimaries(-1),
68     fNparticles(-1) 
69
70 // Copy constructor
71 }
72
73
74 AliMCEvent& AliMCEvent::operator=(const AliMCEvent& mcEvnt)
75 {
76     // assignment operator
77     if (this!=&mcEvnt) { 
78         AliVEvent::operator=(mcEvnt); 
79     }
80   
81     return *this; 
82 }
83
84 void AliMCEvent::ConnectTreeE (TTree* tree)
85 {
86     // Connect the event header tree
87     tree->SetBranchAddress("Header", &fHeader);
88 }
89
90 void AliMCEvent::ConnectTreeK (TTree* tree)
91 {
92     // Connect the kinematics tree to the stack
93     fStack = fHeader->Stack();
94     fStack->ConnectTree(tree);
95     //
96     // Load the event
97     fStack->GetEvent();
98     fNparticles = fStack->GetNtrack();
99     fNprimaries = fStack->GetNprimary();
100     AliInfo(Form("AliMCEvent: Number of particles: %5d (all) %5d (primaries)\n", 
101                  fNparticles, fNprimaries));
102  
103     // This is a cache for the TParticles converted to MCParticles on user request
104     if (fMCParticleMap) {
105         fMCParticleMap->Clear();
106         if (fNparticles>0) fMCParticleMap->Expand(fNparticles);}
107     else
108         fMCParticleMap = new TRefArray(fNparticles);
109
110 }
111
112 void AliMCEvent::ConnectTreeTR (TTree* tree)
113 {
114     // Connect the track reference tree
115     fTreeTR = tree;
116     
117     if (fTreeTR->GetBranch("AliRun")) {
118         if (fTmpFileTR) {
119             fTmpFileTR->Close();
120             delete fTmpFileTR;
121         }
122         // This is an old format with one branch per detector not in synch with TreeK
123         ReorderAndExpandTreeTR();
124     } else {
125         // New format 
126         fTreeTR->SetBranchAddress("TrackReferences", &fTRBuffer);
127     }
128 }
129
130 Int_t AliMCEvent::GetParticleAndTR(Int_t i, TParticle*& particle, TClonesArray*& trefs)
131 {
132     // Retrieve entry i
133     if (i < 0 || i >= fNparticles) {
134         AliWarning(Form("AliMCEventHandler::GetEntry: Index out of range"));
135         particle = 0;
136         trefs    = 0;
137         return (-1);
138     }
139     particle = fStack->Particle(i);
140     if (fTreeTR) {
141         fTreeTR->GetEntry(fStack->TreeKEntry(i));
142         trefs    = fTRBuffer;
143         return trefs->GetEntries();
144     } else {
145         trefs = 0;
146         return -1;
147     }
148 }
149
150
151 void AliMCEvent::Clean()
152 {
153     // Clean-up before new trees are connected
154     
155     if (fHeader) {
156         delete fHeader;
157         fHeader = 0;
158     }
159     
160     delete fStack;
161
162     // Clear TR
163     if (fTRBuffer) {
164         fTRBuffer->Clear();
165         delete fTRBuffer;
166         fTRBuffer = 0;
167     }
168 }
169
170 void AliMCEvent::FinishEvent()
171 {
172     // Clean-up after event
173      fStack->Reset(0);
174      fMCParticles->Clear();
175      fTrackReferences->Clear();
176 }
177
178
179
180 void AliMCEvent::DrawCheck(Int_t i, Int_t search)
181 {
182     //
183     // Simple event display for debugging
184     if (!fTreeTR) {
185         AliWarning("No Track Reference information available");
186         return;
187     } 
188     
189     if (i > -1 && i < fNparticles) {
190         fTreeTR->GetEntry(fStack->TreeKEntry(i));
191     } else {
192         AliWarning("AliMCEvent::GetEntry: Index out of range");
193     }
194     
195     Int_t nh = fTRBuffer->GetEntries();
196     
197     
198     if (search) {
199         while(nh <= search && i < fNparticles - 1) {
200             i++;
201             fTreeTR->GetEntry(fStack->TreeKEntry(i));
202             nh =  fTRBuffer->GetEntries();
203         }
204         printf("Found Hits at %5d\n", i);
205     }
206     TParticle* particle = fStack->Particle(i);
207     
208     TH2F*    h = new TH2F("", "", 100, -500, 500, 100, -500, 500);
209     Float_t x0 = particle->Vx();
210     Float_t y0 = particle->Vy();
211
212     Float_t x1 = particle->Vx() + particle->Px() * 50.;
213     Float_t y1 = particle->Vy() + particle->Py() * 50.;
214     
215     TArrow*  a = new TArrow(x0, y0, x1, y1, 0.01);
216     h->Draw();
217     a->SetLineColor(2);
218     
219     a->Draw();
220     
221     for (Int_t ih = 0; ih < nh; ih++) {
222         AliTrackReference* ref = (AliTrackReference*) fTRBuffer->At(ih);
223         TMarker* m = new TMarker(ref->X(), ref->Y(), 20);
224         m->Draw();
225         m->SetMarkerSize(0.4);
226         
227     }
228 }
229
230
231 void AliMCEvent::ReorderAndExpandTreeTR()
232 {
233 //
234 //  Reorder and expand the track reference tree in order to match the kinematics tree.
235 //  Copy the information from different branches into one
236 //
237 //  TreeTR
238
239     fTmpFileTR = new TFile("TrackRefsTmp.root", "recreate");
240     fTmpTreeTR = new TTree("TreeTR", "TrackReferences");
241     if (!fTRBuffer)  fTRBuffer = new TClonesArray("AliTrackReference", 100);
242     fTmpTreeTR->Branch("TrackReferences", "TClonesArray", &fTRBuffer, 32000, 0);
243     
244
245 //
246 //  Activate the used branches only. Otherwisw we get a bad memory leak.
247     fTreeTR->SetBranchStatus("*",        0);
248     fTreeTR->SetBranchStatus("AliRun.*", 1);
249     fTreeTR->SetBranchStatus("ITS.*",    1);
250     fTreeTR->SetBranchStatus("TPC.*",    1);
251     fTreeTR->SetBranchStatus("TRD.*",    1);
252     fTreeTR->SetBranchStatus("TOF.*",    1);
253     fTreeTR->SetBranchStatus("FRAME.*",  1);
254     fTreeTR->SetBranchStatus("MUON.*",   1);
255 //
256 //  Connect the active branches
257     TClonesArray* trefs[7];
258     for (Int_t i = 0; i < 7; i++) trefs[i] = 0;
259     if (fTreeTR){
260         // make branch for central track references
261         if (fTreeTR->GetBranch("AliRun")) fTreeTR->SetBranchAddress("AliRun", &trefs[0]);
262         if (fTreeTR->GetBranch("ITS"))    fTreeTR->SetBranchAddress("ITS",    &trefs[1]);
263         if (fTreeTR->GetBranch("TPC"))    fTreeTR->SetBranchAddress("TPC",    &trefs[2]);
264         if (fTreeTR->GetBranch("TRD"))    fTreeTR->SetBranchAddress("TRD",    &trefs[3]);
265         if (fTreeTR->GetBranch("TOF"))    fTreeTR->SetBranchAddress("TOF",    &trefs[4]);
266         if (fTreeTR->GetBranch("FRAME"))  fTreeTR->SetBranchAddress("FRAME",  &trefs[5]);
267         if (fTreeTR->GetBranch("MUON"))   fTreeTR->SetBranchAddress("MUON",   &trefs[6]);
268     }
269
270     Int_t np = fStack->GetNprimary();
271     Int_t nt = fTreeTR->GetEntries();
272     
273     //
274     // Loop over tracks and find the secondaries with the help of the kine tree
275     Int_t ifills = 0;
276     Int_t it     = 0;
277     Int_t itlast = 0;
278     TParticle* part;
279
280     for (Int_t ip = np - 1; ip > -1; ip--) {
281         part = fStack->Particle(ip);
282 //      printf("Particle %5d %5d %5d %5d %5d %5d \n", 
283 //             ip, part->GetPdgCode(), part->GetFirstMother(), part->GetFirstDaughter(), 
284 //             part->GetLastDaughter(), part->TestBit(kTransportBit));
285
286         // Determine range of secondaries produced by this primary during transport     
287         Int_t dau1  = part->GetFirstDaughter();
288         if (dau1 < np) continue;  // This particle has no secondaries produced during transport
289         Int_t dau2  = -1;
290         if (dau1 > -1) {
291             Int_t inext = ip - 1;
292             while (dau2 < 0) {
293                 if (inext >= 0) {
294                     part = fStack->Particle(inext);
295                     dau2 =  part->GetFirstDaughter();
296                     if (dau2 == -1 || dau2 < np) {
297                         dau2 = -1;
298                     } else {
299                         dau2--;
300                     }
301                 } else {
302                     dau2 = fStack->GetNtrack() - 1;
303                 }
304                 inext--;
305             } // find upper bound
306         }  // dau2 < 0
307         
308
309 //      printf("Check (1) %5d %5d %5d %5d %5d \n", ip, np, it, dau1, dau2);
310 //
311 // Loop over reference hits and find secondary label
312 // First the tricky part: find the entry in treeTR than contains the hits or
313 // make sure that no hits exist.
314 //
315         Bool_t hasHits   = kFALSE;
316         Bool_t isOutside = kFALSE;
317
318         it = itlast;
319         while (!hasHits && !isOutside && it < nt) {
320             fTreeTR->GetEntry(it++);
321             for (Int_t ib = 0; ib < 7; ib++) {
322                 if (!trefs[ib]) continue;
323                 Int_t nh = trefs[ib]->GetEntries();
324                 for (Int_t ih = 0; ih < nh; ih++) {
325                     AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
326                     Int_t label = tr->Label();
327                     if (label >= dau1 && label <= dau2) {
328                         hasHits = kTRUE;
329                         itlast = it - 1;
330                         break;
331                     }
332                     if (label > dau2 || label < ip) {
333                         isOutside = kTRUE;
334                         itlast = it - 1;
335                         break;
336                     }
337                 } // hits
338                 if (hasHits || isOutside) break;
339             } // branches
340         } // entries
341
342         if (!hasHits) {
343             // Write empty entries
344             for (Int_t id = dau1; (id <= dau2); id++) {
345                 fTmpTreeTR->Fill();
346                 ifills++;
347             } 
348         } else {
349             // Collect all hits
350             fTreeTR->GetEntry(itlast);
351             for (Int_t id = dau1; (id <= dau2) && (dau1 > -1); id++) {
352                 for (Int_t ib = 0; ib < 7; ib++) {
353                     if (!trefs[ib]) continue;
354                     Int_t nh = trefs[ib]->GetEntries();
355                     for (Int_t ih = 0; ih < nh; ih++) {
356                         AliTrackReference* tr = (AliTrackReference*) trefs[ib]->At(ih);
357                         Int_t label = tr->Label();
358                         // Skip primaries
359                         if (label == ip) continue;
360                         if (label > dau2 || label < dau1) 
361                             printf("AliMCEventHandler::Track Reference Label out of range !: %5d %5d %5d %5d \n", 
362                                    itlast, label, dau1, dau2);
363                         if (label == id) {
364                             // secondary found
365                             tr->SetDetectorId(ib-1);
366                             Int_t nref =  fTRBuffer->GetEntriesFast();
367                             TClonesArray &lref = *fTRBuffer;
368                             new(lref[nref]) AliTrackReference(*tr);
369                         }
370                     } // hits
371                 } // branches
372                 fTmpTreeTR->Fill();
373                 fTRBuffer->Clear();
374                 ifills++;
375             } // daughters
376         } // has hits
377     } // tracks
378
379     //
380     // Now loop again and write the primaries
381     //
382     it = nt - 1;
383     for (Int_t ip = 0; ip < np; ip++) {
384         Int_t labmax = -1;
385         while (labmax < ip && it > -1) {
386             fTreeTR->GetEntry(it--);
387             for (Int_t ib = 0; ib < 7; ib++) {
388                 if (!trefs[ib]) continue;
389                 Int_t nh = trefs[ib]->GetEntries();
390                 // 
391                 // Loop over reference hits and find primary labels
392                 for (Int_t ih = 0; ih < nh; ih++) {
393                     AliTrackReference* tr = (AliTrackReference*)  trefs[ib]->At(ih);
394                     Int_t label = tr->Label();
395                     if (label < np && label > labmax) {
396                         labmax = label;
397                     }
398                     
399                     if (label == ip) {
400                         tr->SetDetectorId(ib-1);
401                         Int_t nref = fTRBuffer->GetEntriesFast();
402                         TClonesArray &lref = *fTRBuffer;
403                         new(lref[nref]) AliTrackReference(*tr);
404                     }
405                 } // hits
406             } // branches
407         } // entries
408         it++;
409         fTmpTreeTR->Fill();
410         fTRBuffer->Clear();
411         ifills++;
412     } // tracks
413     // Check
414
415
416     // Clean-up
417     delete fTreeTR; fTreeTR = 0;
418     
419     for (Int_t ib = 0; ib < 7; ib++) {
420         if (trefs[ib]) {
421             trefs[ib]->Clear();
422             delete trefs[ib];
423             trefs[ib] = 0;
424         }
425     }
426
427     if (ifills != fStack->GetNtrack()) 
428         printf("AliMCEvent:Number of entries in TreeTR (%5d) unequal to TreeK (%5d) \n", 
429                ifills, fStack->GetNtrack());
430
431     fTmpTreeTR->Write();
432     fTreeTR = fTmpTreeTR;
433 }
434
435 AliMCParticle* AliMCEvent::GetTrack(Int_t i) const
436 {
437     // Get MC Particle i
438     AliMCParticle *mcParticle = 0;
439     TParticle     *particle   = 0;
440     TClonesArray  *trefs      = 0;
441     Int_t          ntref      = 0;
442     TRefArray     *rarray     = 0;
443     
444     // Out of range check
445     if (i < 0 || i >= fNparticles) {
446         AliWarning(Form("AliMCEvent::GetEntry: Index out of range"));
447         mcParticle = 0;
448         return (mcParticle);
449     }
450     
451
452
453     //
454     // First check of the MC Particle has been already cached
455     if(!fMCParticleMap->At(i)) {
456         // Get particle from the stack
457         particle   = fStack->Particle(i);
458         // Get track references from Tree TR
459         if (fTreeTR) {
460             fTreeTR->GetEntry(fStack->TreeKEntry(i));
461             trefs     = fTRBuffer;
462             ntref     = trefs->GetEntriesFast();
463             rarray    = new TRefArray(ntref);
464             Int_t nen = fTrackReferences->GetEntriesFast();
465             for (Int_t j = 0; j < ntref; j++) {
466                 // Save the track references in a TClonesArray
467                 AliTrackReference* ref = dynamic_cast<AliTrackReference*>((*fTRBuffer)[j]);
468                 // Save the pointer in a TRefArray
469                 new ((*fTrackReferences)[nen]) AliTrackReference(*ref);
470                 rarray->AddAt((*fTrackReferences)[nen], j);
471                 nen++;
472             } // loop over track references for entry i
473         } // if TreeTR available
474
475         Int_t nentries = fMCParticles->GetEntriesFast();
476         new ((*fMCParticles)[nentries]) AliMCParticle(particle, rarray);
477         mcParticle = dynamic_cast<AliMCParticle*>((*fMCParticles)[nentries]);
478         fMCParticleMap->AddAt(mcParticle, i);
479     } else {
480         mcParticle = dynamic_cast<AliMCParticle*>(fMCParticleMap->At(i));
481     }
482
483     return mcParticle;
484 }
485
486  AliGenEventHeader* AliMCEvent::GenEventHeader() {return (fHeader->GenEventHeader());}
487
488 ClassImp(AliMCEvent)