Skip to content

Commit 935d875

Browse files
author
Aleksandra Mulewicz
committed
PWGJE: add truth-level jet tables and reco-truth matching
1 parent 72def59 commit 935d875

1 file changed

Lines changed: 219 additions & 1 deletion

File tree

‎PWGJE/TableProducer/quarkGluonJetsProducer.cxx‎

Lines changed: 219 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
#include <Framework/InitContext.h>
4040
#include <Framework/runDataProcessing.h>
4141

42+
#include <TDatabasePDG.h>
4243
#include <TPDGCode.h>
4344

4445
#include <Rtypes.h>
@@ -47,6 +48,8 @@
4748
#include <cstdint>
4849
#include <limits>
4950
#include <string>
51+
#include <tuple>
52+
#include <unordered_map>
5053
#include <vector>
5154

5255
namespace o2::aod
@@ -117,6 +120,7 @@ DECLARE_SOA_COLUMN(TofNSigmaPi, tofNSigmaPi, float);
117120
DECLARE_SOA_COLUMN(TofNSigmaKa, tofNSigmaKa, float);
118121
DECLARE_SOA_COLUMN(TofNSigmaPr, tofNSigmaPr, float);
119122
DECLARE_SOA_COLUMN(TofBeta, tofBeta, float);
123+
DECLARE_SOA_COLUMN(HasTruthMatch, hasTruthMatch, bool);
120124
DECLARE_SOA_COLUMN(TruthPdg, truthPdg, int32_t);
121125
DECLARE_SOA_COLUMN(TruthPt, truthPt, float);
122126
DECLARE_SOA_COLUMN(TruthEta, truthEta, float);
@@ -148,11 +152,96 @@ DECLARE_SOA_TABLE(QGMLConstituents, "AOD", "QGMLCONSTS",
148152
qgmlconst::TofNSigmaKa,
149153
qgmlconst::TofNSigmaPr,
150154
qgmlconst::TofBeta,
155+
qgmlconst::HasTruthMatch,
151156
qgmlconst::TruthPdg,
152157
qgmlconst::TruthPt,
153158
qgmlconst::TruthEta,
154159
qgmlconst::TruthPhi,
155160
qgmlconst::IsPhysicalPrimary);
161+
162+
namespace qgmltruthjet
163+
{
164+
DECLARE_SOA_COLUMN(SourceMcCollisionId, sourceMcCollisionId, int64_t);
165+
DECLARE_SOA_COLUMN(SourceJetId, sourceJetId, int64_t);
166+
DECLARE_SOA_COLUMN(FlavorLabel, flavorLabel, int32_t);
167+
DECLARE_SOA_COLUMN(LeadingPartonPdg, leadingPartonPdg, int32_t);
168+
DECLARE_SOA_COLUMN(LeadingPartonPt, leadingPartonPt, float);
169+
DECLARE_SOA_COLUMN(LeadingPartonDeltaR, leadingPartonDeltaR, float);
170+
DECLARE_SOA_COLUMN(SubleadingPartonPdg, subleadingPartonPdg, int32_t);
171+
DECLARE_SOA_COLUMN(SubleadingPartonPt, subleadingPartonPt, float);
172+
DECLARE_SOA_COLUMN(SubleadingPartonDeltaR, subleadingPartonDeltaR, float);
173+
DECLARE_SOA_COLUMN(NPartonsInCone, nPartonsInCone, int32_t);
174+
DECLARE_SOA_COLUMN(LabelAmbiguous, labelAmbiguous, bool);
175+
DECLARE_SOA_COLUMN(Pt, pt, float);
176+
DECLARE_SOA_COLUMN(Eta, eta, float);
177+
DECLARE_SOA_COLUMN(Phi, phi, float);
178+
DECLARE_SOA_COLUMN(Area, area, float);
179+
DECLARE_SOA_COLUMN(Radius, radius, float);
180+
DECLARE_SOA_COLUMN(NConstituents, nConstituents, int32_t);
181+
DECLARE_SOA_COLUMN(NPions, nPions, int32_t);
182+
DECLARE_SOA_COLUMN(NKaons, nKaons, int32_t);
183+
DECLARE_SOA_COLUMN(NProtons, nProtons, int32_t);
184+
} // namespace qgmltruthjet
185+
DECLARE_SOA_TABLE(QGMLTruthJets, "AOD", "QGMLTJET",
186+
o2::soa::Index<>,
187+
qgmltruthjet::SourceMcCollisionId,
188+
qgmltruthjet::SourceJetId,
189+
qgmltruthjet::FlavorLabel,
190+
qgmltruthjet::LeadingPartonPdg,
191+
qgmltruthjet::LeadingPartonPt,
192+
qgmltruthjet::LeadingPartonDeltaR,
193+
qgmltruthjet::SubleadingPartonPdg,
194+
qgmltruthjet::SubleadingPartonPt,
195+
qgmltruthjet::SubleadingPartonDeltaR,
196+
qgmltruthjet::NPartonsInCone,
197+
qgmltruthjet::LabelAmbiguous,
198+
qgmltruthjet::Pt,
199+
qgmltruthjet::Eta,
200+
qgmltruthjet::Phi,
201+
qgmltruthjet::Area,
202+
qgmltruthjet::Radius,
203+
qgmltruthjet::NConstituents,
204+
qgmltruthjet::NPions,
205+
qgmltruthjet::NKaons,
206+
qgmltruthjet::NProtons);
207+
208+
namespace qgmltruthconst
209+
{
210+
DECLARE_SOA_INDEX_COLUMN(QGMLTruthJet, qGMLTruthJet);
211+
DECLARE_SOA_COLUMN(Pt, pt, float);
212+
DECLARE_SOA_COLUMN(Eta, eta, float);
213+
DECLARE_SOA_COLUMN(Phi, phi, float);
214+
DECLARE_SOA_COLUMN(DeltaEta, deltaEta, float);
215+
DECLARE_SOA_COLUMN(DeltaPhi, deltaPhi, float);
216+
DECLARE_SOA_COLUMN(DeltaR, deltaR, float);
217+
DECLARE_SOA_COLUMN(PtFraction, ptFraction, float);
218+
DECLARE_SOA_COLUMN(Pdg, pdg, int32_t);
219+
DECLARE_SOA_COLUMN(Charge, charge, float);
220+
DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, bool);
221+
} // namespace qgmltruthconst
222+
DECLARE_SOA_TABLE(QGMLTruthConstituents, "AOD", "QGMLTCONST",
223+
qgmltruthconst::QGMLTruthJetId,
224+
qgmltruthconst::Pt,
225+
qgmltruthconst::Eta,
226+
qgmltruthconst::Phi,
227+
qgmltruthconst::DeltaEta,
228+
qgmltruthconst::DeltaPhi,
229+
qgmltruthconst::DeltaR,
230+
qgmltruthconst::PtFraction,
231+
qgmltruthconst::Pdg,
232+
qgmltruthconst::Charge,
233+
qgmltruthconst::IsPhysicalPrimary);
234+
235+
namespace qgmlmatch
236+
{
237+
DECLARE_SOA_INDEX_COLUMN(QGMLJet, qGMLJet);
238+
DECLARE_SOA_INDEX_COLUMN(QGMLTruthJet, qGMLTruthJet);
239+
DECLARE_SOA_COLUMN(DeltaR, deltaR, float);
240+
} // namespace qgmlmatch
241+
DECLARE_SOA_TABLE(QGMLJetMatches, "AOD", "QGMLMATCH",
242+
qgmlmatch::QGMLJetId,
243+
qgmlmatch::QGMLTruthJetId,
244+
qgmlmatch::DeltaR);
156245
} // namespace o2::aod
157246

