[u/mrichter/AliRoot.git] / PWG2 / RESONANCES / AliRsnReader.cxx

/**************************************************************************
 * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
 *                                                                        *
 * Author: The ALICE Off-line Project.                                    *
 * Contributors are mentioned in the code where appropriate.              *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, provided that the above copyright notice appears in all   *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purpose. It is          *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/
 
//-------------------------------------------------------------------------
//                     Class AliRsnReader
//
//   Reader for conversion of ESD or Kinematics output into AliRsnEvent
//   .....
//   .....
//   .....
//   .....
//   .....
//   .....
//   .....
// 
// author: A. Pulvirenti             (email: alberto.pulvirenti@ct.infn.it)
//-------------------------------------------------------------------------

#include <Riostream.h>

#include <TH1.h>
#include <TH3.h>
#include <TFile.h>
#include <TTree.h>
#include <TArrayF.h>
#include <TParticle.h>
#include <TRandom.h>

#include "AliRun.h"
#include "AliESD.h"
#include "AliStack.h"
#include "AliHeader.h"
#include "AliESDtrack.h"
#include "AliRunLoader.h"

#include "AliRsnDaughter.h"
#include "AliRsnEvent.h"
#include "AliRsnReader.h"

ClassImp(AliRsnReader)

//--------------------------------------------------------------------------------------------------------
AliRsnReader::AliRsnReader() :
  TObject(),
  fPIDMethod(kESDPID),
  fPtLimit4PID(4.0),
  fProbThreshold(0.0),
  fMaxRadius(3.0),
  fUseKineInfo(kFALSE),
  fEvents(0x0)
{
//
// Constructor.
// Initializes all working parameters to default values:
// - ESD particle identification
// - rejection of non-ITS-refitted tracks
// - maximum distance allowed from primary vertex = 3.0 cm (beam pipe)
//
	Int_t i;
	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = 1.0;
}
//--------------------------------------------------------------------------------------------------------
AliRsnReader::AliRsnReader(const AliRsnReader &copy) : 
  TObject(copy),
  fPIDMethod(copy.fPIDMethod),
  fPtLimit4PID(copy.fPtLimit4PID),
  fProbThreshold(copy.fProbThreshold),
  fMaxRadius(copy.fMaxRadius),
  fUseKineInfo(copy.fUseKineInfo),
  fEvents(0x0)
{
//
// Copy constructor.
// Initializes all working parameters to same values of another simlar object.
// TTree data member is not created.
//
	Int_t i;
	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
}
//--------------------------------------------------------------------------------------------------------
AliRsnReader& AliRsnReader::operator=(const AliRsnReader &copy)
{
//
// Assignment operator.
// Initializes all working parameters to same values of another simlar object.
// TTree data member is not created.
//
	fPIDMethod = copy.fPIDMethod;
	Int_t i;
	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
	fPtLimit4PID = copy.fPtLimit4PID;
	fProbThreshold = copy.fProbThreshold;
	fMaxRadius = copy.fMaxRadius;
	
	fUseKineInfo = copy.fUseKineInfo;

	fEvents = 0;
	
	return (*this);
}
//--------------------------------------------------------------------------------------------------------
void AliRsnReader::Clear(Option_t *option)
{
//
// Clear collection of filenames.
// If requested with the option "DELETELIST", 
// the collection object is also deleted.
//
	TString opt(option);
	
	if (!opt.CompareTo("TREE", TString::kIgnoreCase)) {
		fEvents->Reset();
		if (!opt.CompareTo("DELTREE", TString::kIgnoreCase)) {
			delete fEvents;
			fEvents = 0;
		}
	}
}
//--------------------------------------------------------------------------------------------------------
Double_t* AliRsnReader::GetPIDprobabilities(AliRsnDaughter track) const
{
//
// Computes PID probabilites starting from priors and weights
//
	Double_t *prob = new Double_t[AliPID::kSPECIES];
	
	Int_t i;
	
	// step 1 - compute the normalization factor
	Double_t sum = 0.0;
	for (i = 0; i < AliPID::kSPECIES; i++) {
		prob[i] = fPrior[i] * track.GetPIDweight(i);
		sum += prob[i];
	}
	if (sum <= 0.0) return 0;
	
	// step 2 - normalize PID weights by the relative prior probability
	for (Int_t i = 0; i < AliPID::kSPECIES; i++) {
		prob[i] /= sum;
	}
	
	return prob;
}
//--------------------------------------------------------------------------------------------------------
void AliRsnReader::Identify(AliRsnDaughter &track)
{
//
// Identifies a track according to method selected
//
	Bool_t doESDPID = (fPIDMethod == kESDPID);
	Bool_t keepRecSign = (fPIDMethod == kPerfectPIDwithRecSign);
	
	if (doESDPID) {
		// when doing ESD PID it is supposed that charge sign
		// comes from ESD track and it is not modified
		Double_t pt = track.GetPt();
		if (pt <= fPtLimit4PID) {
			Double_t *prob = GetPIDprobabilities(track);
			if (!prob) track.SetPDG(0);
			Int_t idx[AliPID::kSPECIES];
			TMath::Sort(static_cast<Int_t>(AliPID::kSPECIES), prob, idx);
			Int_t imax = idx[0];
			Double_t maxprob = prob[imax];
			if (maxprob >= fProbThreshold) {
				track.SetPDG((UShort_t)AliPID::ParticleCode(imax));
			}
			delete [] prob;
		}
		else {
			track.SetPDG(0);
		}
	}
	else {
		Short_t truePDG = track.GetTruePDG();
		track.SetPDG((UShort_t)TMath::Abs(truePDG));
		if (!keepRecSign) {
			if (TMath::Abs(truePDG) <= 13) {
				if (truePDG > 0) track.SetSign(-1); else track.SetSign(1);
			}
			else {
				if (truePDG > 0) track.SetSign(1); else track.SetSign(-1);
			}
		}
	}
}
//--------------------------------------------------------------------------------------------------------
TTree* AliRsnReader::ReadTracks(const char *path, Option_t *option)
{
//
// Reads AliESD in a given path, with and "experimental" method.
// When using this method, no Kinematics information is assumed
// and particle identification is always done with the bayesian method.
// No Kinematics informations are stored.
// Allowed options (actually):
// - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
// - "F" : reject 'fake' tracks (negative label)
//
	fPIDMethod = kESDPID;

	TTree *events = new TTree("selection", "AliRsnEvents");
	TTree::SetBranchStyle(1);
	AliRsnEvent *event = new AliRsnEvent;
	TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
	branch->SetAutoDelete(kFALSE);
	
	// get path
	TString strPath(path);
	if (strPath.Length() < 1) return 0;
	strPath += '/';
	
	// evaluate options
	TString opt(option);
	opt.ToUpper();
	Bool_t checkITSrefit = (opt.Contains("R"));
	Bool_t rejectFakes = (opt.Contains("F"));
	
	// opens ESD file
	TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
	if (!fileESD) return 0;
	if (fileESD->IsOpen() == kFALSE) return 0;
	TTree* treeESD = (TTree*)fileESD->Get("esdTree");
	AliESD *esd = 0;
	treeESD->SetBranchAddress("ESD", &esd);
	Int_t nev = (Int_t)treeESD->GetEntries();
	
	// loop on events
	Int_t i, nSelTracks = 0; 
	Double_t vertex[3];
	for (i = 0; i < nev; i++) {
	
		// message
//		cout << "\rEvent " << i << flush;
		treeESD->GetEntry(i);
		
		// primary vertex
		vertex[0] = esd->GetVertex()->GetXv();
		vertex[1] = esd->GetVertex()->GetYv();
		vertex[2] = esd->GetVertex()->GetZv();
		
		// create new AliRsnEvent object
		event->Clear("DELETE");
		event->Init();
				
		// get number of tracks
		Int_t ntracks = esd->GetNumberOfTracks();
		if (!ntracks) continue;
		
		// store tracks from ESD
		Int_t index, label;
		Double_t vtot, v[3];
		for (index = 0; index < ntracks; index++) {
			
			// get track
			AliESDtrack *esdTrack = esd->GetTrack(index);
			
			// check against vertex constraint
			esdTrack->GetXYZ(v);
			vtot  = (v[0] - vertex[0])*(v[0] - vertex[0]);
			vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
			vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
			vtot  = TMath::Sqrt(vtot);
			if (vtot > fMaxRadius) continue;
			
			// check for ITS refit
			if (checkITSrefit) {
				if (!(esdTrack->GetStatus() & AliESDtrack::kITSrefit)) continue;
			}
			
			// check for fakes
			label = esdTrack->GetLabel();
			if (rejectFakes && (label < 0)) continue;
			
			// create AliRsnDaughter (and make Bayesian PID)
			AliRsnDaughter track;
			if (!track.Adopt(esdTrack, checkITSrefit)) continue;
			track.SetIndex(index);
			track.SetLabel(label);
			Identify(track);
			
			// store in TClonesArray
			event->AddTrack(track);
			nSelTracks++;
		}
		
		// compute total multiplicity
		event->ComputeMultiplicity();
	
		// link to events tree and fill
		events->Fill();
	}
	
	fileESD->Close();
	return events;
}
//--------------------------------------------------------------------------------------------------------
TTree* AliRsnReader::ReadTracksAndParticles(const char *path, Option_t *option)
{
//
// Reads AliESD in a given path, getting also informations from Kinematics.
// In this case, the PID method used is the one selected with apposite setter.
// Allowed options (actually):
// - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
// - "F" : reject 'fake' tracks (negative label)
// - "M" : use 'true' momentum instead of reconstructed one
//
	TTree *events = new TTree("selection", "AliRsnEvents");
	TTree::SetBranchStyle(1);
	AliRsnEvent *event = new AliRsnEvent;
	TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
	branch->SetAutoDelete(kFALSE);
	
	// get path
	TString strPath(path);
	if (strPath.Length() < 1) return 0;
	
	// evaluate options
	TString opt(option);
	opt.ToUpper();
	Bool_t checkITSrefit = (opt.Contains("R"));
	Bool_t rejectFakes = (opt.Contains("F"));
	Bool_t copyMomentum = (opt.Contains("M"));
	
	// opens ESD file
	TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
	if (!fileESD) return 0;
	if (fileESD->IsOpen() == kFALSE) return 0;
	TTree* treeESD = (TTree*)fileESD->Get("esdTree");
	AliESD *esd = 0;
	treeESD->SetBranchAddress("ESD", &esd);
	Int_t nevRec = (Int_t)treeESD->GetEntries();
	
	// initialize run loader
	AliRunLoader *runLoader = OpenRunLoader(path);
	if (!runLoader) return 0;
	Int_t nevSim = runLoader->GetNumberOfEvents();
	
	// check number of reconstructed and simulated events
	if ( (nevSim != 0 && nevRec != 0) && (nevSim != nevRec) ) {
		cerr << "Count mismatch: sim = " << nevSim << " -- rec = " << nevRec << endl;
		return 0;
	}
	else if (nevSim == 0 && nevRec == 0) {
		cerr << "Count error: sim = rec = 0" << endl;
		return 0;
	}
	
	// loop on events
	Int_t i /*, procType*/, ntracks, nSelTracks = 0; 
	Double_t vertex[3];
	for (i = 0; i < nevRec; i++) {
	
		// get event
//		cout << "\rEvent " << i << " " << flush;
		treeESD->GetEntry(i);
		runLoader->GetEvent(i);
				
		// reject event if it is diffractive
		//procType = ((AliGenPythiaEventHeader*)runLoader->GetHeader()->GenEventHeader())->ProcessType();
		//if (procType == 92 || procType == 93 || procType == 94) {
		//	cout << "Skipping diffractive event" << endl;
		//	continue;
		//}
		
		// get particle stack
		AliStack *stack = runLoader->Stack();
				
		// primary vertex
		vertex[0] = esd->GetVertex()->GetXv();
		vertex[1] = esd->GetVertex()->GetYv();
		vertex[2] = esd->GetVertex()->GetZv();
		
		// multiplicity
		ntracks = esd->GetNumberOfTracks();
		if (!ntracks) {
			Warning("ReadTracksAndParticles", "Event %d has no tracks!", i);
			continue;
		}
		
		// create new AliRsnEvent object
		event->Clear("DELETE");
		event->Init();
				
		// store tracks from ESD
		Int_t index, label;
		Double_t vtot, v[3];
		for (index = 0; index < ntracks; index++) {
			
			// get track
			AliESDtrack *esdTrack = esd->GetTrack(index);
			
			// check against vertex constraint
			esdTrack->GetXYZ(v);
			vtot  = (v[0] - vertex[0])*(v[0] - vertex[0]);
			vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
			vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
			vtot  = TMath::Sqrt(vtot);
			if (vtot > fMaxRadius) continue;
			
			// check for fakes
			label = esdTrack->GetLabel();
			if (rejectFakes) {
				if (label < 0) continue;
			}
			
			// create AliRsnDaughter (and make Bayesian PID)
			AliRsnDaughter track;
			if (!track.Adopt(esdTrack, checkITSrefit)) continue;
			track.SetIndex(index);
			
			// retrieve particle and get Kine info
			TParticle *part = stack->Particle(TMath::Abs(label));
			track.SetTruePDG(part->GetPdgCode());
			Int_t mother = part->GetFirstMother();
			track.SetMother(mother);
			if (mother >= 0) {
				TParticle *mum = stack->Particle(mother);
				track.SetMotherPDG(mum->GetPdgCode());
			}
			if (copyMomentum) track.SetPxPyPz(part->Px(), part->Py(), part->Pz());
			
			// identification
			Identify(track);
			
			// store in TClonesArray
//			track.Print();
			event->AddTrack(track);
			nSelTracks++;
		}
		
		// compute total multiplicity
		event->ComputeMultiplicity();
	
		// link to events tree and fill
		events->Fill();
	}
	
	runLoader->UnloadKinematics();
	delete runLoader;
	fileESD->Close();
	
	return events;
}
//--------------------------------------------------------------------------------------------------------
void AliRsnReader::SetPriorProbabilities(Double_t *prior)
{
//
// Set prior probabilities to be used in case of ESD PID.
//
	if (!prior) return;
	
	Int_t i = 0;
	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = prior[i];
}
//--------------------------------------------------------------------------------------------------------
void AliRsnReader::SetPriorProbability(AliPID::EParticleType type, Double_t value)
{
//
// Sets prior probability referred to a single particle species.
//
	if (type >= AliPID::kElectron && type < AliPID::kPhoton) fPrior[type] = value;
}
//--------------------------------------------------------------------------------------------------------
AliPID::EParticleType AliRsnReader::FindType(Int_t pdg)
{
//
// Finds the enum value corresponding to a PDG code
//
	pdg = TMath::Abs(pdg);
	switch (pdg) {
		case   11: return AliPID::kElectron; break;
		case   13: return AliPID::kMuon; break;
		case  211: return AliPID::kPion; break;
		case  321: return AliPID::kKaon; break;
		case 2212: return AliPID::kProton; break;
		default  : return AliPID::kPhoton;
	}
}
//--------------------------------------------------------------------------------------------------------
AliRunLoader* AliRsnReader::OpenRunLoader(const char *path)
{
//
// Open the Run loader with events in a given path
//
	// clear gALICE
	if (gAlice) {
		delete gAlice;
		gAlice = 0;
	}
	
	// initialize run loader
	TString name(path);
	name += "/galice.root";
	AliRunLoader *runLoader = AliRunLoader::Open(name.Data());
	if (runLoader) {
		runLoader->LoadgAlice();
		gAlice = runLoader->GetAliRun();
		runLoader->LoadKinematics();
		runLoader->LoadHeader();
	}
	
	return runLoader;
}
Commit	Line	Data
b35947a8	1	/**************************************************************************
	2	* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
	3	* *
	4	* Author: The ALICE Off-line Project. *
	5	* Contributors are mentioned in the code where appropriate. *
	6	* *
	7	* Permission to use, copy, modify and distribute this software and its *
	8	* documentation strictly for non-commercial purposes is hereby granted *
	9	* without fee, provided that the above copyright notice appears in all *
	10	* copies and that both the copyright notice and this permission notice *
	11	* appear in the supporting documentation. The authors make no claims *
	12	* about the suitability of this software for any purpose. It is *
	13	* provided "as is" without express or implied warranty. *
	14	**************************************************************************/
	15
	16	//-------------------------------------------------------------------------
	17	// Class AliRsnReader
	18	//
	19	// Reader for conversion of ESD or Kinematics output into AliRsnEvent
