[u/mrichter/AliRoot.git] / PWGJE / EMCALJetTasks / AliJetConstituentTagCopier.cxx

// $Id: AliJetConstituentTagCopier.cxx  $
//
// Copy tags from particle level constituent to detector level
//
// Author: S. Aiola

#include "AliJetConstituentTagCopier.h"

#include <TClonesArray.h>
#include <TMath.h>
#include <TLorentzVector.h>

#include "AliNamedArrayI.h"
#include "AliVCluster.h"
#include "AliVParticle.h"
#include "AliEmcalParticle.h"
#include "AliLog.h"

ClassImp(AliJetConstituentTagCopier)

//________________________________________________________________________
AliJetConstituentTagCopier::AliJetConstituentTagCopier() : 
  AliAnalysisTaskEmcal("AliJetConstituentTagCopier", kFALSE),
  fMCParticlesName(),
  fCleanBeforeCopy(kFALSE),
  fMCLabelShift(0),
  fMCParticles(0),
  fMCParticlesMap(0)
{
  // Default constructor.
}

//________________________________________________________________________
AliJetConstituentTagCopier::AliJetConstituentTagCopier(const char *name) : 
  AliAnalysisTaskEmcal(name, kFALSE),
  fMCParticlesName("MCParticles"),
  fCleanBeforeCopy(kFALSE),
  fMCLabelShift(0),
  fMCParticles(0),
  fMCParticlesMap(0)
{
  // Standard constructor.
}

//________________________________________________________________________
AliJetConstituentTagCopier::~AliJetConstituentTagCopier()
{
  // Destructor
}

//________________________________________________________________________
void AliJetConstituentTagCopier::ExecOnce()
{
  // Execute once.

  if (!fMCParticles && !fMCParticlesName.IsNull()) {
    fMCParticles = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fMCParticlesName));
    if (!fMCParticles) {
      AliError(Form("%s: Could not retrieve MC particles %s!", GetName(), fMCParticlesName.Data()));
      return;
    }
    else if (!fMCParticles->GetClass()->GetBaseClass("AliVParticle")) {
      AliError(Form("%s: Collection %s does not contain AliVParticle objects!", GetName(), fMCParticlesName.Data())); 
      fMCParticles = 0;
      return;
    }
  }

  if (fMCParticles && !fMCParticlesMap) {
    fMCParticlesMap = dynamic_cast<AliNamedArrayI*>(InputEvent()->FindListObject(fMCParticlesName + "_Map"));
    // this is needed to map the MC labels with the indexes of the MC particle collection
    // if teh map is not given, the MC labels are assumed to be consistent with the indexes (which is not the case if AliEmcalMCTrackSelector is used)
    if (!fMCParticlesMap) {
      AliWarning(Form("%s: Could not retrieve map for MC particles %s! Will assume MC labels consistent with indexes...", GetName(), fMCParticlesName.Data())); 
      fMCParticlesMap = new AliNamedArrayI("tracksMap",9999);
      for (Int_t i = 0; i < 9999; i++) {
	fMCParticlesMap->AddAt(i,i);
      }
    }
  }

  AliAnalysisTaskEmcal::ExecOnce();
}

//________________________________________________________________________
Bool_t AliJetConstituentTagCopier::Run()
{
  if (fTracks) {
    if (fTracks->GetClass()->GetBaseClass("AliVParticle"))
      DoTrackLoop(fTracks);
    else if (fTracks->GetClass()->GetBaseClass("AliEmcalParticle"))
      DoEmcalParticleLoop(fTracks);
    else 
      AliError(Form("%s: Object type not recognized in collection %s. Nothing will be done.", GetName(), fTracks->GetName()));
  }

  if (fCaloClusters) {
    if (fCaloClusters->GetClass()->GetBaseClass("AliVCluster"))
      DoClusterLoop(fCaloClusters);
    else if (fCaloClusters->GetClass()->GetBaseClass("AliEmcalParticle"))
      DoEmcalParticleLoop(fCaloClusters);
    else 
      AliError(Form("%s: Object type not recognized in collection %s. Nothing will be done.", GetName(), fCaloClusters->GetName()));
  }

  return kTRUE;
}

