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