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