2f769150	20	// .....
	21	// .....
	22	// .....
	23	// .....
	24	// .....
	25	// .....
	26	// .....
b35947a8	27	//
	28	// author: A. Pulvirenti (email: alberto.pulvirenti@ct.infn.it)
	29	//-------------------------------------------------------------------------
	30
	31	#include <Riostream.h>
	32
	33	#include <TH1.h>
	34	#include <TH3.h>
	35	#include <TFile.h>
2f769150	36	#include <TTree.h>
b35947a8	37	#include <TArrayF.h>
	38	#include <TParticle.h>
	39	#include <TRandom.h>
b35947a8	40
	41	#include "AliRun.h"
	42	#include "AliESD.h"
	43	#include "AliStack.h"
	44	#include "AliHeader.h"
	45	#include "AliESDtrack.h"
	46	#include "AliRunLoader.h"
b35947a8	47
	48	#include "AliRsnDaughter.h"
	49	#include "AliRsnEvent.h"
	50	#include "AliRsnReader.h"
	51
	52	ClassImp(AliRsnReader)
	53
	54	//--------------------------------------------------------------------------------------------------------
c37c6481	55	AliRsnReader::AliRsnReader() :
	56	TObject(),
	57	fPIDMethod(kESDPID),
	58	fPtLimit4PID(4.0),
	59	fProbThreshold(0.0),
	60	fMaxRadius(3.0),
	61	fUseKineInfo(kFALSE),
	62	fEvents(0x0)
