Modified macros for phi analysis: improved pseudorapidity cut flexibility (A.Knospe)
[u/mrichter/AliRoot.git] / PWGLF / RESONANCES / AliRsnEvent.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-1999, 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 //  This class works as generic interface to an event.
19 //  Its main purpose is to provide a unique reference which includes all the
20 //  facilities available in the AliVEvent generic base class, plus all info
21 //  which could be needed during analysis, which are not in AliVEvent but
22 //  need to be accessed from ESD or AOD objects, usually in different ways.
23 //  When MC is available, it is properly taken into account.
24 //
25 //  authors: A. Pulvirenti (alberto.pulvirenti@ct.infn.it)
26 //           M. Vala (martin.vala@cern.ch)
27 //
28 ////////////////////////////////////////////////////////////////////////////////
29
30 #include <Riostream.h>
31 #include <TArrayF.h>
32
33 #include "AliGenEventHeader.h"
34
35 #include "AliRsnCutSet.h"
36 #include "AliRsnEvent.h"
37
38 ClassImp(AliRsnEvent)
39
40 //_____________________________________________________________________________
41 AliRsnEvent::AliRsnEvent(AliVEvent *ref, AliVEvent *refMC) :
42    fRef(ref),
43    fRefMC(refMC),
44    fLeading(-1),
45    fPID(0x0),
46    fAODList(0x0)
47 {
48 //
49 // Default constructor.
50 //
51 }
52
53 //_____________________________________________________________________________
54 AliRsnEvent::AliRsnEvent(const AliRsnEvent &event) :
55    TObject(event),
56    fRef(event.fRef),
57    fRefMC(event.fRefMC),
58    fLeading(event.fLeading),
59    fPID(event.fPID),
60    fAODList(event.fAODList)
61 {
62 //
63 // Copy constructor.
64 //
65 }
66
67 //_____________________________________________________________________________
68 AliRsnEvent &AliRsnEvent::operator= (const AliRsnEvent &event)
69 {
70 //
71 // Works in the same way as the copy constructor.
72 //
73
74    TObject::operator=(event);
75    if (this == &event)
76       return *this;
77    fRef             = event.fRef;
78    fRefMC           = event.fRefMC;
79    fLeading         = event.fLeading;
80    fPID             = event.fPID;
81    fAODList         = event.fAODList;
82
83    return (*this);
84 }
85
86 //_____________________________________________________________________________
87 AliRsnEvent::~AliRsnEvent()
88 {
89 //
90 // Destructor (does nothing since there are not owned pointers)
91 //
92 }
93
94 //_____________________________________________________________________________
95 void AliRsnEvent::SetDaughter(AliRsnDaughter &out, Int_t index, Bool_t fromMC)
96 {
97 //
98 // Assigns to the first argument the reference to the i-th track in the ref event.
99 // What assignment method to be used will depend on the index and on the type of input.
100 // If the last argument is kTRUE and an MC is referenced, then both fRef and fRefMC will
101 // point to the MC particle (pure MC analysis)
102 //
103
104    // by default the daughter is reset
105    // and assigned the used index
106    out.Reset();
107    out.SetRsnID(index);
108    out.SetOwnerEvent(this);
109
110    // check input type
111    if (!InputOK()) return;
112    Bool_t inputESD = IsESD();
113
114    // if it is pure MC, the index tells what particle
115    // to be read in the stack of MC particles, otherwise
116    // it is converted into a real collection index
117    if (fromMC) {
118       out.SetLabel(index);
119       Bool_t ok = (inputESD ? SetMCInfoESD(out) : SetMCInfoAOD(out));
120       if (ok) {
121          out.SetGood();
122          out.SetRef(out.GetRefMC());
123       }
124    } else {
125       Int_t trueIndex;
126       AliRsnDaughter::ERefType type;
127       if (!ConvertAbsoluteIndex(index, trueIndex, type)) {
128          AliError(Form("Failed to convert absolute index %d", index));
129          return;
130       }
131       switch (type) {
132          case AliRsnDaughter::kTrack:
133             if (inputESD) SetDaughterESDtrack(out, trueIndex); else SetDaughterAODtrack(out, trueIndex);
134             break;
135          case AliRsnDaughter::kV0:
136             if (inputESD) SetDaughterESDv0(out, trueIndex); else SetDaughterAODv0(out, trueIndex);
137             break;
138          case AliRsnDaughter::kCascade:
139             if (inputESD) SetDaughterESDcascade(out, trueIndex); else SetDaughterAODcascade(out, trueIndex);
140             break;
141          default:
142             AliError("Unrecognized daughter type");
143             return;
144       }
145    }
146 }
147
148 //_____________________________________________________________________________
149 AliRsnDaughter AliRsnEvent::GetDaughter(Int_t i, Bool_t fromMC)
150 {
151 //
152 // Returns a daughter set using same criteria as SetDaughter
153 //
154
155    AliRsnDaughter d;
156    SetDaughter(d, i, fromMC);
157    return d;
158 }
159
160 //_____________________________________________________________________________
161 void AliRsnEvent::SetDaughterESDtrack(AliRsnDaughter &out, Int_t i)
162 {
163 //
164 // Setup the first argument to the track identified by the index.
165 // When available, adds the MC information and references.
166 // ---
167 // Version #1: ESD tracks
168 //
169
170    AliESDEvent *esd = (AliESDEvent *)fRef;
171
172    if (i >= 0 && i < esd->GetNumberOfTracks()) {
173       AliESDtrack *track = esd->GetTrack(i);
174       if (track) {
175          out.SetRef(track);
176          out.SetGood();
177          // if MC is present, assign label and retrieve corresponding particle
178          if (fRefMC) {
179             out.SetLabel(TMath::Abs(track->GetLabel()));
180             if (!SetMCInfoESD(out)) {
181                AliWarning("Failed assignment of MC info");
182             }
183          }
184       } else {
185          AliWarning("Null track");
186       }
187    } else {
188       AliWarning(Form("Overflow: required index = %d, max = %d", i, esd->GetNumberOfTracks()));
189    }
190 }
191
192 //_____________________________________________________________________________
193 void AliRsnEvent::SetDaughterAODtrack(AliRsnDaughter &out, Int_t i)
194 {
195 //
196 // Setup the first argument to the track identified by the index.
197 // When available, adds the MC information and references.
198 // ---
199 // Version #2: AOD tracks
200 //
201
202    AliAODEvent *aod = (AliAODEvent *)fRef;
203
204    if (i >= 0 && i < aod->GetNumberOfTracks()) {
205       AliAODTrack *track = aod->GetTrack(i);
206       if (track) {
207          out.SetRef(track);
208          out.SetGood();
209          // if MC is present, assign label and retrieve corresponding particle
210          if (fRefMC) {
211             out.SetLabel(TMath::Abs(track->GetLabel()));
212             if (!SetMCInfoAOD(out)) {
213                AliWarning("Failed assignment of MC info");
214             }
215          }
216       } else {
217          AliWarning("Null track");
218       }
219    } else {
220       AliWarning(Form("Overflow: required index = %d, max = %d", i, aod->GetNumberOfTracks()));
221    }
222 }
223
224 //_____________________________________________________________________________
225 void AliRsnEvent::SetDaughterESDv0(AliRsnDaughter &out, Int_t i)
226 {
227 //
228 // Setup the first argument to the track identified by the index.
229 // When available, adds the MC information and references.
230 // ---
231 // Version #3: ESD v0
232 //
233
234    if (i >= 0 && i < fRef->GetNumberOfV0s()) {
235       AliESDEvent *ev = GetRefESD();
236       AliESDv0    *v0 = ev->GetV0(i);
237       if (v0) {
238          out.SetRef(v0);
239          out.SetGood();
240          // if MC is present, retrieve the label of V0 from those of daughters
241          if (fRefMC) {
242             AliMCEvent  *mc = (AliMCEvent *)fRefMC;
243             AliESDtrack *tp = ev->GetTrack(v0->GetPindex());
244             AliESDtrack *tn = ev->GetTrack(v0->GetNindex());
245             if (mc && tp && tn) {
246                Int_t lp = TMath::Abs(tp->GetLabel());
247                Int_t ln = TMath::Abs(tn->GetLabel());
248                TParticle *pp = mc->Stack()->Particle(lp);
249                TParticle *pn = mc->Stack()->Particle(ln);
250                if (pp && pn) {
251                   // if their first mothers are the same, the V0 is true
252                   // otherwise label remains '-1' --> fake V0
253                   if (pp->GetFirstMother() == pn->GetFirstMother() && pp->GetFirstMother() >= 0) {
254                      out.SetLabel(pp->GetFirstMother());
255                      //patch for k0s/k0l
256                      TParticle *mom = mc->Stack()->Particle(pn->GetFirstMother());
257                      if(mom->GetPdgCode() == 310) {
258                         //take the mother of the k0s which is a k0 (311)
259                         out.SetLabel(mom->GetFirstMother());
260                      }
261                      SetMCInfoESD(out);
262                   }
263                }
264             }
265          }
266       }
267    }
268 }
269
270 //_____________________________________________________________________________
271 void AliRsnEvent::SetDaughterAODv0(AliRsnDaughter &out, Int_t i)
272 {
273 //
274 // Setup the first argument to the track identified by the index.
275 // When available, adds the MC information and references.
276 // ---
277 // Version #4: AOD v0
278 //
279
280    if (i >= 0 && i < fRef->GetNumberOfV0s()) {
281       AliAODEvent *ev = (AliAODEvent *)fRef;
282       AliAODv0    *v0 = ev->GetV0(i);
283       if (v0) {
284          out.SetRef(v0);
285          out.SetGood();
286          if (fRefMC) {
287             AliAODEvent  *mc = (AliAODEvent *)fRefMC;
288             TClonesArray *mcArray = (TClonesArray *)mc->GetList()->FindObject(AliAODMCParticle::StdBranchName());
289             AliAODTrack  *tp  = (AliAODTrack *)v0->GetDaughter(0);
290             AliAODTrack  *tn  = (AliAODTrack *)v0->GetDaughter(1);
291             if (mcArray && tp && tn) {
292                Int_t lp = TMath::Abs(tp->GetLabel());
293                Int_t ln = TMath::Abs(tn->GetLabel());
294                AliAODMCParticle *pp = (AliAODMCParticle *)mcArray->At(lp);
295                AliAODMCParticle *pn = (AliAODMCParticle *)mcArray->At(ln);
296                if (pp && pn) {
297                   // if their first mothers are the same, the V0 is true
298                   // otherwise label remains '-1' --> fake V0
299                   if (pp->GetMother() == pn->GetMother() && pp->GetMother() >= 0) {
300                      out.SetLabel(pp->GetMother());
301                      SetMCInfoAOD(out);
302                   }
303                }
304             }
305          }
306       }
307    }
308 }
309
310 //_____________________________________________________________________________
311 void AliRsnEvent::SetDaughterESDcascade(AliRsnDaughter &out, Int_t i)
312 {
313 //
314 // Setup the first argument to the track identified by the index.
315 // When available, adds the MC information and references.
316 // ---
317 // Version #3: ESD cascade
318 //
319
320    if (i >= 0 && i < fRef->GetNumberOfCascades()) {
321       AliESDEvent   *ev   = GetRefESD();
322       AliESDcascade *casc = ev->GetCascade(i);
323       if (casc) {
324          out.SetRef(casc);
325          out.SetGood();
326          if (fRefMC) {
327
328          }
329       }
330    }
331 }
332
333 //_____________________________________________________________________________
334 void AliRsnEvent::SetDaughterAODcascade(AliRsnDaughter &out, Int_t i)
335 {
336 //
337 // Setup the first argument to the track identified by the index.
338 // When available, adds the MC information and references.
339 // ---
340 // Version #4: AOD cascade
341 //
342
343    if (i >= 0 && i < fRef->GetNumberOfCascades()) {
344       AliAODEvent *ev = GetRefAOD();
345       AliAODv0    *casc = ev->GetCascade(i);
346       if (casc) {
347          out.SetRef(casc);
348          out.SetGood();
349          if (fRefMC) {
350
351          }
352       }
353    }
354 }
355
356 //_____________________________________________________________________________
357 Bool_t AliRsnEvent::SetMCInfoESD(AliRsnDaughter &out)
358 {
359 //
360 // Using the label assigned to the daughter, searches for the MC informations:
361 // -- MC reference
362 // -- mother
363 //
364
365    // if label makes no sense --> failed
366    Int_t label = out.GetLabel();
367    if (label < 0 || !fRefMC) return kFALSE;
368
369    // get number of particles
370    Int_t nMC = fRefMC->GetNumberOfTracks();
371
372    // if label too large --> failed
373    if (label >= nMC) {
374       AliWarning(Form("Stack overflow: track label = %d -- stack maximum = %d", label, nMC));
375       return kFALSE;
376    }
377
378    // retrieve particle
379    AliMCEvent    *mc = (AliMCEvent *)fRefMC;
380    AliMCParticle *mcPart = (AliMCParticle *)mc->GetTrack(label);
381
382    // if particle = NULL --> failed
383    if (!mcPart) {
384       AliWarning(Form("Stack discontinuity: label %d refers to a NULL object", label));
385       return kFALSE;
386    }
387    // otherwise --> success
388    out.SetRefMC(mcPart);
389
390    // if the particle is not primary, find the mother and get its PDG
391    Int_t imum = mcPart->Particle()->GetFirstMother();
392    if (imum >= 0 && imum < nMC) {
393       AliMCParticle *mcMother = (AliMCParticle *)mc->GetTrack(imum);
394       if (mcMother) {
395          out.SetMotherPDG(mcMother->Particle()->GetPdgCode());
396       } else {
397          AliWarning(Form("Stack discontinuity: label mother %d refers to a NULL object", imum));
398       }
399    } else {
400       AliWarning(Form("Stack overflow: mother label = %d -- stack maximum = %d", imum, nMC));
401    }
402
403    return kTRUE;
404 }
405
406 //_____________________________________________________________________________
407 Bool_t AliRsnEvent::SetMCInfoAOD(AliRsnDaughter &out)
408 {
409 //
410 // Using the label assigned to the daughter, searches for the MC informations:
411 // -- MC reference
412 // -- mother
413 //
414
415    // if label makes no sense --> failed
416    Int_t label = out.GetLabel();
417    if (label < 0 || !fRefMC) return kFALSE;
418
419    // retrieve particle
420    AliAODEvent  *mc = (AliAODEvent *)fRefMC;
421    TClonesArray *mcArray = (TClonesArray *)mc->GetList()->FindObject(AliAODMCParticle::StdBranchName());
422
423    // get number of particles
424    Int_t nMC = mcArray->GetEntriesFast();
425
426    // if label too large --> failed
427    if (label >= nMC) {
428       AliWarning(Form("Stack overflow: track label = %d -- stack maximum = %d", label, nMC));
429       return kFALSE;
430    }
431
432    // if particle = NULL --> failed
433    AliAODMCParticle *mcPart = (AliAODMCParticle *)mcArray->At(label);
434    if (!mcPart) {
435       AliWarning(Form("Stack discontinuity: label %d refers to a NULL object", label));
436       return kFALSE;
437    }
438    // otherwise --> success
439    out.SetRefMC(mcPart);
440
441    // if the particle is not primary, find the mother and get its PDG
442    Int_t imum = mcPart->GetMother();
443    if (imum >= 0 && imum < nMC) {
444       AliAODMCParticle *mcMother = (AliAODMCParticle *)mcArray->At(imum);
445       if (mcMother) {
446          out.SetMotherPDG(mcMother->GetPdgCode());
447       } else {
448          AliWarning(Form("Stack discontinuity: label mother %d refers to a NULL object", imum));
449       }
450    } else if (imum >= nMC) {
451       AliWarning(Form("Stack overflow: mother label = %d -- stack maximum = %d", imum, nMC));
452    }
453
454    return kTRUE;
455 }
456
457 //_____________________________________________________________________________
458 Bool_t AliRsnEvent::ConvertAbsoluteIndex(Int_t index, Int_t &realIndex, AliRsnDaughter::ERefType &type)
459 {
460 //
461 // Using the phylosophy of the absolute index, which loops over
462 // all tracks, V0s and cascades, returns the result of a check
463 // on it (first argument) based on this criterion:
464 // 1) if the absolute ID is smaller than number of tracks,
465 //    return itself and the type 'track'
466 // 2) if the absolute ID is larger than number of tracks, subtract it
467 //    and if the result is smaller than number of V0s,
468 //    return the corresponding V0 index and type
469 // 3) if the absolute ID is larger than number of tracks + V0s, subtract them
470 //    and if the result is smaller than number of cascades,
471 //    return the corresponding cascade index and type
472 // The results of this check are stored in the reference arguments, while the outcome of
473 // the function is kTRUE if one of these checks was successful, otherwise it is kFALSE,
474 // meaning that the absolute index reached the end.
475 //
476
477    Int_t nTracks   = fRef->GetNumberOfTracks();
478    Int_t nV0s      = fRef->GetNumberOfV0s();
479    Int_t nCascades = fRef->GetNumberOfCascades();
480
481    if (index < nTracks) {
482       realIndex = index;
483       type = AliRsnDaughter::kTrack;
484       return kTRUE;
485    } else if (index >= nTracks && index < nTracks + nV0s) {
486       realIndex = index - nTracks;
487       type = AliRsnDaughter::kV0;
488       return kTRUE;
489    } else if (index >= nTracks + nV0s && index < nTracks + nV0s + nCascades) {
490       realIndex = index - nTracks - nV0s;
491       type = AliRsnDaughter::kCascade;
492       return kTRUE;
493    }
494
495    realIndex = -1;
496    type = AliRsnDaughter::kNoType;
497    return kFALSE;
498 }
499
500 //_____________________________________________________________________________
501 Int_t AliRsnEvent::ConvertRealIndex(Int_t index, AliRsnDaughter::ERefType type)
502 {
503 //
504 // Translates a pair made by index + object type into the corresponding
505 // absolute index, which is set to -1 in case the real index overflows.
506 //
507
508    Int_t nTracks   = fRef->GetNumberOfTracks();
509    Int_t nV0s      = fRef->GetNumberOfV0s();
510    Int_t nCascades = fRef->GetNumberOfCascades();
511
512    switch (type) {
513       case AliRsnDaughter::kTrack:
514          if (index >= 0 && index < nTracks)
515             return index;
516          else
517             return -1;
518       case AliRsnDaughter::kV0:
519          if (index >= 0 && index < nV0s)
520             return nTracks + index;
521          else
522             return -1;
523       case AliRsnDaughter::kCascade:
524          if (index >= 0 && index < nCascades)
525             return nTracks + nV0s + index;
526          else
527             return -1;
528       default:
529          return -1;
530    }
531 }
532
533 //_____________________________________________________________________________
534 Int_t AliRsnEvent::SelectLeadingParticle(AliRsnCutSet *cuts)
535 {
536 //
537 // Searches the collection of all particles with given PID type and charge,
538 // and returns the one with largest momentum, provided that it is greater than 1st argument.
539 // If one specifies AliRsnPID::kUnknown as type or AliRsnDaughter::kNoPID as method,
540 // the check is done over all particles irrespectively of their PID.
541 // If the sign argument is '+' or '-', the check is done over the particles of that charge,
542 // otherwise it is done irrespectively of the charge.
543 //
544
545    // check input type
546    Bool_t inputESD = IsESD();
547    if (!inputESD && !IsAOD()) {
548       AliError("Need to process ESD or AOD input");
549       return -1;
550    }
551
552    Double_t ptMax = 0.0;
553    Int_t i, nTracks = fRef->GetNumberOfTracks();
554
555    fLeading = -1;
556    AliRsnDaughter leading;
557
558    for (i = 0; i < nTracks; i++) {
559       if (inputESD)
560          SetDaughterESDtrack(leading, i);
561       else
562          SetDaughterAODtrack(leading, i);
563       if (!leading.IsOK()) {
564          AliDebugClass(1, Form("Failed assignment of track %d", i));
565          continue;
566       }
567       if (cuts && !cuts->IsSelected(&leading)) {
568          AliDebugClass(1, Form("Track %d didn't pass cuts", i));
569          continue;
570       }
571       // check if it has largest momentum
572       if (leading.GetRef()->Pt() > ptMax) {
573          ptMax = leading.GetRef()->Pt();
574          fLeading = i;
575       }
576    }
577
578    return fLeading;
579 }