PWG2/EBYE/AliAnalysisTaskBF.cxx

   1 #include "TChain.h"
   2 #include "TTree.h"
   3 #include "TCanvas.h"
   4 #include "TLorentzVector.h"
   5 #include "TGraphErrors.h"
   6
   7 #include "AliAnalysisTask.h"
   8 #include "AliAnalysisManager.h"
   9
  10 #include "AliESDEvent.h"
  11 #include "AliESDInputHandler.h"
  12 #include "AliAODEvent.h"
  13 #include "AliAODTrack.h"
  14 #include "AliAODInputHandler.h"
  15 #include "AliMCEventHandler.h"
  16 #include "AliMCEvent.h"
  17 #include "AliStack.h"
  18
  19 #include "AliBalance.h"
  20
  21 #include "AliAnalysisTaskBF.h"
  22
  23 // Analysis task for the BF code
  24 // Authors: Panos Cristakoglou@cern.ch
  25
  26 ClassImp(AliAnalysisTaskBF)
  27
  28 //________________________________________________________________________
  29 AliAnalysisTaskBF::AliAnalysisTaskBF(const char *name)
  30   : AliAnalysisTask(name, ""), fESD(0), fAOD(0), fMC(0), fBalance(0) {
  31   // Constructor
  32
  33   // Define input and output slots here
  34   // Input slot #0 works with a TChain
  35   DefineInput(0, TChain::Class());
  36   // Output slot #0 writes into a TH1 container
  37   DefineOutput(0, AliBalance::Class());
  38 }
  39
  40 //________________________________________________________________________
  41 void AliAnalysisTaskBF::ConnectInputData(Option_t *) {
  42   // Connect ESD or AOD here
  43   // Called once
  44   TString gAnalysisLevel = fBalance->GetAnalysisLevel();
  45
  46   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
  47   if (!tree) {
  48     Printf("ERROR: Could not read chain from input slot 0");
  49   }
  50   else {
  51     if(gAnalysisLevel == "ESD") {
  52       AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  53
  54       if (!esdH) {
  55         Printf("ERROR: Could not get ESDInputHandler");
  56       }
  57       else
  58         fESD = esdH->GetEvent();
  59     }//ESD
  60     else if(gAnalysisLevel == "AOD") {
  61       AliAODInputHandler *aodH = dynamic_cast<AliAODInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
  62
  63       if (!aodH) {
  64         Printf("ERROR: Could not get AODInputHandler");
  65       } else
  66         fAOD = aodH->GetEvent();
  67     }//AOD
  68     else if(gAnalysisLevel == "MC") {
  69       AliMCEventHandler* mcH = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
  70       if (!mcH) {
  71         Printf("ERROR: Could not retrieve MC event handler");
  72       }
  73       else
  74         fMC = mcH->MCEvent();
  75     }//MC
  76     else
  77       Printf("Wrong analysis type: Only ESD, AOD and MC types are allowed!");
  78   }
  79 }
  80
  81 //________________________________________________________________________
  82 void AliAnalysisTaskBF::CreateOutputObjects() {
  83   // Create histograms
  84   // Called once
  85
  86 }
  87
  88 //________________________________________________________________________
  89 void AliAnalysisTaskBF::Exec(Option_t *) {
  90   // Main loop
  91   // Called for each event
  92   TString gAnalysisLevel = fBalance->GetAnalysisLevel();
  93
  94   TObjArray *array = new TObjArray();
  95
  96   //ESD analysis
  97   if(gAnalysisLevel == "ESD") {
  98     if (!fESD) {
  99       Printf("ERROR: fESD not available");
 100       return;
 101     }
 102
 103     Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
 104     for (Int_t iTracks = 0; iTracks < fESD->GetNumberOfTracks(); iTracks++) {
 105       AliESDtrack* track = fESD->GetTrack(iTracks);
 106       if (!track) {
 107         Printf("ERROR: Could not receive track %d", iTracks);
 108         continue;
 109       }
 110       array->Add(track);
 111     } //track loop
 112   }//ESD analysis
 113   //AOD analysis
 114   else if(gAnalysisLevel == "AOD") {
 115     if (!fAOD) {
 116       Printf("ERROR: fAOD not available");
 117       return;
 118     }
 119
 120     Printf("There are %d tracks in this event", fAOD->GetNumberOfTracks());
 121     for (Int_t iTracks = 0; iTracks < fAOD->GetNumberOfTracks(); iTracks++) {
 122       AliAODTrack* track = fAOD->GetTrack(iTracks);
 123       if (!track) {
 124         Printf("ERROR: Could not receive track %d", iTracks);
 125         continue;
 126       }
 127       array->Add(track);
 128     } //track loop
 129   }//AOD analysis
 130   //MC analysis
 131   else if(gAnalysisLevel == "MC") {
 132     if (!fMC) {
 133       Printf("ERROR: fMC not available");
 134       return;
 135     }
 136
 137     Printf("There are %d tracks in this event", fMC->GetNumberOfPrimaries());
 138     for (Int_t iTracks = 0; iTracks < fMC->GetNumberOfPrimaries(); iTracks++) {
 139       AliMCParticle* track = dynamic_cast<AliMCParticle *>(fMC->GetTrack(iTracks));
 140       if (!track) {
 141         Printf("ERROR: Could not receive particle %d", iTracks);
 142         continue;
 143       }
 144       array->Add(track);
 145     } //track loop
 146   }//MC analysis
 147
 148   fBalance->CalculateBalance(array);
 149
 150   delete array;
 151
 152   // Post output data.
 153   PostData(0, fBalance);
 154 }
 155
 156 //________________________________________________________________________
 157 void AliAnalysisTaskBF::Terminate(Option_t *) {
 158   // Draw result to the screen
 159   // Called once at the end of the query
 160
 161   fBalance = dynamic_cast<AliBalance*> (GetOutputData(0));
 162   if (!fBalance) {
 163     Printf("ERROR: fBalance not available");
 164     return;
 165   }
 166
 167   TGraphErrors *gr = fBalance->DrawBalance();
 168   gr->SetMarkerStyle(20);
 169   gr->Draw("AP");
 170
 171   fBalance->PrintResults();
 172 }