2f769150	63	{
b35947a8	64	//
	65	// Constructor.
	66	// Initializes all working parameters to default values:
	67	// - ESD particle identification
	68	// - rejection of non-ITS-refitted tracks
	69	// - maximum distance allowed from primary vertex = 3.0 cm (beam pipe)
	70	//
b35947a8	71	Int_t i;
b35947a8	72	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = 1.0;
b35947a8	73	}
b35947a8	74	//--------------------------------------------------------------------------------------------------------
c37c6481	75	AliRsnReader::AliRsnReader(const AliRsnReader &copy) :
	76	TObject(copy),
	77	fPIDMethod(copy.fPIDMethod),
	78	fPtLimit4PID(copy.fPtLimit4PID),
	79	fProbThreshold(copy.fProbThreshold),
	80	fMaxRadius(copy.fMaxRadius),
	81	fUseKineInfo(copy.fUseKineInfo),
	82	fEvents(0x0)
2f769150	83	{
b35947a8	84	//
	85	// Copy constructor.
	86	// Initializes all working parameters to same values of another simlar object.
	87	// TTree data member is not created.
	88	//
2f769150	89	Int_t i;
2f769150	90	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
2f769150	91	}
	92	//--------------------------------------------------------------------------------------------------------
	93	AliRsnReader& AliRsnReader::operator=(const AliRsnReader &copy)
