This is the initial commit for a pilot nanoAOD implementation

[u/mrichter/AliRoot.git] / PWG / DevNanoAOD / AliAnalysisTaskNanoAODFilter.cxx

diff --git a/PWG/DevNanoAOD/AliAnalysisTaskNanoAODFilter.cxx b/PWG/DevNanoAOD/AliAnalysisTaskNanoAODFilter.cxx
new file mode 100644 (file)
index 0000000..ec8c73d
--- /dev/null
+++ b/PWG/DevNanoAOD/AliAnalysisTaskNanoAODFilter.cxx
@@ -0,0 +1,201 @@
+/**************************************************************************
+ * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
+ *                                                                        *
+ * Author: The ALICE Off-line Project.                                    *
+ * Contributors are mentioned in the code where appropriate.              *
+ *                                                                        *
+ * Permission to use, copy, modify and distribute this software and its   *
+ * documentation strictly for non-commercial purposes is hereby granted   *
+ * without fee, provided that the above copyright notice appears in all   *
+ * copies and that both the copyright notice and this permission notice   *
+ * appear in the supporting documentation. The authors make no claims     *
+ * about the suitability of this software for any purpose. It is          *
+ * provided "as is" without express or implied warranty.                  *
+ **************************************************************************/
+
+/* $Id$ */
+
+// Generic task to produce nanoAOD
+// Author: Michele Floris, michele.floris@cern.ch
+
+#include "AliAnalysisTaskNanoAODFilter.h"
+
+#include "Riostream.h"
+#include "TChain.h"
+#include "TTree.h"
+#include "TH1F.h"
+#include "TH2F.h"
+#include "TCanvas.h"
+#include "TList.h"
+
+#include "AliAnalysisTaskSE.h"
+#include "AliAnalysisManager.h"
+#include "AliStack.h"
+#include "AliESDtrackCuts.h"
+#include "AliESDEvent.h"
+#include "AliESDInputHandler.h"
+#include "AliAODEvent.h"
+#include "AliAODTrack.h"
+#include "AliMCEvent.h"
+#include "AliESDtrack.h"
+#include "AliAODHandler.h"
+#include "AliNanoAODReplicator.h"
+#include "AliNanoAODTrackMapping.h"
+
+ClassImp(AliAnalysisTaskNanoAODFilter)
+
+
+//________________________________________________________________________
+AliAnalysisTaskNanoAODFilter::AliAnalysisTaskNanoAODFilter() // All data members should be initialised here
+:AliAnalysisTaskSE(),
+  fTrkrep(0),
+  fVarList(""),
+  fVarListHead(""),
+  fEvtCuts(0),
+  fTrkCuts(0),
+  fSetter(0),
+  fSaveCutsFlag(0)
+{
+  // Dummy constructor ALWAYS needed for I/O.
+}
+
+//________________________________________________________________________
+AliAnalysisTaskNanoAODFilter::AliAnalysisTaskNanoAODFilter(const char *name, Bool_t saveCutsFlag) // All data members should be initialised here
+  :AliAnalysisTaskSE(name),
+   fTrkrep(0),
+   fVarList(""),
+   fVarListHead(""),
+   fEvtCuts(0),
+   fTrkCuts(0),
+   fSetter(0),
+   fSaveCutsFlag(saveCutsFlag)
+     
+{
+  // Constructor
+  if(fSaveCutsFlag) {
+    DefineOutput(1, AliAnalysisCuts::Class());
+    DefineOutput(2, AliAnalysisCuts::Class());
+  }
+
+
+}
+
+//________________________________________________________________________
+AliAnalysisTaskNanoAODFilter::~AliAnalysisTaskNanoAODFilter()
+{
+  // Destructor. Clean-up the output list, but not the histograms that are put inside
+  // (the list is owner and will clean-up these histograms). Protect in PROOF case.
+}
+
+//________________________________________________________________________
+void AliAnalysisTaskNanoAODFilter::UserCreateOutputObjects()
+{
+  // Create histograms
+  // Called once (on the worker node)
+      
+  if(fSaveCutsFlag) {
+    PostData(1, fEvtCuts); 
+    PostData(2, fTrkCuts); 
+  }
+}
+
+void AliAnalysisTaskNanoAODFilter::AddFilteredAOD(const char* aodfilename, const char* title)
+{
+  // The replicator is added to the extension
+
+  AliAODHandler *aodH = (AliAODHandler*)((AliAnalysisManager::GetAnalysisManager())->GetOutputEventHandler());
+  if (!aodH) AliFatal("No AOD handler");
+  AliAODExtension* ext = aodH->AddFilteredAOD(aodfilename,title);
+    
+  if (!ext){
+    AliFatal("Cannot get extension");
+  }
+  
+
+  AliNanoAODReplicator * rep = new AliNanoAODReplicator("NanoAODReplicator",
+                                                             "remove non interesting tracks, "
+                                                             "writes special tracks array tracks",
+                                                             fVarList,
+                                                             fTrkCuts,
+                                                             fMCMode);
+
+     
+  cout<<"rep: "<<rep<<endl;
+  rep->SetCustomSetter(fSetter);
+  std::cout << "SETTER: " << fSetter << " " << rep->GetCustomSetter() << std::endl;
+  
+  ext->DropUnspecifiedBranches(); // all branches not part of a FilterBranch call (below) will be dropped
+      
+  ext->FilterBranch("tracks",rep);
+  ext->FilterBranch("vertices",rep);  
+  ext->FilterBranch("header",rep);  
+            
+  if ( fMCMode > 0 ) 
+    {
+      // MC branches will be copied (if present), as they are, but only
+      // for events with at least one muon. 
+      // For events w/o muon, mcparticles array will be empty and mcheader will be dummy
+      // (e.g. strlen(GetGeneratorName())==0)
+      
+      ext->FilterBranch("mcparticles",rep);
+      ext->FilterBranch("mcHeader",rep);
+    }
+}
+
+void AliAnalysisTaskNanoAODFilter::Init()
+{
+
+  // Initialization
+  AddFilteredAOD("AliAOD.NanoAOD.root", "NanoAODTracksEvents");
+
+}
+
+
+
+//________________________________________________________________________
+void AliAnalysisTaskNanoAODFilter::UserExec(Option_t *) 
+{
+  // Main loop
+  Long64_t ientry = Entry();
+  if(fDebug)printf("Nano AOD Filter: Analysing event # %5d\n", (Int_t) ientry);
+
+  AliAODEvent *lAODevent=(AliAODEvent*)InputEvent();
+  
+    
+  AliAODHandler* handler = dynamic_cast<AliAODHandler*>(AliAnalysisManager::GetAnalysisManager()->GetOutputEventHandler());
+
+  if(fEvtCuts && !fEvtCuts->IsSelected(lAODevent)) return;// FIXME: should event cuts be called here or in the branch replicator? Do we get duplicated events if we skip here (arrays not reset in the branch replicator?)
+
+  if ( handler ){
+    AliAODExtension *extNanoAOD = handler->GetFilteredAOD("AliAOD.NanoAOD.root");
+   if ( extNanoAOD ) {                         
+     extNanoAOD->SetEvent(lAODevent);
+     extNanoAOD->SelectEvent();
+     extNanoAOD->FinishEvent();
+   }
+  }
+
+
+  if(fSaveCutsFlag) {
+    PostData(1, fEvtCuts); 
+    PostData(2, fTrkCuts); 
+  }
+}
+
+
+//________________________________________________________________________
+void AliAnalysisTaskNanoAODFilter::Terminate(Option_t *) 
+{
+  // print some debug info
+
+}
+
+void AliAnalysisTaskNanoAODFilter::FinishTaskOutput() {
+
+  // We save here the user info //
+
+  AliAODHandler* handler = dynamic_cast<AliAODHandler*>(AliAnalysisManager::GetAnalysisManager()->GetOutputEventHandler());
+  AliAODExtension *extNanoAOD = handler->GetFilteredAOD("AliAOD.NanoAOD.root");
+  extNanoAOD->GetTree()->GetUserInfo()->Add(AliNanoAODTrackMapping::GetInstance());
+
+}