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