add qa check

[u/mrichter/AliRoot.git] / PWG / EMCAL / AliEmcalMCTrackSelector.cxx

diff --git a/PWG/EMCAL/AliEmcalMCTrackSelector.cxx b/PWG/EMCAL/AliEmcalMCTrackSelector.cxx
index d6e15f68b76be21a1487078464fd035c98d2bd04..be4eacb88d796646e90c67380cbb8969c4ae8710 100644 (file)
--- a/PWG/EMCAL/AliEmcalMCTrackSelector.cxx
+++ b/PWG/EMCAL/AliEmcalMCTrackSelector.cxx
@@ -7,12 +7,15 @@
 #include "AliEmcalMCTrackSelector.h"
 
 #include <TClonesArray.h>
-#include <TH1I.h>
 
+#include "AliNamedArrayI.h"
 #include "AliAnalysisManager.h"
+#include "AliVEventHandler.h"
 #include "AliVEvent.h"
 #include "AliMCEvent.h"
 #include "AliVParticle.h"
+#include "AliAODMCParticle.h"
+#include "AliMCParticle.h"
 #include "AliLog.h"
 
 ClassImp(AliEmcalMCTrackSelector)
@@ -23,8 +26,11 @@ AliEmcalMCTrackSelector::AliEmcalMCTrackSelector() :
   fTracksOutName("PicoTracks"),
   fRejectNK(kFALSE),
   fChargedMC(kFALSE),
-  fInit(kFALSE),
   fTracksMapName(""),
