From e0cc5442f700567ae94aa4abfcaf69329141379a Mon Sep 17 00:00:00 2001 From: nburmaso Date: Mon, 14 Sep 2026 17:05:22 +0300 Subject: [PATCH] Introduce MC tables to upcCandProducerBarrel --- PWGUD/TableProducer/upcCandProducerBarrel.cxx | 206 ++++++++++++++++-- 1 file changed, 193 insertions(+), 13 deletions(-) diff --git a/PWGUD/TableProducer/upcCandProducerBarrel.cxx b/PWGUD/TableProducer/upcCandProducerBarrel.cxx index 66e75d700f2..4d2a75be18a 100644 --- a/PWGUD/TableProducer/upcCandProducerBarrel.cxx +++ b/PWGUD/TableProducer/upcCandProducerBarrel.cxx @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -39,14 +40,15 @@ #include #include -#include - #include #include #include #include #include #include +#include +#include +#include #include using namespace o2; @@ -59,19 +61,37 @@ constexpr int NBCsPerOrbit = o2::constants::lhc::LHCMaxBunches; constexpr int MaxStoredDistance = 15; constexpr float MaxFITTime = 30.f; +constexpr int NoSignalMCRef = -1; +constexpr int NoRecoCollision = -1; +constexpr int MultipleRecoTracks = -2; + +constexpr uint16_t MCTrackFakeMask = 1u << 15; + +enum MCParticleFlag : uint8_t { + kNoFlags = 0, + kPhysicalPrimary = 1u << 0, // generated physical primary + kHasRecoTrack = 1u << 1, // has at least one non-fake reconstructed track + kHighMultiplicityCollision = 1u << 2, // matched track belongs to a collision with extra contributors + kHasPVContributor = 1u << 3, // has a track contributing to the PV + kHasAmbiguousTrack = 1u << 4 // has an ambiguous track +}; + +using SignalRef = std::pair; using DistanceMap = std::vector>; DistanceMap buildMinimumDistanceMap(std::vector const& tfIDs, int64_t bcSOR, int64_t nBCsPerTF, - std::vector& activeBCs, uint32_t maxDistance) + std::vector& activeBCs, uint32_t maxDistance, + bool skipFirstExactMatch = false) { const uint32_t overflowDistance = maxDistance + 1; DistanceMap distances(tfIDs.size(), std::vector(nBCsPerTF, overflowDistance)); if (activeBCs.empty()) { return distances; } + // Keep duplicate BCs so collisions at the same BC have distance zero std::sort(activeBCs.begin(), activeBCs.end()); - activeBCs.erase(std::unique(activeBCs.begin(), activeBCs.end()), activeBCs.end()); + // The forward-only cursor makes the scan linear after sorting size_t nextIndex = 0; for (size_t iTF = 0; iTF < tfIDs.size(); ++iTF) { const int64_t tfStartBC = bcSOR + tfIDs[iTF] * nBCsPerTF; @@ -80,9 +100,14 @@ DistanceMap buildMinimumDistanceMap(std::vector const& tfIDs, while (nextIndex < activeBCs.size() && activeBCs[nextIndex] < currentBC) { ++nextIndex; } + size_t rightIndex = nextIndex; + // MC maps contain the queried collision itself, so skip one exact self-match + if (skipFirstExactMatch && rightIndex < activeBCs.size() && activeBCs[rightIndex] == currentBC) { + ++rightIndex; + } int64_t distance = overflowDistance; - if (nextIndex < activeBCs.size()) { - distance = std::min(distance, activeBCs[nextIndex] - currentBC); + if (rightIndex < activeBCs.size()) { + distance = std::min(distance, activeBCs[rightIndex] - currentBC); } if (nextIndex > 0) { distance = std::min(distance, currentBC - activeBCs[nextIndex - 1]); @@ -143,6 +168,13 @@ DECLARE_SOA_COLUMN(Sign, sign, std::vector); DECLARE_SOA_COLUMN(MinimumDistanceFT0, minimumDistanceFT0, int8_t); DECLARE_SOA_COLUMN(MinimumDistanceFV0, minimumDistanceFV0, int8_t); DECLARE_SOA_COLUMN(MinimumDistanceFDD, minimumDistanceFDD, int8_t); +DECLARE_SOA_COLUMN(Pdg, pdg, std::vector); +DECLARE_SOA_COLUMN(MinimumDistanceMCCol, minimumDistanceMCCol, int8_t); +DECLARE_SOA_COLUMN(RecoFlags, recoFlags, std::vector); // See MCParticleFlag for the bit definitions +DECLARE_SOA_COLUMN(RecoCollisionIds, recoCollisionIds, std::vector); // -1: no collision, -2: multiple tracks +DECLARE_SOA_COLUMN(ParticleIds, particleIds, std::vector); +DECLARE_SOA_COLUMN(Mult, mult, int32_t); // number of physical-primary MC particles in the collision +DECLARE_SOA_COLUMN(IsSignal, isSignal, bool); // whether the collision belongs to the configured signal source } // namespace o2::aod::upc_cand_prod_bar namespace o2::aod @@ -173,10 +205,39 @@ DECLARE_SOA_TABLE(UPCBarrelCands, "AOD", "UPCBARRELCANDS", upc_cand_prod_bar::MinimumDistanceFT0, upc_cand_prod_bar::MinimumDistanceFV0, upc_cand_prod_bar::MinimumDistanceFDD); + +DECLARE_SOA_TABLE(UPCBarrelMcColls, "AOD", "UPCBMCCOLL", + upc_cand_prod_bar::TfId, + upc_cand_prod_bar::GlobalBC, + upc_cand_prod_bar::PosX, + upc_cand_prod_bar::PosY, + upc_cand_prod_bar::PosZ, + upc_cand_prod_bar::Px, + upc_cand_prod_bar::Py, + upc_cand_prod_bar::Pz, + upc_cand_prod_bar::Pdg, + upc_cand_prod_bar::Mult, + upc_cand_prod_bar::MinimumDistanceMCCol, + upc_cand_prod_bar::RecoFlags, + upc_cand_prod_bar::RecoCollisionIds, + upc_cand_prod_bar::IsSignal); + +namespace upc_cand_prod_bar +{ +// Array index into UPCBarrelMcColls +DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL_CUSTOM(McCollision, mcCollision, int32_t, UPCBarrelMcColls, "UPCBMCColls", ""); +} // namespace upc_cand_prod_bar + +// Both IDs are -1 if the track has no valid stored signal reference +DECLARE_SOA_TABLE(UPCBarrelMCLabels, "AOD", "UPCBMCLABEL", + upc_cand_prod_bar::McCollisionIds, + upc_cand_prod_bar::ParticleIds); } // namespace o2::aod struct UpcCandProducerBarrel { Produces selectedCandidates; + Produces selectedMCCollisions; + Produces selectedCandidateMCLabels; HistogramRegistry registry{"registry", {}}; Service ccdb{}; @@ -197,13 +258,21 @@ struct UpcCandProducerBarrel { Configurable vetoFT0{"vetoFT0", 1, "FT0 veto: 0=off, 1=on"}; Configurable vetoFV0{"vetoFV0", 0, "FV0 veto: 0=off, 1=on"}; Configurable vetoFDD{"vetoFDD", 0, "FDD veto: 0=off, 1=on"}; + Configurable signalSourceId{"signalSourceId", 1, "MC source ID stored as embedded signal"}; using CollisionsWithSels = soa::Join; using BCsWithSels = soa::Join; using TracksWithPID = soa::Join; + using TracksWithPIDMC = soa::Join; Preslice tracksPerCollision = aod::track::collisionId; + Preslice mcTracksPerCollision = aod::track::collisionId; + PresliceUnsorted tracksPerMCParticle = aod::mctracklabel::mcParticleId; + Preslice particlesPerMCCollision = aod::mcparticle::mcCollisionId; RCTFlagsChecker rctChecker{kFDDBad, kFT0Bad, kFV0Bad, kITSBad, kITSLimAccMCRepr, kTPCBadTracking, kTPCLimAccMCRepr, kTPCBadPID, kTOFBad, kTOFLimAccMCRepr, kCcdbObjectLoaded}; @@ -250,13 +319,19 @@ struct UpcCandProducerBarrel { cachedRunNumber = runNumber; } - void process(CollisionsWithSels const& collisions, TracksWithPID const& tracks, BCsWithSels const& bcs, aod::FT0s const& ft0s, aod::FV0As const& fv0s, aod::FDDs const& fdds) + template + void processImpl(CollisionsT const& collisions, TracksT const& tracks, PresliceT const& tracksPerColl, + BCsWithSels const& bcs, aod::FT0s const& ft0s, aod::FV0As const& fv0s, + aod::FDDs const& fdds, aod::McCollisions const* mcCollisions = nullptr, + aod::McParticles const* mcParticles = nullptr, + aod::AmbiguousTracks const* ambiguousTracks = nullptr) { updateRunInfo(bcs.begin().runNumber()); std::vector tfIDs; std::vector tfPassesRCT; std::vector localTFIndex; localTFIndex.reserve(bcs.size()); + // Cache the local TF slot for each BC row for collision and MC lookups for (const auto& bc : bcs) { const auto tfID = static_cast((static_cast(bc.globalBC()) - bcSOR) / nBCsPerTF); if (tfIDs.empty() || tfID != tfIDs.back()) { @@ -267,12 +342,84 @@ struct UpcCandProducerBarrel { } registry.fill(HIST("hProcessedTFs"), 0.5, static_cast(tfIDs.size())); + std::unordered_map signalParticleRefs; + if constexpr (isMC) { + std::unordered_set ambiguousTrackIds; + ambiguousTrackIds.reserve(ambiguousTracks->size()); + for (const auto& ambiguousTrack : *ambiguousTracks) { + ambiguousTrackIds.insert(ambiguousTrack.trackId()); + } + + std::vector bcsWithMCCollision; + bcsWithMCCollision.reserve(mcCollisions->size()); + // Use collisions from every source and retain same-BC pile-up + for (const auto& mcCollision : *mcCollisions) { + bcsWithMCCollision.push_back(static_cast(mcCollision.template bc_as().globalBC())); + } + const auto minDistanceMC = buildMinimumDistanceMap(tfIDs, bcSOR, nBCsPerTF, bcsWithMCCollision, MaxStoredDistance, true); + + for (const auto& mcCollision : *mcCollisions) { + const auto bc = mcCollision.template bc_as(); + const auto globalBC = bc.globalBC(); + const auto iTF = localTFIndex[bc.globalIndex()]; + const auto bcInTF = (static_cast(globalBC) - bcSOR) % nBCsPerTF; + const auto particles = mcParticles->sliceBy(particlesPerMCCollision, mcCollision.globalIndex()); + const bool isSignal = mcCollision.getSourceId() == signalSourceId.value; + + std::vector px, py, pz; + std::vector pdg, recoCollisionIds; + std::vector recoFlags; + int32_t mult = 0; + // Output rows follow input MC collisions, making this the compact collision label + const auto newMCCollisionId = static_cast(selectedMCCollisions.lastIndex() + 1); + for (const auto& particle : particles) { + mult += particle.isPhysicalPrimary(); + // Background rows keep only collision metadata and physical-primary multiplicity + if (!isSignal) { + continue; + } + const auto newMCParticleId = static_cast(pdg.size()); + px.push_back(particle.px()); + py.push_back(particle.py()); + pz.push_back(particle.pz()); + pdg.push_back(particle.pdgCode()); + + uint8_t flags = particle.isPhysicalPrimary() ? kPhysicalPrimary : kNoFlags; + int32_t recoCollisionId = NoRecoCollision; + int nRecoTracks = 0; + // Inspect only reconstructed tracks matched to this signal particle + for (const auto& track : tracks.sliceBy(tracksPerMCParticle, particle.globalIndex())) { + if ((track.mcMask() & MCTrackFakeMask) != 0u) { + continue; + } + ++nRecoTracks; + flags |= kHasRecoTrack; + flags |= track.isPVContributor() ? kHasPVContributor : kNoFlags; + flags |= ambiguousTrackIds.contains(track.globalIndex()) ? kHasAmbiguousTrack : kNoFlags; + if (track.has_collision()) { + recoCollisionId = track.collisionId(); + flags |= collisions.iteratorAt(recoCollisionId).numContrib() > nTracks ? kHighMultiplicityCollision : kNoFlags; + } + } + // A single collision label is meaningful only for exactly one reconstructed track + if (nRecoTracks > 1) { + recoCollisionId = MultipleRecoTracks; + } + recoFlags.push_back(flags); + recoCollisionIds.push_back(recoCollisionId); + // Remap original particle IDs to output index for candidate labels + signalParticleRefs.emplace(particle.globalIndex(), SignalRef{newMCCollisionId, newMCParticleId}); + } + selectedMCCollisions(tfIDs[iTF], globalBC, mcCollision.posX(), mcCollision.posY(), mcCollision.posZ(), + px, py, pz, pdg, mult, minDistanceMC[iTF][bcInTF], + recoFlags, recoCollisionIds, isSignal); + } + } + auto hTVX = registry.get(HIST("hTVX")); auto hTVXRCT = registry.get(HIST("hTVXRCT")); - std::vector bcsWithFT0; - std::vector bcsWithFV0; - std::vector bcsWithFDD; + std::vector bcsWithFT0, bcsWithFV0, bcsWithFDD; bcsWithFT0.reserve(ft0s.size()); bcsWithFV0.reserve(fv0s.size()); bcsWithFDD.reserve(fdds.size()); @@ -284,7 +431,7 @@ struct UpcCandProducerBarrel { if (ft0.timeA() < MaxFITTime || ft0.timeC() < MaxFITTime) { bcsWithFT0.push_back(gbc); } - if (!collidingBCs[bcInOrbit] || !TESTBIT(ft0.triggerMask(), o2::fit::Triggers::bitVertex)) { + if (!collidingBCs[bcInOrbit] || !(ft0.triggerMask() & (1U << o2::fit::Triggers::bitVertex))) { continue; } hTVX->Fill(bcInOrbit); @@ -337,7 +484,7 @@ struct UpcCandProducerBarrel { if (collision.numContrib() != nTracks || !collision.selection_bit(aod::evsel::kNoTimeFrameBorder) || !rctChecker(collision)) { continue; } - auto bc = collision.bc_as(); + auto bc = collision.template bc_as(); auto gbc = bc.globalBC(); const auto iTF = localTFIndex[bc.globalIndex()]; const int64_t bcInTF = (static_cast(gbc) - bcSOR) % nBCsPerTF; @@ -354,7 +501,9 @@ struct UpcCandProducerBarrel { std::vector itsClusterMap; std::vector nClusters; std::vector sign; - for (const auto& track : tracks.sliceBy(tracksPerCollision, collision.globalIndex())) { + std::vector newMCCollisionIds; + std::vector newParticleIds; + for (const auto& track : tracks.sliceBy(tracksPerColl, collision.globalIndex())) { if (!track.isPVContributor() || !track.hasITS() || !track.hasTPC() || std::abs(track.eta()) > maxAbsEta.value || track.pt() < minPt.value) { continue; } @@ -373,6 +522,17 @@ struct UpcCandProducerBarrel { itsClusterMap.push_back(track.itsClusterMap()); nClusters.push_back(track.tpcNClsFound()); sign.push_back(track.sign()); + if constexpr (isMC) { + // Fake, background, and unmatched particles keep the -1 label + SignalRef label{NoSignalMCRef, NoSignalMCRef}; + if (track.has_mcParticle() && (track.mcMask() & MCTrackFakeMask) == 0u) { + if (const auto ref = signalParticleRefs.find(track.mcParticleId()); ref != signalParticleRefs.end()) { + label = ref->second; + } + } + newMCCollisionIds.push_back(label.first); + newParticleIds.push_back(label.second); + } } if (px.size() != static_cast(nTracks)) { continue; @@ -383,8 +543,28 @@ struct UpcCandProducerBarrel { std::min(distFT0, MaxStoredDistance + 1), std::min(distFV0, MaxStoredDistance + 1), std::min(distFDD, MaxStoredDistance + 1)); + if constexpr (isMC) { + selectedCandidateMCLabels(newMCCollisionIds, newParticleIds); + } } } + + void processData(CollisionsWithSels const& collisions, TracksWithPID const& tracks, BCsWithSels const& bcs, + aod::FT0s const& ft0s, aod::FV0As const& fv0s, aod::FDDs const& fdds) + { + processImpl(collisions, tracks, tracksPerCollision, bcs, ft0s, fv0s, fdds); + } + PROCESS_SWITCH(UpcCandProducerBarrel, processData, "Process reconstructed data", true); + + void processMC(CollisionsWithSels const& collisions, TracksWithPIDMC const& tracks, BCsWithSels const& bcs, + aod::FT0s const& ft0s, aod::FV0As const& fv0s, aod::FDDs const& fdds, + aod::McCollisions const& mcCollisions, aod::McParticles const& mcParticles, + aod::AmbiguousTracks const& ambiguousTracks) + { + processImpl(collisions, tracks, mcTracksPerCollision, bcs, ft0s, fv0s, fdds, + &mcCollisions, &mcParticles, &ambiguousTracks); + } + PROCESS_SWITCH(UpcCandProducerBarrel, processMC, "Process reconstructed MC and write compact MC truth", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& context) { return WorkflowSpec{adaptAnalysisTask(context)}; }