Skip to content

Commit f9d02c6

Browse files
[PWGJE] isolation fixes & MC prompt photon information fixes (#16308)
Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
1 parent 6700ceb commit f9d02c6

1 file changed

Lines changed: 177 additions & 37 deletions

File tree

PWGJE/Tasks/gammaJetTreeProducer.cxx

Lines changed: 177 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
#include <string>
5656
#include <type_traits>
5757
#include <unordered_map>
58+
#include <unordered_set>
5859
#include <utility>
5960
#include <vector>
6061

@@ -120,9 +121,11 @@ struct GammaJetTreeProducer {
120121
std::vector<float> trackEta;
121122
std::vector<float> trackPhi;
122123
std::vector<float> trackPt;
124+
std::vector<int> trackSourceIndex;
123125
std::vector<float> mcParticleEta;
124126
std::vector<float> mcParticlePhi;
125127
std::vector<float> mcParticlePt;
128+
std::vector<int> mcParticleSourceIndex;
126129
TKDTree<int, float>* trackTree = nullptr;
127130
TKDTree<int, float>* mcParticleTree = nullptr;
128131

@@ -265,33 +268,57 @@ struct GammaJetTreeProducer {
265268
curMcParticleTreeColIndex = collision.globalIndex();
266269
}
267270

268-
trackEta.clear();
269-
trackPhi.clear();
270-
trackPt.clear();
271-
mcParticleEta.clear();
272-
mcParticlePhi.clear();
273-
mcParticlePt.clear();
274-
275271
// if the track type is aod::JetTracks, we need to build the kd tree for the tracks
276272
if constexpr (std::is_same_v<typename std::decay_t<T>, aod::JetTracks>) {
273+
trackEta.clear();
274+
trackPhi.clear();
275+
trackPt.clear();
276+
trackSourceIndex.clear();
277+
const float maxMatchingDistance = std::max(static_cast<float>(isoR), static_cast<float>(perpConeJetR));
278+
const float additionalMargin = 0.05f;
277279
for (const auto& track : objects) {
278280
if (!isTrackSelected(track)) {
279281
continue;
280282
}
283+
const float phi = RecoDecay::constrainAngle(track.phi(), 0.0);
284+
const int originalIndex = static_cast<int>(trackPt.size());
281285
trackEta.push_back(track.eta());
282-
trackPhi.push_back(track.phi());
286+
trackPhi.push_back(phi);
283287
trackPt.push_back(track.pt());
288+
trackSourceIndex.push_back(originalIndex);
289+
if (phi <= (maxMatchingDistance + additionalMargin)) {
290+
trackEta.push_back(track.eta());
291+
trackPhi.push_back(phi + o2::constants::math::TwoPI);
292+
trackPt.push_back(track.pt());
293+
trackSourceIndex.push_back(originalIndex);
294+
}
295+
if (phi >= (o2::constants::math::TwoPI - (maxMatchingDistance + additionalMargin))) {
296+
trackEta.push_back(track.eta());
297+
trackPhi.push_back(phi - o2::constants::math::TwoPI);
298+
trackPt.push_back(track.pt());
299+
trackSourceIndex.push_back(originalIndex);
300+
}
284301
}
285302
if (trackEta.size() > 0) {
286303
delete trackTree;
287304
trackTree = new TKDTree<int, float>(trackEta.size(), 2, 1);
288305
trackTree->SetData(0, trackEta.data());
289306
trackTree->SetData(1, trackPhi.data());
290307
trackTree->Build();
308+
} else {
309+
// Keep tree state consistent with the current event content.
310+
delete trackTree;
311+
trackTree = nullptr;
291312
}
292313
}
293314
// if the track type is aod::JetParticles, we need to build the kd tree for the mc particles
294315
if constexpr (std::is_same_v<typename std::decay_t<T>, aod::JetParticles>) {
316+
mcParticleEta.clear();
317+
mcParticlePhi.clear();
318+
mcParticlePt.clear();
319+
mcParticleSourceIndex.clear();
320+
const float maxMatchingDistance = std::max(static_cast<float>(isoR), static_cast<float>(perpConeJetR));
321+
const float additionalMargin = 0.05f;
295322
for (const auto& particle : objects) {
296323
if (!particle.isPhysicalPrimary()) {
297324
continue;
@@ -302,16 +329,35 @@ struct GammaJetTreeProducer {
302329
if (particle.pt() < trackMinPt) {
303330
continue;
304331
}
332+
const float phi = RecoDecay::constrainAngle(particle.phi(), 0.0);
333+
const int originalIndex = static_cast<int>(mcParticlePt.size());
305334
mcParticleEta.push_back(particle.eta());
306-
mcParticlePhi.push_back(particle.phi());
335+
mcParticlePhi.push_back(phi);
307336
mcParticlePt.push_back(particle.pt());
337+
mcParticleSourceIndex.push_back(originalIndex);
338+
if (phi <= (maxMatchingDistance + additionalMargin)) {
339+
mcParticleEta.push_back(particle.eta());
340+
mcParticlePhi.push_back(phi + o2::constants::math::TwoPI);
341+
mcParticlePt.push_back(particle.pt());
342+
mcParticleSourceIndex.push_back(originalIndex);
343+
}
344+
if (phi >= (o2::constants::math::TwoPI - (maxMatchingDistance + additionalMargin))) {
345+
mcParticleEta.push_back(particle.eta());
346+
mcParticlePhi.push_back(phi - o2::constants::math::TwoPI);
347+
mcParticlePt.push_back(particle.pt());
348+
mcParticleSourceIndex.push_back(originalIndex);
349+
}
308350
}
309351
if (mcParticleEta.size() > 0) {
310352
delete mcParticleTree;
311353
mcParticleTree = new TKDTree<int, float>(mcParticleEta.size(), 2, 1);
312354
mcParticleTree->SetData(0, mcParticleEta.data());
313355
mcParticleTree->SetData(1, mcParticlePhi.data());
314356
mcParticleTree->Build();
357+
} else {
358+
// Keep tree state consistent with the current event content.
359+
delete mcParticleTree;
360+
mcParticleTree = nullptr;
315361
}
316362
}
317363
}
@@ -350,7 +396,7 @@ struct GammaJetTreeProducer {
350396
{
351397
mHistograms.fill(HIST("eventQA"), 0);
352398

353-
if (collision.posZ() > mVertexCut) {
399+
if (std::abs(collision.posZ()) > mVertexCut) {
354400
return false;
355401
}
356402
mHistograms.fill(HIST("eventQA"), 1);
@@ -385,6 +431,33 @@ struct GammaJetTreeProducer {
385431
return std::abs(pdg->GetParticle(particle.pdgCode())->Charge()) >= 1.;
386432
}
387433

434+
/// \brief Sums pt values once per unique source index from KD-tree query indices
435+
/// \param indices KD-tree result indices
436+
/// \param sourceIndexMap Maps KD-tree entry index to original object index
437+
/// \param ptValues pt values of original objects
438+
/// \param uniqueSourceIndices Set used to deduplicate source indices
439+
/// \return Sum of unique pt values for the given indices
440+
double sumUniquePtFromIndices(const std::vector<int>& indices,
441+
const std::vector<int>& sourceIndexMap,
442+
const std::vector<float>& ptValues,
443+
std::unordered_set<int>& uniqueSourceIndices)
444+
{
445+
double ptSum = 0;
446+
for (const auto& index : indices) {
447+
if (index < 0 || static_cast<size_t>(index) >= sourceIndexMap.size()) {
448+
continue;
449+
}
450+
const int sourceIndex = sourceIndexMap[index];
451+
if (sourceIndex < 0 || static_cast<size_t>(sourceIndex) >= ptValues.size()) {
452+
continue;
453+
}
454+
if (uniqueSourceIndices.insert(sourceIndex).second) {
455+
ptSum += ptValues[sourceIndex];
456+
}
457+
}
458+
return ptSum;
459+
}
460+
388461
/// \brief Calculates the charged particle isolation in a cone of given size using a pre-built kd tree
389462
/// \param particle The particle to calculate the isolation for
390463
/// \param radius The cone radius
@@ -394,25 +467,22 @@ struct GammaJetTreeProducer {
394467
double ch_iso_in_cone(const T& particle, float radius = 0.4, bool mcGenIso = false)
395468
{
396469
double iso = 0;
397-
float point[2] = {particle.eta(), particle.phi()};
470+
float point[2] = {particle.eta(), RecoDecay::constrainAngle(particle.phi(), 0.0)};
398471
std::vector<int> indices;
472+
std::unordered_set<int> uniqueSourceIndices;
399473

400474
if (!mcGenIso) {
401475
if (trackTree) {
402476
trackTree->FindInRange(point, radius, indices);
403-
for (const auto& index : indices) {
404-
iso += trackPt[index];
405-
}
477+
iso += sumUniquePtFromIndices(indices, trackSourceIndex, trackPt, uniqueSourceIndices);
406478
} else {
407479
LOG(error) << "Track tree not found";
408480
return 0;
409481
}
410482
} else {
411483
if (mcParticleTree) {
412484
mcParticleTree->FindInRange(point, radius, indices);
413-
for (const auto& index : indices) {
414-
iso += mcParticlePt[index];
415-
}
485+
iso += sumUniquePtFromIndices(indices, mcParticleSourceIndex, mcParticlePt, uniqueSourceIndices);
416486
} else {
417487
LOG(error) << "MC particle tree not found";
418488
return 0;
@@ -435,14 +505,16 @@ struct GammaJetTreeProducer {
435505
double cPhi = TVector2::Phi_0_2pi(object.phi());
436506

437507
// rotate cone left by 90 degrees
438-
float cPhiLeft = cPhi - o2::constants::math::PIHalf;
439-
float cPhiRight = cPhi + o2::constants::math::PIHalf;
508+
float cPhiLeft = RecoDecay::constrainAngle(cPhi - o2::constants::math::PIHalf, 0.0);
509+
float cPhiRight = RecoDecay::constrainAngle(cPhi + o2::constants::math::PIHalf, 0.0);
440510

441511
float pointLeft[2] = {object.eta(), cPhiLeft};
442512
float pointRight[2] = {object.eta(), cPhiRight};
443513

444514
std::vector<int> indicesLeft;
445515
std::vector<int> indicesRight;
516+
std::unordered_set<int> uniqueSourceIndicesLeft;
517+
std::unordered_set<int> uniqueSourceIndicesRight;
446518

447519
if (!mcGenIso) {
448520
if (trackTree) {
@@ -453,12 +525,8 @@ struct GammaJetTreeProducer {
453525
return 0;
454526
}
455527

456-
for (const auto& index : indicesLeft) {
457-
ptSumLeft += trackPt[index];
458-
}
459-
for (const auto& index : indicesRight) {
460-
ptSumRight += trackPt[index];
461-
}
528+
ptSumLeft += sumUniquePtFromIndices(indicesLeft, trackSourceIndex, trackPt, uniqueSourceIndicesLeft);
529+
ptSumRight += sumUniquePtFromIndices(indicesRight, trackSourceIndex, trackPt, uniqueSourceIndicesRight);
462530
} else {
463531
if (mcParticleTree) {
464532
mcParticleTree->FindInRange(pointLeft, radius, indicesLeft);
@@ -467,12 +535,8 @@ struct GammaJetTreeProducer {
467535
LOG(error) << "MC particle tree not found";
468536
return 0;
469537
}
470-
for (const auto& index : indicesLeft) {
471-
ptSumLeft += mcParticlePt[index];
472-
}
473-
for (const auto& index : indicesRight) {
474-
ptSumRight += mcParticlePt[index];
475-
}
538+
ptSumLeft += sumUniquePtFromIndices(indicesLeft, mcParticleSourceIndex, mcParticlePt, uniqueSourceIndicesLeft);
539+
ptSumRight += sumUniquePtFromIndices(indicesRight, mcParticleSourceIndex, mcParticlePt, uniqueSourceIndicesRight);
476540
}
477541

478542
float rho = (ptSumLeft + ptSumRight) / (o2::constants::math::TwoPI * radius * radius);
@@ -494,6 +558,46 @@ struct GammaJetTreeProducer {
494558
}
495559
}
496560

561+
/// \brief Prints the mother chain of a given MC particle.
562+
/// \param particle The particle to print the mother chain for
563+
template <typename T>
564+
void printMotherChain(const T& particle) const
565+
{
566+
T current = particle;
567+
int depth = 0;
568+
bool continueTracing = true;
569+
while (continueTracing) {
570+
LOG(info) << "Level " << depth
571+
<< " | PDG: " << current.pdgCode()
572+
<< " | E: " << current.energy()
573+
<< " | pt: " << current.pt()
574+
<< " | eta: " << current.eta()
575+
<< " | phi: " << current.phi()
576+
<< " | genStatusCode: " << current.getGenStatusCode()
577+
<< " | globalIndex: " << current.globalIndex();
578+
auto mothers = current.template mothers_as<aod::JMcParticles>();
579+
if (mothers.size() == 0) {
580+
break;
581+
}
582+
// Handle case with potentially duplicate mothers
583+
int selectedMother = -1;
584+
if (mothers.size() == 1) {
585+
selectedMother = 0;
586+
} else if (mothers.size() == 2 &&
587+
mothers[0].globalIndex() == mothers[1].globalIndex() &&
588+
mothers[0].pdgCode() == mothers[1].pdgCode() &&
589+
mothers[0].getGenStatusCode() == mothers[1].getGenStatusCode()) {
590+
selectedMother = 0;
591+
}
592+
if (selectedMother == -1) {
593+
LOG(info) << "Branching in mother chain or unclear mothers, stopping trace.";
594+
break;
595+
}
596+
current = mothers[selectedMother];
597+
++depth;
598+
}
599+
}
600+
497601
/// \brief Finds the top-most copy of a particle in the decay chain (following carbon copies)
498602
/// \param particle The particle to start from
499603
/// \return The top-most copy of the particle
@@ -504,10 +608,36 @@ struct GammaJetTreeProducer {
504608
T currentParticle = particle;
505609
int pdgCode = particle.pdgCode();
506610
auto mothers = particle.template mothers_as<aod::JMcParticles>();
507-
while (iUp > 0 && mothers.size() == 1 && mothers[0].globalIndex() > 0 && mothers[0].pdgCode() == pdgCode) {
508-
iUp = mothers[0].globalIndex();
509-
currentParticle = mothers[0];
611+
// sometimes the particle will have two identical mothers, not sure why, but we need to cover this case
612+
// if there are two mothers, but they are not identical, we will stop
613+
bool twoMothersIdentical = false;
614+
int selectedMother = -1;
615+
if (mothers.size() == 1) {
616+
selectedMother = 0;
617+
} else if (mothers.size() == 2) {
618+
twoMothersIdentical = (mothers[0].globalIndex() == mothers[1].globalIndex() &&
619+
mothers[0].pdgCode() == mothers[1].pdgCode() &&
620+
mothers[0].getGenStatusCode() == mothers[1].getGenStatusCode());
621+
if (twoMothersIdentical) {
622+
selectedMother = 0;
623+
}
624+
}
625+
while (iUp > 0 && selectedMother >= 0 && mothers[selectedMother].globalIndex() > 0 && mothers[selectedMother].pdgCode() == pdgCode) {
626+
iUp = mothers[selectedMother].globalIndex();
627+
currentParticle = mothers[selectedMother];
510628
mothers = currentParticle.template mothers_as<aod::JMcParticles>();
629+
selectedMother = -1;
630+
twoMothersIdentical = false;
631+
if (mothers.size() == 1) {
632+
selectedMother = 0;
633+
} else if (mothers.size() == 2) {
634+
twoMothersIdentical = (mothers[0].globalIndex() == mothers[1].globalIndex() &&
635+
mothers[0].pdgCode() == mothers[1].pdgCode() &&
636+
mothers[0].getGenStatusCode() == mothers[1].getGenStatusCode());
637+
if (twoMothersIdentical) {
638+
selectedMother = 0;
639+
}
640+
}
511641
}
512642
return currentParticle;
513643
}
@@ -889,7 +1019,7 @@ struct GammaJetTreeProducer {
8891019
// check that mother has exactly two daughters which are e+ and e-
8901020
if (daughters.size() == 2) {
8911021
LOG(debug) << "Got the daughters";
892-
if ((daughters.iteratorAt(0).pdgCode() == PDG_t::kElectron && daughters.iteratorAt(1).pdgCode() == PDG_t::kPositron) || (daughters.iteratorAt(0).pdgCode() == -PDG_t::kPositron && daughters.iteratorAt(1).pdgCode() == PDG_t::kElectron)) {
1022+
if ((daughters.iteratorAt(0).pdgCode() == PDG_t::kElectron && daughters.iteratorAt(1).pdgCode() == PDG_t::kPositron) || (daughters.iteratorAt(0).pdgCode() == PDG_t::kPositron && daughters.iteratorAt(1).pdgCode() == PDG_t::kElectron)) {
8931023
SETBIT(origin, static_cast<uint16_t>(gjanalysis::ClusterOrigin::kConvertedPhoton));
8941024
LOG(debug) << "Cluster is a converted photon";
8951025
}
@@ -932,7 +1062,7 @@ struct GammaJetTreeProducer {
9321062
int nRecCollisions = 0;
9331063
mHistograms.fill(HIST("mcCollisionsWithRecCollisions"), 0);
9341064
for (auto const& collision : collisions) {
935-
if (collision.posZ() > mVertexCut) {
1065+
if (std::abs(collision.posZ()) > mVertexCut) {
9361066
continue;
9371067
}
9381068
if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits, true, true, rctLabel)) {
@@ -1203,7 +1333,7 @@ struct GammaJetTreeProducer {
12031333
if (jet.pt() < jetPtMin) {
12041334
return;
12051335
}
1206-
for (auto& jetConstituent : jet.template tracks_as<U>()) {
1336+
for (const auto& jetConstituent : jet.template tracks_as<U>()) {
12071337
fastjetutilities::fillTracks(jetConstituent, jetConstituents, jetConstituent.globalIndex());
12081338
}
12091339

@@ -1305,11 +1435,21 @@ struct GammaJetTreeProducer {
13051435
if (particle.pdgCode() != PDG_t::kPi0 && particle.pdgCode() != PDG_t::kGamma) {
13061436
continue;
13071437
}
1438+
13081439
// check the origin of the particle
13091440
uint16_t origin = getMCParticleOrigin(particle);
13101441
double mcIsolation = ch_iso_in_cone(particle, isoR, true);
13111442
mcParticlesTable(storedColIndex, particle.energy(), particle.eta(), particle.phi(), particle.pt(), particle.pdgCode(), mcIsolation, origin);
13121443

1444+
// DEBUGGING for photons. If it is a photon, print the origin and then print the chain
1445+
// leaving this in here
1446+
// if (particle.pdgCode() == PDG_t::kGamma) {
1447+
// LOG(info) << "Particle PDG: " << particle.pdgCode() << " isPhysicalPrimary: " << particle.isPhysicalPrimary() << " getGenStatusCode: " << particle.getGenStatusCode();
1448+
// LOG(info) << "Origin: " << "kPromptPhoton (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kPromptPhoton))) << "), kDirectPromptPhoton (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kDirectPromptPhoton))) << "), kFragmentationPhoton (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kFragmentationPhoton))) << "), kDecayPhoton (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kDecayPhoton))) << "), kDecayPhotonPi0 (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kDecayPhotonPi0))) << "), kDecayPhotonEta (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kDecayPhotonEta))) << "), kDecayPhotonOther (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kDecayPhotonOther))) << "), kPi0 (" << (origin & (1 << static_cast<uint16_t>(gjanalysis::ParticleOrigin::kPi0))) << ")";
1449+
// LOG(info) << "Mother chain: ";
1450+
// printMotherChain(particle);
1451+
// }
1452+
13131453
// fill mc gen trigger particle histograms
13141454
mHistograms.fill(HIST("mcGenTrigger_E"), particle.energy());
13151455
mHistograms.fill(HIST("mcGenTrigger_Eta"), particle.eta());

0 commit comments

Comments
 (0)