PWG/EMCAL/AliEmcalMCTrackSelector.cxx

   1 //
   2 // Class to select particles in MC events.
   3 //
   4 // Author: S. Aiola
   5
   6 #include "AliEmcalMCTrackSelector.h"
   7
   8 #include <TClonesArray.h>
   9
  10 #include "AliVEvent.h"
  11 #include "AliMCEvent.h"
  12 #include "AliAODMCParticle.h"
  13 #include "AliMCParticle.h"
  14 #include "AliStack.h"
  15 #include "AliNamedArrayI.h"
  16
  17 #include "AliLog.h"
  18
  19 ClassImp(AliEmcalMCTrackSelector)
  20
  21 //________________________________________________________________________
  22 AliEmcalMCTrackSelector::AliEmcalMCTrackSelector() :
  23   AliAnalysisTaskSE("AliEmcalMCTrackSelector"),
  24   fParticlesOutName("MCParticlesSelected"),
  25   fOnlyPhysPrim(kTRUE),
  26   fRejectNK(kFALSE),
  27   fChargedMC(kFALSE),
  28   fOnlyHIJING(kFALSE),
  29   fEtaMax(1),
  30   fParticlesMapName(""),
  31   fInit(kFALSE),
  32   fParticlesIn(0),
  33   fParticlesOut(0),
  34   fParticlesMap(0),
  35   fEvent(0),
  36   fMC(0),
  37   fIsESD(kFALSE),
  38   fDisabled(kFALSE)
  39 {
  40   // Constructor.
  41 }
  42
  43 //________________________________________________________________________
  44 AliEmcalMCTrackSelector::AliEmcalMCTrackSelector(const char *name) :
  45   AliAnalysisTaskSE(name),
  46   fParticlesOutName("MCParticlesSelected"),
  47   fOnlyPhysPrim(kTRUE),
  48   fRejectNK(kFALSE),
  49   fChargedMC(kFALSE),
  50   fOnlyHIJING(kFALSE),
  51   fEtaMax(1),
  52   fParticlesMapName(""),
  53   fInit(kFALSE),
  54   fParticlesIn(0),
  55   fParticlesOut(0),
  56   fParticlesMap(0),
  57   fEvent(0),
  58   fMC(0),
  59   fIsESD(kFALSE),
  60   fDisabled(kFALSE)
  61 {
  62   // Constructor.
  63 }
  64
  65 //________________________________________________________________________
  66 AliEmcalMCTrackSelector::~AliEmcalMCTrackSelector()
  67 {
  68   // Destructor.
  69 }
  70
  71 //________________________________________________________________________
  72 void AliEmcalMCTrackSelector::UserCreateOutputObjects()
  73 {
  74   // Create my user objects.
  75 }
  76
  77 //________________________________________________________________________
  78 void AliEmcalMCTrackSelector::UserExec(Option_t *)
  79 {
  80   // Main loop, called for each event.
  81
  82   if (fDisabled) return;
  83
  84   if (!fInit) {
  85     fEvent = InputEvent();
  86     if (!fEvent) {
  87       AliError("Could not retrieve event! Returning");
  88       return;
  89     }
  90
  91     if (fEvent->InheritsFrom("AliESDEvent")) fIsESD = kTRUE;
  92     else fIsESD = kFALSE;
  93
  94     TObject *obj = fEvent->FindListObject(fParticlesOutName);
  95     if (obj) { // the output array is already present in the array!
  96       AliError(Form("The output array %s is already present in the event! Task will be disabled.", fParticlesOutName.Data()));
  97       fDisabled = kTRUE;
  98       return;
  99     }
 100     else {  // copy the array from the standard ESD/AOD collections, and filter if requested
 101
 102       fParticlesOut = new TClonesArray("AliAODMCParticle");  // the output will always be of AliAODMCParticle, regardless of the input
 103       fParticlesOut->SetName(fParticlesOutName);
 104       fEvent->AddObject(fParticlesOut);
 105
 106       fParticlesMapName = fParticlesOutName;
 107       fParticlesMapName += "_Map";
 108
 109       if (fEvent->FindListObject(fParticlesMapName)) {
 110         AliError(Form("The output array map %s is already present in the event! Task will be disabled.", fParticlesMapName.Data()));
 111         fDisabled = kTRUE;
 112         return;
 113       }
 114       else {
 115         fParticlesMap = new AliNamedArrayI(fParticlesMapName, 99999);
 116         fEvent->AddObject(fParticlesMap);
 117       }
 118
 119       if (!fIsESD) {
 120         fParticlesIn = static_cast<TClonesArray*>(InputEvent()->FindListObject(AliAODMCParticle::StdBranchName()));
 121         if (!fParticlesIn) {
 122           AliError("Could not retrieve AOD MC particles! Task will be disabled.");
 123           fDisabled = kTRUE;
 124           return;
 125         }
 126         TClass *cl = fParticlesIn->GetClass();
 127         if (!cl->GetBaseClass("AliAODMCParticle")) {
 128           AliError(Form("%s: Collection %s does not contain AliAODMCParticle! Task will be disabled.", GetName(), AliAODMCParticle::StdBranchName()));
 129           fDisabled = kTRUE;
 130           fParticlesIn = 0;
 131           return;
 132         }
 133       }
 134     }
 135
 136     fMC = MCEvent();
 137     if (!fMC) {
 138       AliError("Could not retrieve MC event! Returning");
 139       fDisabled = kTRUE;
 140       return;
 141     }
 142
 143     fInit = kTRUE;
 144   }
 145
 146   if (fIsESD) ConvertMCParticles();
 147   else CopyMCParticles();
 148 }
 149
 150 //________________________________________________________________________
 151 void AliEmcalMCTrackSelector::ConvertMCParticles()
 152 {
 153   // Convert MC particles in MC AOD particles.
 154
 155   // clear container (normally a null operation as the event should clean it already)
 156   fParticlesOut->Delete();
 157
 158   // clear particles map
 159   fParticlesMap->Clear();
 160
 161   const Int_t Nparticles = fMC->GetNumberOfTracks();
 162   const Int_t nprim = fMC->GetNumberOfPrimaries();
 163
 164   if (fParticlesMap->GetSize() <= Nparticles) fParticlesMap->Set(Nparticles*2);
 165
 166   // loop over particles
 167   for (Int_t iPart = 0, nacc = 0; iPart < Nparticles; iPart++) {
 168
 169     fParticlesMap->AddAt(-1, iPart);
 170
 171     AliMCParticle* part = static_cast<AliMCParticle*>(fMC->GetTrack(iPart));
 172
 173     if (!part) continue;
 174
 175     if (fEtaMax > 0. && TMath::Abs(part->Eta()) > fEtaMax) continue;
 176
 177     if (fRejectNK && (part->PdgCode() == 130 || part->PdgCode() == 2112)) continue;
 178
 179     if (fChargedMC && part->Charge() == 0) continue;
 180
 181     Int_t genIndex = part->GetGeneratorIndex();
 182     if (fOnlyHIJING && genIndex != 0) continue;
 183
 184     Bool_t isPhysPrim = fMC->IsPhysicalPrimary(iPart);
 185     if (fOnlyPhysPrim && !isPhysPrim) continue;
 186
 187     fParticlesMap->AddAt(nacc, iPart);
 188
 189     Int_t flag = 0;
 190     if (iPart < nprim) flag |= AliAODMCParticle::kPrimary;
 191     if (isPhysPrim) flag |= AliAODMCParticle::kPhysicalPrim;
 192     if (fMC->IsSecondaryFromWeakDecay(iPart)) flag |= AliAODMCParticle::kSecondaryFromWeakDecay;
 193     if (fMC->IsSecondaryFromMaterial(iPart)) flag |= AliAODMCParticle::kSecondaryFromMaterial;
 194
 195     AliAODMCParticle *aodPart = new ((*fParticlesOut)[nacc]) AliAODMCParticle(part, iPart, flag);
 196     aodPart->SetGeneratorIndex(genIndex);
 197     aodPart->SetStatus(part->Particle()->GetStatusCode());
 198     aodPart->SetMCProcessCode(part->Particle()->GetUniqueID());
 199
 200     nacc++;
 201   }
 202 }
 203
 204 //________________________________________________________________________
 205 void AliEmcalMCTrackSelector::CopyMCParticles()
 206 {
 207   // Convert standard MC AOD particles in a new array, and filter if requested.
 208
 209   if (!fParticlesIn) return;
 210
 211   // clear container (normally a null operation as the event should clean it already)
 212   fParticlesOut->Delete();
 213
 214   // clear particles map
 215   fParticlesMap->Clear();
 216
 217   const Int_t Nparticles = fParticlesIn->GetEntriesFast();
 218
 219   if (fParticlesMap->GetSize() <= Nparticles) fParticlesMap->Set(Nparticles*2);
 220
 221   // loop over particles
 222   for (Int_t iPart = 0, nacc = 0; iPart < Nparticles; iPart++) {
 223
 224     AliAODMCParticle* part = static_cast<AliAODMCParticle*>(fParticlesIn->At(iPart));
 225
 226     fParticlesMap->AddAt(-1, iPart);
 227
 228     if (!part) continue;
 229
 230     if (fEtaMax > 0. && TMath::Abs(part->Eta()) > fEtaMax) continue;
 231
 232     if (fRejectNK && (part->PdgCode() == 130 || part->PdgCode() == 2112)) continue;
 233
 234     if (fChargedMC && part->Charge() == 0) continue;
 235
 236     if (fOnlyHIJING && (part->GetGeneratorIndex() != 0)) continue;
 237
 238     if (fOnlyPhysPrim && !part->IsPhysicalPrimary()) continue;
 239
 240     fParticlesMap->AddAt(nacc, iPart);
 241
 242     AliAODMCParticle *newPart = new ((*fParticlesOut)[nacc]) AliAODMCParticle(*part);
 243     newPart->SetGeneratorIndex(part->GetGeneratorIndex());
 244
 245     nacc++;
 246   }
 247 }