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


/* $Id$ */

#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 (!fOutputUserAOD) {
    Printf("ERROR: fOutputUserAOD not available\n");
    return;
  }

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

  for (Int_t iEv=0; iEv<fBufferSize; iEv++) {
    if (!fInputAOD[iEv]) {
      Printf("ERROR: fInputAOD[%d] not available\n",iEv);
      continue;
    }

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

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

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