//________________________________________________________________________
void AliJetConstituentTagCopier::DoClusterLoop(TClonesArray *array)
{
  if (fCleanBeforeCopy) {
    for (Int_t i = 0; i < array->GetEntries(); i++) {
      AliVCluster *cluster = static_cast<AliVCluster*>(array->At(i));
      if (!cluster) {
	AliError(Form("%s: Could not get cluster %d", GetName(), i));
	continue;
      }
      if (!AcceptCluster(cluster))
	continue;
      Int_t mcLabel = cluster->GetLabel();
      if (mcLabel > 0)
	cluster->SetBit(TObject::kBitMask, kFALSE);
    }
  }

  if (!fMCParticles)
    return;
  
  Double_t totalEnergy = 0;
  for (Int_t i = 0; i < array->GetEntries(); i++) {
    AliVCluster *cluster = static_cast<AliVCluster*>(array->At(i));
    if (!cluster) {
      AliError(Form("%s: Could not get cluster %d", GetName(), i));
      continue;
    }
    if (!AcceptCluster(cluster))
      continue;
    Int_t mcLabel = cluster->GetLabel();
    if (mcLabel > fMCLabelShift)
      mcLabel -= fMCLabelShift;
    if (mcLabel > 0) {
      TLorentzVector vect;
      cluster->GetMomentum(vect, fVertex);
      AliDebug(2, Form("Cluster %d, pt = %f, eta = %f, phi = %f, label = %d",
		       i, cluster->E(), vect.Eta(), vect.Phi(), mcLabel));
      totalEnergy += cluster->E();
      Int_t index = -1;
      if (mcLabel < fMCParticlesMap->GetSize())
	index = fMCParticlesMap->At(mcLabel);
      if (index < 0)
	continue;
      AliVParticle *part = static_cast<AliVParticle*>(fMCParticles->At(index));
      if (!part) {
	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	continue;
      }      
      AliDebug(2, Form("Matched with particle %d, pt = %f, eta = %f, phi = %f", 
		       index, part->E(), part->Eta(), part->Phi()));
      UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
      cluster->SetBit(bits);
    }
  }

  AliDebug(2, Form("Total energy of MC clusters = %f", 
		   totalEnergy));
}

//________________________________________________________________________
void AliJetConstituentTagCopier::DoTrackLoop(TClonesArray *array)
{
  if (fCleanBeforeCopy) {
    for (Int_t i = 0; i < array->GetEntries(); i++) {
      AliVParticle *track = static_cast<AliVParticle*>(array->At(i));
      if (!track) {
	AliError(Form("%s: Could not get track %d", GetName(), i));
	continue;
      }
      if (!AcceptTrack(track))
      continue;
      Int_t mcLabel = TMath::Abs(track->GetLabel());
      if (mcLabel != 0) 
	track->SetBit(TObject::kBitMask, kFALSE);
    }
  }

  if (!fMCParticles)
    return;

  for (Int_t i = 0; i < array->GetEntries(); i++) {
    AliVParticle *track = static_cast<AliVParticle*>(array->At(i));
    if (!track) {
      AliError(Form("%s: Could not get track %d", GetName(), i));
      continue;
    }
    if (!AcceptTrack(track))
      continue;
    Int_t mcLabel = TMath::Abs(track->GetLabel());
    if (mcLabel > fMCLabelShift)
      mcLabel -= fMCLabelShift;
    if (mcLabel > 0) {
      Int_t index = -1;
      if (mcLabel < fMCParticlesMap->GetSize())
	index = fMCParticlesMap->At(mcLabel);
      if (index < 0)
	continue;
      AliVParticle *part = static_cast<AliVParticle*>(fMCParticles->At(index));
      if (!part) {
	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	continue;
      }
      AliDebug(3, Form("Track %d, pt = %f, eta = %f, phi = %f, label = %d is matched with particle %d, pt = %f, eta = %f, phi = %f", 
		       i, track->Pt(), track->Eta(), track->Phi(), mcLabel, index, part->Pt(), part->Eta(), part->Phi()));
      UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
      track->SetBit(bits);
    }
  }
}

