[u/mrichter/AliRoot.git] / PWG / muon / AliAnalysisTaskCreateMixedDimuons.cxx

#include "TChain.h"
#include "TTree.h"

#include "AliAnalysisTaskME.h"
#include "AliAnalysisManager.h"

#include "AliAODEvent.h"
#include "AliAODTrack.h"
#include "AliAODVertex.h"
#include "AliMultiEventInputHandler.h"
#include "AliAODHandler.h"

#include "AliAnalysisTaskCreateMixedDimuons.h"
#include "AliEventPoolMuon.h"
#include "TDatabasePDG.h"
#include "TRandom.h"

// Example of an analysis task creating aod events filled with mixed muon pairs
//
// Authors Alessandro De Falco and Antonio Uras, INFN Cagliari
// alessandro.de.falco@ca.infn.it  antonio.uras@ca.infn.it

#define AliAnalysisTaskCreateMixedDimuons_CXX

ClassImp(AliAnalysisTaskCreateMixedDimuons)

//=================================================================================

AliAnalysisTaskCreateMixedDimuons::AliAnalysisTaskCreateMixedDimuons(const char *name) 
: AliAnalysisTaskME(name),
  fBufferSize(0),
  fOutputUserHandler(0x0),
  fOutputUserAOD(0X0),
  fOutputUserAODTree(0X0),
  fPoolMuon(0X0),
  fDebug(0X0)
{

  // Constructor

  // Default input and output containers
  DefineInput (0, TChain::Class());

  // User-defined input and output containers
  DefineOutput(1, TTree::Class());

  // ---------------------------------------------------------------

  fDebug = kFALSE;
  fBufferSize = 0;

  RequireFreshBuffer();

  for (Int_t i=0; i<100; i++) fInputAOD[i] = 0;

  fOutputUserHandler = 0;
  fOutputUserAOD     = 0;
  fOutputUserAODTree = 0;  

}

//=================================================================================

void AliAnalysisTaskCreateMixedDimuons::ConnectInputData(Option_t *) {

  // Connect ESD or AOD here
  // Called once

  printf("-> AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");

  TTree* tree = (TTree*) GetInputData(0);

  if (!tree) {
    Printf("ERROR: Could not read chain from input slot 0");
  } 
  else {

    fInputHandler = (AliMultiEventInputHandler*) AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
    fPoolMuon     = (AliEventPoolMuon*)          AliAnalysisManager::GetAnalysisManager()->GetEventPool();

    if (!fInputHandler) {
      Printf("ERROR: Could not get AliMultiAODInputHandler");
    }
    else {  
      fBufferSize = fInputHandler->GetBufferSize();
      if (fBufferSize>100) {
        printf("\n*** WARNING AliAnalysisTaskCreateMixedDimuons::ConnectInputData -> Trying to set fBufferSize>100, forcing fBufferSize=100 ***\n\n");
        fBufferSize = 100;
      }
      for (Int_t i=0; i<fBufferSize; i++) fInputAOD[i] = (AliAODEvent*) fInputHandler->GetEvent(i);
    }
  }

  printf("<- AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");

}

//=================================================================================

void AliAnalysisTaskCreateMixedDimuons::UserCreateOutputObjects() {

  // Here the user-defined output containers should be created!!!
  // Called once

  fOutputUserHandler = new AliAODHandler();

  fOutputUserHandler -> Init("");

  fOutputUserAOD = fOutputUserHandler -> GetAOD();

  fOutputUserAODTree = fOutputUserHandler -> GetTree();

}

//=================================================================================

