[u/mrichter/AliRoot.git] / PWG / EMCAL / AliEmcalEsdTrackFilterTask.cxx

// $Id$
//
// Task to filter Esd tracks and propagate to Emcal surface.
//
// Author: C.Loizides

#include "AliEmcalEsdTrackFilterTask.h"
#include <TClonesArray.h>
#include <TRandom3.h>
#include <TGeoGlobalMagField.h>
#include <AliAnalysisManager.h>
#include <AliEMCALRecoUtils.h>
#include <AliESDEvent.h>
#include <AliESDtrackCuts.h>
#include <AliMagF.h>
#include <AliTrackerBase.h>


ClassImp(AliEmcalEsdTrackFilterTask)

//________________________________________________________________________
AliEmcalEsdTrackFilterTask::AliEmcalEsdTrackFilterTask() : 
  AliAnalysisTaskSE("AliEmcalEsdTrackFilterTask"),
  fEsdTrackCuts(0),
  fDoSpdVtxCon(0),
  fHybridTrackCuts(0),
  fTracksName(),
  fIncludeNoITS(kTRUE),
  fDoPropagation(kFALSE),
  fDist(440),
  fTrackEfficiency(1),
  fIsMC(kFALSE),
  fEsdEv(0),
  fTracks(0)
{
  // Constructor.
}

//________________________________________________________________________
AliEmcalEsdTrackFilterTask::AliEmcalEsdTrackFilterTask(const char *name) : 
  AliAnalysisTaskSE(name),
  fEsdTrackCuts(0),
  fDoSpdVtxCon(0),
  fHybridTrackCuts(0),
  fTracksName("EsdTracksOut"),
  fIncludeNoITS(kTRUE),
  fDoPropagation(kFALSE),
  fDist(440),
  fTrackEfficiency(1),
  fIsMC(kFALSE),
  fEsdEv(0),
  fTracks(0)
{
  // Constructor.

  if (!name)
    return;

  SetName(name);

  fBranchNames = "ESD:AliESDHeader.,AliESDRun.,SPDVertex.,Tracks";
}

//________________________________________________________________________
AliEmcalEsdTrackFilterTask::~AliEmcalEsdTrackFilterTask()
{
  //Destructor

  delete fEsdTrackCuts;
}

//________________________________________________________________________
void AliEmcalEsdTrackFilterTask::UserCreateOutputObjects()
{
  // Create histograms.

  fTracks = new TClonesArray("AliESDtrack");
  fTracks->SetName(fTracksName);

  if (fDoSpdVtxCon) {
    if (!fEsdTrackCuts) {
      AliInfo("No track cuts given, creating default (standard only TPC) cuts");
      fEsdTrackCuts = AliESDtrackCuts::GetStandardTPCOnlyTrackCuts();
      fEsdTrackCuts->SetPtRange(0.15,1e3);
    } 
  } else {
    AliWarning("No track cuts given, but maybe this is indeed intended?");
  }
}

