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