158247
using namespace o2;
@@ -162,6 +251,11 @@ using namespace o2::framework;
162251
using MCDJetEvents = aod::JetCollisionsMCD;
163252
using ChargedMCDJets = soa::Join<aod::ChargedMCDetectorLevelJets,
164253
aod::ChargedMCDetectorLevelJetConstituents>;
254+
using ChargedMCPJets = soa::Join<aod::ChargedMCParticleLevelJets,
255+
aod::ChargedMCParticleLevelJetConstituents>;
256+
using MatchedChargedMCDJets = soa::Join<aod::ChargedMCDetectorLevelJets,
257+
aod::ChargedMCDetectorLevelJetConstituents,
258+
aod::ChargedMCDetectorLevelJetsMatchedToChargedMCParticleLevelJets>;
165259
using JetTrackRefs = soa::Join<aod::JetTracks, aod::JTrackPIs>;
166260

167261
using HadronTracksMC = soa::Join<aod::Tracks,
@@ -181,6 +275,9 @@ using HadronTracksMC = soa::Join<aod::Tracks,
181275
struct QuarkGluonJetsProducer {
182276
Produces<aod::QGMLJets> qgMLJets;
183277
Produces<aod::QGMLConstituents> qgMLConstituents;
278+
Produces<aod::QGMLTruthJets> qgMLTruthJets;
279+
Produces<aod::QGMLTruthConstituents> qgMLTruthConstituents;
280+
Produces<aod::QGMLJetMatches> qgMLJetMatches;
184281

185282
static constexpr double JetRadiusTolerance = 1.e-3;
186283

@@ -220,6 +317,11 @@ struct QuarkGluonJetsProducer {
220317
Configurable<bool> dropAmbiguousLabels{"dropAmbiguousLabels", false, "drop ambiguous labels"};
221318

222319
Preslice<aod::JetParticles> mcParticlesPerCollision = aod::jmcparticle::mcCollisionId;
320+
Preslice<ChargedMCPJets> truthJetsPerCollision = aod::jet::mcCollisionId;
321+
322+
Configurable<double> minTruthJetPt{"minTruthJetPt", 10.0, "minimum truth jet pT"};
323+
Configurable<double> maxTruthJetPt{"maxTruthJetPt", 1000.0, "maximum truth jet pT"};
324+
Configurable<bool> truthPhysicalPrimariesOnly{"truthPhysicalPrimariesOnly", true, "retain physical primaries only"};
223325

224326
enum JetFlavor { UnknownJet = 0,
225327
QuarkJet = 1,
@@ -340,7 +442,8 @@ struct QuarkGluonJetsProducer {
340442
}
341443

342444
void process(MCDJetEvents::iterator const& collision,
343-
ChargedMCDJets const& detectorLevelJets,
445+
MatchedChargedMCDJets const& detectorLevelJets,
446+
ChargedMCPJets const& truthJets,
344447
JetTrackRefs const&,
345448
HadronTracksMC const&,
346449
aod::JetParticles const& jetParticles,
@@ -358,6 +461,9 @@ struct QuarkGluonJetsProducer {
358461
return;
359462
}
360463

464+
const int64_t mcCollisionId = collision.mcCollisionId();
465+
std::vector<std::tuple<int64_t, int64_t, float>> pendingMatches;
466+
361467
for (auto const& jet : detectorLevelJets) {
362468
const double jetRadius = static_cast<double>(jet.r()) / 100.0;
363469
if (std::abs(jetRadius - static_cast<double>(rJet)) > JetRadiusTolerance) {
@@ -427,6 +533,32 @@ struct QuarkGluonJetsProducer {
427533

428534
const auto qgMLJetId = qgMLJets.lastIndex();
429535

536+
auto geoMatches = jet.matchedJetGeo_as<ChargedMCPJets>();
537+
auto ptMatches = jet.matchedJetPt_as<ChargedMCPJets>();
538+
float bestDR = std::numeric_limits<float>::infinity();
539+
int64_t bestTruthSourceId = -1;
540+
for (auto const& candidate : geoMatches) {
541+
bool sharedMatch = false;
542+
for (auto const& ptCandidate : ptMatches) {
543+
if (candidate.globalIndex() == ptCandidate.globalIndex()) {
544+
sharedMatch = true;
545+
break;
546+
}
547+
}
548+
if (!sharedMatch || candidate.mcCollisionId() != mcCollisionId) {
549+
continue;
550+
}
551+
const auto dPhi = RecoDecay::constrainAngle(candidate.phi() - jet.phi(), -PI);
552+
const float dR = std::hypot(candidate.eta() - jet.eta(), dPhi);
553+
if (dR < bestDR) {
554+
bestDR = dR;
555+
bestTruthSourceId = candidate.globalIndex();
556+
}
557+
}
558+
if (bestTruthSourceId >= 0) {
559+
pendingMatches.emplace_back(qgMLJetId, bestTruthSourceId, bestDR);
560+
}
561+
430562
for (auto const& jtrack : jet.tracks_as<JetTrackRefs>()) {
431563
auto track = jtrack.track_as<HadronTracksMC>();
432564

@@ -477,13 +609,99 @@ struct QuarkGluonJetsProducer {
477609
track.hasTOF() ? static_cast<float>(track.tofNSigmaKa()) : missingTOF,
478610
track.hasTOF() ? static_cast<float>(track.tofNSigmaPr()) : missingTOF,
479611
track.hasTOF() ? static_cast<float>(track.beta()) : missingTOF,
612+
track.has_mcParticle(),
480613
truthPdg,
481614
truthPt,
482615
truthEta,
483616
truthPhi,
484617
isPhysicalPrimary);
485618
}
486619
}
620+
621+
auto truthJetsInCollision = truthJets.sliceBy(truthJetsPerCollision, mcCollisionId);
622+
std::unordered_map<int64_t, int64_t> truthOutputIndex;
623+
for (auto const& truthJet : truthJetsInCollision) {
624+
const double radius = static_cast<double>(truthJet.r()) / 100.0;
625+
if (std::abs(radius - static_cast<double>(rJet)) > JetRadiusTolerance) {
626+
continue;
627+
}
628+
if (isppRefAnalysis && std::abs(truthJet.eta()) > cfgEtaJetMax) {
629+
continue;
630+
}
631+
if (!isppRefAnalysis && std::abs(truthJet.eta()) + radius > maxEta - deltaEtaEdge) {
632+
continue;
633+
}
634+
if (truthJet.pt() < minTruthJetPt || truthJet.pt() > maxTruthJetPt) {
635+
continue;
636+
}
637+
if (applyAreaCut && truthJet.area() / (PI * radius * radius) < minNormalizedJetArea) {
638+
continue;
639+
}
640+
641+
struct SelectedTruthParticle {
642+
float pt, eta, phi, charge;
643+
int32_t pdg;
644+
bool primary;
645+
};
646+
std::vector<SelectedTruthParticle> selected;
647+
int32_t nPions = 0, nKaons = 0, nProtons = 0;
648+
for (auto const& particle : truthJet.tracks_as<aod::JetParticles>()) {
649+
if (truthPhysicalPrimariesOnly && !particle.isPhysicalPrimary()) {
650+
continue;
651+
}
652+
if (particle.pt() < minPt || particle.pt() > maxPt ||
653+
particle.eta() < minEta || particle.eta() > maxEta) {
654+
continue;
655+
}
656+
auto* pdgInfo = TDatabasePDG::Instance()->GetParticle(particle.pdgCode());
657+
if (!pdgInfo) {
658+
continue;
659+
}
660+
const float charge = static_cast<float>(pdgInfo->Charge() / 3.0);
661+
if (std::abs(charge) < 0.5f) {
662+
continue;
663+
}
664+
const int apdg = std::abs(particle.pdgCode());
665+
nPions += (apdg == 211);
666+
nKaons += (apdg == 321);
667+
nProtons += (apdg == 2212);
668+
selected.push_back({particle.pt(), particle.eta(), particle.phi(), charge,
669+
particle.pdgCode(), particle.isPhysicalPrimary()});
670+
}
671+
if (selected.empty()) {
672+
continue;
673+
}
674+
auto collisionParticles = jetParticles.sliceBy(mcParticlesPerCollision, mcCollisionId);
675+
auto flavor = getJetFlavorTag(truthJet, collisionParticles);
676+
if (dropUnknownLabels && flavor.flavor == UnknownJet) {
677+
continue;
678+
}
679+
if (dropAmbiguousLabels && flavor.ambiguous) {
680+
continue;
681+
}
682+
qgMLTruthJets(mcCollisionId, truthJet.globalIndex(),
683+
flavor.flavor, flavor.leadingPartonPdg, flavor.leadingPartonPt,
684+
flavor.leadingPartonDeltaR, flavor.subleadingPartonPdg,
685+
flavor.subleadingPartonPt, flavor.subleadingPartonDeltaR,
686+
flavor.nPartonsInCone, flavor.ambiguous,
687+
truthJet.pt(), truthJet.eta(), truthJet.phi(), truthJet.area(),
688+
radius, static_cast<int32_t>(selected.size()), nPions, nKaons, nProtons);
689+
const auto truthOutputId = qgMLTruthJets.lastIndex();
690+
truthOutputIndex.emplace(truthJet.globalIndex(), truthOutputId);
691+
for (auto const& part : selected) {
692+
const float dEta = part.eta - truthJet.eta();
693+
const float dPhi = RecoDecay::constrainAngle(part.phi - truthJet.phi(), -PI);
694+
qgMLTruthConstituents(truthOutputId, part.pt, part.eta, part.phi,
695+
dEta, dPhi, std::hypot(dEta, dPhi),
696+
truthJet.pt() > 0 ? part.pt / truthJet.pt() : 0.f,
697+
part.pdg, part.charge, part.primary);
698+
}
699+
}
700+
for (auto const& [recoOutputId, sourceTruthId, dR] : pendingMatches) {
701+
if (auto it = truthOutputIndex.find(sourceTruthId); it != truthOutputIndex.end()) {
702+
qgMLJetMatches(recoOutputId, it->second, dR);
703+
}
704+
}
487705
}
488706
};
489707

0 commit comments

Comments
 (0)