[u/mrichter/AliRoot.git] / PWG0 / dNdEta / makeCorrection.C

//
// Script to make correction maps for dndeta measurements using the
// dNdEtaCorrection class.
// 

makeCorrection(Char_t* dataDir, Int_t nRuns=20) {

  Char_t str[256];

  gSystem->Load("../libPWG0base.so");

  // ########################################################
  // selection of esd tracks
  AliESDtrackCuts* esdTrackCuts = new AliESDtrackCuts();    
  esdTrackCuts->DefineHistograms(1);
  
  esdTrackCuts->SetMinNClustersTPC(50);
  esdTrackCuts->SetMaxChi2PerClusterTPC(3.5);
  esdTrackCuts->SetMaxCovDiagonalElements(2,2,0.5,0.5,2);
  esdTrackCuts->SetRequireTPCRefit(kTRUE);
  
  esdTrackCuts->SetMinNsigmaToVertex(3);
  esdTrackCuts->SetAcceptKingDaughters(kFALSE);

  AliLog::SetClassDebugLevel("AliESDtrackCuts",1);

  // ########################################################
  // definition of dNdEta correction object

  AlidNdEtaCorrection* dNdEtaMap = new AlidNdEtaCorrection("dndeta_correction", "dndeta_correction");

  // ########################################################
  // get the data dir  
  Char_t execDir[256];
  sprintf(execDir,gSystem->pwd());
  TSystemDirectory* baseDir = new TSystemDirectory(".",dataDir);
  TList* dirList            = baseDir->GetListOfFiles();
  Int_t nDirs               = dirList->GetEntries();
  // go back to the dir where this script is executed
  gSystem->cd(execDir);

  // ########################################################
  // definition of used pointers
  TFile* esdFile;
  TTree* esdTree;
  TBranch* esdBranch;

  AliESD* esd =0;

  // ########################################################
  // loop over runs
  Int_t nRunCounter = 0;
  for (Int_t r=0; r<nDirs; r++) {

    TSystemFile* presentDir = (TSystemFile*)dirList->At(r);
    if (!presentDir || !presentDir->IsDirectory())
      continue;
    // check that the files are there
    TString currentDataDir;
    currentDataDir.Form("%s/%s",dataDir, presentDir->GetName());
    cout << "Processing directory " << currentDataDir.Data() << endl;
    if ((!gSystem->Which(currentDataDir,"galice.root")) ||
          (!gSystem->Which(currentDataDir,"AliESDs.root"))) 
      continue;
    
    if (nRunCounter++ >= nRuns)
      break;    
    
    cout << "run #" << nRunCounter << endl;

    // #########################################################
    // setup galice and runloader
    if (gAlice) {
      delete gAlice->GetRunLoader();
      delete gAlice;
      gAlice=0;
    }

    sprintf(str,"%s/galice.root",currentDataDir.Data());
    AliRunLoader* runLoader = AliRunLoader::Open(str);

    runLoader->LoadgAlice();
    gAlice = runLoader->GetAliRun();
    runLoader->LoadKinematics();
    runLoader->LoadHeader();

    // #########################################################
    // open esd file and get the tree

    sprintf(str,"%s/AliESDs.root",currentDataDir.Data());
    // close it first to avoid memory leak
    if (esdFile)
      if (esdFile->IsOpen())
        esdFile->Close();

    esdFile = TFile::Open(str);
    esdTree = (TTree*)esdFile->Get("esdTree");
    if (!esdTree)
      continue;
    esdBranch = esdTree->GetBranch("ESD");
    esdBranch->SetAddress(&esd);
    if (!esdBranch)
      continue;

    // ########################################################
    // Magnetic field
    AliTracker::SetFieldMap(gAlice->Field(),kTRUE); // kTRUE means uniform magnetic field
    //AliKalmanTrack::SetConvConst(1000/0.299792458/5.);

    // ########################################################
    // getting number of events
    Int_t nEvents = (Int_t)runLoader->GetNumberOfEvents();
    Int_t nESDEvents = esdBranch->GetEntries();

    Int_t nEventsTriggered = 0;
    Int_t nEventsAll       = 0;
    
    if (nEvents!=nESDEvents) {
      cout << " Different number of events from runloader and esdtree!!!" << nEvents << " / " << nESDEvents << endl;
      return;
    }
   // ########################################################
    // loop over number of events
    cout << " looping over events..." << endl;
    for(Int_t i=0; i<nEvents; i++) {

      esdBranch->GetEntry(i);
      runLoader->GetEvent(i);              

      // ########################################################
      // get the EDS vertex
      AliESDVertex* vtxESD = esd->GetVertex();

      Double_t vtx[3];
      Double_t vtx_res[3];
      vtxESD->GetXYZ(vtx);          
    
      vtx_res[0] = vtxESD->GetXRes();
      vtx_res[1] = vtxESD->GetYRes();
      vtx_res[2] = vtxESD->GetZRes();

      Bool_t vertexReconstructed = kTRUE;

      // the vertex should be reconstructed
      if (strcmp(vtxESD->GetName(),"default")==0) 
	vertexReconstructed = kFALSE;

      // the resolution should be reasonable???
      if (vtx_res[2]==0 || vtx_res[2]>0.01)
	vertexReconstructed = kFALSE;

      // ########################################################
      // get the trigger info
      
      Bool_t triggered = kFALSE;

      // MB should be 
      // ITS_SPD_GFO_L0  : 32
      // VZERO_OR_LEFT   : 1
      // VZERO_OR_RIGHT  : 2

      ULong64_t triggerMask = esd->GetTriggerMask();

      nEventsAll++;      

      if (triggerMask&32 && ((triggerMask&1) || (triggerMask&2))) {
	triggered = kTRUE;       
	nEventsTriggered++;
      }

      // ########################################################
      // get the MC vertex
      AliGenPythiaEventHeader* genHeader =
        (AliGenPythiaEventHeader*)runLoader->GetHeader()->GenEventHeader();
      
      TArrayF vtxMC(3);
      genHeader->PrimaryVertex(vtxMC);
      Double_t vtx[3];
      vtx[0] = vtxMC[0];
      vtx[1] = vtxMC[1];
      vtx[2] = vtxMC[2];

      // ########################################################
      // loop over mc particles
      AliStack* particleStack = runLoader->Stack();
      Int_t nPrim    = particleStack->GetNprimary();

      for (Int_t i_mc=0; i_mc<nPrim; i_mc++) {
      
	TParticle* part = particleStack->Particle(i_mc);      
	if (!part || strcmp(part->GetName(),"XXX")==0) 
	  continue;
      
	TParticlePDG* pdgPart = part->GetPDG();

	Bool_t prim = kFALSE;
	// check if it's a primary particle - is there a better way ???
	if ((part->GetFirstDaughter() >= nPrim) || (part->GetFirstDaughter()==-1)) {
	  if (TMath::Abs(pdgPart->PdgCode())>10 && pdgPart->PdgCode()!=21 && strcmp(pdgPart->ParticleClass(),"Unknown")!=0)
	    prim = kTRUE;
	}
	if (!prim)
	  continue;

	if (pdgPart->Charge()==0)
	  continue;	

	Float_t eta = part->Eta();
	Float_t pt  = part->Pt();

	if (prim) {

	  dNdEtaMap->FillParticleAllEvents(eta, pt);		  

	  if (triggered)
	    dNdEtaMap->FillParticleWhenEventTriggered(eta, pt);
	  
	  if (vertexReconstructed)
	    dNdEtaMap->FillParticle(vtx[2], eta, 1.);	
	}
	
      }// end of mc particle
      
      // ########################################################
      // loop over esd tracks      
      Int_t nGoodTracks = 0;

      Int_t nTracks = esd->GetNumberOfTracks();            
      for (Int_t t=0; t<nTracks; t++) {
	AliESDtrack* esdTrack = esd->GetTrack(t);      
	
	// cut the esd track?
	if (!esdTrackCuts->AcceptTrack(esdTrack))
	  continue;
	
	nGoodTracks++;	
	
	Double_t p[3];
	esdTrack->GetPxPyPz(p);
	Float_t momentum = TMath::Sqrt(TMath::Power(p[0],2) + TMath::Power(p[1],2) + TMath::Power(p[2],2));
	
	Float_t eta = -100.;
	if((momentum != TMath::Abs(p[2]))&&(momentum != 0))
	  eta = 0.5*TMath::Log((momentum + p[2])/(momentum - p[2]));
	
	// using the eta of the mc particle
	Int_t label = TMath::Abs(esdTrack->GetLabel());
	if (label<0) {
	  cout << " cannot find corresponding mc part !!! " << label << endl;
	  continue;
	}
	TParticle* mcPart = particleStack->Particle(label);	
	eta = mcPart->Eta();
	Float_t pt = mcPart->Pt();
	
	if (vertexReconstructed)
	  dNdEtaMap->FillParticleWhenMeasuredTrack(vtx[2], eta, pt);	
	
      } // end of track loop
      
      dNdEtaMap->FillEvent(vtxMC[2], nGoodTracks);
      
      if (vertexReconstructed)
	dNdEtaMap->FillEventWithReconstructedVertex(vtxMC[2], nGoodTracks);

    } // end  of event loop
  } // end of run loop
  
  dNdEtaMap->Finish(nEventsAll, nEventsTriggered);  

  TFile* fout = new TFile("correction_map.root","RECREATE");
  
  esdTrackCuts->SaveHistograms("esd_track_cuts");
  dNdEtaMap->SaveHistograms();
  
  fout->Write();
  fout->Close();
  
  dNdEtaMap->DrawHistograms();

}
Commit	Line	Data
75ec0f41	1	//
	2	// Script to make correction maps for dndeta measurements using the
	3	// dNdEtaCorrection class.
	4	//
	5
	6	makeCorrection(Char_t* dataDir, Int_t nRuns=20) {
	7
	8	Char_t str[256];
	9
745e836a	10	gSystem->Load("../libPWG0base.so");
75ec0f41	11
	12	// ########################################################
	13	// selection of esd tracks
745e836a	14	AliESDtrackCuts* esdTrackCuts = new AliESDtrackCuts();
75ec0f41	15	esdTrackCuts->DefineHistograms(1);
	16
	17	esdTrackCuts->SetMinNClustersTPC(50);
	18	esdTrackCuts->SetMaxChi2PerClusterTPC(3.5);
	19	esdTrackCuts->SetMaxCovDiagonalElements(2,2,0.5,0.5,2);
	20	esdTrackCuts->SetRequireTPCRefit(kTRUE);
	21
	22	esdTrackCuts->SetMinNsigmaToVertex(3);
	23	esdTrackCuts->SetAcceptKingDaughters(kFALSE);
	24
745e836a	25	AliLog::SetClassDebugLevel("AliESDtrackCuts",1);
75ec0f41	26
	27	// ########################################################
	28	// definition of dNdEta correction object
	29
8b3563f4	30	AlidNdEtaCorrection* dNdEtaMap = new AlidNdEtaCorrection("dndeta_correction", "dndeta_correction");
75ec0f41	31
	32	// ########################################################
	33	// get the data dir
	34	Char_t execDir[256];
	35	sprintf(execDir,gSystem->pwd());
	36	TSystemDirectory* baseDir = new TSystemDirectory(".",dataDir);
	37	TList* dirList = baseDir->GetListOfFiles();
	38	Int_t nDirs = dirList->GetEntries();
	39	// go back to the dir where this script is executed
	40	gSystem->cd(execDir);
	41
	42	// ########################################################
	43	// definition of used pointers
	44	TFile* esdFile;
	45	TTree* esdTree;
	46	TBranch* esdBranch;
	47
	48	AliESD* esd =0;
	49
	50	// ########################################################
	51	// loop over runs
	52	Int_t nRunCounter = 0;
	53	for (Int_t r=0; r<nDirs; r++) {
	54
	55	TSystemFile* presentDir = (TSystemFile*)dirList->At(r);
	56	if (!presentDir \|\| !presentDir->IsDirectory())
	57	continue;
	58	// check that the files are there
	59	TString currentDataDir;
	60	currentDataDir.Form("%s/%s",dataDir, presentDir->GetName());
	61	cout << "Processing directory " << currentDataDir.Data() << endl;
	62	if ((!gSystem->Which(currentDataDir,"galice.root")) \|\|
	63	(!gSystem->Which(currentDataDir,"AliESDs.root")))
	64	continue;
	65
	66	if (nRunCounter++ >= nRuns)
	67	break;
	68
	69	cout << "run #" << nRunCounter << endl;
	70
	71	// #########################################################
	72	// setup galice and runloader
	73	if (gAlice) {
	74	delete gAlice->GetRunLoader();
	75	delete gAlice;
	76	gAlice=0;
	77	}
	78
	79	sprintf(str,"%s/galice.root",currentDataDir.Data());
	80	AliRunLoader* runLoader = AliRunLoader::Open(str);
	81
	82	runLoader->LoadgAlice();
	83	gAlice = runLoader->GetAliRun();
	84	runLoader->LoadKinematics();
	85	runLoader->LoadHeader();
	86
	87	// #########################################################
	88	// open esd file and get the tree
	89
	90	sprintf(str,"%s/AliESDs.root",currentDataDir.Data());
	91	// close it first to avoid memory leak
	92	if (esdFile)
	93	if (esdFile->IsOpen())
	94	esdFile->Close();
95
96	esdFile = TFile::Open(str);
97	esdTree = (TTree*)esdFile->Get("esdTree");
98	if (!esdTree)
99	continue;
100	esdBranch = esdTree->GetBranch("ESD");
101	esdBranch->SetAddress(&esd);
102	if (!esdBranch)
103	continue;
104
105	// ########################################################
106	// Magnetic field
107	AliTracker::SetFieldMap(gAlice->Field(),kTRUE); // kTRUE means uniform magnetic field
108	//AliKalmanTrack::SetConvConst(1000/0.299792458/5.);
109
110	// ########################################################
111	// getting number of events
112	Int_t nEvents = (Int_t)runLoader->GetNumberOfEvents();
113	Int_t nESDEvents = esdBranch->GetEntries();
55e05544	114
	115	Int_t nEventsTriggered = 0;
	116	Int_t nEventsAll = 0;
75ec0f41	117
	118	if (nEvents!=nESDEvents) {
	119	cout << " Different number of events from runloader and esdtree!!!" << nEvents << " / " << nESDEvents << endl;
	120	return;
	121	}
55e05544	122	// ########################################################
75ec0f41	123	// loop over number of events
	124	cout << " looping over events..." << endl;
	125	for(Int_t i=0; i<nEvents; i++) {
	126
	127	esdBranch->GetEntry(i);
	128	runLoader->GetEvent(i);
	129
	130	// ########################################################
	131	// get the EDS vertex
	132	AliESDVertex* vtxESD = esd->GetVertex();
	133
	134	Double_t vtx[3];
	135	Double_t vtx_res[3];
	136	vtxESD->GetXYZ(vtx);
	137
	138	vtx_res[0] = vtxESD->GetXRes();
	139	vtx_res[1] = vtxESD->GetYRes();
	140	vtx_res[2] = vtxESD->GetZRes();
	141
55e05544	142	Bool_t vertexReconstructed = kTRUE;
745e836a	143
75ec0f41	144	// the vertex should be reconstructed
75ec0f41	145	if (strcmp(vtxESD->GetName(),"default")==0)
55e05544	146	vertexReconstructed = kFALSE;
75ec0f41	147
75ec0f41	148	// the resolution should be reasonable???
745e836a	149	if (vtx_res[2]==0 \|\| vtx_res[2]>0.01)
55e05544	150	vertexReconstructed = kFALSE;
	151
	152	// ########################################################
	153	// get the trigger info
	154
	155	Bool_t triggered = kFALSE;
	156
	157	// MB should be
	158	// ITS_SPD_GFO_L0 : 32
	159	// VZERO_OR_LEFT : 1
	160	// VZERO_OR_RIGHT : 2
	161
	162	ULong64_t triggerMask = esd->GetTriggerMask();
	163
	164	nEventsAll++;
	165
	166	if (triggerMask&32 && ((triggerMask&1) \|\| (triggerMask&2))) {
	167	triggered = kTRUE;
	168	nEventsTriggered++;
	169	}
75ec0f41	170
	171	// ########################################################
	172	// get the MC vertex
	173	AliGenPythiaEventHeader* genHeader =
	174	(AliGenPythiaEventHeader*)runLoader->GetHeader()->GenEventHeader();
	175
	176	TArrayF vtxMC(3);
	177	genHeader->PrimaryVertex(vtxMC);
	178	Double_t vtx[3];
	179	vtx[0] = vtxMC[0];
	180	vtx[1] = vtxMC[1];
	181	vtx[2] = vtxMC[2];
	182
75ec0f41	183	// ########################################################
	184	// loop over mc particles
	185	AliStack* particleStack = runLoader->Stack();
	186	Int_t nPrim = particleStack->GetNprimary();
	187
	188	for (Int_t i_mc=0; i_mc<nPrim; i_mc++) {
	189
	190	TParticle* part = particleStack->Particle(i_mc);
	191	if (!part \|\| strcmp(part->GetName(),"XXX")==0)
	192	continue;
	193
	194	TParticlePDG* pdgPart = part->GetPDG();
	195
	196	Bool_t prim = kFALSE;
	197	// check if it's a primary particle - is there a better way ???
	198	if ((part->GetFirstDaughter() >= nPrim) \|\| (part->GetFirstDaughter()==-1)) {
	199	if (TMath::Abs(pdgPart->PdgCode())>10 && pdgPart->PdgCode()!=21 && strcmp(pdgPart->ParticleClass(),"Unknown")!=0)
	200	prim = kTRUE;
	201	}
	202	if (!prim)
	203	continue;
	204
	205	if (pdgPart->Charge()==0)
	206	continue;
	207
	208	Float_t eta = part->Eta();
55e05544	209	Float_t pt = part->Pt();
75ec0f41	210
745e836a	211	if (prim) {
55e05544	212
	213	dNdEtaMap->FillParticleAllEvents(eta, pt);
	214
	215	if (triggered)
	216	dNdEtaMap->FillParticleWhenEventTriggered(eta, pt);
745e836a	217
55e05544	218	if (vertexReconstructed)
55e05544	219	dNdEtaMap->FillParticle(vtx[2], eta, 1.);
745e836a	220	}
75ec0f41	221
75ec0f41	222	}// end of mc particle
55e05544	223
75ec0f41	224	// ########################################################
75ec0f41	225	// loop over esd tracks
55e05544	226	Int_t nGoodTracks = 0;
55e05544	227
75ec0f41	228	Int_t nTracks = esd->GetNumberOfTracks();
	229	for (Int_t t=0; t<nTracks; t++) {
	230	AliESDtrack* esdTrack = esd->GetTrack(t);
	231
	232	// cut the esd track?
	233	if (!esdTrackCuts->AcceptTrack(esdTrack))
	234	continue;
	235
55e05544	236	nGoodTracks++;
	237
	238	Double_t p[3];
	239	esdTrack->GetPxPyPz(p);
	240	Float_t momentum = TMath::Sqrt(TMath::Power(p[0],2) + TMath::Power(p[1],2) + TMath::Power(p[2],2));
	241
	242	Float_t eta = -100.;
	243	if((momentum != TMath::Abs(p[2]))&&(momentum != 0))
	244	eta = 0.5*TMath::Log((momentum + p[2])/(momentum - p[2]));
	245
75ec0f41	246	// using the eta of the mc particle
	247	Int_t label = TMath::Abs(esdTrack->GetLabel());
	248	if (label<0) {
	249	cout << " cannot find corresponding mc part !!! " << label << endl;
	250	continue;
	251	}
	252	TParticle* mcPart = particleStack->Particle(label);
	253	eta = mcPart->Eta();
55e05544	254	Float_t pt = mcPart->Pt();
	255
	256	if (vertexReconstructed)
	257	dNdEtaMap->FillParticleWhenMeasuredTrack(vtx[2], eta, pt);
	258
75ec0f41	259	} // end of track loop
55e05544	260
	261	dNdEtaMap->FillEvent(vtxMC[2], nGoodTracks);
	262
	263	if (vertexReconstructed)
	264	dNdEtaMap->FillEventWithReconstructedVertex(vtxMC[2], nGoodTracks);
	265
75ec0f41	266	} // end of event loop
75ec0f41	267	} // end of run loop
55e05544	268
55e05544	269	dNdEtaMap->Finish(nEventsAll, nEventsTriggered);
75ec0f41	270
	271	TFile* fout = new TFile("correction_map.root","RECREATE");
	272
	273	esdTrackCuts->SaveHistograms("esd_track_cuts");
	274	dNdEtaMap->SaveHistograms();
	275
	276	fout->Write();
	277	fout->Close();
	278
	279	dNdEtaMap->DrawHistograms();
	280
	281	}