b35947a8	94	{
2f769150	95	//
	96	// Assignment operator.
	97	// Initializes all working parameters to same values of another simlar object.
	98	// TTree data member is not created.
	99	//
b35947a8	100	fPIDMethod = copy.fPIDMethod;
	101	Int_t i;
	102	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = copy.fPrior[i];
	103	fPtLimit4PID = copy.fPtLimit4PID;
	104	fProbThreshold = copy.fProbThreshold;
	105	fMaxRadius = copy.fMaxRadius;
30a3fb9a	106
30a3fb9a	107	fUseKineInfo = copy.fUseKineInfo;
b35947a8	108
b35947a8	109	fEvents = 0;
2f769150	110
2f769150	111	return (*this);
b35947a8	112	}
	113	//--------------------------------------------------------------------------------------------------------
	114	void AliRsnReader::Clear(Option_t *option)
2f769150	115	{
b35947a8	116	//
	117	// Clear collection of filenames.
	118	// If requested with the option "DELETELIST",
	119	// the collection object is also deleted.
	120	//
b35947a8	121	TString opt(option);
	122
	123	if (!opt.CompareTo("TREE", TString::kIgnoreCase)) {
	124	fEvents->Reset();
	125	if (!opt.CompareTo("DELTREE", TString::kIgnoreCase)) {
	126	delete fEvents;
	127	fEvents = 0;
	128	}
	129	}
b35947a8	130	}
b35947a8	131	//--------------------------------------------------------------------------------------------------------
2f769150	132	Double_t* AliRsnReader::GetPIDprobabilities(AliRsnDaughter track) const
2f769150	133	{
b35947a8	134	//
	135	// Computes PID probabilites starting from priors and weights
	136	//
b35947a8	137	Double_t *prob = new Double_t[AliPID::kSPECIES];
	138
	139	Int_t i;
	140
	141	// step 1 - compute the normalization factor
	142	Double_t sum = 0.0;
	143	for (i = 0; i < AliPID::kSPECIES; i++) {
	144	prob[i] = fPrior[i] * track.GetPIDweight(i);
	145	sum += prob[i];
	146	}
	147	if (sum <= 0.0) return 0;
	148
	149	// step 2 - normalize PID weights by the relative prior probability
	150	for (Int_t i = 0; i < AliPID::kSPECIES; i++) {
	151	prob[i] /= sum;
	152	}
	153
	154	return prob;
	155	}
	156	//--------------------------------------------------------------------------------------------------------
	157	void AliRsnReader::Identify(AliRsnDaughter &track)
