Code updates with bugfixes; new PID class; filtered AOD config macro

[u/mrichter/AliRoot.git] / PWG3 / dielectron / AliAnalysisTaskDielectronFilter.cxx

diff --git a/PWG3/dielectron/AliAnalysisTaskDielectronFilter.cxx b/PWG3/dielectron/AliAnalysisTaskDielectronFilter.cxx
index 89826cf79ece354901704dc2afde2bb4771259d1..e5d4dab8500dc04a7b736e4cc31800f57d7e45d8 100644 (file)
--- a/PWG3/dielectron/AliAnalysisTaskDielectronFilter.cxx
+++ b/PWG3/dielectron/AliAnalysisTaskDielectronFilter.cxx
@@ -25,17 +25,22 @@
 #include <AliAODHandler.h>
 #include <AliAnalysisManager.h>
 #include <AliVEvent.h>
+#include <AliInputEventHandler.h>
+#include <AliESDInputHandler.h>
 
 #include "AliDielectron.h"
+#include "AliDielectronMC.h"
 #include "AliDielectronHistos.h"
+#include "AliDielectronVarManager.h"
 #include "AliAnalysisTaskDielectronFilter.h"
 
 ClassImp(AliAnalysisTaskDielectronFilter)
 
 //_________________________________________________________________________________
 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter() :
-  AliAnalysisTaskSE(),
-  fDielectron(0)
+    AliAnalysisTaskSE(),
+                      fDielectron(0),
+                                  fSelectPhysics(kTRUE)
 {
   //
   // Constructor
@@ -44,8 +49,9 @@ AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter() :
 
 //_________________________________________________________________________________
 AliAnalysisTaskDielectronFilter::AliAnalysisTaskDielectronFilter(const char *name) :
-  AliAnalysisTaskSE(name),
-  fDielectron(0)
+    AliAnalysisTaskSE(name),
+                      fDielectron(0),
+                                  fSelectPhysics(kTRUE)
 {
   //
   // Constructor
@@ -69,6 +75,7 @@ void AliAnalysisTaskDielectronFilter::Init()
     Error("Init","Dielectron framework class required. Please create and instance with proper cuts and set it via 'SetDielectron' before executing this task!!!");
     return;
   }
+  fDielectron->Init();
   
   aodH->AddFilteredAOD("AliAOD.Dielectron.root", "DielectronEvents");
 //   AddAODBranch("AliDielectronCandidates",fDielectron->GetPairArraysPointer(),"deltaAOD.Dielectron.root");
@@ -80,18 +87,69 @@ void AliAnalysisTaskDielectronFilter::UserExec(Option_t *)
   //
   // Main loop. Called for every event
   //
-
+  
   if (!fDielectron) return;
   
+  AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
+  AliESDInputHandler *esdHandler=0x0;
+  if ( (esdHandler=dynamic_cast<AliESDInputHandler*>(man->GetInputEventHandler())) && esdHandler->GetESDpid() ){
+    AliDielectronVarManager::SetESDpid(esdHandler->GetESDpid());
+  } else {
+    //load esd pid bethe bloch parameters depending on the existance of the MC handler
+    // yes: MC parameters
+    // no:  data parameters
+    if (!AliDielectronVarManager::GetESDpid()){
+      if (AliDielectronMC::Instance()->HasMC()) {
+        AliDielectronVarManager::InitESDpid();
+      } else {
+        AliDielectronVarManager::InitESDpid(1);
+      }
+    }
+  }
+  // Was event selected ?
+  AliInputEventHandler* inputHandler = (AliInputEventHandler*) (man->GetInputEventHandler());
+  Bool_t isSelected = kTRUE;
+  if( fSelectPhysics && inputHandler && inputHandler->GetEventSelection() ) {
+    isSelected = inputHandler->IsEventSelected();
+  }
+  
+  if (!isSelected) return;
+  
   //bz for AliKF
   Double_t bz = InputEvent()->GetMagneticField();
   AliKFParticle::SetField( bz );
   
   fDielectron->Process(InputEvent());
-
+  
   if(fDielectron->HasCandidates()){
+    //If input event is an AliESDevent
+    // replace the references of the legs with the AOD references
+    if(man->GetInputEventHandler()->IsA()==AliESDInputHandler::Class()){
+      AliAODEvent *aod = ((AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler()))->GetAOD();
+      TObjArray *obj = 0x0;
+      AliAODTrack *leg1 = 0x0;
+      AliAODTrack *leg2 = 0x0;
+      for(Int_t i=0; i < 10; i++ ){
+        obj = (TObjArray*)(*(fDielectron->GetPairArraysPointer()))->UncheckedAt(i);
+        if(!obj) continue;
+        for(int j=0;j<obj->GetEntriesFast();j++)
+        {
+          AliDielectronPair *pairObj = (AliDielectronPair*)obj->UncheckedAt(j);
+          Int_t id1 = ((AliESDtrack*)pairObj->GetFirstDaughter())->GetID();
+          Int_t id2 = ((AliESDtrack*)pairObj->GetSecondDaughter())->GetID();
+          for(Int_t it=0;it<aod->GetNumberOfTracks();it++){
+            if(aod->GetTrack(it)->GetID() == id1) leg1 = aod->GetTrack(it);
+            if(aod->GetTrack(it)->GetID() == id2) leg2 = aod->GetTrack(it);
+          }
+          if(!leg1 || !leg2) continue;
+          pairObj->SetRefFirstDaughter(leg1);
+          pairObj->SetRefSecondDaughter(leg2);
+        }
+      }
+    }
+      
     AliAODExtension *extDielectron = dynamic_cast<AliAODHandler*>
-      ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Dielectron.root");
+        ((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler())->GetFilteredAOD("AliAOD.Dielectron.root");
     extDielectron->SelectEvent();
     //see if dielectron candidate branch exists, if not create is
     TTree *t=extDielectron->GetTree();