void AliAnalysisTaskCreateMixedDimuons::UserExec(Option_t *) {

  if (!fOutputUserAOD) {
    Printf("ERROR: fOutputUserAOD not available\n");
    return;
  }

  printf("Calling USER EXEC\n\n");

  for (Int_t iEv=0; iEv<fBufferSize; iEv++) {
    for (Int_t jEv=0; jEv<iEv; jEv++) {

      Int_t nTracksEv[2]  = {0};
      Int_t nFWMuonsEv[2] = {0};
      
      nTracksEv[0] = fInputAOD[iEv]->GetNTracks();
      nTracksEv[1] = fInputAOD[jEv]->GetNTracks();
      
      for (Int_t i=0; i<nTracksEv[0]; i++) if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[0]++;
      for (Int_t i=0; i<nTracksEv[1]; i++) if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[1]++;
      
      // Muon track mixing to fill a mass spectrum
      
      if (nFWMuonsEv[0] && nFWMuonsEv[1]) {
	
	Int_t rndMuonTrack[2] = {0};
	rndMuonTrack[0] = gRandom->Integer(nFWMuonsEv[0]);
	rndMuonTrack[1] = gRandom->Integer(nFWMuonsEv[1]);
	
	Int_t nFWMUonsAdded = 0;
	Int_t nPosTracksAdded = 0;
	Int_t nNegTracksAdded = 0;
	
	AliAODVertex *vertex = new AliAODVertex();
	vertex -> SetX(0.0);
	vertex -> SetY(0.0);
	vertex -> SetZ(fPoolMuon->GetMeanPrimaryVertexZ());
	
	Int_t muonCounter[2] = {0};
	
	// adding tracks and vertex to the output event...
	
	for (Int_t i=0; i<nTracksEv[0]; i++) {
	  if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) {
	    if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p    pt = %f     uniqueID = %d\n",
			       iEv,i,fInputAOD[iEv]->GetTrack(i),fInputAOD[iEv]->GetTrack(i)->Pt(),
			       fInputAOD[iEv]->GetTrack(i)->GetUniqueID());
	    if (muonCounter[0]==rndMuonTrack[0]) {
	      fOutputUserAOD->AddTrack(fInputAOD[iEv]->GetTrack(i));
	      nFWMUonsAdded++;
	      if (fInputAOD[iEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
	      else nNegTracksAdded++;
	    }
	    muonCounter[0]++;
	  }
	}
	
	for (Int_t i=0; i<nTracksEv[1]; i++) {
	  if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) {
	    if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p    pt = %f     uniqueID = %d\n",
			       jEv,i,fInputAOD[jEv]->GetTrack(i),fInputAOD[jEv]->GetTrack(i)->Pt(),
			       fInputAOD[jEv]->GetTrack(i)->GetUniqueID());
	    if (muonCounter[1]==rndMuonTrack[1]) {
	      fOutputUserAOD->AddTrack(fInputAOD[jEv]->GetTrack(i));
	      nFWMUonsAdded++;
	      if (fInputAOD[jEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
	      else nNegTracksAdded++;
	    }
	    muonCounter[1]++;	
	  }
	}
	
	fOutputUserAOD->AddVertex(vertex);
	
	// ... done!
	
	if (fDebug) {
	  for (Int_t i=0; i<nFWMUonsAdded; i++) {
	    AliAODTrack *tr = (AliAODTrack*) fOutputUserAOD->GetTrack(i);
	    printf("fOutputUserAOD->GetTrack(%d) = %p    pt = %f\n",i,tr,tr->Pt());
	  }
	}
	
	fOutputUserAOD->GetHeader()->SetRefMultiplicity(nFWMUonsAdded); 
	fOutputUserAOD->GetHeader()->SetRefMultiplicityPos(nPosTracksAdded);
	fOutputUserAOD->GetHeader()->SetRefMultiplicityNeg(nNegTracksAdded);
	
	fOutputUserHandler -> FinishEvent();
	
      }

      PostData(1, fOutputUserAODTree);
      
    }
    
  }
  
}      

//===================================================================================

void AliAnalysisTaskCreateMixedDimuons::Terminate(Option_t *) {

  // Called once at the end of the query

  printf("\n\nCalling TERMINATE \n\n\n");

  fOutputUserHandler -> Terminate();

}