//________________________________________________________________________
void AliJetConstituentTagCopier::DoEmcalParticleLoop(TClonesArray *array)
{
  if (fCleanBeforeCopy) {
    for (Int_t i = 0; i < array->GetEntries(); i++) {
      AliEmcalParticle *emcpart = static_cast<AliEmcalParticle*>(array->At(i));
      if (!emcpart) {
	AliError(Form("%s: Could not get EmcalParticle %d", GetName(), i));
	continue;
      }
      if (!AcceptEmcalPart(emcpart))
	continue;
      AliVCluster *cluster = emcpart->GetCluster();
      AliVParticle *track = emcpart->GetTrack();
      Int_t mcLabel = 0;
      if (cluster)
	mcLabel = cluster->GetLabel();
      else if (track)
	mcLabel = TMath::Abs(track->GetLabel());
      if (mcLabel != 0) 
	emcpart->SetBit(TObject::kBitMask, kFALSE);
    }
  }

  if (!fMCParticles)
    return;

  for (Int_t i = 0; i < array->GetEntries(); i++) {
    AliEmcalParticle *emcpart = static_cast<AliEmcalParticle*>(array->At(i));
    if (!emcpart) {
      AliError(Form("%s: Could not get EmcalParticle %d", GetName(), i));
      continue;
    }
    if (!AcceptEmcalPart(emcpart))
      continue;
    AliVCluster *cluster = emcpart->GetCluster();
    AliVParticle *track = emcpart->GetTrack();
    Int_t mcLabel = 0;
    if (cluster)
      mcLabel = cluster->GetLabel();
    else if (track)
      mcLabel = TMath::Abs(track->GetLabel());
    if (mcLabel > fMCLabelShift)
      mcLabel -= fMCLabelShift;
    if (mcLabel > 0) {
      Int_t index = -1;
      if (mcLabel < fMCParticlesMap->GetSize())
	index = fMCParticlesMap->At(mcLabel);
      if (index < 0)
	continue;
      AliVParticle *part = static_cast<AliVParticle*>(fMCParticles->At(index));
      if (!part) {
	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	continue;
      }
      UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
      emcpart->SetBit(bits);
    }
  }
}
Commit	Line	Data
8afc7c8d	1	// $Id: AliJetConstituentTagCopier.cxx $
	2	//
	3	// Copy tags from particle level constituent to detector level
	4	//
	5	// Author: S. Aiola
	6
	7	#include "AliJetConstituentTagCopier.h"
	8
	9	#include <TClonesArray.h>
5ce8ae64	10	#include <TMath.h>
4358e58a	11	#include <TLorentzVector.h>
8afc7c8d	12
5be3857d	13	#include "AliNamedArrayI.h"
8afc7c8d	14	#include "AliVCluster.h"
	15	#include "AliVParticle.h"
	16	#include "AliEmcalParticle.h"
	17	#include "AliLog.h"
	18
	19	ClassImp(AliJetConstituentTagCopier)
	20
	21	//________________________________________________________________________
	22	AliJetConstituentTagCopier::AliJetConstituentTagCopier() :
	23	AliAnalysisTaskEmcal("AliJetConstituentTagCopier", kFALSE),
	24	fMCParticlesName(),
