[u/mrichter/AliRoot.git] / PWG3 / 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();
    fBufferSize = fInputHandler->GetBufferSize();
    if (fBufferSize>100) {
      printf("\n*** WARNING AliAnalysisTaskCreateMixedDimuons::ConnectInputData -> Trying to set fBufferSize>100, forcing fBufferSize=100 ***\n\n");
      fBufferSize = 100;
    }

    if (!fInputHandler) Printf("ERROR: Could not get AliMultiAODInputHandler");
    else 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 (!fInputAOD) {
    Printf("ERROR: fInputAOD not available\n");
    return;
  }

  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();
80	fBufferSize = fInputHandler->GetBufferSize();
81	if (fBufferSize>100) {
82	printf("\n* WARNING AliAnalysisTaskCreateMixedDimuons::ConnectInputData -> Trying to set fBufferSize>100, forcing fBufferSize=100 *\n\n");
83	fBufferSize = 100;
84	}
85
86	if (!fInputHandler) Printf("ERROR: Could not get AliMultiAODInputHandler");
87	else for (Int_t i=0; i<fBufferSize; i++) fInputAOD[i] = (AliAODEvent*) fInputHandler->GetEvent(i);
88	}
89
90	printf("<- AliAnalysisTaskCreateMixedDimuons::ConnectInputData\n");
91
92	}
93
94	//=================================================================================
95
96	void AliAnalysisTaskCreateMixedDimuons::UserCreateOutputObjects() {
97
98	// Here the user-defined output containers should be created!!!
99	// Called once
100
101	fOutputUserHandler = new AliAODHandler();
102
103	fOutputUserHandler -> Init("");
104
105	fOutputUserAOD = fOutputUserHandler -> GetAOD();
106
107	fOutputUserAODTree = fOutputUserHandler -> GetTree();
108
109	}
110
111	//=================================================================================
112
113	void AliAnalysisTaskCreateMixedDimuons::UserExec(Option_t *) {
114
115	if (!fInputAOD) {
116	Printf("ERROR: fInputAOD not available\n");
117	return;
118	}
119
120	if (!fOutputUserAOD) {
121	Printf("ERROR: fOutputUserAOD not available\n");
122	return;
123	}
124
125	printf("Calling USER EXEC\n\n");
126
127	for (Int_t iEv=0; iEv<fBufferSize; iEv++) {
128	for (Int_t jEv=0; jEv<iEv; jEv++) {
129
130	Int_t nTracksEv[2] = {0};
131	Int_t nFWMuonsEv[2] = {0};
132
133	nTracksEv[0] = fInputAOD[iEv]->GetNTracks();
134	nTracksEv[1] = fInputAOD[jEv]->GetNTracks();
135
136	for (Int_t i=0; i<nTracksEv[0]; i++) if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[0]++;
137	for (Int_t i=0; i<nTracksEv[1]; i++) if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) nFWMuonsEv[1]++;
138
139	// Muon track mixing to fill a mass spectrum
140
141	if (nFWMuonsEv[0] && nFWMuonsEv[1]) {
142
143	Int_t rndMuonTrack[2] = {0};
144	rndMuonTrack[0] = gRandom->Integer(nFWMuonsEv[0]);
145	rndMuonTrack[1] = gRandom->Integer(nFWMuonsEv[1]);
146
147	Int_t nFWMUonsAdded = 0;
148	Int_t nPosTracksAdded = 0;
149	Int_t nNegTracksAdded = 0;
150
151	AliAODVertex *vertex = new AliAODVertex();
152	vertex -> SetX(0.0);
153	vertex -> SetY(0.0);
154	vertex -> SetZ(fPoolMuon->GetMeanPrimaryVertexZ());
155
156	Int_t muonCounter[2] = {0};
157
158	// adding tracks and vertex to the output event...
159
160	for (Int_t i=0; i<nTracksEv[0]; i++) {
161	if(fInputAOD[iEv]->GetTrack(i)->IsMuonTrack()) {
162	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
163	iEv,i,fInputAOD[iEv]->GetTrack(i),fInputAOD[iEv]->GetTrack(i)->Pt(),
164	fInputAOD[iEv]->GetTrack(i)->GetUniqueID());
165	if (muonCounter[0]==rndMuonTrack[0]) {
166	fOutputUserAOD->AddTrack(fInputAOD[iEv]->GetTrack(i));
167	nFWMUonsAdded++;
168	if (fInputAOD[iEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
169	else nNegTracksAdded++;
170	}
171	muonCounter[0]++;
172	}
173	}
174
175	for (Int_t i=0; i<nTracksEv[1]; i++) {
176	if(fInputAOD[jEv]->GetTrack(i)->IsMuonTrack()) {
177	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
178	jEv,i,fInputAOD[jEv]->GetTrack(i),fInputAOD[jEv]->GetTrack(i)->Pt(),
179	fInputAOD[jEv]->GetTrack(i)->GetUniqueID());
180	if (muonCounter[1]==rndMuonTrack[1]) {
181	fOutputUserAOD->AddTrack(fInputAOD[jEv]->GetTrack(i));
182	nFWMUonsAdded++;
183	if (fInputAOD[jEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
184	else nNegTracksAdded++;
185	}
186	muonCounter[1]++;
187	}
188	}
189
190	fOutputUserAOD->AddVertex(vertex);
191
192	// ... done!
193
194	if (fDebug) {
195	for (Int_t i=0; i<nFWMUonsAdded; i++) {
196	AliAODTrack tr = (AliAODTrack) fOutputUserAOD->GetTrack(i);
197	printf("fOutputUserAOD->GetTrack(%d) = %p pt = %f\n",i,tr,tr->Pt());
198	}
199	}
200
201	fOutputUserAOD->GetHeader()->SetRefMultiplicity(nFWMUonsAdded);
202	fOutputUserAOD->GetHeader()->SetRefMultiplicityPos(nPosTracksAdded);
203	fOutputUserAOD->GetHeader()->SetRefMultiplicityNeg(nNegTracksAdded);
204
205	fOutputUserHandler -> FinishEvent();
206
207	}
208
209	PostData(1, fOutputUserAODTree);
210
211	}
212
213	}
214
215	}
216
217	//===================================================================================
218
219	void AliAnalysisTaskCreateMixedDimuons::Terminate(Option_t *) {
220
221	// Called once at the end of the query
222
223	printf("\n\nCalling TERMINATE \n\n\n");
224
225	fOutputUserHandler -> Terminate();
226
227	}
228
229	//===================================================================================