2f769150	158	{
b35947a8	159	//
	160	// Identifies a track according to method selected
	161	//
b35947a8	162	Bool_t doESDPID = (fPIDMethod == kESDPID);
	163	Bool_t keepRecSign = (fPIDMethod == kPerfectPIDwithRecSign);
	164
	165	if (doESDPID) {
	166	// when doing ESD PID it is supposed that charge sign
	167	// comes from ESD track and it is not modified
	168	Double_t pt = track.GetPt();
	169	if (pt <= fPtLimit4PID) {
	170	Double_t *prob = GetPIDprobabilities(track);
	171	if (!prob) track.SetPDG(0);
	172	Int_t idx[AliPID::kSPECIES];
e5d53b83	173	TMath::Sort(static_cast<Int_t>(AliPID::kSPECIES), prob, idx);
b35947a8	174	Int_t imax = idx[0];
	175	Double_t maxprob = prob[imax];
	176	if (maxprob >= fProbThreshold) {
	177	track.SetPDG((UShort_t)AliPID::ParticleCode(imax));
	178	}
	179	delete [] prob;
	180	}
	181	else {
	182	track.SetPDG(0);
	183	}
	184	}
	185	else {
	186	Short_t truePDG = track.GetTruePDG();
	187	track.SetPDG((UShort_t)TMath::Abs(truePDG));
	188	if (!keepRecSign) {
	189	if (TMath::Abs(truePDG) <= 13) {
	190	if (truePDG > 0) track.SetSign(-1); else track.SetSign(1);
	191	}
	192	else {
	193	if (truePDG > 0) track.SetSign(1); else track.SetSign(-1);
	194	}
	195	}
	196	}
	197	}
	198	//--------------------------------------------------------------------------------------------------------
