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