//________________________________________________________________________
void AliEmcalEsdTrackFilterTask::UserExec(Option_t *) 
{
  // Main loop, called for each event.

  fEsdEv = dynamic_cast<AliESDEvent*>(InputEvent());
  if (!fEsdEv) {
    AliError("Task works only on ESD events, returning");
    return;
  }

  AliAnalysisManager *am = AliAnalysisManager::GetAnalysisManager();
  if (!am) {
    AliError("Manager zero, returning");
    return;
  }

  // add tracks to event if not yet there
  fTracks->Delete();
  if (!(InputEvent()->FindListObject(fTracksName)))
    InputEvent()->AddObject(fTracks);

  if (!fHybridTrackCuts) { // constrain TPC tracks to SPD vertex if fDoSpdVtxCon==kTRUE
    am->LoadBranch("AliESDRun.");
    am->LoadBranch("AliESDHeader.");
    am->LoadBranch("Tracks");

    if (fDoSpdVtxCon) {
      if (!TGeoGlobalMagField::Instance()->GetField()) { // construct field map
        fEsdEv->InitMagneticField();
      }
      am->LoadBranch("SPDVertex.");
      const AliESDVertex *vtxSPD = fEsdEv->GetPrimaryVertexSPD();
      if (!vtxSPD) {
        AliError("No SPD vertex, returning");
        return;
      }
      Int_t ntr = fEsdEv->GetNumberOfTracks();
      for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
        AliESDtrack *etrack = fEsdEv->GetTrack(i);
        if (!etrack)
          continue;

        if (!fEsdTrackCuts->AcceptTrack(etrack))
          continue;

        AliESDtrack *ntrack = AliESDtrackCuts::GetTPCOnlyTrack(fEsdEv,etrack->GetID());
        if (!ntrack)
          continue;
        if (ntrack->Pt()<=0) {
          delete ntrack;
          continue;
        }
        Double_t bfield[3] = {0,0,0};
        ntrack->GetBxByBz(bfield);
        AliExternalTrackParam exParam;
        Bool_t relate = ntrack->RelateToVertexBxByBz(vtxSPD,bfield,kVeryBig,&exParam);
        if (!relate) {
          delete ntrack;
          continue;
        }
        // set the constraint parameters to the track
        ntrack->Set(exParam.GetX(),exParam.GetAlpha(),exParam.GetParameter(),exParam.GetCovariance());
        if (ntrack->Pt()<=0) {
          delete ntrack;
          continue;
        }

	if (fTrackEfficiency < 1) {
	  Double_t r = gRandom->Rndm();
	  if (fTrackEfficiency < r)
	    continue;
	}

	if (fDoPropagation) 	
	  AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(ntrack,fDist);
        new ((*fTracks)[ntrnew++]) AliESDtrack(*ntrack);
        delete ntrack;
      }
    } else { /* no spd vtx constraint */
      Int_t ntr = fEsdEv->GetNumberOfTracks();
      for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
        AliESDtrack *etrack = fEsdEv->GetTrack(i);
        if (!etrack)
          continue;

	if ((fEsdTrackCuts!=0) && !fEsdTrackCuts->AcceptTrack(etrack))
          continue;
	
	if (fTrackEfficiency < 1) {
	  Double_t r = gRandom->Rndm();
	  if (fTrackEfficiency < r)
	    continue;
	}

        AliESDtrack *ntrack = new ((*fTracks)[ntrnew++]) AliESDtrack(*etrack);
	if (fDoPropagation) 	
	  AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(ntrack,fDist);
      }
    }

  } else { // use hybrid track cuts

    am->LoadBranch("Tracks");
    Int_t ntr = fEsdEv->GetNumberOfTracks();
    for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
      AliESDtrack *etrack = fEsdEv->GetTrack(i);
      if (!etrack) 
	continue;

      if (fEsdTrackCuts->AcceptTrack(etrack)) {
	if (fTrackEfficiency < 1) {
	  Double_t r = gRandom->Rndm();
	  if (fTrackEfficiency < r)
	    continue;
	}
        AliESDtrack *newTrack = new ((*fTracks)[ntrnew]) AliESDtrack(*etrack);
	if (fDoPropagation) 
	  AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(newTrack,fDist);
        newTrack->SetBit(BIT(22),0); 
        newTrack->SetBit(BIT(23),0);
	if (!fIsMC) newTrack->SetLabel(0);
        ++ntrnew;
      } else if (fHybridTrackCuts->AcceptTrack(etrack)) {
	if (!etrack->GetConstrainedParam())
	  continue;
	UInt_t status = etrack->GetStatus();
	if (!fIncludeNoITS && ((status&AliESDtrack::kITSrefit)==0))
	  continue;

	if (fTrackEfficiency < 1) {
	  Double_t r = gRandom->Rndm();
	  if (fTrackEfficiency < r)
	    continue;
	}
	AliESDtrack *newTrack = new ((*fTracks)[ntrnew]) AliESDtrack(*etrack);
	if (fDoPropagation) 	
	  AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(newTrack,fDist);
	const AliExternalTrackParam* constrainParam = etrack->GetConstrainedParam();
	newTrack->Set(constrainParam->GetX(),
		      constrainParam->GetAlpha(),
		      constrainParam->GetParameter(),
		      constrainParam->GetCovariance());
	if ((status&AliESDtrack::kITSrefit)==0) {
	  newTrack->SetBit(BIT(22),0); //type 2
	  newTrack->SetBit(BIT(23),1);
	} else {
	  newTrack->SetBit(BIT(22),1); //type 1
	  newTrack->SetBit(BIT(23),0);
	}
	if (!fIsMC) newTrack->SetLabel(0);
	++ntrnew;
      }
    }
  }
}
Commit	Line	Data
cc2b5fb5	1	// $Id$
b0a53615	2	//
cdd4e0e2	3	// Task to filter Esd tracks and propagate to Emcal surface.
b0a53615	4	//
cd231d42	5	// Author: C.Loizides
cc2b5fb5	6
33a8056d	7	#include "AliEmcalEsdTrackFilterTask.h"
cc2b5fb5	8	#include <TClonesArray.h>
d47bb504	9	#include <TRandom3.h>
cc2b5fb5	10	#include <TGeoGlobalMagField.h>
fe75ca27	11	#include <AliAnalysisManager.h>
33a8056d	12	#include <AliEMCALRecoUtils.h>
fe75ca27 CL	13	#include <AliESDEvent.h>
	14	#include <AliESDtrackCuts.h>
	15	#include <AliMagF.h>
	16	#include <AliTrackerBase.h>
