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