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