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(const char *name, const char *title)
  76   :AliTRDrecoTask(name, title)
  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.5)
  88 {
  89   //
  90   // Default constructor
  91   //
  92
  93   fReconstructor = new AliTRDReconstructor();
  94   fReconstructor->SetRecoParam(AliTRDrecoParam::GetLowFluxParam());
  95   memset(fdEdx, 0, 10*sizeof(Float_t));
  96   memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
  97
  98   DefineInput(1, TObjArray::Class());
  99   DefineInput(2, TObjArray::Class());
 100   DefineOutput(1, TTree::Class());
 101 }
 102
 103
 104 //________________________________________________________________________
 105 AliTRDpidRefMaker::~AliTRDpidRefMaker()
 106 {
 107   if(fPIDdataArray) delete fPIDdataArray;
 108   if(fReconstructor) delete fReconstructor;
 109 }
 110
 111 //________________________________________________________________________
 112 Bool_t AliTRDpidRefMaker::CheckQuality(AliTRDseedV1* /*trklt*/)
 113 {
 114 // Place holder for checking tracklet quality for PID.
 115   return kTRUE;
 116 }
 117
 118
 119 //________________________________________________________________________
 120 Float_t* AliTRDpidRefMaker::CookdEdx(AliTRDseedV1 *trklt)
 121 {
 122   trklt->CookdEdx(AliTRDpidUtil::kNNslices);
 123   memcpy(fdEdx, trklt->GetdEdx(), AliTRDpidUtil::kNNslices*sizeof(Float_t));
 124   return fdEdx;
 125 }
 126
 127
 128 //________________________________________________________________________
 129 void AliTRDpidRefMaker::ConnectInputData(Option_t *opt)
 130 {
 131   AliTRDrecoTask::ConnectInputData(opt);
 132   fV0s  = dynamic_cast<TObjArray*>(GetInputData(1));
 133   fInfo = dynamic_cast<TObjArray*>(GetInputData(2));
 134 }
 135
 136 //________________________________________________________________________
 137 void AliTRDpidRefMaker::CreateOutputObjects()
 138 {
 139   // Create histograms
 140   // Called once
 141
 142   OpenFile(0, "RECREATE");
 143   fContainer = new TObjArray();
 144   fContainer->SetName(Form("Moni%s", GetName()));
 145
 146   TH2 *h2 = new TH2I("hPDG","Particle abundance", AliPID::kSPECIES, -0.5, 4.5, AliTRDCalPID::kNMom, -0.5, AliTRDCalPID::kNMom-0.5);
 147   TAxis *ax = h2->GetXaxis();
 148   ax->SetNdivisions(505);
 149   ax->SetTitle("Particle species");
 150   for(Int_t is=AliPID::kSPECIES; is--;) ax->SetBinLabel(is+1, AliPID::ParticleShortName(is));
 151   h2->GetYaxis()->SetTitle("P bins");
 152   h2->GetYaxis()->SetNdivisions(511);
 153
 154   fContainer->AddAt(h2, 0);
 155
 156   fData = new TTree(GetName(), Form("Reference data for %s", GetName()));
 157   fData->Branch("s", &fPIDbin, "s/b");
 158   fData->Branch("data", &fPIDdataArray);
 159 }
 160
 161 //________________________________________________________________________
 162 void AliTRDpidRefMaker::Exec(Option_t *)
 163 {
 164   // Main loop
 165   // Called for each event
 166
 167   AliTRDtrackInfo     *track = NULL;
 168   AliTRDtrackV1    *trackTRD = NULL;
 169   AliTrackReference     *ref = NULL;
 170   const AliTRDtrackInfo::AliESDinfo *infoESD = NULL;
 171   for(Int_t itrk=0; itrk<fTracks->GetEntriesFast(); itrk++){
 172     track = (AliTRDtrackInfo*)fTracks->UncheckedAt(itrk);
 173     if(!track->HasESDtrack()) continue;
 174     trackTRD = track->GetTrack();
 175     infoESD  = track->GetESDinfo();
 176     Double32_t *infoPID = infoESD->GetSliceIter();
 177     Int_t n = infoESD->GetNSlices() - AliTRDgeometry::kNlayer;
 178     Double32_t *p = &infoPID[n];
 179     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]));
 180
 181     ULong_t status = track->GetStatus();
 182     if(!(status&AliESDtrack::kTPCout)) continue;
 183
 184     // fill the pid information
 185     SetRefPID(fRefPID, track, fPID);
 186
 187     // prepare PID data array
 188     if(!fPIDdataArray){
 189       fPIDdataArray = new AliTRDpidRefDataArray();
 190     } else fPIDdataArray->Reset();
 191
 192     // fill PID information
 193     for(Int_t ily = 0; ily < AliTRDgeometry::kNlayer; ily++){
 194
 195       // fill P & dE/dx information
 196       if(HasFriends()){ // from TRD track
 197         if(!trackTRD) continue;
 198         AliTRDseedV1 *trackletTRD(NULL);
 199         trackTRD -> SetReconstructor(fReconstructor);
 200         if(!(trackletTRD = trackTRD -> GetTracklet(ily))) continue;
 201         if(!CheckQuality(trackletTRD)) continue;
 202         if(!CookdEdx(trackletTRD)) continue;
 203
 204         // fill momentum information
 205         fP = 0.;
 206         switch(fRefP){
 207         case kMC:
 208           if(!(ref = track->GetTrackRef(trackletTRD))) continue;
 209           fP = ref->P();
 210           break;
 211         case kRec:
 212           fP = trackletTRD->GetMomentum();
 213           break;
 214         default: continue;
 215         }
 216       } else { // from ESD track
 217         // fill momentum information
 218         switch(fRefP){
 219         case kMC:
 220           if(!(ref = track->GetTrackRef(ily))) continue;
 221           fP = ref->P();
 222           break;
 223         case kRec:
 224           fP = p[ily];
 225           break;
 226         default: continue;
 227         }
 228         Double32_t *it = &infoPID[ily*AliTRDCalPID::kNSlicesNN];
 229         for(Int_t is=AliTRDCalPID::kNSlicesNN; is--; it++) fdEdx[is] = (*it);
 230       }
 231
 232       // momentum threshold
 233       if(fP < fPthreshold) continue;
 234
 235       // store information
 236       fPIDdataArray->PushBack(ily, AliTRDpidUtil::GetMomentumBin(fP), fdEdx);
 237     }
 238
 239     Fill();
 240   }
 241
 242   PostData(0, fContainer);
 243   PostData(1, fData);
 244 }
 245
 246
 247 //________________________________________________________________________
 248 void AliTRDpidRefMaker::Fill()
 249 {
 250 // Fill data tree
 251
 252   if(!fPIDdataArray->fNtracklets) return;
 253   fPIDbin = TMath::LocMax(AliPID::kSPECIES, fPID); // get particle type
 254 // Fill data tree
 255   fData->Fill();
 256
 257
 258   // fill monitor
 259   for(Int_t ily=fPIDdataArray->fNtracklets; ily--;){
 260     Int_t pBin = fPIDdataArray->fData[ily].fPLbin & 0xf;
 261     ((TH2*)fContainer->At(0))->Fill(fPIDbin, pBin);
 262   }
 263 }
 264
 265 //________________________________________________________________________
 266 void AliTRDpidRefMaker::LinkPIDdata()
 267 {
 268 // Link data tree to data members
 269   fData->SetBranchAddress("s", &fPIDbin);
 270   fData->SetBranchAddress("data", &fPIDdataArray);
 271 }
 272
 273 //________________________________________________________________________
 274 void AliTRDpidRefMaker::SetRefPID(ETRDpidRefMakerSource select, void *source, Float_t *pid)
 275 {
 276 // Fill the reference PID values "pid" from "source" object
 277 // according to the option "select". Possible options are
 278 // - kV0  - v0 based PID
 279 // - kMC  - MC truth [default]
 280 // - kRec - outside detectors
 281
 282
 283   memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
 284   switch(select){
 285   case kV0:
 286     {
 287       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
 288       if(!track){
 289         AliError("No trackInfo found");
 290         return;
 291       }
 292
 293       //Get V0 PID decisions from the AliTRDv0Info for all particle species (implemented so far : electrons from conversions, pions from K0s and protons from Lambdas) :
 294       AliTRDv0Info v0;
 295       for(Int_t is=AliPID::kSPECIES; is--;) fPID[is] = v0.GetV0PID(is, track);
 296     }
 297     break;
 298   case kMC:
 299     if(!HasMCdata()){
 300       AliError("Could not retrive reference PID from MC");
 301       return;
 302     }
 303     {
 304       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
 305       switch(track->GetPDG()){
 306       case kElectron:
 307       case kPositron:
 308         fPID[AliPID::kElectron] = 1.;
 309         break;
 310       case kMuonPlus:
 311       case kMuonMinus:
 312         fPID[AliPID::kMuon] = 1.;
 313         break;
 314       case kPiPlus:
 315       case kPiMinus:
 316         fPID[AliPID::kPion] = 1.;
 317         break;
 318       case kKPlus:
 319       case kKMinus:
 320         fPID[AliPID::kKaon] = 1.;
 321         break;
 322       case kProton:
 323       case kProtonBar:
 324         fPID[AliPID::kProton] = 1.;
 325         break;
 326       }
 327     }
 328     break;
 329   case kRec:
 330     {
 331       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
 332       AliTRDtrackV1 *trackTRD = track->GetTrack();
 333       trackTRD -> SetReconstructor(fReconstructor);
 334       //fReconstructor -> SetOption("nn");
 335       trackTRD -> CookPID();
 336       for(Int_t iPart = 0; iPart < AliPID::kSPECIES; iPart++){
 337         pid[iPart] = trackTRD -> GetPID(iPart);
 338         AliDebug(4, Form("PDG is (in V0info) %d %f", iPart, pid[iPart]));
 339       }
 340     }
 341     break;
 342   default:
 343     AliWarning("PID reference source not implemented");
 344     return;
 345   }
 346   AliDebug(4, Form("Ref PID [%] : %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f]"
 347     ,AliPID::ParticleShortName(0), 1.e2*fPID[0]
 348     ,AliPID::ParticleShortName(1), 1.e2*fPID[1]
 349     ,AliPID::ParticleShortName(2), 1.e2*fPID[2]
 350     ,AliPID::ParticleShortName(3), 1.e2*fPID[3]
 351     ,AliPID::ParticleShortName(4), 1.e2*fPID[4]
 352   ));
 353 }
 354
 355 //________________________________________________________________________
 356 void AliTRDpidRefMaker::SetAbundance(Float_t train)
 357 {
 358 // Split data sample between trainning and testing
 359
 360   if(train<0. || train >1.){
 361     AliWarning("The input data should be in the interval [0, 1]");
 362     return;
 363   }
 364
 365   fFreq = train;
 366 }
 367