//===================================================================================
Commit	Line	Data
866de0ba	1	#include "TChain.h"
	2	#include "TTree.h"
	3
	4	#include "AliAnalysisTaskME.h"
	5	#include "AliAnalysisManager.h"
	6
	7	#include "AliAODEvent.h"
	8	#include "AliAODTrack.h"
	9	#include "AliAODVertex.h"
	10	#include "AliMultiEventInputHandler.h"
	11	#include "AliAODHandler.h"
	12
	13	#include "AliAnalysisTaskCreateMixedDimuons.h"
	14	#include "AliEventPoolMuon.h"
	15	#include "TDatabasePDG.h"
	16	#include "TRandom.h"
	17
	18	// Example of an analysis task creating aod events filled with mixed muon pairs
	19	//
	20	// Authors Alessandro De Falco and Antonio Uras, INFN Cagliari
	21	// alessandro.de.falco@ca.infn.it antonio.uras@ca.infn.it
	22
	23	#define AliAnalysisTaskCreateMixedDimuons_CXX
	24
	25	ClassImp(AliAnalysisTaskCreateMixedDimuons)
	26
	27	//=================================================================================
	28
	29	AliAnalysisTaskCreateMixedDimuons::AliAnalysisTaskCreateMixedDimuons(const char *name)
	30	: AliAnalysisTaskME(name),
	31	fBufferSize(0),
	32	fOutputUserHandler(0x0),
	33	fOutputUserAOD(0X0),
	34	fOutputUserAODTree(0X0),
	35	fPoolMuon(0X0),
	36	fDebug(0X0)
	37	{
	38
	39	// Constructor
	40
	41	// Default input and output containers
	42	DefineInput (0, TChain::Class());
	43
	44	// User-defined input and output containers
	45	DefineOutput(1, TTree::Class());
	46
	47	// ---------------------------------------------------------------
	48
	49	fDebug = kFALSE;
	50	fBufferSize = 0;
	51
	52	RequireFreshBuffer();
	53
	54	for (Int_t i=0; i<100; i++) fInputAOD[i] = 0;
	55
	56	fOutputUserHandler = 0;
	57	fOutputUserAOD = 0;
	58	fOutputUserAODTree = 0;
	59
	60	}
	61
	62	//=================================================================================
	63
	64	void AliAnalysisTaskCreateMixedDimuons::ConnectInputData(Option_t *) {
65
66	// Connect ESD or AOD here
67	// Called once
68
69	printf("-> AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");
70
71	TTree* tree = (TTree*) GetInputData(0);
72
73	if (!tree) {
74	Printf("ERROR: Could not read chain from input slot 0");
75	}
76	else {
77
78	fInputHandler = (AliMultiEventInputHandler*) AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
79	fPoolMuon = (AliEventPoolMuon*) AliAnalysisManager::GetAnalysisManager()->GetEventPool();
866de0ba	80
e46b3d21	81	if (!fInputHandler) {
	82	Printf("ERROR: Could not get AliMultiAODInputHandler");
	83	}
	84	else {
	85	fBufferSize = fInputHandler->GetBufferSize();
	86	if (fBufferSize>100) {
	87	printf("\n* WARNING AliAnalysisTaskCreateMixedDimuons::ConnectInputData -> Trying to set fBufferSize>100, forcing fBufferSize=100 *\n\n");
	88	fBufferSize = 100;
	89	}
	90	for (Int_t i=0; i<fBufferSize; i++) fInputAOD[i] = (AliAODEvent*) fInputHandler->GetEvent(i);
	91	}
866de0ba	92	}
	93
	94	printf("<- AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");
	95
	96	}
	97
	98	//=================================================================================
	99
	100	void AliAnalysisTaskCreateMixedDimuons::UserCreateOutputObjects() {
	101
	102	// Here the user-defined output containers should be created!!!
	103	// Called once
	104
	105	fOutputUserHandler = new AliAODHandler();
	106
	107	fOutputUserHandler -> Init("");
	108
	109	fOutputUserAOD = fOutputUserHandler -> GetAOD();
	110
	111	fOutputUserAODTree = fOutputUserHandler -> GetTree();
	112
	113	}
	114
	115	//=================================================================================
	116
	117	void AliAnalysisTaskCreateMixedDimuons::UserExec(Option_t *) {
	118
866de0ba	119	if (!fOutputUserAOD) {
	120	Printf("ERROR: fOutputUserAOD not available\n");
	121	return;
	122	}
	123
	124	printf("Calling USER EXEC\n\n");
	125
	126	for (Int_t iEv=0; iEv<fBufferSize; iEv++) {
	127	for (Int_t jEv=0; jEv<iEv; jEv++) {
	128
	129	Int_t nTracksEv[2] = {0};
	130	Int_t nFWMuonsEv[2] = {0};
	131
	132	nTracksEv[0] = fInputAOD[iEv]->GetNTracks();
	133	nTracksEv[1] = fInputAOD[jEv]->GetNTracks();
	134
	135	for (Int_t i=0; i<nTracksEv[0]; i++) if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[0]++;
	136	for (Int_t i=0; i<nTracksEv[1]; i++) if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[1]++;
	137
	138	// Muon track mixing to fill a mass spectrum
	139
	140	if (nFWMuonsEv[0] && nFWMuonsEv[1]) {
	141
	142	Int_t rndMuonTrack[2] = {0};
	143	rndMuonTrack[0] = gRandom->Integer(nFWMuonsEv[0]);
	144	rndMuonTrack[1] = gRandom->Integer(nFWMuonsEv[1]);
	145
	146	Int_t nFWMUonsAdded = 0;
	147	Int_t nPosTracksAdded = 0;
	148	Int_t nNegTracksAdded = 0;
	149
	150	AliAODVertex *vertex = new AliAODVertex();
	151	vertex -> SetX(0.0);
	152	vertex -> SetY(0.0);
	153	vertex -> SetZ(fPoolMuon->GetMeanPrimaryVertexZ());
	154
	155	Int_t muonCounter[2] = {0};
	156
	157	// adding tracks and vertex to the output event...
	158
	159	for (Int_t i=0; i<nTracksEv[0]; i++) {
	160	if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) {
	161	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
	162	iEv,i,fInputAOD[iEv]->GetTrack(i),fInputAOD[iEv]->GetTrack(i)->Pt(),
	163	fInputAOD[iEv]->GetTrack(i)->GetUniqueID());
	164	if (muonCounter[0]==rndMuonTrack[0]) {
	165	fOutputUserAOD->AddTrack(fInputAOD[iEv]->GetTrack(i));
	166	nFWMUonsAdded++;
	167	if (fInputAOD[iEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
	168	else nNegTracksAdded++;
	169	}
	170	muonCounter[0]++;
	171	}
	172	}
	173
	174	for (Int_t i=0; i<nTracksEv[1]; i++) {
	175	if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) {
	176	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
	177	jEv,i,fInputAOD[jEv]->GetTrack(i),fInputAOD[jEv]->GetTrack(i)->Pt(),
	178	fInputAOD[jEv]->GetTrack(i)->GetUniqueID());
	179	if (muonCounter[1]==rndMuonTrack[1]) {
	180	fOutputUserAOD->AddTrack(fInputAOD[jEv]->GetTrack(i));
	181	nFWMUonsAdded++;
	182	if (fInputAOD[jEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
183	else nNegTracksAdded++;
184	}
185	muonCounter[1]++;
186	}
187	}
188
189	fOutputUserAOD->AddVertex(vertex);
190
191	// ... done!
192
193	if (fDebug) {
194	for (Int_t i=0; i<nFWMUonsAdded; i++) {
195	AliAODTrack tr = (AliAODTrack) fOutputUserAOD->GetTrack(i);
196	printf("fOutputUserAOD->GetTrack(%d) = %p pt = %f\n",i,tr,tr->Pt());
197	}
198	}
199
200	fOutputUserAOD->GetHeader()->SetRefMultiplicity(nFWMUonsAdded);
201	fOutputUserAOD->GetHeader()->SetRefMultiplicityPos(nPosTracksAdded);
202	fOutputUserAOD->GetHeader()->SetRefMultiplicityNeg(nNegTracksAdded);
203
204	fOutputUserHandler -> FinishEvent();
205
206	}
207
208	PostData(1, fOutputUserAODTree);
209
210	}
211
212	}
213
214	}
215
216	//===================================================================================
217
218	void AliAnalysisTaskCreateMixedDimuons::Terminate(Option_t *) {
219
220	// Called once at the end of the query
221
222	printf("\n\nCalling TERMINATE \n\n\n");
223
224	fOutputUserHandler -> Terminate();
225
226	}
227
228	//===================================================================================