[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]->GetNumberOfTracks();
      nTracksEv[1] = fInputAOD[jEv]->GetNumberOfTracks();
      // Check if both events contain std aod tracks
      if (!dynamic_cast<AliAODTrack*>(fInputAOD[iEv]->GetTrack(0)) || dynamic_cast<AliAODTrack*>(fInputAOD[jEv]->GetTrack(0))){
        AliFatal("Not a standard AOD");
      }

      for (Int_t i=0; i<nTracksEv[0]; i++) if(((AliAODTrack*)fInputAOD[iEv]->GetTrack(i))->IsMuonTrack()) nFWMuonsEv[0]++;
      for (Int_t i=0; i<nTracksEv[1]; i++) if(((AliAODTrack*)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(((AliAODTrack*)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((AliAODTrack*)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(((AliAODTrack*)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((AliAODTrack*)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 = dynamic_cast<AliAODTrack*>( fOutputUserAOD->GetTrack(i));
	    if(!tr) AliFatal("Not a standard AOD");
	    printf("fOutputUserAOD->GetTrack(%d) = %p    pt = %f\n",i,tr,tr->Pt());
	  }
	}
	

        AliAODHeader * header = dynamic_cast<AliAODHeader*>(fOutputUserAOD->GetHeader());
        if(!header) AliFatal("Not a standard AOD");
	header->SetRefMultiplicity(nFWMUonsAdded); 
	header->SetRefMultiplicityPos(nPosTracksAdded);
	header->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
4640c275	132	nTracksEv[0] = fInputAOD[iEv]->GetNumberOfTracks();
4640c275	133	nTracksEv[1] = fInputAOD[jEv]->GetNumberOfTracks();
f15c1f69	134	// Check if both events contain std aod tracks
	135	if (!dynamic_cast<AliAODTrack>(fInputAOD[iEv]->GetTrack(0)) \|\| dynamic_cast<AliAODTrack>(fInputAOD[jEv]->GetTrack(0))){
	136	AliFatal("Not a standard AOD");
	137	}
	138
	139	for (Int_t i=0; i<nTracksEv[0]; i++) if(((AliAODTrack*)fInputAOD[iEv]->GetTrack(i))->IsMuonTrack()) nFWMuonsEv[0]++;
	140	for (Int_t i=0; i<nTracksEv[1]; i++) if(((AliAODTrack*)fInputAOD[jEv]->GetTrack(i))->IsMuonTrack()) nFWMuonsEv[1]++;
866de0ba	141
	142	// Muon track mixing to fill a mass spectrum
	143
	144	if (nFWMuonsEv[0] && nFWMuonsEv[1]) {
	145
	146	Int_t rndMuonTrack[2] = {0};
	147	rndMuonTrack[0] = gRandom->Integer(nFWMuonsEv[0]);
	148	rndMuonTrack[1] = gRandom->Integer(nFWMuonsEv[1]);
	149
	150	Int_t nFWMUonsAdded = 0;
	151	Int_t nPosTracksAdded = 0;
	152	Int_t nNegTracksAdded = 0;
	153
	154	AliAODVertex *vertex = new AliAODVertex();
	155	vertex -> SetX(0.0);
	156	vertex -> SetY(0.0);
	157	vertex -> SetZ(fPoolMuon->GetMeanPrimaryVertexZ());
	158
	159	Int_t muonCounter[2] = {0};
	160
	161	// adding tracks and vertex to the output event...
	162
	163	for (Int_t i=0; i<nTracksEv[0]; i++) {
f15c1f69	164	if(((AliAODTrack*)fInputAOD[iEv]->GetTrack(i))->IsMuonTrack()) {
866de0ba	165	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
	166	iEv,i,fInputAOD[iEv]->GetTrack(i),fInputAOD[iEv]->GetTrack(i)->Pt(),
	167	fInputAOD[iEv]->GetTrack(i)->GetUniqueID());
	168	if (muonCounter[0]==rndMuonTrack[0]) {
f15c1f69	169	fOutputUserAOD->AddTrack((AliAODTrack*)fInputAOD[iEv]->GetTrack(i));
866de0ba	170	nFWMUonsAdded++;
	171	if (fInputAOD[iEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
	172	else nNegTracksAdded++;
	173	}
	174	muonCounter[0]++;
	175	}
	176	}
	177
	178	for (Int_t i=0; i<nTracksEv[1]; i++) {
f15c1f69	179	if(((AliAODTrack*)fInputAOD[jEv]->GetTrack(i))->IsMuonTrack()) {
866de0ba	180	if (fDebug) printf("fInputAOD[%d]->GetTrack(%d) = %p pt = %f uniqueID = %d\n",
	181	jEv,i,fInputAOD[jEv]->GetTrack(i),fInputAOD[jEv]->GetTrack(i)->Pt(),
	182	fInputAOD[jEv]->GetTrack(i)->GetUniqueID());
	183	if (muonCounter[1]==rndMuonTrack[1]) {
f15c1f69	184	fOutputUserAOD->AddTrack((AliAODTrack*)fInputAOD[jEv]->GetTrack(i));
866de0ba	185	nFWMUonsAdded++;
	186	if (fInputAOD[jEv]->GetTrack(i)->Charge()>0) nPosTracksAdded++;
	187	else nNegTracksAdded++;
	188	}
	189	muonCounter[1]++;
	190	}
	191	}
	192
	193	fOutputUserAOD->AddVertex(vertex);
	194
	195	// ... done!
	196
	197	if (fDebug) {
	198	for (Int_t i=0; i<nFWMUonsAdded; i++) {
f15c1f69	199	AliAODTrack tr = dynamic_cast<AliAODTrack>( fOutputUserAOD->GetTrack(i));
f15c1f69	200	if(!tr) AliFatal("Not a standard AOD");
866de0ba	201	printf("fOutputUserAOD->GetTrack(%d) = %p pt = %f\n",i,tr,tr->Pt());
	202	}
	203	}
	204
0a918d8d	205
	206	AliAODHeader * header = dynamic_cast<AliAODHeader*>(fOutputUserAOD->GetHeader());
	207	if(!header) AliFatal("Not a standard AOD");
	208	header->SetRefMultiplicity(nFWMUonsAdded);
	209	header->SetRefMultiplicityPos(nPosTracksAdded);
	210	header->SetRefMultiplicityNeg(nNegTracksAdded);
866de0ba	211
	212	fOutputUserHandler -> FinishEvent();
	213
	214	}
	215
	216	PostData(1, fOutputUserAODTree);
	217
	218	}
	219
	220	}
	221
	222	}
	223
	224	//===================================================================================
	225
	226	void AliAnalysisTaskCreateMixedDimuons::Terminate(Option_t *) {
	227
	228	// Called once at the end of the query
	229
	230	printf("\n\nCalling TERMINATE \n\n\n");
	231
	232	fOutputUserHandler -> Terminate();
	233
	234	}
	235
	236	//===================================================================================