486883571fb9535069b8e6a89afd4e19978f346e
[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnReader.cxx
1 //
2 // Class AliRsnReader
3 //
4 // This is the universal converter from any kind of source event
5 // (i.e. ESD, standard AOD, MC) into the internal non-standard
6 // AOD format used by RSN package.
7 // ---
8 // This class reads all tracks in an input event and converts them
9 // into AliRsnDaughters, and computes preliminarily the PID probabilities
10 // by doing the Bayesian combination of a set of assigned prior probabilities
11 // with the PID weights defined in each track.
12 // ---
13 // When filling the output event (AliRsnEvent), some arrays of indexes
14 // are created in order to organize tracks according to their PID and charge,
15 // which will then be used in further analysis steps.
16 //
17 // author: A. Pulvirenti (alberto.pulvirenti@ct.infn.it)
18 //
19
20 #include <Riostream.h>
21 #include <TString.h>
22
23 #include "AliLog.h"
24
25 #include "AliVEvent.h"
26
27 #include "AliESDEvent.h"
28 #include "AliESDtrack.h"
29 #include "AliESDVertex.h"
30 #include "AliESDtrackCuts.h"
31
32 #include "AliAODEvent.h"
33 #include "AliAODTrack.h"
34 #include "AliAODVertex.h"
35
36 #include "AliStack.h"
37 #include "AliMCEvent.h"
38 #include "AliMCParticle.h"
39 #include "AliGenEventHeader.h"
40
41 #include "AliRsnMCInfo.h"
42 #include "AliRsnDaughter.h"
43 #include "AliRsnEvent.h"
44 #include "AliRsnPIDDefESD.h"
45 #include "AliRsnCut.h"
46
47 #include "AliRsnReader.h"
48
49 ClassImp(AliRsnReader)
50
51 //_____________________________________________________________________________
52 AliRsnReader::AliRsnReader() :
53   fCheckSplit(kFALSE),
54   fRejectFakes(kFALSE),
55   fTPCOnly(kFALSE),
56   fUseESDTrackCuts(kFALSE),
57   fUseRsnTrackCuts(kFALSE),
58   fPIDDef(),
59   fITSClusters(0),
60   fTPCClusters(0),
61   fTRDClusters(0),
62   fTrackRefs(0),
63   fTrackRefsITS(0),
64   fTrackRefsTPC(0),
65   fESDTrackCuts(),
66   fRsnTrackCuts("primaryCuts")
67 {
68 //
69 // Constructor.
70 //
71 }
72
73 //_____________________________________________________________________________
74 Bool_t AliRsnReader::AreSplitted(AliESDtrack *track1, AliESDtrack *track2)
75 {
76 //
77 // Checks if two tracks are splitted.
78 // Currently, two split tracks are the ones which
79 // have equal GEANT labels.
80 // On real data, this criterion will have to change
81 // into an "experimental" one.
82 //
83
84   Int_t lab1 = TMath::Abs(track1->GetLabel());
85   Int_t lab2 = TMath::Abs(track2->GetLabel());
86
87   return (lab1 == lab2);
88 }
89
90 //_____________________________________________________________________________
91 Bool_t AliRsnReader::ResolveSplit(AliESDtrack *track1, AliESDtrack *track2)
92 {
93 //
94 // Resolves a split of two tracks.
95 // Only two values can be returned:
96 //  kTRUE  --> accept "track1" and reject "track2"
97 //  kFALSE --> accept "track2" and reject "track1"
98 //
99
100   Double_t chiSq1 = track1->GetConstrainedChi2();
101   Double_t chiSq2 = track2->GetConstrainedChi2();
102
103   if (chiSq1 < chiSq2) return 1; else return 2;
104 }
105
106 //_____________________________________________________________________________
107 void AliRsnReader::SetTPCOnly(Bool_t doit)
108 {
109 //
110 // Set the flag for TPC only.
111 // If this is true, exclude all other detectors from PID
112 //
113
114   fTPCOnly = doit;
115
116   if (fTPCOnly) {
117     fPIDDef.ExcludeAll();
118     fPIDDef.IncludeDet(AliRsnPIDDefESD::kTPC);
119     fPIDDef.SetDivValue(AliRsnPIDDefESD::kTPC, 0.0);
120   }
121 }
122
123 //_____________________________________________________________________________
124 Bool_t AliRsnReader::ConvertTrack(AliRsnDaughter *daughter, AliESDtrack *esdTrack)
125 {
126 //
127 // Translates an ESD track into a RSN daughter
128 //
129
130   if (!esdTrack) {
131     AliError("Passed NULL object: nothing done");
132     return kFALSE;
133   }
134
135   Double_t p[3], v[3], pid[AliRsnPID::kSpecies];
136
137   // copy values which don't need treatment:
138   // momentum, vertex, chi2, flags, charge, number of TPC and ITS clusters
139   esdTrack->GetPxPyPz(p);
140   daughter->SetP(p[0], p[1], p[2]);
141   esdTrack->GetXYZ(v);
142   daughter->SetV(v[0], v[1], v[2]);
143   daughter->SetChi2(esdTrack->GetConstrainedChi2());
144   daughter->SetFlags(esdTrack->GetStatus());
145   daughter->SetCharge((Short_t)esdTrack->Charge());
146   daughter->SetNumberOfITSClusters(esdTrack->GetITSclusters(0x0));
147   daughter->SetNumberOfTPCClusters(esdTrack->GetTPCclusters(0x0));
148
149   // define the kink index:
150   //  0 = no kink
151   //  1 = kink daughter
152   // -1 = kink mother
153   Int_t i, ik[3];
154   for (i = 0; i < 3; i++) ik[i] = esdTrack->GetKinkIndex(i);
155   if (ik[0] < 0 || ik[1] < 0 || ik[2] < 0) daughter->SetKinkMother();
156   else if (ik[0] > 0 || ik[1] > 0 || ik[2] > 0) daughter->SetKinkDaughter();
157   else daughter->SetNoKink();
158
159   // store PID weights according to definition
160   // check: if the sum of all weights is null, the adoption fails
161   Double_t sum = 0.0;
162   fPIDDef.ComputeWeights(esdTrack, pid);
163   for (i = 0; i < AliRsnPID::kSpecies; i++)
164   {
165     daughter->SetPIDWeight(i, pid[i]);
166     sum += pid[i];
167   }
168   if (sum <= 0.0) return kFALSE;
169
170   // calculate N sigma to vertex
171   AliESDtrackCuts trkCut;
172   daughter->SetNSigmaToVertex(trkCut.GetSigmaToVertex(esdTrack));
173
174   return kTRUE;
175 }
176
177 //_____________________________________________________________________________
178 Bool_t AliRsnReader::ConvertTrack(AliRsnDaughter *daughter, AliAODTrack *aodTrack)
179 {
180 //
181 // Translates an AOD track into a RSN daughter
182 //
183
184   if (!aodTrack)
185   {
186     AliError("Passed NULL object: nothing done");
187     return kFALSE;
188   }
189
190   // copy momentum  and vertex
191   daughter->SetP(aodTrack->Px(), aodTrack->Py(), aodTrack->Pz());
192   daughter->SetV(aodTrack->Xv(), aodTrack->Yv(), aodTrack->Zv());
193
194   // chi2
195   daughter->SetChi2(aodTrack->Chi2perNDF());
196
197   // copy PID weights
198   Int_t i;
199   for (i = 0; i < 5; i++) daughter->SetPIDWeight(i, aodTrack->PID()[i]);
200
201   // copy flags
202   daughter->SetFlags(aodTrack->GetStatus());
203
204   // copy sign
205   daughter->SetCharge(aodTrack->Charge());
206
207   return kTRUE;
208 }
209
210 //_____________________________________________________________________________
211 Bool_t AliRsnReader::ConvertTrack(AliRsnDaughter *daughter, TParticle *particle)
212 {
213 //
214 // Translates an MC particle into a RSN daughter
215 //
216
217   if (!particle) return kFALSE;
218
219   // copy other MC info (mother PDG code cannot be retrieved here)
220   daughter->InitMCInfo(particle);
221
222   // copy momentum  and vertex
223   daughter->SetP(particle->Px(), particle->Py(), particle->Pz());
224   daughter->SetV(particle->Vx(), particle->Vy(), particle->Vz());
225
226   // recognize charge sign from PDG code sign
227   Int_t pdg = particle->GetPdgCode();
228   Int_t absPDG = TMath::Abs(pdg);
229   if (absPDG == 11 || absPDG == 13)
230   {
231     if (pdg > 0) daughter->SetCharge(-1); else daughter->SetCharge(1);
232   }
233   else if (absPDG == 211 || absPDG == 321 || absPDG == 2212)
234   {
235     if (pdg > 0) daughter->SetCharge(1); else daughter->SetCharge(-1);
236   }
237   else
238   {
239     // when trying to "adopt" a neutral track (photon, neutron, etc.)
240     // for the moment a "failed" message is returned
241     return kFALSE;
242   }
243
244   // flags and PID weights make no sense with MC tracks
245   daughter->SetFlags(0);
246   for (pdg = 0; pdg < AliRsnPID::kSpecies; pdg++) daughter->SetPIDWeight(pdg, 0.0);
247   daughter->SetPIDWeight(AliRsnPID::InternalType(absPDG), 1.0);
248
249   return kTRUE;
250 }
251
252 //_____________________________________________________________________________
253 Bool_t AliRsnReader::Fill
254 (AliRsnEvent *rsn, AliVEvent *event, AliMCEvent *mc)
255 {
256 //
257 // According to the class type of event and the selected source
258 // recalls one of the private reading methods to fill the RsnEvent
259 // passed by reference as first argument.
260 //
261
262   Bool_t success = kFALSE;
263   TString str(event->ClassName());
264
265   if (!str.CompareTo("AliESDEvent")) {
266     AliDebug(1, "Reading an ESD event");
267     success = FillFromESD(rsn, (AliESDEvent*)event, mc);
268   }
269   else if (!str.CompareTo("AliAODEvent")) {
270     AliDebug(1, "Reading an AOD event");
271     success = FillFromAOD(rsn, (AliAODEvent*)event, mc);
272   }
273   else if (!str.CompareTo("AliMCEvent")) {
274     AliDebug(1, "Reading an MC event");
275     success = FillFromMC(rsn, (AliMCEvent*)event);
276   }
277   else {
278     AliError(Form("ClassName '%s' not recognized as possible source data: aborting.", str.Data()));
279     return kFALSE;
280   }
281
282   return success;
283 }
284
285 //_____________________________________________________________________________
286 Bool_t AliRsnReader::FillFromESD(AliRsnEvent *rsn, AliESDEvent *esd, AliMCEvent *mc)
287 {
288 //
289 // Filler from an ESD event.
290 // Stores all tracks (if a filter is defined, it will store
291 // only the ones which survive the cuts).
292 // If a reference MC event is provided, it is used to store
293 // the MC informations for each track (true PDG code,
294 // GEANT label of mother, PDG code of mother, if any).
295 // When this is used, the 'source' flag of the output
296 // AliRsnEvent object will be set to 'kESD'.
297 //
298
299   // retrieve stack (if possible)
300   AliStack *stack = 0x0;
301   if (mc) stack = mc->Stack();
302
303   // set MC primary vertex if present
304   TArrayF fvertex(3);
305   Double_t mcvx = 0., mcvy = 0., mcvz = 0.;
306   if (mc) {
307     mc->GenEventHeader()->PrimaryVertex(fvertex);
308     mcvx = (Double_t)fvertex[0];
309     mcvy = (Double_t)fvertex[1];
310     mcvz = (Double_t)fvertex[2];
311     rsn->SetPrimaryVertexMC(mcvx, mcvy, mcvz);
312   }
313
314   // get number of tracks
315   Int_t ntracks = esd->GetNumberOfTracks();
316   if (!ntracks) return kFALSE;
317
318   // if required with the flag, scans the event
319   // and searches all split tracks (= 2 tracks with the same label);
320   // for each pair of split tracks, only the better (best chi2) is kept
321   // and the other is rejected: this info is stored into a Boolean array
322   Int_t i, i1, i2;
323   Bool_t *accept = new Bool_t[ntracks];
324   for (i = 0; i < ntracks; i++) accept[i] = kTRUE;
325   if (fCheckSplit) {
326     for (i1 = 0; i1 < ntracks; i1++) {
327       AliESDtrack *track1 = esd->GetTrack(i1);
328       for (i2 = i1 + 1; i2 < ntracks; i2++) {
329         AliESDtrack *track2 = esd->GetTrack(i2);
330         if (AreSplitted(track1, track2)) {
331           if (ResolveSplit(track1, track2)) accept[i2] = kFALSE;
332           else accept[i1] = kFALSE;
333         }
334       }
335     }
336   }
337
338   // get primary vertex
339   Double_t vertex[3];
340   if (!fTPCOnly) {
341     // when taking vertex from ESD event there are two options:
342     // if a vertex with tracks was successfully reconstructed,
343     // it is used for computing DCA;
344     // otherwise, the one computed with SPD is used.
345     // This is known from the "Status" parameter of the vertex itself.
346     const AliESDVertex *v = esd->GetPrimaryVertex();
347     if (!v->GetStatus()) v = esd->GetPrimaryVertexSPD();
348
349     // get primary vertex
350     vertex[0] = (Double_t)v->GetXv();
351     vertex[1] = (Double_t)v->GetYv();
352     vertex[2] = (Double_t)v->GetZv();
353   }
354   else {
355     const AliESDVertex *v = esd->GetPrimaryVertexTPC();
356
357     // get primary vertex
358     vertex[0] = (Double_t)v->GetXv();
359     vertex[1] = (Double_t)v->GetYv();
360     vertex[2] = (Double_t)v->GetZv();
361   }
362   rsn->SetPrimaryVertex(vertex[0], vertex[1], vertex[2]);
363
364   // store tracks from ESD
365   Float_t p[2], cov[3];
366   Int_t   index, label, labmum;
367   Bool_t  check;
368   AliRsnDaughter temp;
369   AliESDtrack *esdTrack = 0x0, esdTrackTmp;
370   for (index = 0; index < ntracks; index++) {
371     // skip track recognized as the worse one in a splitted pair
372     if (!accept[index]) {
373       AliDebug(10, Form("Rejecting split track #%d in this event", index));
374       continue;
375     }
376     // get ESD track
377     esdTrack = esd->GetTrack(index);
378     // check for ESD track cuts (if required)
379     if (fUseESDTrackCuts && (!fESDTrackCuts.AcceptTrack(esdTrack))) continue;
380     // check for fake tracks
381     label = esdTrack->GetLabel();
382     if (fRejectFakes && (label < 0)) continue;
383     // check for minimum number of clusters in ITS, TPC and TRD
384     if (fITSClusters > 0 && (esdTrack->GetITSclusters(0x0) < fITSClusters)) if (!fTPCOnly) continue;
385     if (fTPCClusters > 0 && (esdTrack->GetTPCclusters(0x0) < fTPCClusters)) continue;
386     if (fTRDClusters > 0 && (esdTrack->GetTRDclusters(0x0) < fTRDClusters)) if (!fTPCOnly) continue;
387     // switch to TPC data if required
388     if (fTPCOnly) {
389       esdTrack->GetImpactParametersTPC(p, cov);
390       if (p[0] == 0. && p[1] == 0.) {
391         esdTrack->RelateToVertexTPC(esd->GetPrimaryVertexTPC(), esd->GetMagneticField(), kVeryBig);
392       }
393       check = esdTrack->FillTPCOnlyTrack(esdTrackTmp);
394       if (check) esdTrack = &esdTrackTmp; else continue;
395     }
396     // try to get information from this track
397     // output value tells if this was successful or not
398     check = ConvertTrack(&temp, esdTrack);
399     if (!check) {
400       AliDebug(10, Form("Failed adopting track #%d", index));
401       continue;
402     }
403
404     // if stack is present, copy MC info
405     if (stack) {
406       TParticle *part = stack->Particle(TMath::Abs(label));
407       if (part) {
408         temp.InitMCInfo(part);
409         labmum = part->GetFirstMother();
410         if (labmum >= 0) {
411           TParticle *mum = stack->Particle(labmum);
412           temp.GetMCInfo()->SetMotherPDG(mum->GetPdgCode());
413         }
414       }
415     }
416
417     // set index and label
418     temp.SetIndex(index);
419     temp.SetLabel(label);
420
421     // check this object against the Rsn cuts (if required)
422     if (fUseRsnTrackCuts) {
423       if (!fRsnTrackCuts.IsSelected(AliRsnCut::kParticle, &temp)) continue;
424     }
425
426     // try to add track to collection and returns an error in case of problems
427     AliRsnDaughter *ptr = rsn->AddTrack(temp);
428     if (!ptr) AliWarning(Form("Failed storing track#%d", index));
429   }
430
431   // compute total multiplicity
432   rsn->MakeComputations();
433   if (rsn->GetMultiplicity() <= 0) {
434       AliDebug(1, "Zero Multiplicity in this event");
435       return kFALSE;
436   }
437
438   // sort tracks w.r. to Pt (from largest to smallest)
439   //rsn->SortTracks();
440
441   // correct tracks for primary vertex
442   //rsn->CorrectTracks();
443
444   return kTRUE;
445 }
446
447 //_____________________________________________________________________________
448 Bool_t AliRsnReader::FillFromAOD(AliRsnEvent *rsn, AliAODEvent *aod, AliMCEvent *mc)
449 {
450 //
451 // Filler from an AOD event.
452 // Stores all tracks (if a filter is defined, it will store
453 // only the ones which survive the cuts).
454 // If a reference MC event is provided, it is used to store
455 // the MC informations for each track (true PDG code,
456 // GEANT label of mother, PDG code of mother, if any).
457 // When this is used, the 'source' flag of the output
458 // AliRsnEvent object will be set to 'kAOD'.
459 //
460
461     // retrieve stack (if possible)
462     AliStack *stack = 0x0;
463     if (mc) stack = mc->Stack();
464
465     // get number of tracks
466     Int_t ntracks = aod->GetNTracks();
467     if (!ntracks) return kFALSE;
468
469     // get primary vertex
470     Double_t vertex[3];
471     vertex[0] = aod->GetPrimaryVertex()->GetX();
472     vertex[1] = aod->GetPrimaryVertex()->GetY();
473     vertex[2] = aod->GetPrimaryVertex()->GetZ();
474     rsn->SetPrimaryVertex(vertex[0], vertex[1], vertex[2]);
475
476     // set MC primary vertex if present
477     TArrayF fvertex(3);
478     Double_t mcvx = 0., mcvy = 0., mcvz = 0.;
479     if (mc) {
480       mc->GenEventHeader()->PrimaryVertex(fvertex);
481       mcvx = (Double_t)fvertex[0];
482       mcvy = (Double_t)fvertex[1];
483       mcvz = (Double_t)fvertex[2];
484       rsn->SetPrimaryVertexMC(mcvx, mcvy, mcvz);
485     }
486
487     // store tracks from ESD
488     Int_t  index, label, labmum;
489     Bool_t check;
490     AliAODTrack *aodTrack = 0;
491     AliRsnDaughter temp;
492     TObjArrayIter iter(aod->GetTracks());
493     while ((aodTrack = (AliAODTrack*)iter.Next()))
494     {
495         // retrieve index
496         index = aod->GetTracks()->IndexOf(aodTrack);
497         label = aodTrack->GetLabel();
498         if (fRejectFakes && (label < 0)) continue;
499         // copy ESD track data into RsnDaughter
500         // if unsuccessful, this track is skipped
501         check = ConvertTrack(&temp, aodTrack);
502         if (!check) continue;
503         // if stack is present, copy MC info
504         if (stack)
505         {
506             TParticle *part = stack->Particle(TMath::Abs(label));
507             if (part)
508             {
509                 temp.InitMCInfo(part);
510                 labmum = part->GetFirstMother();
511                 if (labmum >= 0)
512                 {
513                     TParticle *mum = stack->Particle(labmum);
514                     temp.GetMCInfo()->SetMotherPDG(mum->GetPdgCode());
515                 }
516             }
517         }
518         // set index and label and add this object to the output container
519         temp.SetIndex(index);
520         temp.SetLabel(label);
521
522         // check this object against the Rsn cuts (if required)
523         if (fUseRsnTrackCuts) {
524           if (!fRsnTrackCuts.IsSelected(AliRsnCut::kParticle, &temp)) continue;
525         }
526
527         AliRsnDaughter *ptr = rsn->AddTrack(temp);
528         // if problems occurred while storin, that pointer is NULL
529         if (!ptr) AliWarning(Form("Failed storing track#%d"));
530     }
531
532     // compute total multiplicity
533     rsn->MakeComputations();
534     if (rsn->GetMultiplicity() <= 0)
535     {
536         AliDebug(1, "Zero multiplicity in this event");
537         return kFALSE;
538     }
539
540     // correct tracks for primary vertex
541     rsn->CorrectTracks();
542
543     return kTRUE;
544 }
545
546 //_____________________________________________________________________________
547 Bool_t AliRsnReader::FillFromMC(AliRsnEvent *rsn, AliMCEvent *mc)
548 {
549 //
550 // Filler from an ESD event.
551 // Stores all tracks which generate at least one
552 // TrackReference (a point in a sensitive volume).
553 // In this case, the MC info is stored by default and
554 // perfect particle identification is the unique available.
555 // When this is used, the 'source' flag of the output
556 // AliRsnEvent object will be set to 'kMC'.
557 //
558
559   // get number of tracks
560   Int_t ntracks = mc->GetNumberOfTracks();
561   if (!ntracks) return kFALSE;
562
563   AliStack *stack = mc->Stack();
564
565   // get primary vertex
566   TArrayF fvertex(3);
567   Double_t vx, vy, vz;
568   mc->GenEventHeader()->PrimaryVertex(fvertex);
569   vx = (Double_t)fvertex[0];
570   vy = (Double_t)fvertex[1];
571   vz = (Double_t)fvertex[2];
572   rsn->SetPrimaryVertex(vx, vy, vz);
573   rsn->SetPrimaryVertexMC(vx, vy, vz);
574
575   // store tracks from MC
576   Int_t  i, index, labmum, nRef, nRefITS, nRefTPC, detectorID;
577   AliRsnDaughter temp;
578   for (index = 0; index < ntracks; index++)
579   {
580     // get MC track & take index and label
581     AliMCParticle *mcTrack = mc->GetTrack(index);
582     temp.SetIndex(index);
583     temp.SetLabel(mcTrack->Label());
584
585     // check particle track references
586     nRef = mcTrack->GetNumberOfTrackReferences();
587     for (i = 0, nRefITS = 0, nRefTPC = 0; i < nRef; i++) {
588       AliTrackReference *trackRef = mcTrack->GetTrackReference(i);
589       if (!trackRef) continue;
590       detectorID = trackRef->DetectorId();
591       switch (detectorID) {
592         case AliTrackReference::kITS  : nRefITS++; break;
593         case AliTrackReference::kTPC  : nRefTPC++; break;
594         default: break;
595       }
596     }
597     if (fTrackRefs > 0 && nRef < fTrackRefs) continue;
598     if (fTrackRefsITS > 0 && nRefITS < fTrackRefsITS) continue;
599     if (fTrackRefsTPC > 0 && nRefTPC < fTrackRefsTPC) continue;
600
601     // try to insert in the RsnDaughter its data
602     TParticle *mcPart = mcTrack->Particle();
603     if (!ConvertTrack(&temp, mcPart)) continue;
604
605     // assign mother label and PDG code (if any)
606     labmum = temp.GetMCInfo()->Mother();
607     if (labmum >= 0) {
608       TParticle *mum = stack->Particle(labmum);
609       temp.GetMCInfo()->SetMotherPDG(mum->GetPdgCode());
610     }
611
612     // check this object against the Rsn cuts (if required)
613     if (fUseRsnTrackCuts) {
614       if (!fRsnTrackCuts.IsSelected(AliRsnCut::kParticle, &temp)) continue;
615     }
616
617     AliRsnDaughter *ptr = rsn->AddTrack(temp);
618     if (!ptr) AliWarning(Form("Failed storing track #%d", index));
619   }
620
621   // compute total multiplicity
622   rsn->MakeComputations();
623   if (rsn->GetMultiplicity() <= 0)
624   {
625     AliDebug(1, "Zero multiplicity in this event");
626     return kFALSE;
627   }
628
629   // sort tracks w.r. to Pt (from largest to smallest)
630   rsn->SortTracks();
631
632   // correct tracks for primary vertex
633   rsn->CorrectTracks();
634
635   return kTRUE;
636 }