+  fEtaMax(1),
+  fInit(kFALSE),
+  fEsdMode(kFALSE),
+  fTracksIn(0),
   fTracksOut(0),
   fTracksMap(0)
 {
@@ -37,8 +43,11 @@ AliEmcalMCTrackSelector::AliEmcalMCTrackSelector(const char *name) :
   fTracksOutName("PicoTracks"),
   fRejectNK(kFALSE),
   fChargedMC(kFALSE),
-  fInit(kFALSE),
   fTracksMapName(""),
+  fEtaMax(1),
+  fInit(kFALSE),
+  fEsdMode(kFALSE),
+  fTracksIn(0),
   fTracksOut(0),
   fTracksMap(0)
 {
@@ -56,63 +65,93 @@ void AliEmcalMCTrackSelector::UserCreateOutputObjects()
 {
   // Create my user objects.
 
-  fTracksOut = new TClonesArray("AliMCParticle");
+  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
+  if (!mgr) {
+    AliFatal("No analysis manager!");
+    return;
+  }  
+  
+  AliVEventHandler *handler = mgr->GetInputEventHandler();
+  if (!handler) {
+    AliFatal("No input handler!");
+    return;
+  }  
+
+  if (handler->InheritsFrom("AliESDInputHandler") || handler->InheritsFrom("AliDummyHandler"))
+    fEsdMode = kTRUE;
+  else
+    fEsdMode = kFALSE;
+
+  if (fEsdMode)
+    fTracksOut = new TClonesArray("AliMCParticle");
+  else
+    fTracksOut = new TClonesArray("AliAODMCParticle");
   fTracksOut->SetName(fTracksOutName);
 
   fTracksMapName = fTracksOutName;
   fTracksMapName += "_Map";
-  fTracksMap = new TH1I(fTracksMapName, fTracksMapName, 1000, 0, 1);
+  fTracksMap = new AliNamedArrayI(fTracksMapName, 99999);
 }
 
 //________________________________________________________________________
 void AliEmcalMCTrackSelector::UserExec(Option_t *) 
 {
   // Main loop, called for each event.
-
-  AliMCEvent *mcevent = MCEvent();
-  AliVEvent *event = InputEvent();
-
-  if (!event) {
+  if (!InputEvent()) {
     AliError("Could not retrieve event! Returning");
     return;
   }
 
-  if (!mcevent) {
+  if (fEsdMode && !MCEvent()) {
     AliError("Could not retrieve MC event! Returning");
     return;
   }
 
   if (!fInit) {
     // add tracks to event if not yet there
-    if (!(event->FindListObject(fTracksOutName))) 
-      event->AddObject(fTracksOut);
-
-    if (!(event->FindListObject(fTracksMapName)))
-      event->AddObject(fTracksMap);
+    if (!(InputEvent()->FindListObject(fTracksOutName))) 
+      InputEvent()->AddObject(fTracksOut);
+
+    if (!(InputEvent()->FindListObject(fTracksMapName)))
+      InputEvent()->AddObject(fTracksMap);
+
+    if (!fEsdMode) {
+      fTracksIn = static_cast<TClonesArray*>(InputEvent()->FindListObject(AliAODMCParticle::StdBranchName()));
+      if (!fTracksIn) {
+       AliError("Could not retrieve AOD MC particles! Returning");
+       return;
+      }
+      TClass *cl = fTracksIn->GetClass();
+      if (!cl->GetBaseClass("AliAODMCParticle")) {
+       AliError(Form("%s: Collection %s does not contain AliAODMCParticle!", GetName(), AliAODMCParticle::StdBranchName())); 
+       fTracksIn = 0;
+       return;
+      }
+    }
 
     fInit = kTRUE;
   }
 
   // clear container (normally a null operation as the event should clean it already)
   fTracksOut->Delete();
+  fTracksMap->Clear();
 
-  const Int_t Ntracks = mcevent->GetNumberOfTracks();
-  new (fTracksMap) TH1I(fTracksMapName, fTracksMapName, Ntracks-2, 0, 1);  // Ntracks - 2, we use also over- and uner-flow bins
+  const Int_t Ntracks = GetNumberOfTracks();
 
   // loop over tracks
   for (Int_t iTracks = 0, nacc = 0; iTracks < Ntracks; ++iTracks) {
 
-    fTracksMap->SetBinContent(iTracks, -1);
+    if (iTracks >= fTracksMap->GetSize())
+      fTracksMap->Set(iTracks*2);
 
-    if (!mcevent->IsPhysicalPrimary(iTracks))
-      continue;
+    fTracksMap->AddAt(-1, iTracks);
 
-    AliMCParticle* track = dynamic_cast<AliMCParticle*>(mcevent->GetTrack(iTracks));
+    AliVParticle* track = GetTrack(iTracks);
 
     if (!track)
       continue;
 
-    if (TMath::Abs(track->Eta()) > 1) 
+    if (TMath::Abs(track->Eta()) > fEtaMax) 
       continue;
     
     Int_t pdgCode = track->PdgCode();
@@ -120,9 +159,50 @@ void AliEmcalMCTrackSelector::UserExec(Option_t *)
     
     if (fChargedMC && track->Charge() == 0) continue;
 
-    fTracksMap->SetBinContent(iTracks, nacc);
-    new ((*fTracksOut)[nacc]) AliMCParticle(track->Particle(), 0, track->Label());
+    fTracksMap->AddAt(nacc, iTracks);
+
+    AddTrack(track, nacc);
 
     ++nacc;
   }
 }
+
+//________________________________________________________________________
+Int_t AliEmcalMCTrackSelector::GetNumberOfTracks() const
+{
+  if (fEsdMode)
+    return MCEvent()->GetNumberOfTracks();
+  else
+    return fTracksIn->GetEntries();
+}
+
+//________________________________________________________________________
+AliVParticle* AliEmcalMCTrackSelector::GetTrack(Int_t i)
+{
+  if (fEsdMode) {
+    if (!MCEvent()->IsPhysicalPrimary(i))
+      return 0;
+
+    return MCEvent()->GetTrack(i);
+  }
+  else {
+    AliAODMCParticle *part = static_cast<AliAODMCParticle*>(fTracksIn->At(i));
+    if (!part->IsPhysicalPrimary()) 
+      return 0;
+    
+    return part;
+  }
+}
+
+//________________________________________________________________________
+void AliEmcalMCTrackSelector::AddTrack(AliVParticle *track, Int_t nacc)
+{
+  if (fEsdMode) {
+    AliMCParticle *part = static_cast<AliMCParticle*>(track);
+    new ((*fTracksOut)[nacc]) AliMCParticle(part->Particle(), 0, part->Label());
+  }
+  else {
+    AliAODMCParticle *part = static_cast<AliAODMCParticle*>(track);
+    new ((*fTracksOut)[nacc]) AliAODMCParticle(*part);
+  }
+}