PWGPP/TRD/AliTRDpidRefMaker.cxx

   1 #include "TPDGCode.h"
   2 #include "TFile.h"
   3 #include "TTree.h"
   4 #include "TEventList.h"
   5 #include "TObjArray.h"
   6 #include "TH2.h"
   7
   8 #include "AliLog.h"
   9 #include "AliESDtrack.h"
  10 #include "AliTrackReference.h"
  11 #include "AliAnalysisManager.h"
  12
  13 #include "AliTRDReconstructor.h"
  14 #include "AliTRDtrackV1.h"
  15 #include "AliTRDseedV1.h"
  16 #include "AliTRDpidRefMaker.h"
  17 #include "AliTRDinfoGen.h"
  18 #include "info/AliTRDeventInfo.h"
  19 #include "info/AliTRDv0Info.h"
  20 #include "info/AliTRDpidInfo.h"
  21
  22
  23 // Defines and implements basic functionality for building reference data for TRD PID.
  24 //
  25 // Here is the list of functionality provided by this class
  26 // 1.
  27 // 2.
  28 //
  29 // Authors:
  30 //   Alex Bercuci <A.Bercuci@gsi.de>
  31 //   Alex Wilk    <wilka@uni-muenster.de>
  32 //   Markus Fasel <mfasel@gsi.de>
  33 //   Markus Heide <mheide@uni-muenster.de>
  34 //
  35
  36 ClassImp(AliTRDpidRefMaker)
  37
  38 //________________________________________________________________________
  39 AliTRDpidRefMaker::AliTRDpidRefMaker()
  40   :AliTRDrecoTask()
  41   ,fV0s(NULL)
  42   ,fData(NULL)
  43   ,fInfo(NULL)
  44   ,fPIDdataArray(NULL)
  45   ,fRefPID(kV0)
  46   ,fRefP(kRec)
  47   ,fFreq(1.)
  48   ,fP(-1.)
  49   ,fPthreshold(0.)
  50 {
  51   //
  52   // Default constructor
  53   //
  54   SetNameTitle("PIDrefMaker", "PID Reference Maker");
  55 }
  56
  57 //________________________________________________________________________
  58 AliTRDpidRefMaker::AliTRDpidRefMaker(const char *name, const char *title)
  59   :AliTRDrecoTask(name, title)
  60   ,fV0s(NULL)
  61   ,fData(NULL)
  62   ,fInfo(NULL)
  63   ,fPIDdataArray(NULL)
  64   ,fRefPID(kV0)
  65   ,fRefP(kRec)
  66   ,fFreq(1.)
  67   ,fP(-1.)
  68   ,fPthreshold(0.5)
  69 {
  70   //
  71   // Default constructor
  72   //
  73
  74   memset(fdEdx, 0, AliTRDpidUtil::kNNslices*sizeof(Float_t));
  75   memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
  76
  77   DefineInput(3, TObjArray::Class()); // v0 list
  78   DefineInput(4, TObjArray::Class()); // pid info list
  79   DefineOutput(2, TTree::Class());
  80 }
  81
  82
  83 //________________________________________________________________________
  84 AliTRDpidRefMaker::~AliTRDpidRefMaker()
  85 {
  86   if(fPIDdataArray) delete fPIDdataArray;
  87 }
  88
  89 //________________________________________________________________________
  90 Bool_t AliTRDpidRefMaker::CheckQuality(AliTRDseedV1* /*trklt*/)
  91 {
  92 // Place holder for checking tracklet quality for PID.
  93   return kTRUE;
  94 }
  95
  96
  97 //________________________________________________________________________
  98 Float_t* AliTRDpidRefMaker::CookdEdx(AliTRDseedV1 *trklt)
  99 {
 100   trklt->CookdEdx(AliTRDpidUtil::kNNslices);
 101   memcpy(fdEdx, trklt->GetdEdx(), AliTRDpidUtil::kNNslices*sizeof(Float_t));
 102   return fdEdx;
 103 }
 104
 105
 106 //________________________________________________________________________
 107 void AliTRDpidRefMaker::UserCreateOutputObjects()
 108 {
 109   // Create histograms
 110   // Called once
 111
 112   fContainer = new TObjArray();
 113   fContainer->SetName(Form("Moni%s", GetName()));
 114
 115   TH2 *h2 = new TH2I("hPDG","Particle abundance", AliPID::kSPECIES, -0.5, 4.5, AliTRDCalPID::kNMom, -0.5, AliTRDCalPID::kNMom-0.5);
 116   TAxis *ax = h2->GetXaxis();
 117   ax->SetNdivisions(505);
 118   ax->SetTitle("Particle species");
 119   for(Int_t is=AliPID::kSPECIES; is--;) ax->SetBinLabel(is+1, AliPID::ParticleShortName(is));
 120   h2->GetYaxis()->SetTitle("P bins");
 121   h2->GetYaxis()->SetNdivisions(511);
 122   fContainer->AddAt(h2, 0);
 123   PostData(1, fContainer);
 124
 125   OpenFile(2);
 126   fData = new TTree("RefPID", "Reference data for PID");
 127   fData->Branch("data", &fPIDdataArray);
 128   PostData(2, fData);
 129 }
 130
 131 //________________________________________________________________________
 132 void AliTRDpidRefMaker::UserExec(Option_t *)
 133 {
 134   // Main loop
 135   // Called for each event
 136   Int_t ev((Int_t)AliAnalysisManager::GetAnalysisManager()->GetCurrentEntry());
 137   if(!(fTracks = dynamic_cast<TObjArray*>(GetInputData(1)))){
 138     AliDebug(3, Form("Missing tracks container in ev %d", ev));
 139     return;
 140   }
 141   if(!(fEvent = dynamic_cast<AliTRDeventInfo*>(GetInputData(2)))){
 142     AliDebug(3, Form("Missing Event Info container in ev %d", ev));
 143     return;
 144   }
 145   if(!(fV0s    = dynamic_cast<TObjArray*>(GetInputData(3)))){
 146     AliDebug(3, Form("Missing v0 container in ev %d", ev));
 147     return;
 148   }
 149   if(!(fInfo   = dynamic_cast<TObjArray*>(GetInputData(4)))){
 150     AliDebug(3, Form("Missing pid info container in ev %d", ev));
 151     return;
 152   }
 153
 154   AliDebug(1, Form("Entries: Ev[%d] Tracks[%d] V0[%d] PID[%d]", ev, fTracks->GetEntriesFast(), fV0s->GetEntriesFast(), fInfo->GetEntriesFast()));
 155   AliTRDtrackInfo     *track = NULL;
 156   //AliTRDtrackV1    *trackTRD = NULL;
 157   AliTrackReference     *ref = NULL;
 158   const AliTRDtrackInfo::AliESDinfo *infoESD = NULL;
 159   for(Int_t itrk=0; itrk<fTracks->GetEntriesFast(); itrk++){
 160     track = (AliTRDtrackInfo*)fTracks->UncheckedAt(itrk);
 161     if(!track->HasESDtrack()) continue;
 162     //trackTRD = track->GetTrack();
 163     infoESD  = track->GetESDinfo();
 164     Double32_t *infoPID = infoESD->GetSliceIter();
 165     Int_t n = infoESD->GetNSlices() - AliTRDgeometry::kNlayer;
 166     if(n==0){
 167       AliWarning(Form("dEdx info missing in ESD track %d", itrk));
 168       continue;
 169     }
 170     Double32_t *p = &infoPID[n];
 171     AliDebug(4, Form("n[%d] p[GeV/c]{%6.2f %6.2f %6.2f %6.2f %6.2f %6.2f}", n, p[0], p[1], p[2], p[3], p[4], p[5]));
 172
 173     ULong_t status = track->GetStatus();
 174     if(!(status&AliESDtrack::kTRDpid)) continue;
 175
 176     // fill the pid information
 177     SetRefPID(fRefPID, track, infoESD, fPID);
 178     // get particle type
 179     Int_t idx(TMath::Max(Long64_t(0), TMath::LocMax(AliPID::kSPECIES, fPID)));
 180     if(fPID[idx]<1.e-5) continue;
 181
 182     // prepare PID data array
 183     if(!fPIDdataArray){
 184       fPIDdataArray = new AliTRDpidInfo();
 185     } else fPIDdataArray->Reset();
 186     fPIDdataArray->SetPID(idx);
 187
 188     // fill PID information
 189     for(Int_t ily = 0; ily < AliTRDgeometry::kNlayer; ily++){
 190
 191       // fill P & dE/dx information
 192       switch(fRefP){
 193       case kMC:
 194         if(!(ref = track->GetTrackRef(ily))) continue;
 195         fP = ref->P();
 196         break;
 197       case kRec:
 198         fP = p[ily];
 199         break;
 200       default:
 201         continue;
 202       }
 203       Double32_t *it = &infoPID[ily*AliTRDCalPID::kNSlicesNN];
 204       for(Int_t is=AliTRDCalPID::kNSlicesNN; is--; it++) fdEdx[is] = (*it);
 205
 206       // momentum threshold
 207       if(fP < fPthreshold) continue;
 208
 209       // store information
 210       fPIDdataArray->PushBack(ily, AliTRDpidUtil::GetMomentumBin(fP), fdEdx);
 211     }
 212
 213     Fill();
 214   }
 215 }
 216
 217
 218 //________________________________________________________________________
 219 void AliTRDpidRefMaker::Fill()
 220 {
 221 // Fill data tree
 222
 223   if(!fPIDdataArray->GetNtracklets()) return;
 224   // Fill data tree
 225   fData->Fill();
 226
 227
 228   // fill monitor
 229   for(Int_t itrklt=fPIDdataArray->GetNtracklets(); itrklt--;){
 230     Int_t pBin = fPIDdataArray->GetData(itrklt)->Momentum();
 231     ((TH2*)fContainer->At(0))->Fill(fPIDdataArray->GetPID(), pBin);
 232   }
 233 }
 234
 235 //________________________________________________________________________
 236 void AliTRDpidRefMaker::LinkPIDdata()
 237 {
 238 // Link data tree to data members
 239   fData->SetBranchAddress("data", &fPIDdataArray);
 240 }
 241
 242 //________________________________________________________________________
 243 void AliTRDpidRefMaker::SetRefPID(ETRDpidRefMakerSource select, AliTRDtrackInfo *track, const AliTRDtrackInfo::AliESDinfo *infoESD, Float_t *pid)
 244 {
 245 // Fill the reference PID values "pid" from "source" object
 246 // according to the option "select". Possible options are
 247 // - kV0  - v0 based PID
 248 // - kMC  - MC truth [default]
 249 // - kRec - outside detectors
 250
 251   if(!track){
 252     AliError("No trackInfo found");
 253     return;
 254   }
 255   memset(pid, 0, AliPID::kSPECIES*sizeof(Float_t));
 256   switch(select){
 257   case kV0:
 258     {
 259       //Get V0 PID decisions for all particle species (implemented so far : electrons from conversions, pions from K0s and protons from Lambdas) :
 260       if(!infoESD->HasV0()) return;
 261       const Int_t *v0pid=infoESD->GetV0pid();
 262       for(Int_t is=AliPID::kSPECIES; is--;){ pid[is] = (Float_t)v0pid[is];}
 263     }
 264     break;
 265   case kMC:
 266     if(!HasMCdata()){
 267       AliError("Could not retrive reference PID from MC");
 268       return;
 269     }
 270     switch(track->GetPDG()){
 271     case kElectron:
 272     case kPositron:
 273       pid[AliPID::kElectron] = 1.;
 274       break;
 275     case kMuonPlus:
 276     case kMuonMinus:
 277       pid[AliPID::kMuon] = 1.;
 278       break;
 279     case kPiPlus:
 280     case kPiMinus:
 281       pid[AliPID::kPion] = 1.;
 282       break;
 283     case kKPlus:
 284     case kKMinus:
 285       pid[AliPID::kKaon] = 1.;
 286       break;
 287     case kProton:
 288     case kProtonBar:
 289       pid[AliPID::kProton] = 1.;
 290       break;
 291     }
 292     break;
 293   case kRec:
 294     {
 295       AliTRDtrackV1 *trackTRD = track->GetTrack();
 296       trackTRD -> SetReconstructor(AliTRDinfoGen::Reconstructor());
 297       //fReconstructor -> SetOption("nn");
 298       trackTRD -> CookPID();
 299       for(Int_t iPart = 0; iPart < AliPID::kSPECIES; iPart++){
 300         pid[iPart] = trackTRD -> GetPID(iPart);
 301         AliDebug(4, Form("PDG is (in V0info) %d %f", iPart, pid[iPart]));
 302       }
 303     }
 304     break;
 305   default:
 306     AliWarning("PID reference source not implemented");
 307     return;
 308   }
 309   AliDebug(4, Form("Ref PID : %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f]"
 310     ,AliPID::ParticleShortName(0), 1.e2*pid[0]
 311     ,AliPID::ParticleShortName(1), 1.e2*pid[1]
 312     ,AliPID::ParticleShortName(2), 1.e2*pid[2]
 313     ,AliPID::ParticleShortName(3), 1.e2*pid[3]
 314     ,AliPID::ParticleShortName(4), 1.e2*pid[4]
 315   ));
 316 }
 317
 318 //________________________________________________________________________
 319 void AliTRDpidRefMaker::SetAbundance(Float_t train)
 320 {
 321 // Split data sample between trainning and testing
 322
 323   if(train<0. || train >1.){
 324     AliWarning("The input data should be in the interval [0, 1]");
 325     return;
 326   }
 327
 328   fFreq = train;
 329 }
 330