cc2b5fb5	17
d47bb504	18
cdd4e0e2	19	ClassImp(AliEmcalEsdTrackFilterTask)
cc2b5fb5	20
cc2b5fb5	21	//________________________________________________________________________
cdd4e0e2 CL	22	AliEmcalEsdTrackFilterTask::AliEmcalEsdTrackFilterTask() :
cdd4e0e2 CL	23	AliAnalysisTaskSE("AliEmcalEsdTrackFilterTask"),
cc2b5fb5	24	fEsdTrackCuts(0),
4d1e3842	25	fDoSpdVtxCon(0),
57184cc6	26	fHybridTrackCuts(0),
cc2b5fb5	27	fTracksName(),
6231f841	28	fIncludeNoITS(kTRUE),
35fb503f	29	fDoPropagation(kFALSE),
af2f024d	30	fDist(440),
d47bb504	31	fTrackEfficiency(1),
1737c13e	32	fIsMC(kFALSE),
cc2b5fb5	33	fEsdEv(0),
	34	fTracks(0)
	35	{
a93ef37c	36	// Constructor.
cc2b5fb5	37	}
	38
	39	//________________________________________________________________________
cdd4e0e2	40	AliEmcalEsdTrackFilterTask::AliEmcalEsdTrackFilterTask(const char *name) :
cc2b5fb5	41	AliAnalysisTaskSE(name),
cc2b5fb5	42	fEsdTrackCuts(0),
4d1e3842	43	fDoSpdVtxCon(0),
57184cc6	44	fHybridTrackCuts(0),
cdd4e0e2	45	fTracksName("EsdTracksOut"),
6231f841	46	fIncludeNoITS(kTRUE),
35fb503f	47	fDoPropagation(kFALSE),
af2f024d	48	fDist(440),
d47bb504	49	fTrackEfficiency(1),
1737c13e	50	fIsMC(kFALSE),
cc2b5fb5	51	fEsdEv(0),
	52	fTracks(0)
	53	{
a93ef37c	54	// Constructor.
27fde6c7	55
	56	if (!name)
	57	return;
	58
	59	SetName(name);
	60
a93ef37c	61	fBranchNames = "ESD:AliESDHeader.,AliESDRun.,SPDVertex.,Tracks";
cc2b5fb5	62	}
	63
	64	//________________________________________________________________________
cdd4e0e2	65	AliEmcalEsdTrackFilterTask::~AliEmcalEsdTrackFilterTask()
cc2b5fb5	66	{
	67	//Destructor
	68
	69	delete fEsdTrackCuts;
	70	}
	71
	72	//________________________________________________________________________