f0c10df2	25	fCleanBeforeCopy(kFALSE),
a7477843	26	fMCLabelShift(0),
8afc7c8d	27	fMCParticles(0),
	28	fMCParticlesMap(0)
	29	{
	30	// Default constructor.
	31	}
	32
	33	//________________________________________________________________________
	34	AliJetConstituentTagCopier::AliJetConstituentTagCopier(const char *name) :
	35	AliAnalysisTaskEmcal(name, kFALSE),
	36	fMCParticlesName("MCParticles"),
f0c10df2	37	fCleanBeforeCopy(kFALSE),
a7477843	38	fMCLabelShift(0),
8afc7c8d	39	fMCParticles(0),
	40	fMCParticlesMap(0)
	41	{
	42	// Standard constructor.
	43	}
	44
	45	//________________________________________________________________________
	46	AliJetConstituentTagCopier::~AliJetConstituentTagCopier()
	47	{
	48	// Destructor
	49	}
	50
	51	//________________________________________________________________________
	52	void AliJetConstituentTagCopier::ExecOnce()
	53	{
	54	// Execute once.
	55
f0c10df2	56	if (!fMCParticles && !fMCParticlesName.IsNull()) {
8afc7c8d	57	fMCParticles = dynamic_cast<TClonesArray*>(InputEvent()->FindListObject(fMCParticlesName));
	58	if (!fMCParticles) {
	59	AliError(Form("%s: Could not retrieve MC particles %s!", GetName(), fMCParticlesName.Data()));
	60	return;
	61	}
	62	else if (!fMCParticles->GetClass()->GetBaseClass("AliVParticle")) {
	63	AliError(Form("%s: Collection %s does not contain AliVParticle objects!", GetName(), fMCParticlesName.Data()));
	64	fMCParticles = 0;
	65	return;
	66	}
	67	}
	68
f0c10df2	69	if (fMCParticles && !fMCParticlesMap) {
5be3857d	70	fMCParticlesMap = dynamic_cast<AliNamedArrayI*>(InputEvent()->FindListObject(fMCParticlesName + "_Map"));
8afc7c8d	71	// this is needed to map the MC labels with the indexes of the MC particle collection
4358e58a	72	// if teh map is not given, the MC labels are assumed to be consistent with the indexes (which is not the case if AliEmcalMCTrackSelector is used)
8afc7c8d	73	if (!fMCParticlesMap) {
8afc7c8d	74	AliWarning(Form("%s: Could not retrieve map for MC particles %s! Will assume MC labels consistent with indexes...", GetName(), fMCParticlesName.Data()));
5be3857d	75	fMCParticlesMap = new AliNamedArrayI("tracksMap",9999);
8afc7c8d	76	for (Int_t i = 0; i < 9999; i++) {
5be3857d	77	fMCParticlesMap->AddAt(i,i);
8afc7c8d	78	}
	79	}
	80	}
	81
	82	AliAnalysisTaskEmcal::ExecOnce();
	83	}
	84
	85	//________________________________________________________________________
	86	Bool_t AliJetConstituentTagCopier::Run()
	87	{
	88	if (fTracks) {
	89	if (fTracks->GetClass()->GetBaseClass("AliVParticle"))
	90	DoTrackLoop(fTracks);
	91	else if (fTracks->GetClass()->GetBaseClass("AliEmcalParticle"))
	92	DoEmcalParticleLoop(fTracks);
	93	else
	94	AliError(Form("%s: Object type not recognized in collection %s. Nothing will be done.", GetName(), fTracks->GetName()));
	95	}
	96
	97	if (fCaloClusters) {
	98	if (fCaloClusters->GetClass()->GetBaseClass("AliVCluster"))
	99	DoClusterLoop(fCaloClusters);
	100	else if (fCaloClusters->GetClass()->GetBaseClass("AliEmcalParticle"))
	101	DoEmcalParticleLoop(fCaloClusters);
	102	else
	103	AliError(Form("%s: Object type not recognized in collection %s. Nothing will be done.", GetName(), fCaloClusters->GetName()));
	104	}
	105
	106	return kTRUE;
	107	}
	108
	109	//________________________________________________________________________
	110	void AliJetConstituentTagCopier::DoClusterLoop(TClonesArray *array)
	111	{
f0c10df2	112	if (fCleanBeforeCopy) {
	113	for (Int_t i = 0; i < array->GetEntries(); i++) {
	114	AliVCluster cluster = static_cast<AliVCluster>(array->At(i));
	115	if (!cluster) {
	116	AliError(Form("%s: Could not get cluster %d", GetName(), i));
	117	continue;
	118	}
	119	if (!AcceptCluster(cluster))
	120	continue;
	121	Int_t mcLabel = cluster->GetLabel();
	122	if (mcLabel > 0)
	123	cluster->SetBit(TObject::kBitMask, kFALSE);
	124	}
	125	}
	126
	127	if (!fMCParticles)
	128	return;
	129
4358e58a	130	Double_t totalEnergy = 0;
8afc7c8d	131	for (Int_t i = 0; i < array->GetEntries(); i++) {
	132	AliVCluster cluster = static_cast<AliVCluster>(array->At(i));
	133	if (!cluster) {
	134	AliError(Form("%s: Could not get cluster %d", GetName(), i));
	135	continue;
	136	}
	137	if (!AcceptCluster(cluster))
	138	continue;
	139	Int_t mcLabel = cluster->GetLabel();
a7477843	140	if (mcLabel > fMCLabelShift)
a7477843	141	mcLabel -= fMCLabelShift;
8afc7c8d	142	if (mcLabel > 0) {
4358e58a	143	TLorentzVector vect;
	144	cluster->GetMomentum(vect, fVertex);
	145	AliDebug(2, Form("Cluster %d, pt = %f, eta = %f, phi = %f, label = %d",
	146	i, cluster->E(), vect.Eta(), vect.Phi(), mcLabel));
	147	totalEnergy += cluster->E();
5be3857d	148	Int_t index = -1;
	149	if (mcLabel < fMCParticlesMap->GetSize())
	150	index = fMCParticlesMap->At(mcLabel);
5ce8ae64	151	if (index < 0)
5ce8ae64	152	continue;
8afc7c8d	153	AliVParticle part = static_cast<AliVParticle>(fMCParticles->At(index));
	154	if (!part) {
	155	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	156	continue;
4358e58a	157	}
	158	AliDebug(2, Form("Matched with particle %d, pt = %f, eta = %f, phi = %f",
	159	index, part->E(), part->Eta(), part->Phi()));
8afc7c8d	160	UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
	161	cluster->SetBit(bits);
	162	}
	163	}
4358e58a	164
	165	AliDebug(2, Form("Total energy of MC clusters = %f",
	166	totalEnergy));
8afc7c8d	167	}
	168
	169	//________________________________________________________________________
	170	void AliJetConstituentTagCopier::DoTrackLoop(TClonesArray *array)
	171	{
f0c10df2	172	if (fCleanBeforeCopy) {
	173	for (Int_t i = 0; i < array->GetEntries(); i++) {
	174	AliVParticle track = static_cast<AliVParticle>(array->At(i));
	175	if (!track) {
	176	AliError(Form("%s: Could not get track %d", GetName(), i));
	177	continue;
	178	}
	179	if (!AcceptTrack(track))
	180	continue;
	181	Int_t mcLabel = TMath::Abs(track->GetLabel());
	182	if (mcLabel != 0)
	183	track->SetBit(TObject::kBitMask, kFALSE);
	184	}
	185	}
	186
	187	if (!fMCParticles)
	188	return;
	189
8afc7c8d	190	for (Int_t i = 0; i < array->GetEntries(); i++) {
	191	AliVParticle track = static_cast<AliVParticle>(array->At(i));
	192	if (!track) {
	193	AliError(Form("%s: Could not get track %d", GetName(), i));
	194	continue;
	195	}
	196	if (!AcceptTrack(track))
	197	continue;
5ce8ae64	198	Int_t mcLabel = TMath::Abs(track->GetLabel());
a7477843	199	if (mcLabel > fMCLabelShift)
	200	mcLabel -= fMCLabelShift;
	201	if (mcLabel > 0) {
5be3857d	202	Int_t index = -1;
	203	if (mcLabel < fMCParticlesMap->GetSize())
	204	index = fMCParticlesMap->At(mcLabel);
5ce8ae64	205	if (index < 0)
5ce8ae64	206	continue;
8afc7c8d	207	AliVParticle part = static_cast<AliVParticle>(fMCParticles->At(index));
	208	if (!part) {
	209	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	210	continue;
	211	}
4358e58a	212	AliDebug(3, Form("Track %d, pt = %f, eta = %f, phi = %f, label = %d is matched with particle %d, pt = %f, eta = %f, phi = %f",
5ce8ae64	213	i, track->Pt(), track->Eta(), track->Phi(), mcLabel, index, part->Pt(), part->Eta(), part->Phi()));
8afc7c8d	214	UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
	215	track->SetBit(bits);
	216	}
	217	}
	218	}
	219
	220	//________________________________________________________________________
	221	void AliJetConstituentTagCopier::DoEmcalParticleLoop(TClonesArray *array)
	222	{
f0c10df2	223	if (fCleanBeforeCopy) {
	224	for (Int_t i = 0; i < array->GetEntries(); i++) {
	225	AliEmcalParticle emcpart = static_cast<AliEmcalParticle>(array->At(i));
	226	if (!emcpart) {
	227	AliError(Form("%s: Could not get EmcalParticle %d", GetName(), i));
	228	continue;
	229	}
	230	if (!AcceptEmcalPart(emcpart))
	231	continue;
	232	AliVCluster *cluster = emcpart->GetCluster();
	233	AliVParticle *track = emcpart->GetTrack();
	234	Int_t mcLabel = 0;
	235	if (cluster)
	236	mcLabel = cluster->GetLabel();
	237	else if (track)
	238	mcLabel = TMath::Abs(track->GetLabel());
	239	if (mcLabel != 0)
	240	emcpart->SetBit(TObject::kBitMask, kFALSE);
	241	}
	242	}
	243
	244	if (!fMCParticles)
	245	return;
	246
8afc7c8d	247	for (Int_t i = 0; i < array->GetEntries(); i++) {
	248	AliEmcalParticle emcpart = static_cast<AliEmcalParticle>(array->At(i));
	249	if (!emcpart) {
	250	AliError(Form("%s: Could not get EmcalParticle %d", GetName(), i));
	251	continue;
	252	}
	253	if (!AcceptEmcalPart(emcpart))
	254	continue;
	255	AliVCluster *cluster = emcpart->GetCluster();
	256	AliVParticle *track = emcpart->GetTrack();
	257	Int_t mcLabel = 0;
	258	if (cluster)
	259	mcLabel = cluster->GetLabel();
	260	else if (track)
5ce8ae64	261	mcLabel = TMath::Abs(track->GetLabel());
a7477843	262	if (mcLabel > fMCLabelShift)
	263	mcLabel -= fMCLabelShift;
	264	if (mcLabel > 0) {
5be3857d	265	Int_t index = -1;
	266	if (mcLabel < fMCParticlesMap->GetSize())
	267	index = fMCParticlesMap->At(mcLabel);
5ce8ae64	268	if (index < 0)
5ce8ae64	269	continue;
8afc7c8d	270	AliVParticle part = static_cast<AliVParticle>(fMCParticles->At(index));
	271	if (!part) {
	272	AliError(Form("%s: Could not get MC particle %d", GetName(), index));
	273	continue;
	274	}
	275	UInt_t bits = (UInt_t)part->TestBits(TObject::kBitMask);
	276	emcpart->SetBit(bits);
	277	}
	278	}
	279	}