b35947a8	199	TTree* AliRsnReader::ReadTracks(const char path, Option_t option)
2f769150	200	{
b35947a8	201	//
	202	// Reads AliESD in a given path, with and "experimental" method.
	203	// When using this method, no Kinematics information is assumed
	204	// and particle identification is always done with the bayesian method.
	205	// No Kinematics informations are stored.
	206	// Allowed options (actually):
	207	// - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
	208	// - "F" : reject 'fake' tracks (negative label)
	209	//
b35947a8	210	fPIDMethod = kESDPID;
	211
	212	TTree *events = new TTree("selection", "AliRsnEvents");
	213	TTree::SetBranchStyle(1);
	214	AliRsnEvent *event = new AliRsnEvent;
	215	TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
	216	branch->SetAutoDelete(kFALSE);
	217
	218	// get path
	219	TString strPath(path);
	220	if (strPath.Length() < 1) return 0;
	221	strPath += '/';
	222
	223	// evaluate options
	224	TString opt(option);
	225	opt.ToUpper();
	226	Bool_t checkITSrefit = (opt.Contains("R"));
	227	Bool_t rejectFakes = (opt.Contains("F"));
	228
	229	// opens ESD file
	230	TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
	231	if (!fileESD) return 0;
	232	if (fileESD->IsOpen() == kFALSE) return 0;
	233	TTree* treeESD = (TTree*)fileESD->Get("esdTree");
	234	AliESD *esd = 0;
	235	treeESD->SetBranchAddress("ESD", &esd);
	236	Int_t nev = (Int_t)treeESD->GetEntries();
	237
	238	// loop on events
	239	Int_t i, nSelTracks = 0;
	240	Double_t vertex[3];
	241	for (i = 0; i < nev; i++) {
	242
	243	// message
30a3fb9a	244	// cout << "\rEvent " << i << flush;
b35947a8	245	treeESD->GetEntry(i);
	246
	247	// primary vertex
	248	vertex[0] = esd->GetVertex()->GetXv();
	249	vertex[1] = esd->GetVertex()->GetYv();
	250	vertex[2] = esd->GetVertex()->GetZv();
	251
	252	// create new AliRsnEvent object
	253	event->Clear("DELETE");
	254	event->Init();
b35947a8	255
	256	// get number of tracks
	257	Int_t ntracks = esd->GetNumberOfTracks();
	258	if (!ntracks) continue;
	259
	260	// store tracks from ESD
	261	Int_t index, label;
	262	Double_t vtot, v[3];
	263	for (index = 0; index < ntracks; index++) {
	264
	265	// get track
	266	AliESDtrack *esdTrack = esd->GetTrack(index);
	267
	268	// check against vertex constraint
	269	esdTrack->GetXYZ(v);
	270	vtot = (v[0] - vertex[0])*(v[0] - vertex[0]);
	271	vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
	272	vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
	273	vtot = TMath::Sqrt(vtot);
	274	if (vtot > fMaxRadius) continue;
	275
	276	// check for ITS refit
	277	if (checkITSrefit) {
	278	if (!(esdTrack->GetStatus() & AliESDtrack::kITSrefit)) continue;
	279	}
	280
	281	// check for fakes
	282	label = esdTrack->GetLabel();
	283	if (rejectFakes && (label < 0)) continue;
	284
	285	// create AliRsnDaughter (and make Bayesian PID)
	286	AliRsnDaughter track;
	287	if (!track.Adopt(esdTrack, checkITSrefit)) continue;
	288	track.SetIndex(index);
	289	track.SetLabel(label);
	290	Identify(track);
	291
	292	// store in TClonesArray
	293	event->AddTrack(track);
	294	nSelTracks++;
	295	}
	296
	297	// compute total multiplicity
30a3fb9a	298	event->ComputeMultiplicity();
b35947a8	299
	300	// link to events tree and fill
	301	events->Fill();
	302	}
	303
	304	fileESD->Close();
	305	return events;
	306	}
	307	//--------------------------------------------------------------------------------------------------------
	308	TTree* AliRsnReader::ReadTracksAndParticles(const char path, Option_t option)
