Update of the Xiaoming code for pp900 first muon analysis. Fixing wanirngs and violti...
[u/mrichter/AliRoot.git] / PWG3 / muon / AliMuonInfoStoreMC.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-2006, 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 /////////////////////////////////////////////////////////////
17 //
18 // class used to extract and store info of MC particle
19 //
20 // Author: X-M. Zhang, zhang@clermont.in2p3.fr
21 //                     zhangxm@iopp.ccnu.edu.cn
22 /////////////////////////////////////////////////////////////
23
24 #include <TParticle.h>
25 #include <TClonesArray.h>
26
27 #include "AliMCEvent.h"
28 #include "AliMCEventHandler.h"
29 #include "AliStack.h"
30 #include "AliAODMCParticle.h"
31 #include "AliESDMuonTrack.h"
32 #include "AliAODTrack.h"
33 #include "AliMuonInfoStoreRD.h"
34 #include "AliMuonInfoStoreMC.h"
35
36 class AliESDEvent;
37
38 ClassImp(AliMuonInfoStoreMC)
39
40 const TString AliMuonInfoStoreMC::fgkStdBranchName("MuonMC");
41 const Int_t   AliMuonInfoStoreMC::fgkNSources = 7;
42
43 //-----------------------------------------------------------------------------
44 AliMuonInfoStoreMC::AliMuonInfoStoreMC() :
45 AliMuonInfoStoreRD(),
46 fIsFull(kFALSE),
47 fLorentzP(),
48 fTrackIndex(-1),
49 fTrackPDGCode(0),
50 fSource(-1),
51 fNParents(0),
52 fOscillation(kFALSE),
53 fWeight(0.)
54 {
55   //
56   // default constructor
57   //
58   for (Int_t i=5; i--;) { fParentIndex[i] = -1; fParentPDGCode[i] = 0; }
59   for (Int_t i=4; i--;) { fQuarkIndex[i]  = -1; fQuarkPDGCode[i]  = 0; }
60 }
61
62 //-----------------------------------------------------------------------------
63 AliMuonInfoStoreMC::AliMuonInfoStoreMC(AliAODTrack *trkAOD, TClonesArray *mcClArr, Bool_t full) :
64 AliMuonInfoStoreRD(trkAOD),
65 fIsFull(full),
66 fLorentzP(),
67 fTrackIndex(-1),
68 fTrackPDGCode(0),
69 fSource(-1),
70 fNParents(0),
71 fOscillation(kFALSE),
72 fWeight(0.)
73 {
74   //
75   // default constructor
76   //
77   for (Int_t i=5; i--;) { fParentIndex[i] = -1; fParentPDGCode[i] = 0; }
78   for (Int_t i=4; i--;) { fQuarkIndex[i]  = -1; fQuarkPDGCode[i]  = 0; }
79
80   AliAODMCParticle *pMC = this->FindTrackRef(trkAOD, mcClArr);
81   if (pMC) this->SetMCInfo(pMC, mcClArr);
82 }
83
84 //-----------------------------------------------------------------------------
85 AliMuonInfoStoreMC::AliMuonInfoStoreMC(AliESDMuonTrack *trkESD, AliMCEventHandler *mcH, Bool_t full) :
86 AliMuonInfoStoreRD(trkESD),
87 fIsFull(full),
88 fLorentzP(),
89 fTrackIndex(-1),
90 fTrackPDGCode(0),
91 fSource(-1),
92 fNParents(0),
93 fOscillation(kFALSE),
94 fWeight(0.)
95 {
96   //
97   // default constructor
98   //
99   for (Int_t i=5; i--;) { fParentIndex[i] = -1; fParentPDGCode[i] = 0; }
100   for (Int_t i=4; i--;) { fQuarkIndex[i]  = -1; fQuarkPDGCode[i]  = 0; }
101
102   TParticle *pMC = this->FindTrackRef(trkESD, mcH);
103   if (pMC) this->SetMCInfo(pMC, mcH);
104 }
105
106 //-----------------------------------------------------------------------------
107 /*AliMuonInfoStoreMC::AliMuonInfoStoreMC(AliESDMuonTrack *trkESD, AliESDEvent *esd, AliMCEventHandler *mcH, Bool_t full) :
108 AliMuonInfoStoreRD(trkESD),
109 fIsFull(full),
110 fLorentzP(),
111 fTrackIndex(-1),
112 fTrackPDGCode(0),
113 fSource(-1),
114 fNParents(0),
115 fOscillation(kFALSE),
116 fWeight(0.)
117 {
118 #include "AliMUONRecoCheck.h"
119 #include "AliMUONVTrackStore.h"
120 #include "AliMUONTrack.h"
121 #include "AliMUONESDInterface.h"
122   //
123   // default constructor
124   //
125   for (Int_t i=5; i--;) { fParentIndex[i] = -1; fParentPDGCode[i] = 0; }
126   for (Int_t i=4; i--;) { fQuarkIndex[i]  = -1; fQuarkPDGCode[i]  = 0; }
127
128   TParticle *pMC = this->FindTrackRef(trkESD, esd, mcH);
129   if (pMC) this->SetMCInfo(pMC, mcH);
130 }*/
131
132 //-----------------------------------------------------------------------------
133 AliMuonInfoStoreMC::AliMuonInfoStoreMC(const AliMuonInfoStoreMC &src) :
134 AliMuonInfoStoreRD(src),
135 fIsFull(src.fIsFull),
136 fLorentzP(src.fLorentzP),
137 fTrackIndex(src.fTrackIndex),
138 fTrackPDGCode(src.fTrackPDGCode),
139 fSource(src.fSource),
140 fNParents(src.fNParents),
141 fOscillation(src.fOscillation),
142 fWeight(src.fWeight)
143 {
144   //
145   // copy constructor
146   //
147   for (Int_t i=5; i--;) {
148     fParentIndex[i]   = src.fParentIndex[i];
149     fParentPDGCode[i] = src.fParentPDGCode[i];
150   }
151   for (Int_t i=4; i--;) {
152     fQuarkIndex[i]    = src.fQuarkIndex[i];
153     fQuarkPDGCode[i]  = src.fQuarkPDGCode[i];
154   }
155 }
156
157 //-----------------------------------------------------------------------------
158 AliMuonInfoStoreMC& AliMuonInfoStoreMC::operator=(const AliMuonInfoStoreMC &src)
159 {
160   //
161   // assignment constructor
162   //
163   if(&src==this) return *this;
164   AliMuonInfoStoreRD::operator=(src);
165
166   fIsFull       = src.fIsFull;
167   fLorentzP     = src.fLorentzP;
168   fTrackIndex   = src.fTrackIndex;
169   fTrackPDGCode = src.fTrackPDGCode;
170   fSource       = src.fSource;
171   fNParents     = src.fNParents;
172   fOscillation  = src.fOscillation;
173   fWeight       = src.fWeight;
174   for (Int_t i=5; i--;) {
175     fParentIndex[i]   = src.fParentIndex[i];
176     fParentPDGCode[i] = src.fParentPDGCode[i];
177   }
178   for (Int_t i=4; i--;) {
179     fQuarkIndex[i]    = src.fQuarkIndex[i];
180     fQuarkPDGCode[i]  = src.fQuarkPDGCode[i];
181   }
182
183   return *this;
184 }
185
186 //-----------------------------------------------------------------------------
187 AliMuonInfoStoreMC::~AliMuonInfoStoreMC()
188 {
189   //
190   // destructor
191   //
192 }
193
194 //-----------------------------------------------------------------------------
195 AliAODMCParticle* AliMuonInfoStoreMC::FindTrackRef(AliAODTrack* const trkAOD, TClonesArray* const mcClArr)
196 {
197   // find MC track ref with AOD base
198
199   AliAODMCParticle *pMC = 0;
200   fTrackIndex = trkAOD->GetLabel();
201   if (fTrackIndex>=0)
202     pMC = (AliAODMCParticle*)mcClArr->At(fTrackIndex);
203   return pMC;
204 }
205
206 //-----------------------------------------------------------------------------
207 TParticle* AliMuonInfoStoreMC::FindTrackRef(AliESDMuonTrack* const trkESD, AliMCEventHandler* const mcH)
208 {
209   // find MC track ref with ESD base
210
211   TParticle *pMCRef = 0;
212   fTrackIndex = trkESD->GetLabel();
213   if (fTrackIndex>=0) pMCRef = mcH->MCEvent()->Stack()->Particle(fTrackIndex);
214   return pMCRef;
215 }
216
217 //-----------------------------------------------------------------------------
218 /*TParticle* AliMuonInfoStoreMC::FindTrackRef(AliESDMuonTrack* const trkESD, AliESDEvent* const esd, AliMCEventHandler* const mcH)
219 {
220   // find MC track ref with ESD trackRef base
221
222   TParticle *pMCRef = 0;
223   AliMUONRecoCheck rc(esd,mcH);
224   AliMUONVTrackStore *trkRefArr = rc.TrackRefs(-1);
225
226   AliMUONTrack trkMuon;
227   AliMUONESDInterface::ESDToMUON(*trkESD, trkMuon, kFALSE);
228
229   Int_t nMatchClusters = 0;
230   AliMUONTrack *trkRef = rc.FindCompatibleTrack(trkMuon, *trkRefArr, nMatchClusters, kFALSE, 10.);
231   if (trkRef) fTrackIndex = trkRef->GetUniqueID();
232   if (fTrackIndex>=0) pMCRef = mcH->MCEvent()->Stack()->Particle(fTrackIndex);
233   return pMCRef;
234 }*/
235
236 //-----------------------------------------------------------------------------
237 void AliMuonInfoStoreMC::SetMCInfo(AliAODMCParticle *pMC, TClonesArray *mcClArr)
238 {
239   // fill track MC info with AOD base
240
241   fLorentzP.SetPxPyPzE(pMC->Px(), pMC->Py(), pMC->Pz(), pMC->E());
242   fTrackPDGCode = pMC->GetPdgCode();
243   if (TMath::Abs(fTrackPDGCode)!=13) {
244     fSource = 4;
245     return;
246   } 
247
248   Int_t lineM = pMC->GetMother();
249   if (lineM<0) {
250     fSource = 2;
251     return;
252   }
253
254   Bool_t isPrimary = ((AliAODMCParticle*)mcClArr->At(lineM))->IsPrimary();
255   if (!isPrimary) {
256     fSource = 3;
257     return;
258   }
259
260   this->FillHistoryParents(lineM, mcClArr);
261   fSource = this->SelectHFMuon();
262   return;
263 }
264
265 //-----------------------------------------------------------------------------
266 void AliMuonInfoStoreMC::SetMCInfo(TParticle *pMC, AliMCEventHandler* const mcH)
267 {
268   // fill track MC info with ESD base
269
270   fLorentzP.SetPxPyPzE(pMC->Px(), pMC->Py(), pMC->Pz(), pMC->Energy());
271   fTrackPDGCode = pMC->GetPdgCode();
272   if (TMath::Abs(fTrackPDGCode)!=13) {
273     fSource = 4;
274     return;
275   }
276
277   Int_t lineM = pMC->GetFirstMother();
278   if (lineM<0) {
279     fSource = 2;
280     return;
281   }
282
283   AliStack *stack = mcH->MCEvent()->Stack();
284   if (lineM>=stack->GetNprimary()) {
285     fSource = 3;
286     return;
287   }
288
289   this->FillHistoryParents(lineM, stack);
290   fSource = this->SelectHFMuon();
291   return;
292 }
293
294 //-----------------------------------------------------------------------------
295 void AliMuonInfoStoreMC::FillHistoryParents(Int_t lineM, TClonesArray *mcClArr)
296 {
297   // find track hadron parents with AOD base
298
299   Int_t countP=-1, pdg=0;
300   Int_t parents[10], parLine[10];
301   AliAODMCParticle *mother = 0;
302   while(lineM>=0){
303     mother = (AliAODMCParticle*)mcClArr->At(lineM);
304     pdg = mother->GetPdgCode();
305     if(pdg==92 || pdg==21 || TMath::Abs(pdg)<10 || IsDiquark(pdg)) break;
306     parents[++countP] = pdg;
307     parLine[countP] = lineM;
308     lineM = mother->GetMother();
309   }
310   for(Int_t i=0; i<=countP; i++){
311     fParentIndex[i] = parLine[countP-i];
312     fParentPDGCode[i] = parents[countP-i];
313   }
314   fNParents = countP + 1;
315
316   if (fIsFull && lineM>=0) this->FillHistoryQuarks(lineM, mcClArr);
317   return;
318 }
319
320 //-----------------------------------------------------------------------------
321 void AliMuonInfoStoreMC::FillHistoryParents(Int_t lineM, AliStack *stack)
322 {
323   // find track hadron parents with ESD base
324
325   Int_t countP=-1, pdg=0;
326   Int_t parents[10], parLine[10];
327   TParticle *mother = 0;
328   while(lineM>=0){
329     mother = stack->Particle(lineM);
330     pdg = mother->GetPdgCode();
331     if(pdg==92 || pdg==21 || TMath::Abs(pdg)<10 || IsDiquark(pdg)) break;
332     parents[++countP] = pdg;
333     parLine[countP] = lineM;
334     lineM = mother->GetFirstMother();
335   }
336   for(Int_t i=0; i<=countP; i++){
337     fParentIndex[i] = parLine[countP-i];
338     fParentPDGCode[i] = parents[countP-i];
339   }
340   fNParents = countP + 1;
341
342   if (fIsFull && lineM>=0) this->FillHistoryQuarks(lineM, stack);
343   return;
344 }
345
346 //-----------------------------------------------------------------------------
347 void AliMuonInfoStoreMC::FillHistoryQuarks(Int_t lineM, TClonesArray* const mcClArr)
348 {
349   // method in $ALICE_ROOT/MUON/AliMUONTrackLight.cxx 
350
351   if (lineM<0) return;
352   Int_t countP=-1, pdg=0;
353   AliAODMCParticle *mother = 0;
354   while(lineM>=0){
355     mother = (AliAODMCParticle*)mcClArr->At(lineM);
356     pdg = mother->GetPdgCode();
357     fQuarkIndex[++countP] = lineM;
358     fQuarkPDGCode[countP] = pdg;
359     lineM = mother->GetMother();
360   }
361
362   // for PYTHIA checking
363   countP = 1;
364   for(Int_t par=0; par<4; par++) {
365     if(TMath::Abs(this->QuarkPDGCode(par))<6) { countP=par; break; }
366   }
367   if(this->QuarkIndex(countP)>-1 && (this->ParentFlavour(0)==4 || this->ParentFlavour(0)==5)) {
368     if(this->ParentFlavour(0)!=TMath::Abs(this->QuarkPDGCode(countP))) {
369       AliWarning(Form("quark flavour of parent and that of quark do not correspond: %d %d --> correcting\n",
370                  this->ParentFlavour(0), TMath::Abs(this->QuarkPDGCode(countP))));
371
372       pdg = this->QuarkPDGCode(countP);
373       Int_t line = this->QuarkIndex(countP);
374       this->ResetQuarkInfo();
375       while(TMath::Abs(pdg)!=this->ParentFlavour(0)) {
376         pdg = ((AliAODMCParticle*)mcClArr->At(++line))->GetPdgCode();
377       }
378       while(line>=0){
379         mother = (AliAODMCParticle*)mcClArr->At(line);
380         pdg = mother->GetPdgCode();
381         fQuarkIndex[countP] = line;
382         fQuarkPDGCode[countP++] = pdg;
383         line = mother->GetMother();
384       }
385     }
386   }
387   return;
388 }
389
390 //-----------------------------------------------------------------------------
391 void AliMuonInfoStoreMC::FillHistoryQuarks(Int_t lineM, AliStack* const stack)
392 {
393   // method in $ALICE_ROOT/MUON/AliMUONTrackLight.cxx 
394
395   if (lineM<0) return;
396   Int_t countP=-1, pdg=0;
397   TParticle *mother = 0;
398   while(lineM>=0){
399     mother = stack->Particle(lineM);
400     pdg = mother->GetPdgCode();
401     fQuarkIndex[++countP] = lineM;
402     fQuarkPDGCode[countP] = pdg;
403     lineM = mother->GetFirstMother();
404   }
405
406   // for PYTHIA checking
407   countP = 1;
408   for(Int_t par=0; par<4; par++) {
409     if(TMath::Abs(this->QuarkPDGCode(par))<6) { countP=par; break; }
410   }
411   if(this->QuarkIndex(countP)>-1 && (this->ParentFlavour(0)==4 || this->ParentFlavour(0)==5)) {
412     if(this->ParentFlavour(0)!=TMath::Abs(this->QuarkPDGCode(countP))) {
413       AliWarning(Form("quark flavour of parent and that of quark do not correspond: %d %d --> correcting\n",
414                  this->ParentFlavour(0), TMath::Abs(this->QuarkPDGCode(countP))));
415
416       pdg = this->QuarkPDGCode(countP);
417       Int_t line = this->QuarkIndex(countP);
418       this->ResetQuarkInfo();
419       while(TMath::Abs(pdg)!=this->ParentFlavour(0)) {
420         pdg = stack->Particle(++line)->GetPdgCode();
421       }
422       while(line>=0){
423         mother = stack->Particle(line);
424         pdg = mother->GetPdgCode();
425         fQuarkIndex[countP] = line;
426         fQuarkPDGCode[countP++] = pdg;
427         line = mother->GetFirstMother();
428       }
429     }
430   }
431   return;
432 }
433
434 //-----------------------------------------------------------------------------
435 Int_t AliMuonInfoStoreMC::SelectHFMuon()
436 {
437   // set info of muon from HF
438
439   Int_t flv = ParentFlavour(0);
440   if (flv!=4 && flv!=5) return 2;
441
442   Bool_t isRes = kFALSE;
443   Int_t i=0, nparents=this->NParents();
444   while (i<nparents && !isRes) {
445     isRes = IsMotherAResonance(i++);
446   }
447
448   if (isRes) return 2;
449   if (flv==5) return 0;
450   else return 1;
451 }
452
453 //-----------------------------------------------------------------------------
454 Bool_t AliMuonInfoStoreMC::IsDiquark(Int_t pdg)
455 {
456   // copy from $ALICE_ROOT/MUON/AliMUONTrackLight.cxx
457   pdg = TMath::Abs(pdg);
458   if(pdg>1000 && (pdg%100)<10) return kTRUE;
459   else return kFALSE;
460 }
461
462 //-----------------------------------------------------------------------------
463 void AliMuonInfoStoreMC::ResetQuarkInfo()
464 {
465   // copy from $ALICE_ROOT/MUON/AliMUONTrackLight.cxx
466   for(Int_t pos=1; pos<4; pos++) {
467     fQuarkIndex[pos] = -1;
468     fQuarkPDGCode[pos] = 0;
469   }
470   return;
471 }
472
473 //-----------------------------------------------------------------------------
474 Int_t AliMuonInfoStoreMC::ParentFlavour(Int_t i) const
475 {
476   // copy from $ALICE_ROOT/MUON/AliMUONTrackLight.cxx
477   Int_t pdg = ParentPDGCode(i);
478   pdg = TMath::Abs(pdg/100);
479   if(pdg>9) pdg /= 10;
480   return pdg;
481 }
482
483 //-----------------------------------------------------------------------------
484 Bool_t AliMuonInfoStoreMC::IsMotherAResonance(Int_t i) const
485 {
486   // copy from $ALICE_ROOT/MUON/AliMUONTrackLight.cxx
487   Int_t pdg = ParentPDGCode(i);
488   Int_t id=pdg%100000;
489   return (!((id-id%10)%110));
490 }