PWG0/esdTrackCuts/AliCutTask.cxx

   1 #include "TChain.h"
   2 #include "TTree.h"
   3 #include "TH1F.h"
   4 #include "TCanvas.h"
   5 #include "TList.h"
   6 #include "TFile.h"
   7 #include "TParticle.h"
   8 #include "TParticlePDG.h"
   9
  10 #include "esdTrackCuts/AliESDtrackCuts.h"
  11 #include "AliPWG0Helper.h"
  12
  13 #include "AliAnalysisTask.h"
  14 #include "AliAnalysisManager.h"
  15 #include "AliESDEvent.h"
  16 #include "AliESDInputHandler.h"
  17 #include "AliESDVertex.h"
  18
  19 #include "AliMCEventHandler.h"
  20 #include "AliMCEvent.h"
  21 #include "AliStack.h"
  22 #include "AliLog.h"
  23
  24 #include "AliCutTask.h"
  25
  26 // simple task that runs the esd track cuts to evaluate the basic plots created during the cuts
  27
  28 ClassImp(AliCutTask)
  29
  30 //________________________________________________________________________
  31 AliCutTask::AliCutTask(const char *name) :
  32   AliAnalysisTask(name, ""), fESD(0), fTrackCuts(0),
  33   fAnalysisMode(AliPWG0Helper::kTPCITS),
  34   fTrackCutsPrimaries(0),
  35   fTrackCutsSecondaries(0), fVertex(0), fTriggerStats(0), fOutput(0),
  36   fPrimStats(0)
  37 {
  38   // Constructor
  39
  40   // Define input and output slots here
  41   DefineInput(0, TChain::Class());
  42   DefineOutput(0, TList::Class());
  43 }
  44
  45 //________________________________________________________________________
  46 void AliCutTask::ConnectInputData(Option_t *)
  47 {
  48   // Connect ESD or AOD here
  49   // Called once
  50
  51   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
  52   if (!tree) {
  53     Printf("ERROR: Could not read chain from input slot 0");
  54   } else {
  55     // Disable all branches and enable only the needed ones
  56     //tree->SetBranchStatus("*", kFALSE);
  57
  58     // old esd
  59     tree->SetBranchStatus("fTriggerMask", kTRUE);
  60     tree->SetBranchStatus("fTracks.*", kTRUE);
  61     tree->SetBranchStatus("fSPDVertex*", kTRUE);
  62
  63     // new esd
  64     tree->SetBranchStatus("TriggerMask", kTRUE);
  65     tree->SetBranchStatus("AliESDHeader", kTRUE);
  66     //AliPWG0Helper::SetBranchStatusRecursive(tree, "Tracks", kTRUE);
  67     //AliPWG0Helper::SetBranchStatusRecursive(tree, "SPDVertex", kTRUE);
  68
  69     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  70
  71     if (!esdH) {
  72       Printf("ERROR: Could not get ESDInputHandler");
  73     } else
  74       fESD = esdH->GetEvent();
  75   }
  76 }
  77
  78 //________________________________________________________________________
  79 void AliCutTask::CreateOutputObjects()
  80 {
  81   // Create histograms
  82   // Called once
  83
  84   fOutput = new TList;
  85   fOutput->SetOwner();
  86
  87   fOutput->Add(fTrackCuts);
  88   if (fTrackCutsPrimaries)
  89     fOutput->Add(fTrackCutsPrimaries);
  90   if (fTrackCutsSecondaries)
  91     fOutput->Add(fTrackCutsSecondaries);
  92
  93   fVertex = new TH1F("fVertex", "fVertex;z vtx (cm);Count", 201, -20, 20);
  94   fOutput->Add(fVertex);
  95
  96   fTriggerStats = new TH1F("fTriggerStats", "fTriggerStats;trigger;Count", 5, -0.5, 4.5);
  97   fTriggerStats->GetXaxis()->SetBinLabel(1, "!MB1 & !MB2");
  98   fTriggerStats->GetXaxis()->SetBinLabel(2, "MB1");
  99   fTriggerStats->GetXaxis()->SetBinLabel(3, "MB2");
 100   fTriggerStats->GetXaxis()->SetBinLabel(4, "ITS_SPD_GFO_L0");
 101   fTriggerStats->GetXaxis()->SetBinLabel(5, "VZERO_OR_LEFT | VZERO_OR_RIGHT");
 102   fOutput->Add(fTriggerStats);
 103
 104   fPrimStats = new TH1F("fPrimStats", "fPrimStats", 8, 0.5, 8.5);
 105   fPrimStats->GetXaxis()->SetBinLabel(1, "Primary accepted once");
 106   fPrimStats->GetXaxis()->SetBinLabel(2, "Primary accepted more than once");
 107   fPrimStats->GetXaxis()->SetBinLabel(3, "Primary accepted more than once (count)");
 108   fPrimStats->GetXaxis()->SetBinLabel(4, "Primary track rejected");
 109   fPrimStats->GetXaxis()->SetBinLabel(5, "Primary track rejected (count)");
 110   fPrimStats->GetXaxis()->SetBinLabel(6, "Primary track rejected, but other track accepted");
 111   fPrimStats->GetXaxis()->SetBinLabel(7, "Primary track rejected, but other track accepted (count)");
 112   fOutput->Add(fPrimStats);
 113
 114   // disable info messages of AliMCEvent (per event)
 115   AliLog::SetClassDebugLevel("AliMCEvent", AliLog::kWarning - AliLog::kDebug + 1);
 116 }
 117
 118 //________________________________________________________________________
 119 void AliCutTask::Exec(Option_t *)
 120 {
 121   // Main loop
 122   // Called for each event
 123
 124   if (!fESD) {
 125     Printf("ERROR: fESD not available");
 126     return;
 127   }
 128
 129   // Post output data.
 130   PostData(0, fOutput);
 131
 132   //fESD->GetVertex()->Print();
 133
 134   if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1) && !AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
 135     fTriggerStats->Fill(0);
 136
 137   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1))
 138     fTriggerStats->Fill(1);
 139
 140   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
 141     fTriggerStats->Fill(2);
 142
 143   if (fESD->GetTriggerMask() & (0x1 << 14))
 144     fTriggerStats->Fill(3);
 145
 146   if (fESD->GetTriggerMask() & 1 || fESD->GetTriggerMask() & 2)
 147     fTriggerStats->Fill(4);
 148
 149   //if (fESD->GetTriggerMask() & (0x1 << 14) == 0)
 150   if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
 151     return;
 152
 153   // get the ESD vertex
 154   const AliESDVertex* vtxESD = AliPWG0Helper::GetVertex(fESD, fAnalysisMode);
 155
 156   if (!vtxESD)
 157   {
 158     Printf("No vertex. Skipping event");
 159     return;
 160   }
 161
 162   Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
 163   Int_t acceptedTracks = fTrackCuts->CountAcceptedTracks(fESD);
 164   Printf("Accepted %d tracks", acceptedTracks);
 165
 166   if (fTrackCutsPrimaries && fTrackCutsSecondaries)
 167   {
 168     AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
 169     if (!eventHandler) {
 170       Printf("ERROR: Could not retrieve MC event handler");
 171       return;
 172     }
 173
 174     AliMCEvent* mcEvent = eventHandler->MCEvent();
 175     if (!mcEvent) {
 176       Printf("ERROR: Could not retrieve MC event");
 177       return;
 178     }
 179
 180     AliStack* stack = mcEvent->Stack();
 181     if (!stack)
 182     {
 183       Printf("ERROR: Stack not available");
 184       return;
 185     }
 186
 187     // count if primaries are counted several times
 188     Int_t max = stack->GetNprimary();
 189     Int_t* primAcc = new Int_t[max];
 190     Int_t* primRej = new Int_t[max];
 191     for (Int_t i=0; i<max; i++)
 192     {
 193       primAcc[i] = 0;
 194       primRej[i] = 0;
 195     }
 196
 197     for (Int_t trackId = 0; trackId < fESD->GetNumberOfTracks(); trackId++)
 198     {
 199       AliESDtrack* track = fESD->GetTrack(trackId);
 200       if (!track)
 201       {
 202         Printf("UNEXPECTED: Could not get track with id %d", trackId);
 203         continue;
 204       }
 205
 206       if (track->GetLabel() == 0)
 207       {
 208         Printf("Track %d has no label. Skipping", trackId);
 209         continue;
 210       }
 211
 212       Int_t label = TMath::Abs(track->GetLabel());
 213       if (stack->IsPhysicalPrimary(label) == kTRUE)
 214       {
 215         if (label >= max)
 216         {
 217           Printf("Warning label %d is higher than number of primaries %d", label, max);
 218           continue;
 219         }
 220
 221         if (fTrackCutsPrimaries->AcceptTrack(track))
 222         {
 223           primAcc[label]++;
 224         }
 225         else
 226           primRej[label]++;
 227
 228       }
 229       else
 230       {
 231         if (!stack->Particle(label)) {
 232           Printf("ERROR: Could not get particle with label %d", label);
 233           continue;
 234         }
 235
 236         // Deuteron is ok, but missing in Pdg particles in root
 237         if (stack->Particle(label)->GetPdgCode() != 10010020) {
 238           if (!stack->Particle(label)->GetPDG()) {
 239             Printf("ERROR: Could not get PDG particle of particle with label %d (pdg code is %d)", label, stack->Particle(label)->GetPdgCode());
 240             stack->Particle(label)->Print();
 241             continue;
 242           }
 243
 244           if (stack->Particle(label)->GetPDG()->Charge() == 0) {
 245             Printf("Particle %d has 0 charge. Skipping.", label);
 246             continue;
 247           }
 248         }
 249
 250         fTrackCutsSecondaries->AcceptTrack(track);
 251       }
 252     }
 253
 254     for (Int_t i=0; i<max; i++) {
 255       if (primAcc[i] == 1) {
 256         fPrimStats->Fill(1);
 257       } else if (primAcc[i] > 1) {
 258         fPrimStats->Fill(2);
 259         fPrimStats->Fill(3, primAcc[i]);
 260       }
 261
 262       if (primRej[i] > 0) {
 263         if (primAcc[i] == 0) {
 264           fPrimStats->Fill(4);
 265           fPrimStats->Fill(5, primRej[i]);
 266         } else if (primAcc[i] > 0) {
 267           fPrimStats->Fill(6);
 268           fPrimStats->Fill(7, primRej[i]);
 269         }
 270       }
 271     }
 272
 273     delete[] primAcc;
 274     delete[] primRej;
 275   }
 276
 277   // get the ESD vertex
 278   fVertex->Fill(vtxESD->GetZv());
 279 }
 280
 281 //________________________________________________________________________
 282 void AliCutTask::Terminate(Option_t *)
 283 {
 284   // Draw result to the screen
 285   // Called once at the end of the query
 286
 287   fOutput = dynamic_cast<TList*> (GetOutputData(0));
 288   if (!fOutput) {
 289     Printf("ERROR: fOutput not available");
 290     return;
 291   }
 292
 293   fTrackCuts = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("AliESDtrackCuts"));
 294   if (!fTrackCuts) {
 295     Printf("ERROR: fTrackCuts not available");
 296     return;
 297   }
 298
 299   fVertex = dynamic_cast<TH1F*> (fOutput->FindObject("fVertex"));
 300   if (!fVertex) {
 301     Printf("ERROR: fVertex not available");
 302     return;
 303   }
 304
 305   fTriggerStats = dynamic_cast<TH1F*> (fOutput->FindObject("fTriggerStats"));
 306   if (!fTriggerStats) {
 307     Printf("ERROR: fTriggerStats not available");
 308     return;
 309   }
 310
 311   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsPrimaries"));
 312   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsSecondaries"));
 313
 314   TFile* file = TFile::Open("trackCuts.root", "RECREATE");
 315
 316   fTrackCuts->SaveHistograms();
 317   fVertex->Write();
 318   fTriggerStats->Write();
 319
 320   if (fTrackCutsPrimaries)
 321     fTrackCutsPrimaries->SaveHistograms();
 322
 323   if (fTrackCutsSecondaries)
 324     fTrackCutsSecondaries->SaveHistograms();
 325
 326   if (fPrimStats)
 327     fPrimStats->Write();
 328
 329   file->Close();
 330
 331   Printf("Writting results to trackCuts.root.");
 332
 333   fTrackCuts->DrawHistograms();
 334
 335   new TCanvas;
 336   fVertex->Draw();
 337
 338   new TCanvas;
 339   fTriggerStats->Draw();
 340 }
 341
 342 void AliCutTask::EnableSecondaryStudy()
 343 {
 344   //
 345   // clones the fTrackCuts
 346   //
 347
 348   if (!fTrackCuts) {
 349     Printf("ERROR: fTrackCuts 0. Do not call this function before SetTrackCuts");
 350     return;
 351   }
 352
 353   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsPrimaries"));
 354   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsSecondaries"));
 355 }