Add dielectron framework to PWG3
[u/mrichter/AliRoot.git] / PWG3 / dielectron / AliDielectronMC.cxx
1 /*************************************************************************
2 * Copyright(c) 1998-2009, 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 AliDielectronMC                #
19 //#       Cut Class for Jpsi->e+e- analysis           #
20 //#                                                   #
21 //#   by WooJin J. Park, GSI / W.J.Park@gsi.de        #
22 //#                                                   #
23 //#####################################################
24
25 #include <AliAnalysisManager.h>
26 #include <AliAODHandler.h>
27 #include <AliESDInputHandler.h>
28 #include <AliMCEventHandler.h>
29 #include <AliMCEvent.h>
30 #include <AliMCParticle.h>
31 #include <AliAODMCParticle.h>
32 #include <AliStack.h>
33 #include <AliESDEvent.h>
34 #include <AliESDtrack.h>
35 #include <AliLog.h>
36
37 #include "AliDielectronMC.h"
38
39 AliDielectronMC* AliDielectronMC::fgInstance=0x0;
40
41 //____________________________________________________________
42 AliDielectronMC* AliDielectronMC::Instance()
43 {
44   //
45   // return pointer to singleton implementation
46   //
47   if (fgInstance) return fgInstance;
48   
49   AnalysisType type=kUNSET;
50   if (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler()->IsA()==AliESDInputHandler::Class()) type=kESD;
51   else if (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler()->IsA()==AliAODHandler::Class()) type=kAOD;
52   
53   fgInstance=new AliDielectronMC(type);
54   
55   return fgInstance;
56 }
57
58 //____________________________________________________________
59 AliDielectronMC::AliDielectronMC(AnalysisType type):
60   fMCEvent(0x0),
61   fStack(0x0),
62   fAnaType(type)
63 {
64   //
65   // default constructor
66   //
67 }
68
69
70 //____________________________________________________________
71 AliDielectronMC::~AliDielectronMC()
72 {
73   //
74   // default destructor
75   //
76   
77 }
78
79 //____________________________________________________________
80 void AliDielectronMC::Initialize()
81 {
82   //
83   // initialize MC class
84   //
85   if (!ConnectMCEvent()) AliError("Initialization of MC object failed!");
86 }
87
88 //____________________________________________________________
89 Int_t AliDielectronMC::GetNMCTracks()
90 {
91   //
92   //  return the number of generated tracks from MC event
93   //
94   if (!fMCEvent){ AliError("No fMCEvent"); return 0; }
95   return fMCEvent->GetNumberOfTracks();
96 }
97
98 //____________________________________________________________
99 Int_t AliDielectronMC::GetNMCTracksFromStack()
100 {
101   //
102   //  return the number of generated tracks from stack
103   //
104   if (!fStack){ AliError("No fStack"); return -999; }
105   return fStack->GetNtrack();
106 }
107
108 //____________________________________________________________
109 AliVParticle* AliDielectronMC::GetMCTrackFromMCEvent(Int_t _itrk)
110 {
111   //
112   // return MC track directly from MC event
113   //
114   if (!fMCEvent){ AliError("No fMCEvent"); return NULL;}
115   AliVParticle * track = fMCEvent->GetTrack(_itrk); //  tracks from MC event
116   return track;
117 }
118
119 //____________________________________________________________
120 Bool_t AliDielectronMC::ConnectMCEvent()
121 {
122   //
123   // connect stack object from the mc handler
124   //
125   AliMCEventHandler* mcHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
126   if (!mcHandler){ AliError("Could not retrive MC event handler!"); return kFALSE; }
127   
128   AliMCEvent* mcEvent = mcHandler->MCEvent();
129   if (!mcEvent){ AliError("Could not retrieve MC event!"); return kFALSE; }
130   fMCEvent = mcEvent;
131   
132   if (!UpdateStack()) return kFALSE;
133   return kTRUE;
134 }
135
136 //____________________________________________________________
137 Bool_t AliDielectronMC::UpdateStack()
138 {
139   //
140   // update stack with new event
141   //
142   if (!fMCEvent){ AliError("No fMCEvent"); return kFALSE;}
143   AliStack* stack = fMCEvent->Stack();
144   if (!stack){ AliError("Could not retrive stack!"); return kFALSE; }
145   fStack = stack;
146   return kTRUE;
147 }
148
149 //____________________________________________________________
150 AliMCParticle* AliDielectronMC::GetMCTrack(AliESDtrack* _track)
151 {
152   //
153   // return MC track
154   //
155   Int_t label = TMath::Abs(_track->GetLabel());
156   AliMCParticle *mctrack = dynamic_cast<AliMCParticle *>(fMCEvent->GetTrack(label));
157   return mctrack;
158 }
159
160 //____________________________________________________________
161 TParticle* AliDielectronMC::GetMCTrackFromStack(AliESDtrack* _track)
162 {
163   //
164   // return MC track from stack
165   //
166   Int_t label = TMath::Abs(_track->GetLabel());
167   if (!fStack) AliWarning("fStack is not available. Update stack first.");
168   TParticle* mcpart = fStack->Particle(label);
169   if (!mcpart) return NULL;
170   return mcpart;
171 }
172
173 //____________________________________________________________
174 AliMCParticle* AliDielectronMC::GetMCTrackMother(AliESDtrack* _track)
175 {
176   //
177   // return MC track mother
178   //
179   AliMCParticle* mcpart = GetMCTrack(_track);
180   if (!mcpart) return NULL;
181 printf("mcpart->GetMother() : %d\n",mcpart->GetMother());
182   AliMCParticle* mcmother = dynamic_cast<AliMCParticle *>(fMCEvent->GetTrack(mcpart->GetMother()));
183   if (!mcmother) return NULL;
184   return mcmother;
185 }
186
187 //____________________________________________________________
188 TParticle* AliDielectronMC::GetMCTrackMotherFromStack(AliESDtrack* _track)
189 {
190   //
191   // return MC track mother from stack
192   //
193   TParticle* mcpart = GetMCTrackFromStack(_track);
194   if ( !mcpart || mcpart->GetFirstMother()<=0 ) return NULL; 
195   TParticle* mcmother = fStack->Particle(mcpart->GetFirstMother());
196   if (!mcmother) return NULL;
197   return mcmother;
198 }
199
200 //____________________________________________________________
201 Int_t AliDielectronMC::GetMCPID(AliESDtrack* _track)
202 {
203   //
204   // return PDG code of the track from the MC truth info
205   //
206   AliMCParticle* mcpart = GetMCTrack(_track);
207   if (!mcpart) return -999;
208   return mcpart->PdgCode();
209 }
210
211 //____________________________________________________________
212 Int_t AliDielectronMC::GetMCPIDFromStack(AliESDtrack* _track)
213 {
214   // 
215   // return MC PDG code from stack
216   //
217   TParticle* mcpart = GetMCTrackFromStack(_track);
218   if (!mcpart) return -999;
219   return mcpart->GetPdgCode();
220 }
221
222 //____________________________________________________________
223 Int_t AliDielectronMC::GetMotherPDG(AliESDtrack* _track)
224 {
225   //
226   // return PDG code of the mother track from the MC truth info
227   //
228   AliMCParticle* mcmother = GetMCTrackMother(_track);
229   if (!mcmother) return -999;
230   return mcmother->PdgCode();
231 }
232
233 //____________________________________________________________
234 Int_t AliDielectronMC::GetMotherPDGFromStack(AliESDtrack* _track)
235 {
236   //
237   // return PDG code of the mother track from stack
238   //
239   TParticle* mcmother = GetMCTrackMotherFromStack(_track);
240   if (!mcmother) return -999;
241   return mcmother->GetPdgCode();
242 }
243
244 //____________________________________________________________
245 Int_t AliDielectronMC::GetMCProcess(AliESDtrack* _track)
246 {
247   //
248   // return process number of the track
249   //
250   AliMCParticle* mcpart = GetMCTrack(_track);
251   if (!mcpart) return -999;
252   return 0;
253 }
254
255 //____________________________________________________________
256 Int_t AliDielectronMC::GetMCProcessFromStack(AliESDtrack* _track)
257 {
258   //
259   // return process number of the track
260   //
261   TParticle* mcpart = GetMCTrackFromStack(_track);
262   if (!mcpart) return -999;
263   return mcpart->GetUniqueID();
264 }
265
266 //____________________________________________________________
267 Int_t AliDielectronMC::GetMCProcessMother(AliESDtrack* _track)
268 {
269   //
270   // return process number of the mother of the track
271   //
272   AliMCParticle* mcmother = GetMCTrackMother(_track);
273   if (!mcmother) return -999;
274   return 0;
275 }
276
277 //____________________________________________________________
278 Int_t AliDielectronMC::GetMCProcessMotherFromStack(AliESDtrack* _track)
279 {
280   //
281   // return process number of the mother of the track
282   //
283   TParticle* mcmother = GetMCTrackMotherFromStack(_track);
284   if (!mcmother) return -999;
285   return mcmother->GetUniqueID();
286 }
287
288 //____________________________________________________________
289 Bool_t AliDielectronMC::IsMCMotherToEE(const AliVParticle *particle, Int_t pdgMother)
290 {
291   //
292   // Check if the Mother 'particle' is of type pdgMother and decays to e+e-
293   //
294   
295   if (!fMCEvent) return kFALSE;
296   
297   if (particle->IsA()==AliMCParticle::Class()){
298     return IsMCMotherToEEesd(static_cast<const AliMCParticle*>(particle),pdgMother);
299   } else if (particle->IsA()==AliAODMCParticle::Class()){
300     return IsMCMotherToEEaod(static_cast<const AliAODMCParticle*>(particle),pdgMother);
301   } else {
302     AliError("Unknown particle type");
303   }
304   return kFALSE;
305   
306 }
307
308 //____________________________________________________________
309 Bool_t AliDielectronMC::IsMCMotherToEEesd(const AliMCParticle *particle, Int_t pdgMother)
310 {
311   //
312   // Check if the Mother 'particle' is of type pdgMother and decays to e+e-
313   // ESD case
314   //
315   
316   //check pdg code
317   if (particle->PdgCode()!=pdgMother) return kFALSE;
318   Int_t ifirst = particle->GetFirstDaughter();
319   Int_t ilast  = particle->GetLastDaughter();
320   
321   //check number of daughters
322   if ((ilast-ifirst)!=1) return kFALSE;
323   AliMCParticle *firstD=static_cast<AliMCParticle*>(GetMCTrackFromMCEvent(ifirst));
324   AliMCParticle *secondD=static_cast<AliMCParticle*>(GetMCTrackFromMCEvent(ilast));
325   
326   if (firstD->Charge()>0){
327     if (firstD->PdgCode()!=-11) return kFALSE;
328     if (secondD->PdgCode()!=11) return kFALSE;
329   }else{
330     if (firstD->PdgCode()!=11) return kFALSE;
331     if (secondD->PdgCode()!=-11) return kFALSE;
332   }
333   
334   return kTRUE;
335 }
336
337 //____________________________________________________________
338 Bool_t AliDielectronMC::IsMCMotherToEEaod(const AliAODMCParticle *particle, Int_t pdgMother)
339 {
340   //
341   // Check if the Mother 'particle' is of type pdgMother and decays to e+e-
342   // AOD case
343   //
344   if (particle->GetPdgCode()!=pdgMother) return kFALSE;
345   if (particle->GetNDaughters()!=2) return kFALSE;
346   
347   Int_t ifirst = particle->GetDaughter(0);
348   Int_t ilast  = particle->GetDaughter(1);
349   
350   //check number of daughters
351   if ((ilast-ifirst)!=1) return kFALSE;
352   AliAODMCParticle *firstD=static_cast<AliAODMCParticle*>(GetMCTrackFromMCEvent(ifirst));
353   AliAODMCParticle *secondD=static_cast<AliAODMCParticle*>(GetMCTrackFromMCEvent(ilast));
354   
355   if (firstD->Charge()>0){
356     if (firstD->GetPdgCode()!=11) return kFALSE;
357     if (secondD->GetPdgCode()!=-11) return kFALSE;
358   }else{
359     if (firstD->GetPdgCode()!=-11) return kFALSE;
360     if (secondD->GetPdgCode()!=11) return kFALSE;
361   }
362   return kTRUE;
363 }
364
365 //____________________________________________________________
366 Bool_t AliDielectronMC::IsMotherPdg(const AliVParticle *particle1, const AliVParticle *particle2, Int_t pdgMother)
367 {
368   //
369   // test if mother of particle 1 and 2 has pdgCode pdgMother and is the same;
370   //
371   if (!fMCEvent) return kFALSE;
372   
373   if (fAnaType==kESD) return IsMotherPdgESD(particle1, particle2, pdgMother);
374   else if (fAnaType==kAOD) return IsMotherPdgAOD(particle1, particle2, pdgMother);
375   
376   return kFALSE;
377 }
378
379 //____________________________________________________________
380 Bool_t AliDielectronMC::IsMotherPdgESD(const AliVParticle *particle1, const AliVParticle *particle2, Int_t pdgMother)
381 {
382   //
383   // test if mother of particle 1 and 2 has pdgCode pdgMother and is the same;
384   // ESD case
385   //
386   
387   AliMCParticle *mcPart1=static_cast<AliMCParticle*>(GetMCTrackFromMCEvent(particle1->GetLabel()));
388   AliMCParticle *mcPart2=static_cast<AliMCParticle*>(GetMCTrackFromMCEvent(particle2->GetLabel()));
389   
390   if (!mcPart1||!mcPart2) return kFALSE;
391   
392   Int_t lblMother1=mcPart1->GetMother();
393   Int_t lblMother2=mcPart2->GetMother();
394   
395   AliMCParticle *mcMother1=static_cast<AliMCParticle*>(GetMCTrackFromMCEvent(lblMother1));
396   if (!mcMother1) return kFALSE;
397   if (lblMother1!=lblMother2) return kFALSE;
398   if (mcMother1->PdgCode()!=pdgMother) return kFALSE;
399   
400   return kTRUE;
401 }
402
403 //____________________________________________________________
404 Bool_t AliDielectronMC::IsMotherPdgAOD(const AliVParticle *particle1, const AliVParticle *particle2, Int_t pdgMother)
405 {
406   //
407   // test if mother of particle 1 and 2 has pdgCode pdgMother and is the same;
408   // AOD case
409   //
410   AliAODMCParticle *mcPart1=static_cast<AliAODMCParticle*>(GetMCTrackFromMCEvent(particle1->GetLabel()));
411   AliAODMCParticle *mcPart2=static_cast<AliAODMCParticle*>(GetMCTrackFromMCEvent(particle2->GetLabel()));
412   
413   if (!mcPart1||!mcPart2) return kFALSE;
414   
415   Int_t lblMother1=mcPart1->GetMother();
416   Int_t lblMother2=mcPart2->GetMother();
417   
418   AliAODMCParticle *mcMother1=static_cast<AliAODMCParticle*>(GetMCTrackFromMCEvent(lblMother1));
419   
420   if (!mcMother1) return kFALSE;
421   if (lblMother1!=lblMother2) return kFALSE;
422   if (mcMother1->GetPdgCode()!=pdgMother) return kFALSE;
423   
424   return kTRUE;
425 }
426