From 935d875b4631fbdf9d5db18ac558fb19327d78a9 Mon Sep 17 00:00:00 2001 From: Aleksandra Mulewicz Date: Wed, 4 Nov 2026 16:23:34 +0100 Subject: [PATCH] PWGJE: add truth-level jet tables and reco-truth matching --- .../TableProducer/quarkGluonJetsProducer.cxx | 220 +++++++++++++++++- 1 file changed, 219 insertions(+), 1 deletion(-) diff --git a/PWGJE/TableProducer/quarkGluonJetsProducer.cxx b/PWGJE/TableProducer/quarkGluonJetsProducer.cxx index 92c2a05d92b..216d58358c7 100644 --- a/PWGJE/TableProducer/quarkGluonJetsProducer.cxx +++ b/PWGJE/TableProducer/quarkGluonJetsProducer.cxx @@ -39,6 +39,7 @@ #include #include +#include #include #include @@ -47,6 +48,8 @@ #include #include #include +#include +#include #include namespace o2::aod @@ -117,6 +120,7 @@ DECLARE_SOA_COLUMN(TofNSigmaPi, tofNSigmaPi, float); DECLARE_SOA_COLUMN(TofNSigmaKa, tofNSigmaKa, float); DECLARE_SOA_COLUMN(TofNSigmaPr, tofNSigmaPr, float); DECLARE_SOA_COLUMN(TofBeta, tofBeta, float); +DECLARE_SOA_COLUMN(HasTruthMatch, hasTruthMatch, bool); DECLARE_SOA_COLUMN(TruthPdg, truthPdg, int32_t); DECLARE_SOA_COLUMN(TruthPt, truthPt, float); DECLARE_SOA_COLUMN(TruthEta, truthEta, float); @@ -148,11 +152,96 @@ DECLARE_SOA_TABLE(QGMLConstituents, "AOD", "QGMLCONSTS", qgmlconst::TofNSigmaKa, qgmlconst::TofNSigmaPr, qgmlconst::TofBeta, + qgmlconst::HasTruthMatch, qgmlconst::TruthPdg, qgmlconst::TruthPt, qgmlconst::TruthEta, qgmlconst::TruthPhi, qgmlconst::IsPhysicalPrimary); + +namespace qgmltruthjet +{ +DECLARE_SOA_COLUMN(SourceMcCollisionId, sourceMcCollisionId, int64_t); +DECLARE_SOA_COLUMN(SourceJetId, sourceJetId, int64_t); +DECLARE_SOA_COLUMN(FlavorLabel, flavorLabel, int32_t); +DECLARE_SOA_COLUMN(LeadingPartonPdg, leadingPartonPdg, int32_t); +DECLARE_SOA_COLUMN(LeadingPartonPt, leadingPartonPt, float); +DECLARE_SOA_COLUMN(LeadingPartonDeltaR, leadingPartonDeltaR, float); +DECLARE_SOA_COLUMN(SubleadingPartonPdg, subleadingPartonPdg, int32_t); +DECLARE_SOA_COLUMN(SubleadingPartonPt, subleadingPartonPt, float); +DECLARE_SOA_COLUMN(SubleadingPartonDeltaR, subleadingPartonDeltaR, float); +DECLARE_SOA_COLUMN(NPartonsInCone, nPartonsInCone, int32_t); +DECLARE_SOA_COLUMN(LabelAmbiguous, labelAmbiguous, bool); +DECLARE_SOA_COLUMN(Pt, pt, float); +DECLARE_SOA_COLUMN(Eta, eta, float); +DECLARE_SOA_COLUMN(Phi, phi, float); +DECLARE_SOA_COLUMN(Area, area, float); +DECLARE_SOA_COLUMN(Radius, radius, float); +DECLARE_SOA_COLUMN(NConstituents, nConstituents, int32_t); +DECLARE_SOA_COLUMN(NPions, nPions, int32_t); +DECLARE_SOA_COLUMN(NKaons, nKaons, int32_t); +DECLARE_SOA_COLUMN(NProtons, nProtons, int32_t); +} // namespace qgmltruthjet +DECLARE_SOA_TABLE(QGMLTruthJets, "AOD", "QGMLTJET", + o2::soa::Index<>, + qgmltruthjet::SourceMcCollisionId, + qgmltruthjet::SourceJetId, + qgmltruthjet::FlavorLabel, + qgmltruthjet::LeadingPartonPdg, + qgmltruthjet::LeadingPartonPt, + qgmltruthjet::LeadingPartonDeltaR, + qgmltruthjet::SubleadingPartonPdg, + qgmltruthjet::SubleadingPartonPt, + qgmltruthjet::SubleadingPartonDeltaR, + qgmltruthjet::NPartonsInCone, + qgmltruthjet::LabelAmbiguous, + qgmltruthjet::Pt, + qgmltruthjet::Eta, + qgmltruthjet::Phi, + qgmltruthjet::Area, + qgmltruthjet::Radius, + qgmltruthjet::NConstituents, + qgmltruthjet::NPions, + qgmltruthjet::NKaons, + qgmltruthjet::NProtons); + +namespace qgmltruthconst +{ +DECLARE_SOA_INDEX_COLUMN(QGMLTruthJet, qGMLTruthJet); +DECLARE_SOA_COLUMN(Pt, pt, float); +DECLARE_SOA_COLUMN(Eta, eta, float); +DECLARE_SOA_COLUMN(Phi, phi, float); +DECLARE_SOA_COLUMN(DeltaEta, deltaEta, float); +DECLARE_SOA_COLUMN(DeltaPhi, deltaPhi, float); +DECLARE_SOA_COLUMN(DeltaR, deltaR, float); +DECLARE_SOA_COLUMN(PtFraction, ptFraction, float); +DECLARE_SOA_COLUMN(Pdg, pdg, int32_t); +DECLARE_SOA_COLUMN(Charge, charge, float); +DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, bool); +} // namespace qgmltruthconst +DECLARE_SOA_TABLE(QGMLTruthConstituents, "AOD", "QGMLTCONST", + qgmltruthconst::QGMLTruthJetId, + qgmltruthconst::Pt, + qgmltruthconst::Eta, + qgmltruthconst::Phi, + qgmltruthconst::DeltaEta, + qgmltruthconst::DeltaPhi, + qgmltruthconst::DeltaR, + qgmltruthconst::PtFraction, + qgmltruthconst::Pdg, + qgmltruthconst::Charge, + qgmltruthconst::IsPhysicalPrimary); + +namespace qgmlmatch +{ +DECLARE_SOA_INDEX_COLUMN(QGMLJet, qGMLJet); +DECLARE_SOA_INDEX_COLUMN(QGMLTruthJet, qGMLTruthJet); +DECLARE_SOA_COLUMN(DeltaR, deltaR, float); +} // namespace qgmlmatch +DECLARE_SOA_TABLE(QGMLJetMatches, "AOD", "QGMLMATCH", + qgmlmatch::QGMLJetId, + qgmlmatch::QGMLTruthJetId, + qgmlmatch::DeltaR); } // namespace o2::aod using namespace o2; @@ -162,6 +251,11 @@ using namespace o2::framework; using MCDJetEvents = aod::JetCollisionsMCD; using ChargedMCDJets = soa::Join; +using ChargedMCPJets = soa::Join; +using MatchedChargedMCDJets = soa::Join; using JetTrackRefs = soa::Join; using HadronTracksMC = soa::Join qgMLJets; Produces qgMLConstituents; + Produces qgMLTruthJets; + Produces qgMLTruthConstituents; + Produces qgMLJetMatches; static constexpr double JetRadiusTolerance = 1.e-3; @@ -220,6 +317,11 @@ struct QuarkGluonJetsProducer { Configurable dropAmbiguousLabels{"dropAmbiguousLabels", false, "drop ambiguous labels"}; Preslice mcParticlesPerCollision = aod::jmcparticle::mcCollisionId; + Preslice truthJetsPerCollision = aod::jet::mcCollisionId; + + Configurable minTruthJetPt{"minTruthJetPt", 10.0, "minimum truth jet pT"}; + Configurable maxTruthJetPt{"maxTruthJetPt", 1000.0, "maximum truth jet pT"}; + Configurable truthPhysicalPrimariesOnly{"truthPhysicalPrimariesOnly", true, "retain physical primaries only"}; enum JetFlavor { UnknownJet = 0, QuarkJet = 1, @@ -340,7 +442,8 @@ struct QuarkGluonJetsProducer { } void process(MCDJetEvents::iterator const& collision, - ChargedMCDJets const& detectorLevelJets, + MatchedChargedMCDJets const& detectorLevelJets, + ChargedMCPJets const& truthJets, JetTrackRefs const&, HadronTracksMC const&, aod::JetParticles const& jetParticles, @@ -358,6 +461,9 @@ struct QuarkGluonJetsProducer { return; } + const int64_t mcCollisionId = collision.mcCollisionId(); + std::vector> pendingMatches; + for (auto const& jet : detectorLevelJets) { const double jetRadius = static_cast(jet.r()) / 100.0; if (std::abs(jetRadius - static_cast(rJet)) > JetRadiusTolerance) { @@ -427,6 +533,32 @@ struct QuarkGluonJetsProducer { const auto qgMLJetId = qgMLJets.lastIndex(); + auto geoMatches = jet.matchedJetGeo_as(); + auto ptMatches = jet.matchedJetPt_as(); + float bestDR = std::numeric_limits::infinity(); + int64_t bestTruthSourceId = -1; + for (auto const& candidate : geoMatches) { + bool sharedMatch = false; + for (auto const& ptCandidate : ptMatches) { + if (candidate.globalIndex() == ptCandidate.globalIndex()) { + sharedMatch = true; + break; + } + } + if (!sharedMatch || candidate.mcCollisionId() != mcCollisionId) { + continue; + } + const auto dPhi = RecoDecay::constrainAngle(candidate.phi() - jet.phi(), -PI); + const float dR = std::hypot(candidate.eta() - jet.eta(), dPhi); + if (dR < bestDR) { + bestDR = dR; + bestTruthSourceId = candidate.globalIndex(); + } + } + if (bestTruthSourceId >= 0) { + pendingMatches.emplace_back(qgMLJetId, bestTruthSourceId, bestDR); + } + for (auto const& jtrack : jet.tracks_as()) { auto track = jtrack.track_as(); @@ -477,6 +609,7 @@ struct QuarkGluonJetsProducer { track.hasTOF() ? static_cast(track.tofNSigmaKa()) : missingTOF, track.hasTOF() ? static_cast(track.tofNSigmaPr()) : missingTOF, track.hasTOF() ? static_cast(track.beta()) : missingTOF, + track.has_mcParticle(), truthPdg, truthPt, truthEta, @@ -484,6 +617,91 @@ struct QuarkGluonJetsProducer { isPhysicalPrimary); } } + + auto truthJetsInCollision = truthJets.sliceBy(truthJetsPerCollision, mcCollisionId); + std::unordered_map truthOutputIndex; + for (auto const& truthJet : truthJetsInCollision) { + const double radius = static_cast(truthJet.r()) / 100.0; + if (std::abs(radius - static_cast(rJet)) > JetRadiusTolerance) { + continue; + } + if (isppRefAnalysis && std::abs(truthJet.eta()) > cfgEtaJetMax) { + continue; + } + if (!isppRefAnalysis && std::abs(truthJet.eta()) + radius > maxEta - deltaEtaEdge) { + continue; + } + if (truthJet.pt() < minTruthJetPt || truthJet.pt() > maxTruthJetPt) { + continue; + } + if (applyAreaCut && truthJet.area() / (PI * radius * radius) < minNormalizedJetArea) { + continue; + } + + struct SelectedTruthParticle { + float pt, eta, phi, charge; + int32_t pdg; + bool primary; + }; + std::vector selected; + int32_t nPions = 0, nKaons = 0, nProtons = 0; + for (auto const& particle : truthJet.tracks_as()) { + if (truthPhysicalPrimariesOnly && !particle.isPhysicalPrimary()) { + continue; + } + if (particle.pt() < minPt || particle.pt() > maxPt || + particle.eta() < minEta || particle.eta() > maxEta) { + continue; + } + auto* pdgInfo = TDatabasePDG::Instance()->GetParticle(particle.pdgCode()); + if (!pdgInfo) { + continue; + } + const float charge = static_cast(pdgInfo->Charge() / 3.0); + if (std::abs(charge) < 0.5f) { + continue; + } + const int apdg = std::abs(particle.pdgCode()); + nPions += (apdg == 211); + nKaons += (apdg == 321); + nProtons += (apdg == 2212); + selected.push_back({particle.pt(), particle.eta(), particle.phi(), charge, + particle.pdgCode(), particle.isPhysicalPrimary()}); + } + if (selected.empty()) { + continue; + } + auto collisionParticles = jetParticles.sliceBy(mcParticlesPerCollision, mcCollisionId); + auto flavor = getJetFlavorTag(truthJet, collisionParticles); + if (dropUnknownLabels && flavor.flavor == UnknownJet) { + continue; + } + if (dropAmbiguousLabels && flavor.ambiguous) { + continue; + } + qgMLTruthJets(mcCollisionId, truthJet.globalIndex(), + flavor.flavor, flavor.leadingPartonPdg, flavor.leadingPartonPt, + flavor.leadingPartonDeltaR, flavor.subleadingPartonPdg, + flavor.subleadingPartonPt, flavor.subleadingPartonDeltaR, + flavor.nPartonsInCone, flavor.ambiguous, + truthJet.pt(), truthJet.eta(), truthJet.phi(), truthJet.area(), + radius, static_cast(selected.size()), nPions, nKaons, nProtons); + const auto truthOutputId = qgMLTruthJets.lastIndex(); + truthOutputIndex.emplace(truthJet.globalIndex(), truthOutputId); + for (auto const& part : selected) { + const float dEta = part.eta - truthJet.eta(); + const float dPhi = RecoDecay::constrainAngle(part.phi - truthJet.phi(), -PI); + qgMLTruthConstituents(truthOutputId, part.pt, part.eta, part.phi, + dEta, dPhi, std::hypot(dEta, dPhi), + truthJet.pt() > 0 ? part.pt / truthJet.pt() : 0.f, + part.pdg, part.charge, part.primary); + } + } + for (auto const& [recoOutputId, sourceTruthId, dR] : pendingMatches) { + if (auto it = truthOutputIndex.find(sourceTruthId); it != truthOutputIndex.end()) { + qgMLJetMatches(recoOutputId, it->second, dR); + } + } } };