2f769150	309	{
b35947a8	310	//
	311	// Reads AliESD in a given path, getting also informations from Kinematics.
	312	// In this case, the PID method used is the one selected with apposite setter.
	313	// Allowed options (actually):
	314	// - "R" : reject tracks which do not have the flag "kITSRefit" set to TRUE
	315	// - "F" : reject 'fake' tracks (negative label)
	316	// - "M" : use 'true' momentum instead of reconstructed one
	317	//
b35947a8	318	TTree *events = new TTree("selection", "AliRsnEvents");
	319	TTree::SetBranchStyle(1);
	320	AliRsnEvent *event = new AliRsnEvent;
	321	TBranch *branch = events->Branch("events", "AliRsnEvent", &event, 32000, 1);
	322	branch->SetAutoDelete(kFALSE);
	323
	324	// get path
	325	TString strPath(path);
	326	if (strPath.Length() < 1) return 0;
	327
	328	// evaluate options
	329	TString opt(option);
	330	opt.ToUpper();
	331	Bool_t checkITSrefit = (opt.Contains("R"));
	332	Bool_t rejectFakes = (opt.Contains("F"));
	333	Bool_t copyMomentum = (opt.Contains("M"));
	334
	335	// opens ESD file
	336	TFile *fileESD = TFile::Open(Form("%s/AliESDs.root", strPath.Data()));
	337	if (!fileESD) return 0;
	338	if (fileESD->IsOpen() == kFALSE) return 0;
	339	TTree* treeESD = (TTree*)fileESD->Get("esdTree");
	340	AliESD *esd = 0;
	341	treeESD->SetBranchAddress("ESD", &esd);
	342	Int_t nevRec = (Int_t)treeESD->GetEntries();
	343
	344	// initialize run loader
	345	AliRunLoader *runLoader = OpenRunLoader(path);
	346	if (!runLoader) return 0;
	347	Int_t nevSim = runLoader->GetNumberOfEvents();
	348
	349	// check number of reconstructed and simulated events
	350	if ( (nevSim != 0 && nevRec != 0) && (nevSim != nevRec) ) {
	351	cerr << "Count mismatch: sim = " << nevSim << " -- rec = " << nevRec << endl;
	352	return 0;
	353	}
	354	else if (nevSim == 0 && nevRec == 0) {
	355	cerr << "Count error: sim = rec = 0" << endl;
	356	return 0;
	357	}
	358
	359	// loop on events
30a3fb9a	360	Int_t i /, procType/, ntracks, nSelTracks = 0;
b35947a8	361	Double_t vertex[3];
	362	for (i = 0; i < nevRec; i++) {
	363
	364	// get event
30a3fb9a	365	// cout << "\rEvent " << i << " " << flush;
b35947a8	366	treeESD->GetEntry(i);
	367	runLoader->GetEvent(i);
	368
	369	// reject event if it is diffractive
30a3fb9a	370	//procType = ((AliGenPythiaEventHeader*)runLoader->GetHeader()->GenEventHeader())->ProcessType();
	371	//if (procType == 92 \|\| procType == 93 \|\| procType == 94) {
	372	// cout << "Skipping diffractive event" << endl;
	373	// continue;
	374	//}
b35947a8	375
	376	// get particle stack
	377	AliStack *stack = runLoader->Stack();
	378
	379	// primary vertex
	380	vertex[0] = esd->GetVertex()->GetXv();
	381	vertex[1] = esd->GetVertex()->GetYv();
	382	vertex[2] = esd->GetVertex()->GetZv();
	383
	384	// multiplicity
	385	ntracks = esd->GetNumberOfTracks();
	386	if (!ntracks) {
	387	Warning("ReadTracksAndParticles", "Event %d has no tracks!", i);
	388	continue;
	389	}
	390
	391	// create new AliRsnEvent object
	392	event->Clear("DELETE");
	393	event->Init();
b35947a8	394
	395	// store tracks from ESD
	396	Int_t index, label;
	397	Double_t vtot, v[3];
	398	for (index = 0; index < ntracks; index++) {
	399
	400	// get track
	401	AliESDtrack *esdTrack = esd->GetTrack(index);
	402
	403	// check against vertex constraint
	404	esdTrack->GetXYZ(v);
	405	vtot = (v[0] - vertex[0])*(v[0] - vertex[0]);
	406	vtot += (v[1] - vertex[1])*(v[1] - vertex[1]);
	407	vtot += (v[2] - vertex[2])*(v[2] - vertex[2]);
	408	vtot = TMath::Sqrt(vtot);
	409	if (vtot > fMaxRadius) continue;
	410
	411	// check for fakes
	412	label = esdTrack->GetLabel();
	413	if (rejectFakes) {
	414	if (label < 0) continue;
	415	}
	416
	417	// create AliRsnDaughter (and make Bayesian PID)
	418	AliRsnDaughter track;
	419	if (!track.Adopt(esdTrack, checkITSrefit)) continue;
	420	track.SetIndex(index);
	421
	422	// retrieve particle and get Kine info
	423	TParticle *part = stack->Particle(TMath::Abs(label));
	424	track.SetTruePDG(part->GetPdgCode());
	425	Int_t mother = part->GetFirstMother();
	426	track.SetMother(mother);
	427	if (mother >= 0) {
	428	TParticle *mum = stack->Particle(mother);
	429	track.SetMotherPDG(mum->GetPdgCode());
	430	}
	431	if (copyMomentum) track.SetPxPyPz(part->Px(), part->Py(), part->Pz());
	432
	433	// identification
	434	Identify(track);
	435
	436	// store in TClonesArray
30a3fb9a	437	// track.Print();
b35947a8	438	event->AddTrack(track);
	439	nSelTracks++;
	440	}
	441
	442	// compute total multiplicity
30a3fb9a	443	event->ComputeMultiplicity();
b35947a8	444
	445	// link to events tree and fill
	446	events->Fill();
	447	}
	448
	449	runLoader->UnloadKinematics();
	450	delete runLoader;
	451	fileESD->Close();
	452
	453	return events;
	454	}
	455	//--------------------------------------------------------------------------------------------------------
