PWG2/AliAnalysisTaskProtons.cxx

   1 #include "TChain.h"
   2 #include "TTree.h"
   3 #include "TList.h"
   4 #include "TH2F.h"
   5 #include "TF1.h"
   6 #include "TCanvas.h"
   7 #include "TLorentzVector.h"
   8
   9 #include "AliAnalysisTask.h"
  10 #include "AliAnalysisManager.h"
  11
  12 #include "AliESDEvent.h"
  13 #include "AliESDInputHandler.h"
  14 #include "AliAODEvent.h"
  15 #include "AliAODInputHandler.h"
  16 #include "AliMCEventHandler.h"
  17 #include "AliMCEvent.h"
  18 #include "AliStack.h"
  19
  20 #include "AliProtonAnalysis.h"
  21 #include "AliAnalysisTaskProtons.h"
  22
  23 // Analysis task creating a the 2d y-p_t spectrum of p and antip
  24 // Author: Panos Cristakoglou
  25
  26 ClassImp(AliAnalysisTaskProtons)
  27
  28 //________________________________________________________________________
  29 AliAnalysisTaskProtons::AliAnalysisTaskProtons()
  30 : AliAnalysisTask(), fESD(0), fAOD(0), fMC(0),fAnalysisType("ESD"),
  31   fList(0), fAnalysis(0),
  32   fElectronFunction(0), fMuonFunction(0),
  33   fPionFunction(0), fKaonFunction(0), fProtonFunction(0),
  34   fFunctionUsed(kFALSE) {
  35   //Dummy constructor
  36 }
  37
  38 //________________________________________________________________________
  39 AliAnalysisTaskProtons::AliAnalysisTaskProtons(const char *name)
  40 : AliAnalysisTask(name, ""), fESD(0), fAOD(0), fMC(0), fAnalysisType("ESD"),
  41   fList(0), fAnalysis(0),
  42   fElectronFunction(0), fMuonFunction(0),
  43   fPionFunction(0), fKaonFunction(0), fProtonFunction(0),
  44   fFunctionUsed(kFALSE) {
  45   // Constructor
  46
  47   // Define input and output slots here
  48   // Input slot #0 works with a TChain
  49   DefineInput(0, TChain::Class());
  50   // Output slot #0 writes into a TList container
  51   DefineOutput(0, TList::Class());
  52 }
  53
  54 //________________________________________________________________________
  55 void AliAnalysisTaskProtons::ConnectInputData(Option_t *) {
  56   // Connect ESD or AOD here
  57   // Called once
  58
  59   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
  60   if (!tree) {
  61     Printf("ERROR: Could not read chain from input slot 0");
  62   } else {
  63     // Disable all branches and enable only the needed ones
  64     // The next two lines are different when data produced as AliESDEvent is read
  65     if(fAnalysisType == "ESD") {
  66       // In train mode branches can be disabled at the level of ESD handler (M.G.)
  67       // tree->SetBranchStatus("*", kFALSE);
  68       tree->SetBranchStatus("*Tracks.*", 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     else if(fAnalysisType == "AOD") {
  78       AliAODInputHandler *aodH = dynamic_cast<AliAODInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  79
  80       if (!aodH) {
  81         Printf("ERROR: Could not get AODInputHandler");
  82       } else
  83         fAOD = aodH->GetEvent();
  84     }
  85     else if(fAnalysisType == "MC") {
  86       AliMCEventHandler* mcH = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
  87       if (!mcH) {
  88         Printf("ERROR: Could not retrieve MC event handler");
  89       }
  90       else
  91         fMC = mcH->MCEvent();
  92     }
  93     else
  94       Printf("Wrong analysis type: Only ESD, AOD and MC types are allowed!");
  95   }
  96 }
  97
  98 //________________________________________________________________________
  99 void AliAnalysisTaskProtons::CreateOutputObjects() {
 100   // Create histograms
 101   // Called once
 102   Double_t partFrac[5] = {0.01, 0.01, 0.85, 0.10, 0.05};
 103
 104   fAnalysis = new AliProtonAnalysis();
 105   fAnalysis->InitHistograms(10,-1.0,1.0,30,0.1,3.1);
 106   if(fAnalysisType == "ESD") {
 107     //Use of TPConly tracks
 108     fAnalysis->UseTPCOnly();
 109
 110     //TPC related cuts
 111     fAnalysis->SetMinTPCClusters(50);
 112     fAnalysis->SetMaxChi2PerTPCCluster(3.5);
 113     fAnalysis->SetMaxCov11(2.0);
 114     fAnalysis->SetMaxCov22(2.0);
 115     fAnalysis->SetMaxCov33(0.5);
 116     fAnalysis->SetMaxCov44(0.5);
 117     fAnalysis->SetMaxCov55(2.0);
 118     fAnalysis->SetMaxSigmaToVertex(2.5);
 119     fAnalysis->SetTPCRefit();
 120
 121     //ITS related cuts - to be used in the case of the analysis of global tracks
 122     //fAnalysis->SetMinITSClusters(5);
 123     //fAnalysis->SetITSRefit();
 124   }
 125   if(fFunctionUsed)
 126     fAnalysis->SetPriorProbabilityFunctions(fElectronFunction,
 127                                             fMuonFunction,
 128                                             fPionFunction,
 129                                             fKaonFunction,
 130                                             fProtonFunction);
 131   else
 132     fAnalysis->SetPriorProbabilities(partFrac);
 133
 134   fList = new TList();
 135   fList->Add(fAnalysis->GetProtonYPtHistogram());
 136   fList->Add(fAnalysis->GetAntiProtonYPtHistogram());
 137 }
 138
 139 //________________________________________________________________________
 140 void AliAnalysisTaskProtons::Exec(Option_t *) {
 141   // Main loop
 142   // Called for each event
 143
 144   if(fAnalysisType == "ESD") {
 145     if (!fESD) {
 146       Printf("ERROR: fESD not available");
 147       return;
 148     }
 149
 150     Printf("Proton ESD analysis task: There are %d tracks in this event", fESD->GetNumberOfTracks());
 151     fAnalysis->Analyze(fESD);
 152   }//ESD analysis
 153
 154   else if(fAnalysisType == "AOD") {
 155     if (!fAOD) {
 156       Printf("ERROR: fAOD not available");
 157       return;
 158     }
 159
 160     Printf("Proton AOD analysis task: There are %d tracks in this event", fAOD->GetNumberOfTracks());
 161     fAnalysis->Analyze(fAOD);
 162   }//AOD analysis
 163
 164   else if(fAnalysisType == "MC") {
 165     if (!fMC) {
 166       Printf("ERROR: Could not retrieve MC event");
 167       return;
 168     }
 169
 170     AliStack* stack = fMC->Stack();
 171     if (!stack) {
 172       Printf("ERROR: Could not retrieve the stack");
 173       return;
 174     }
 175     Printf("Proton MC analysis task: There are %d primaries in this event", stack->GetNprimary());
 176     fAnalysis->Analyze(stack);
 177   }//MC analysis
 178
 179   // Post output data.
 180   PostData(0, fList);
 181 }
 182
 183 //________________________________________________________________________
 184 void AliAnalysisTaskProtons::Terminate(Option_t *) {
 185   // Draw result to the screen
 186   // Called once at the end of the query
 187
 188   fList = dynamic_cast<TList*> (GetOutputData(0));
 189   if (!fList) {
 190     Printf("ERROR: fList not available");
 191     return;
 192   }
 193
 194   TH2F *fHistYPtProtons = (TH2F *)fList->At(0);
 195   TH2F *fHistYPtAntiProtons = (TH2F *)fList->At(1);
 196
 197   TCanvas *c2 = new TCanvas("c2","p-\bar{p}",200,0,800,400);
 198   c2->SetFillColor(10); c2->SetHighLightColor(10); c2->Divide(2,1);
 199
 200   c2->cd(1)->SetLeftMargin(0.15); c2->cd(1)->SetBottomMargin(0.15);
 201   if (fHistYPtProtons) fHistYPtProtons->DrawCopy("colz");
 202   c2->cd(2)->SetLeftMargin(0.15); c2->cd(2)->SetBottomMargin(0.15);
 203   if (fHistYPtAntiProtons) fHistYPtAntiProtons->DrawCopy("colz");
 204 }
 205
 206