PWG/EMCAL/AliEmcalMCTrackSelector.cxx

   1 // $Id$
   2 //
   3 // Class to select tracks in MC events.
   4 //
   5 // Author: S. Aiola
   6
   7 #include "AliEmcalMCTrackSelector.h"
   8
   9 #include <TClonesArray.h>
  10 #include <TH1I.h>
  11
  12 #include "AliAnalysisManager.h"
  13 #include "AliVEventHandler.h"
  14 #include "AliVEvent.h"
  15 #include "AliMCEvent.h"
  16 #include "AliVParticle.h"
  17 #include "AliAODMCParticle.h"
  18 #include "AliMCParticle.h"
  19 #include "AliLog.h"
  20
  21 ClassImp(AliEmcalMCTrackSelector)
  22
  23 //________________________________________________________________________
  24 AliEmcalMCTrackSelector::AliEmcalMCTrackSelector() :
  25   AliAnalysisTaskSE("AliEmcalMCTrackSelector"),
  26   fTracksOutName("PicoTracks"),
  27   fRejectNK(kFALSE),
  28   fChargedMC(kFALSE),
  29   fInit(kFALSE),
  30   fTracksMapName(""),
  31   fEsdMode(kFALSE),
  32   fTracksIn(0),
  33   fTracksOut(0),
  34   fTracksMap(0)
  35 {
  36   // Constructor.
  37 }
  38
  39 //________________________________________________________________________
  40 AliEmcalMCTrackSelector::AliEmcalMCTrackSelector(const char *name) :
  41   AliAnalysisTaskSE(name),
  42   fTracksOutName("PicoTracks"),
  43   fRejectNK(kFALSE),
  44   fChargedMC(kFALSE),
  45   fInit(kFALSE),
  46   fTracksMapName(""),
  47   fEsdMode(kFALSE),
  48   fTracksIn(0),
  49   fTracksOut(0),
  50   fTracksMap(0)
  51 {
  52   // Constructor.
  53 }
  54
  55 //________________________________________________________________________
  56 AliEmcalMCTrackSelector::~AliEmcalMCTrackSelector()
  57 {
  58   // Destructor.
  59 }
  60
  61 //________________________________________________________________________
  62 void AliEmcalMCTrackSelector::UserCreateOutputObjects()
  63 {
  64   // Create my user objects.
  65
  66   AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  67   if (!mgr) {
  68     AliFatal("No analysis manager!");
  69     return;
  70   }
  71
  72   AliVEventHandler *handler = mgr->GetInputEventHandler();
  73   if (!handler) {
  74     AliFatal("No input handler!");
  75     return;
  76   }
  77
  78   if (handler->InheritsFrom("AliESDInputHandler"))
  79     fEsdMode = kTRUE;
  80   else
  81     fEsdMode = kFALSE;
  82
  83   if (fEsdMode)
  84     fTracksOut = new TClonesArray("AliMCParticle");
  85   else
  86     fTracksOut = new TClonesArray("AliAODMCParticle");
  87   fTracksOut->SetName(fTracksOutName);
  88
  89   fTracksMapName = fTracksOutName;
  90   fTracksMapName += "_Map";
  91   fTracksMap = new TH1I(fTracksMapName, fTracksMapName, 1000, 0, 1);
  92 }
  93
  94 //________________________________________________________________________
  95 void AliEmcalMCTrackSelector::UserExec(Option_t *)
  96 {
  97   // Main loop, called for each event.
  98   if (!InputEvent()) {
  99     AliError("Could not retrieve event! Returning");
 100     return;
 101   }
 102
 103   if (fEsdMode && !MCEvent()) {
 104     AliError("Could not retrieve MC event! Returning");
 105     return;
 106   }
 107
 108   if (!fInit) {
 109     // add tracks to event if not yet there
 110     if (!(InputEvent()->FindListObject(fTracksOutName)))
 111       InputEvent()->AddObject(fTracksOut);
 112
 113     if (!(InputEvent()->FindListObject(fTracksMapName)))
 114       InputEvent()->AddObject(fTracksMap);
 115
 116     if (!fEsdMode) {
 117       fTracksIn = static_cast<TClonesArray*>(InputEvent()->FindListObject(AliAODMCParticle::StdBranchName()));
 118       if (!fTracksIn) {
 119         AliError("Could not retrieve AOD MC particles! Returning");
 120         return;
 121       }
 122       TClass *cl = fTracksIn->GetClass();
 123       if (!cl->GetBaseClass("AliAODMCParticle")) {
 124         AliError(Form("%s: Collection %s does not contain AliAODMCParticle!", GetName(), AliAODMCParticle::StdBranchName()));
 125         fTracksIn = 0;
 126         return;
 127       }
 128     }
 129
 130     fInit = kTRUE;
 131   }
 132
 133   // clear container (normally a null operation as the event should clean it already)
 134   fTracksOut->Delete();
 135
 136   const Int_t Ntracks = GetNumberOfTracks();
 137   new (fTracksMap) TH1I(fTracksMapName, fTracksMapName, Ntracks-2, 0, 1);  // Ntracks - 2, we use also over- and uner-flow bins
 138
 139   // loop over tracks
 140   for (Int_t iTracks = 0, nacc = 0; iTracks < Ntracks; ++iTracks) {
 141
 142     fTracksMap->SetBinContent(iTracks, 0);
 143
 144     AliVParticle* track = GetTrack(iTracks);
 145
 146     if (!track)
 147       continue;
 148
 149     if (TMath::Abs(track->Eta()) > 1)
 150       continue;
 151
 152     Int_t pdgCode = track->PdgCode();
 153     if (fRejectNK && (pdgCode == 130 || pdgCode == 2112)) continue;
 154
 155     if (fChargedMC && track->Charge() == 0) continue;
 156
 157     fTracksMap->SetBinContent(iTracks, nacc);
 158
 159     AddTrack(track, nacc);
 160
 161     ++nacc;
 162   }
 163 }
 164
 165 //________________________________________________________________________
 166 Int_t AliEmcalMCTrackSelector::GetNumberOfTracks() const
 167 {
 168   if (fEsdMode)
 169     return MCEvent()->GetNumberOfTracks();
 170   else
 171     return fTracksIn->GetEntries();
 172 }
 173
 174 //________________________________________________________________________
 175 AliVParticle* AliEmcalMCTrackSelector::GetTrack(Int_t i)
 176 {
 177   if (fEsdMode) {
 178     if (!MCEvent()->IsPhysicalPrimary(i))
 179       return 0;
 180
 181     return MCEvent()->GetTrack(i);
 182   }
 183   else {
 184     AliAODMCParticle *part = static_cast<AliAODMCParticle*>(fTracksIn->At(i));
 185     if (!part->IsPhysicalPrimary())
 186       return 0;
 187
 188     return part;
 189   }
 190 }
 191
 192 //________________________________________________________________________
 193 void AliEmcalMCTrackSelector::AddTrack(AliVParticle *track, Int_t nacc)
 194 {
 195   if (fEsdMode) {
 196     AliMCParticle *part = static_cast<AliMCParticle*>(track);
 197     new ((*fTracksOut)[nacc]) AliMCParticle(part->Particle(), 0, part->Label());
 198   }
 199   else {
 200     AliAODMCParticle *part = static_cast<AliAODMCParticle*>(track);
 201     new ((*fTracksOut)[nacc]) AliAODMCParticle(*part);
 202   }
 203 }