b35947a8	456	void AliRsnReader::SetPriorProbabilities(Double_t *prior)
2f769150	457	{
b35947a8	458	//
	459	// Set prior probabilities to be used in case of ESD PID.
	460	//
b35947a8	461	if (!prior) return;
	462
	463	Int_t i = 0;
	464	for (i = 0; i < AliPID::kSPECIES; i++) fPrior[i] = prior[i];
	465	}
	466	//--------------------------------------------------------------------------------------------------------
	467	void AliRsnReader::SetPriorProbability(AliPID::EParticleType type, Double_t value)
2f769150	468	{
b35947a8	469	//
	470	// Sets prior probability referred to a single particle species.
	471	//
b35947a8	472	if (type >= AliPID::kElectron && type < AliPID::kPhoton) fPrior[type] = value;
	473	}
	474	//--------------------------------------------------------------------------------------------------------
b35947a8	475	AliPID::EParticleType AliRsnReader::FindType(Int_t pdg)
2f769150	476	{
b35947a8	477	//
	478	// Finds the enum value corresponding to a PDG code
	479	//
b35947a8	480	pdg = TMath::Abs(pdg);
	481	switch (pdg) {
	482	case 11: return AliPID::kElectron; break;
	483	case 13: return AliPID::kMuon; break;
	484	case 211: return AliPID::kPion; break;
	485	case 321: return AliPID::kKaon; break;
	486	case 2212: return AliPID::kProton; break;
	487	default : return AliPID::kPhoton;
	488	}
	489	}
	490	//--------------------------------------------------------------------------------------------------------
	491	AliRunLoader* AliRsnReader::OpenRunLoader(const char *path)
2f769150	492	{
b35947a8	493	//
	494	// Open the Run loader with events in a given path
	495	//
b35947a8	496	// clear gALICE
	497	if (gAlice) {
	498	delete gAlice;
	499	gAlice = 0;
	500	}
	501
	502	// initialize run loader
	503	TString name(path);
	504	name += "/galice.root";
	505	AliRunLoader *runLoader = AliRunLoader::Open(name.Data());
	506	if (runLoader) {
	507	runLoader->LoadgAlice();
	508	gAlice = runLoader->GetAliRun();
	509	runLoader->LoadKinematics();
	510	runLoader->LoadHeader();
	511	}
	512
	513	return runLoader;
	514	}