cdd4e0e2	73	void AliEmcalEsdTrackFilterTask::UserCreateOutputObjects()
cc2b5fb5	74	{
a93ef37c	75	// Create histograms.
cc2b5fb5	76
	77	fTracks = new TClonesArray("AliESDtrack");
	78	fTracks->SetName(fTracksName);
	79
04410c6f	80	if (fDoSpdVtxCon) {
	81	if (!fEsdTrackCuts) {
	82	AliInfo("No track cuts given, creating default (standard only TPC) cuts");
	83	fEsdTrackCuts = AliESDtrackCuts::GetStandardTPCOnlyTrackCuts();
	84	fEsdTrackCuts->SetPtRange(0.15,1e3);
	85	}
	86	} else {
	87	AliWarning("No track cuts given, but maybe this is indeed intended?");
cc2b5fb5	88	}
	89	}
	90
	91	//________________________________________________________________________
cdd4e0e2	92	void AliEmcalEsdTrackFilterTask::UserExec(Option_t *)
cc2b5fb5	93	{
	94	// Main loop, called for each event.
	95
	96	fEsdEv = dynamic_cast<AliESDEvent*>(InputEvent());
	97	if (!fEsdEv) {
	98	AliError("Task works only on ESD events, returning");
	99	return;
	100	}
	101
	102	AliAnalysisManager *am = AliAnalysisManager::GetAnalysisManager();
	103	if (!am) {
	104	AliError("Manager zero, returning");
	105	return;
	106	}
	107
	108	// add tracks to event if not yet there
71c1dd76	109	fTracks->Delete();
cc2b5fb5	110	if (!(InputEvent()->FindListObject(fTracksName)))
	111	InputEvent()->AddObject(fTracks);
	112
ce2a8490	113	if (!fHybridTrackCuts) { // constrain TPC tracks to SPD vertex if fDoSpdVtxCon==kTRUE
57184cc6	114	am->LoadBranch("AliESDRun.");
57184cc6	115	am->LoadBranch("AliESDHeader.");
57184cc6	116	am->LoadBranch("Tracks");
4d1e3842	117
	118	if (fDoSpdVtxCon) {
	119	if (!TGeoGlobalMagField::Instance()->GetField()) { // construct field map
	120	fEsdEv->InitMagneticField();
57184cc6	121	}
4d1e3842	122	am->LoadBranch("SPDVertex.");
	123	const AliESDVertex *vtxSPD = fEsdEv->GetPrimaryVertexSPD();
	124	if (!vtxSPD) {
	125	AliError("No SPD vertex, returning");
	126	return;
57184cc6	127	}
4d1e3842	128	Int_t ntr = fEsdEv->GetNumberOfTracks();
	129	for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
	130	AliESDtrack *etrack = fEsdEv->GetTrack(i);
	131	if (!etrack)
	132	continue;
d47bb504	133
4d1e3842	134	if (!fEsdTrackCuts->AcceptTrack(etrack))
4d1e3842	135	continue;
d47bb504	136
4d1e3842	137	AliESDtrack *ntrack = AliESDtrackCuts::GetTPCOnlyTrack(fEsdEv,etrack->GetID());
	138	if (!ntrack)
	139	continue;
	140	if (ntrack->Pt()<=0) {
	141	delete ntrack;
	142	continue;
	143	}
	144	Double_t bfield[3] = {0,0,0};
	145	ntrack->GetBxByBz(bfield);
	146	AliExternalTrackParam exParam;
	147	Bool_t relate = ntrack->RelateToVertexBxByBz(vtxSPD,bfield,kVeryBig,&exParam);
	148	if (!relate) {
	149	delete ntrack;
	150	continue;
	151	}
	152	// set the constraint parameters to the track
	153	ntrack->Set(exParam.GetX(),exParam.GetAlpha(),exParam.GetParameter(),exParam.GetCovariance());
	154	if (ntrack->Pt()<=0) {
	155	delete ntrack;
	156	continue;
	157	}
c0440012	158
	159	if (fTrackEfficiency < 1) {
	160	Double_t r = gRandom->Rndm();
	161	if (fTrackEfficiency < r)
	162	continue;
	163	}
	164
04410c6f	165	if (fDoPropagation)
04410c6f	166	AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(ntrack,fDist);
4d1e3842	167	new ((fTracks)[ntrnew++]) AliESDtrack(ntrack);
57184cc6	168	delete ntrack;
57184cc6	169	}
4d1e3842	170	} else { /* no spd vtx constraint */
	171	Int_t ntr = fEsdEv->GetNumberOfTracks();
	172	for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
	173	AliESDtrack *etrack = fEsdEv->GetTrack(i);
	174	if (!etrack)
	175	continue;
c0440012	176
	177	if ((fEsdTrackCuts!=0) && !fEsdTrackCuts->AcceptTrack(etrack))
	178	continue;
	179
d47bb504	180	if (fTrackEfficiency < 1) {
	181	Double_t r = gRandom->Rndm();
	182	if (fTrackEfficiency < r)
	183	continue;
	184	}
c0440012	185
04410c6f	186	AliESDtrack ntrack = new ((fTracks)[ntrnew++]) AliESDtrack(*etrack);
	187	if (fDoPropagation)
	188	AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(ntrack,fDist);
4d1e3842	189	}
cc2b5fb5	190	}
cc2b5fb5	191
10d1f756	192	} else { // use hybrid track cuts
57184cc6	193
	194	am->LoadBranch("Tracks");
	195	Int_t ntr = fEsdEv->GetNumberOfTracks();
	196	for (Int_t i=0, ntrnew=0; i<ntr; ++i) {
	197	AliESDtrack *etrack = fEsdEv->GetTrack(i);
35fb503f CL	198	if (!etrack)
35fb503f CL	199	continue;
d47bb504	200
57184cc6	201	if (fEsdTrackCuts->AcceptTrack(etrack)) {
310dd8e8	202	if (fTrackEfficiency < 1) {
	203	Double_t r = gRandom->Rndm();
	204	if (fTrackEfficiency < r)
	205	continue;
	206	}
b3bef4aa	207	AliESDtrack newTrack = new ((fTracks)[ntrnew]) AliESDtrack(*etrack);
cdd4e0e2 CL	208	if (fDoPropagation)
cdd4e0e2 CL	209	AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(newTrack,fDist);
32c97abe	210	newTrack->SetBit(BIT(22),0);
32c97abe	211	newTrack->SetBit(BIT(23),0);
ef4d8e9a	212	if (!fIsMC) newTrack->SetLabel(0);
57184cc6	213	++ntrnew;
57184cc6	214	} else if (fHybridTrackCuts->AcceptTrack(etrack)) {
3f4c638b	215	if (!etrack->GetConstrainedParam())
3f4c638b	216	continue;
6231f841	217	UInt_t status = etrack->GetStatus();
3f4c638b	218	if (!fIncludeNoITS && ((status&AliESDtrack::kITSrefit)==0))
3f4c638b	219	continue;
c0440012	220
	221	if (fTrackEfficiency < 1) {
	222	Double_t r = gRandom->Rndm();
	223	if (fTrackEfficiency < r)
	224	continue;
	225	}
3f4c638b	226	AliESDtrack newTrack = new ((fTracks)[ntrnew]) AliESDtrack(*etrack);
	227	if (fDoPropagation)
	228	AliEMCALRecoUtils::ExtrapolateTrackToEMCalSurface(newTrack,fDist);
	229	const AliExternalTrackParam* constrainParam = etrack->GetConstrainedParam();
	230	newTrack->Set(constrainParam->GetX(),
	231	constrainParam->GetAlpha(),
	232	constrainParam->GetParameter(),
	233	constrainParam->GetCovariance());
	234	if ((status&AliESDtrack::kITSrefit)==0) {
	235	newTrack->SetBit(BIT(22),0); //type 2
	236	newTrack->SetBit(BIT(23),1);
	237	} else {
	238	newTrack->SetBit(BIT(22),1); //type 1
	239	newTrack->SetBit(BIT(23),0);
af2f024d	240	}
1737c13e	241	if (!fIsMC) newTrack->SetLabel(0);
3f4c638b	242	++ntrnew;
57184cc6	243	}
57184cc6	244	}
cc2b5fb5	245	}
cc2b5fb5	246	}