diff --git a/PWGLF/DataModel/LFOmega2012MlTables.h b/PWGLF/DataModel/LFOmega2012MlTables.h index dd81f009845..e17bde5d569 100644 --- a/PWGLF/DataModel/LFOmega2012MlTables.h +++ b/PWGLF/DataModel/LFOmega2012MlTables.h @@ -16,6 +16,8 @@ /// Mode A (XiK0s): Omega(2012)- -> Xi- K0S. Mode B (Xi1530K): Omega(2012)- -> Xi(1530)0 K- -> Xi- pi+ K-. /// Only the input-object tables (events, cascades, V0s, tracks) are common; every event row belongs to one /// decay mode (column omDecayMode), so the object rows of the two modes never mix either. +/// Event rows are written only for collisions with at least one written candidate. In reconstructed MC the +/// object MC labels and the pruned McParticles rows replace the source MC tables, which are not written. /// /// Cumulative pass bits (each bit requires all lower bits): /// mode A: 1 loose (valid canonical candidate inside the rapidity window), 2 Xi selection, 4 K0s selection, @@ -234,6 +236,38 @@ DECLARE_SOA_TABLE(Omega2012MlGenAudit, "AOD", "OMMLGENAUDIT", o2::soa::Index<>, omega2012ml::OmRecoCollisionId, omega2012ml::OmOriginalMcParticleId, omega2012ml::OmPdg, omega2012ml::OmDaughterPdg1, omega2012ml::OmDaughterPdg2, omega2012ml::OmGenChannel, omega2012ml::OmGenPt, omega2012ml::OmGenY); + +// MC truth of reconstructed MC: daughter labels of the written objects and the pruned source McParticles rows +// (Omega(2012) decay trees and ancestry of the daughter labels). Source rows are DF-local scalar values. +namespace omega2012ml +{ +DECLARE_SOA_INDEX_COLUMN_FULL_CUSTOM(Omega2012MlTrack, omega2012MlTrack, int, Omega2012MlTracks, "OmmlTracks", ""); +DECLARE_SOA_COLUMN(OmMcCollisionRow, omMcCollisionRow, int64_t); //! row of the MC collision in the input McCollisions table +DECLARE_SOA_COLUMN(OmCascadeMcLabels, omCascadeMcLabels, int[3]); //! McParticles rows of the cascade daughters, order of omCascadeIndices; -1 if absent +DECLARE_SOA_COLUMN(OmV0McLabels, omV0McLabels, int[2]); //! McParticles rows of the V0 daughters, order of omV0Indices; -1 if absent +DECLARE_SOA_COLUMN(OmMcLabel, omMcLabel, int); //! McParticles row of the track; -1 if absent +DECLARE_SOA_COLUMN(OmMcRow, omMcRow, int64_t); //! row of the particle in the input McParticles table +DECLARE_SOA_COLUMN(OmMcSelection, omMcSelection, uint8_t); //! 1 Omega(2012) decay tree, 2 ancestry of a written daughter label +DECLARE_SOA_COLUMN(OmNMothers, omNMothers, uint8_t); //! number of mothers in the input row +DECLARE_SOA_COLUMN(OmMotherRows, omMotherRows, int[2]); //! first two McParticles mother rows; -1 if absent +DECLARE_SOA_COLUMN(OmDaughterRows, omDaughterRows, int[2]); //! McParticles daughter slice (first, last) as stored in the input row +DECLARE_SOA_COLUMN(OmE, omE, float); +DECLARE_SOA_COLUMN(OmWeight, omWeight, float); +} // namespace omega2012ml + +DECLARE_SOA_TABLE(Omega2012MlEventMcs, "AOD", "OMMLEVENTMC", + o2::soa::Index<>, omega2012ml::Omega2012MlEventId, omega2012ml::OmMcCollisionRow); +DECLARE_SOA_TABLE(Omega2012MlCascadeMcs, "AOD", "OMMLCASCMC", + o2::soa::Index<>, omega2012ml::Omega2012MlCascadeId, omega2012ml::OmCascadeMcLabels); +DECLARE_SOA_TABLE(Omega2012MlV0Mcs, "AOD", "OMMLV0MC", + o2::soa::Index<>, omega2012ml::Omega2012MlV0Id, omega2012ml::OmV0McLabels); +DECLARE_SOA_TABLE(Omega2012MlTrackMcs, "AOD", "OMMLTRACKMC", + o2::soa::Index<>, omega2012ml::Omega2012MlTrackId, omega2012ml::OmMcLabel); +DECLARE_SOA_TABLE(Omega2012MlMcParticles, "AOD", "OMMLMCPART", + o2::soa::Index<>, omega2012ml::OmMcRow, omega2012ml::OmMcSelection, omega2012ml::OmPdg, + omega2012ml::OmMcCollisionRow, omega2012ml::OmNMothers, omega2012ml::OmMotherRows, + omega2012ml::OmDaughterRows, omega2012ml::OmPx, omega2012ml::OmPy, omega2012ml::OmPz, + omega2012ml::OmE, omega2012ml::OmWeight); } // namespace o2::aod #endif // PWGLF_DATAMODEL_LFOMEGA2012MLTABLES_H_ diff --git a/PWGLF/Tasks/Resonances/omega2012TrainingTable.cxx b/PWGLF/Tasks/Resonances/omega2012TrainingTable.cxx index 18502ccedb1..3d2a397f6d6 100644 --- a/PWGLF/Tasks/Resonances/omega2012TrainingTable.cxx +++ b/PWGLF/Tasks/Resonances/omega2012TrainingTable.cxx @@ -14,9 +14,12 @@ /// \author Bong-Hwi Lim /// /// Mode A (XiK0s): processXiK0s, processXiK0sMC. Mode B (Xi1530K, Xi- pi+ K- with micro tracks): -/// processXi1530KMicro, processXi1530KMCMicro. Generated parents: processMCTrue. +/// processXi1530KMicro, processXi1530KMCMicro. Generated parents: processMCTrue. MC truth of reconstructed MC: processMcTruth. /// Each candidate is written only to the tables of its own mode (PWGLF/DataModel/LFOmega2012MlTables.h); /// the selection is the one of omega2012Analysis.cxx (PWGLF/Core/Omega2012AnalysisCore.h, same JSON keys). +/// Only collisions with written candidates get an event row. In reconstructed MC the daughter labels of the written +/// objects and the pruned McParticles rows (Omega(2012) decay trees, ancestry of the labels) carry the truth, so the +/// output needs no source table. #include "PWGLF/Core/Omega2012AnalysisCore.h" #include "PWGLF/Core/Omega2012MlFeatures.h" @@ -24,6 +27,7 @@ #include "PWGLF/DataModel/LFOmega2012MlTables.h" #include "PWGLF/DataModel/LFResonanceTables.h" +#include #include #include #include @@ -38,13 +42,19 @@ #include +#include #include #include #include #include +#include +#include #include #include #include +#include +#include +#include using namespace o2; using namespace o2::framework; @@ -59,7 +69,7 @@ static_assert(std::extent_v == ml::Xi struct Omega2012TrainingTable { using ResoCollisions = aod::ResoCollisions_001; - using ResoMCCols = soa::Join; + using ResoMCCols = soa::Join; using ResoMicroTracks = aod::ResoMicroTracks_001; using ResoMCMicroTracks = soa::Join; using ResoMCCascades = soa::Join; @@ -73,6 +83,12 @@ struct Omega2012TrainingTable { static constexpr unsigned int BitsPerFloat = 32; static constexpr uint32_t ByteMask = 0xffU; static constexpr int NBuildStatus = 6; + // Pruned McParticles: generations of Omega(2012) descendants and mother steps above a daughter label + static constexpr int McTreeDepth = 3; + static constexpr int MaxAncestorSteps = 8; + static constexpr uint8_t McSelectionTree = 1; + static constexpr uint8_t McSelectionAncestry = 2; + static constexpr std::size_t NStoredMothers = 2; Produces mlEvents; Produces mlCascades; @@ -85,6 +101,11 @@ struct Omega2012TrainingTable { Produces mlXi1530KInputs; Produces mlXi1530KTruth; Produces mlGenAudit; + Produces mlEventMcs; + Produces mlCascadeMcs; + Produces mlV0Mcs; + Produces mlTrackMcs; + Produces mlMcParticles; HistogramRegistry histos{"histos", {}, OutputObjHandlingPolicy::AnalysisObject}; @@ -103,8 +124,26 @@ struct Omega2012TrainingTable { Configurable omega2012MlParityRows{"omega2012MlParityRows", 0, "Log feature hashes of the first N exported candidates of each mode"}; Configurable omega2012MlExportWrongSign{"omega2012MlExportWrongSign", false, "Also export the wrong-sign charge-pattern controls of the Xi1530K mode"}; + // Event row written before the first candidate of the collision + struct PendingEvent { + DecayMode mode = DecayMode::XiK0s; + int64_t recoCollision = -1; + float posZ = 0.f; + float bField = 0.f; + float centrality = 0.f; + float multiplicity = 0.f; + bool recINELgt0 = false; + bool isMC = false; + int64_t mcCollision = -1; + }; + Omega2012AnalysisCore core; + PendingEvent pendingEvent; int64_t mlEventRow = -1; + // MC tables of the current DF, set only by the reconstructed-MC process functions + aod::McParticles const* mcParticlesDf = nullptr; + aod::McTrackLabels const* mcTrackLabelsDf = nullptr; + std::unordered_set mlAncestryRows; std::unordered_map mlCascadeRows; std::unordered_map mlV0Rows; std::unordered_map mlTrackRows; @@ -124,6 +163,9 @@ struct Omega2012TrainingTable { if (doprocessMCTrue && !doprocessXiK0sMC && !doprocessXi1530KMCMicro) { LOG(fatal) << "processMCTrue requires processXiK0sMC or processXi1530KMCMicro"; } + if (doprocessMcTruth != (doprocessXiK0sMC || doprocessXi1530KMCMicro)) { + LOG(fatal) << "processMcTruth must be enabled exactly when processXiK0sMC or processXi1530KMCMicro is enabled"; + } if (omega2012MlParityRows < 0) { LOG(fatal) << "omega2012MlParityRows must not be negative"; } @@ -176,15 +218,84 @@ struct Omega2012TrainingTable { } template - void writeEvent(Collision const& collision, DecayMode mode) + void beginEvent(Collision const& collision, DecayMode mode) { mlCascadeRows.clear(); mlV0Rows.clear(); mlTrackRows.clear(); + mlEventRow = -1; // ResoCollisions_001 carries no run number or BC; the reduced collision row identifies the event within its DF. - mlEvents(static_cast(mode), static_cast(collision.globalIndex()), - collision.posZ(), collision.bMagField(), collision.cent(), collision.multiplicity(), collision.isRecINELgt0()); + pendingEvent = PendingEvent{mode, static_cast(collision.globalIndex()), collision.posZ(), collision.bMagField(), + collision.cent(), collision.multiplicity(), collision.isRecINELgt0(), false, -1}; + if constexpr (requires { collision.originalMcCollisionId(); }) { + pendingEvent.isMC = true; + pendingEvent.mcCollision = static_cast(collision.originalMcCollisionId()); + } + } + + void ensureEventRow() + { + if (mlEventRow >= 0) { + return; + } + mlEvents(static_cast(pendingEvent.mode), pendingEvent.recoCollision, pendingEvent.posZ, pendingEvent.bField, + pendingEvent.centrality, pendingEvent.multiplicity, pendingEvent.recINELgt0); mlEventRow = mlEvents.lastIndex(); + if (pendingEvent.isMC) { + mlEventMcs(mlEventRow, pendingEvent.mcCollision); + } + } + + template + void writeMcParticle(Particle const& particle, uint8_t selection) + { + const auto mothers = particle.mothersIds(); + std::array motherRows{-1, -1}; + for (std::size_t i = 0; i < NStoredMothers && i < mothers.size(); ++i) { + motherRows[i] = mothers[i]; + } + const auto daughters = particle.daughtersIds(); + std::array daughterRows{daughters[0], daughters[1]}; + const auto nMothers = static_cast(std::min(mothers.size(), std::numeric_limits::max())); + mlMcParticles(static_cast(particle.globalIndex()), selection, particle.pdgCode(), static_cast(particle.mcCollisionId()), + nMothers, motherRows.data(), daughterRows.data(), particle.px(), particle.py(), particle.pz(), particle.e(), particle.weight()); + } + + // McParticles row of a source track, -1 without MC tables or label + int mcLabel(int64_t trackId) const + { + if (mcTrackLabelsDf == nullptr) { + return -1; + } + const auto row = trackId - static_cast(mcTrackLabelsDf->offset()); + if (row < 0 || row >= static_cast(mcTrackLabelsDf->size())) { + return -1; + } + return mcTrackLabelsDf->rawIteratorAt(row).mcParticleId(); + } + + // Writes the label particle and its first-mother chain, once per DF + void writeAncestry(int label) + { + if (mcParticlesDf == nullptr) { + return; + } + int64_t row = label; + for (int step = 0; step <= MaxAncestorSteps; ++step) { + const auto local = row - static_cast(mcParticlesDf->offset()); + if (row < 0 || local < 0 || local >= static_cast(mcParticlesDf->size())) { + return; + } + const auto particle = mcParticlesDf->rawIteratorAt(local); + if (mlAncestryRows.insert(row).second) { + writeMcParticle(particle, McSelectionAncestry); + } + const auto mothers = particle.mothersIds(); + if (mothers.empty()) { + return; + } + row = mothers[0]; + } } template @@ -209,6 +320,14 @@ struct Omega2012TrainingTable { xi.nCrossedRowsPos(), xi.nCrossedRowsNeg(), xi.nCrossedRowsBach(), daughterIds.data()); const auto row = mlCascades.lastIndex(); mlCascadeRows.emplace(id, row); + if (mcTrackLabelsDf != nullptr) { + std::array labels{}; + for (std::size_t i = 0; i < labels.size(); ++i) { + labels[i] = mcLabel(daughterIds[i]); + writeAncestry(labels[i]); + } + mlCascadeMcs(row, labels.data()); + } return row; } @@ -231,6 +350,14 @@ struct Omega2012TrainingTable { v0.nCrossedRowsPos(), v0.nCrossedRowsNeg(), daughterIds.data()); const auto row = mlV0s.lastIndex(); mlV0Rows.emplace(id, row); + if (mcTrackLabelsDf != nullptr) { + std::array labels{}; + for (std::size_t i = 0; i < labels.size(); ++i) { + labels[i] = mcLabel(daughterIds[i]); + writeAncestry(labels[i]); + } + mlV0Mcs(row, labels.data()); + } return row; } @@ -246,6 +373,11 @@ struct Omega2012TrainingTable { track.trackSelectionFlags(), track.trackFlags(), track.tpcNClsCrossedRows(), track.itsClusterMap()); const auto row = mlTracks.lastIndex(); mlTrackRows.emplace(id, row); + if (mcTrackLabelsDf != nullptr) { + const auto label = mcLabel(static_cast(track.trackId())); + writeAncestry(label); + mlTrackMcs(row, label); + } return row; } @@ -267,6 +399,7 @@ struct Omega2012TrainingTable { collision.globalIndex(), xi.globalIndex(), v0.globalIndex(), passBits, featureHash(pack.master)); ++parityLoggedXiK0s; } + ensureEventRow(); const auto cascadeRow = writeCascade(xi); const auto v0Row = writeV0(v0); mlXiK0sCandidates(mlEventRow, cascadeRow, v0Row, @@ -308,6 +441,7 @@ struct Omega2012TrainingTable { values.chargePattern, passBits, featureHash(pack.master)); ++parityLoggedXi1530K; } + ensureEventRow(); const auto cascadeRow = writeCascade(xi); const auto pionRow = writeTrack(pion); const auto kaonRow = writeTrack(kaon); @@ -330,7 +464,7 @@ struct Omega2012TrainingTable { template void exportXiK0s(Collision const& collision, Cascades const& cascades, V0s const& v0s) { - writeEvent(collision, DecayMode::XiK0s); + beginEvent(collision, DecayMode::XiK0s); // Selection only: the analysis histograms belong to the Omega(2012) analysis task core.forEachXiK0sCandidate(histos, collision, collision, cascades, v0s, true, nullptr, nullptr, nullptr, [&](auto const& coll, auto const& xi, auto const& v0, XiK0sCandidateValues const& values, uint16_t passBits) { @@ -341,7 +475,7 @@ struct Omega2012TrainingTable { template void exportXi1530K(Collision const& collision, Cascades const& cascades, Tracks const& tracks) { - writeEvent(collision, DecayMode::Xi1530K); + beginEvent(collision, DecayMode::Xi1530K); core.forEachXi1530KCandidate(histos, collision, cascades, tracks, nullptr, true, nullptr, nullptr, nullptr, [&](auto const& coll, auto const& xi, auto const& pion, auto const& kaon, Xi1530KCandidateValues const& values, uint16_t passBits) { @@ -358,12 +492,19 @@ struct Omega2012TrainingTable { } PROCESS_SWITCH(Omega2012TrainingTable, processXiK0s, "Write XiK0s-mode candidates from data", true); - void processXiK0sMC(ResoMCCols::iterator const& collision, ResoMCCascades const& cascades, ResoMCV0s const& v0s) + void processXiK0sMC(ResoMCCols::iterator const& collision, ResoMCCascades const& cascades, ResoMCV0s const& v0s, + aod::McParticles const& mcParticles, aod::McTrackLabels const& mcTrackLabels) { if (!core.passesEventCuts(collision) || !core.passesMCEventCuts(collision)) { return; } + mcParticlesDf = &mcParticles; + mcTrackLabelsDf = &mcTrackLabels; exportXiK0s(collision, cascades, v0s); + // cppcheck-suppress redundantAssignment + mcParticlesDf = nullptr; + // cppcheck-suppress redundantAssignment + mcTrackLabelsDf = nullptr; } PROCESS_SWITCH(Omega2012TrainingTable, processXiK0sMC, "Write XiK0s-mode candidates with truth from reconstructed MC", false); @@ -376,12 +517,19 @@ struct Omega2012TrainingTable { } PROCESS_SWITCH(Omega2012TrainingTable, processXi1530KMicro, "Write Xi1530K-mode candidates from data micro v001 tables", false); - void processXi1530KMCMicro(ResoMCCols::iterator const& collision, ResoMCCascades const& cascades, ResoMCMicroTracks const& tracks) + void processXi1530KMCMicro(ResoMCCols::iterator const& collision, ResoMCCascades const& cascades, ResoMCMicroTracks const& tracks, + aod::McParticles const& mcParticles, aod::McTrackLabels const& mcTrackLabels) { if (!core.passesEventCuts(collision) || !core.passesMCEventCuts(collision)) { return; } + mcParticlesDf = &mcParticles; + mcTrackLabelsDf = &mcTrackLabels; exportXi1530K(collision, cascades, tracks); + // cppcheck-suppress redundantAssignment + mcParticlesDf = nullptr; + // cppcheck-suppress redundantAssignment + mcTrackLabelsDf = nullptr; } PROCESS_SWITCH(Omega2012TrainingTable, processXi1530KMCMicro, "Write Xi1530K-mode candidates with truth from reconstructed MC micro v001 tables", false); @@ -396,6 +544,41 @@ struct Omega2012TrainingTable { }); } PROCESS_SWITCH(Omega2012TrainingTable, processMCTrue, "Write generated Omega(2012) parents of selected reconstructed MC events", false); + + // Once per DF: every Omega(2012) with McTreeDepth generations of descendants. The ancestry set is reset here, so + // the result does not depend on the order of this call relative to the per-collision process functions. + void processMcTruth(aod::McParticles const& mcParticles) + { + mlAncestryRows.clear(); + const auto offset = static_cast(mcParticles.offset()); + const auto nRows = static_cast(mcParticles.size()); + std::unordered_set treeRows; + std::vector> stack; + for (const auto& parent : mcParticles) { + if (std::abs(parent.pdgCode()) != o2::constants::physics::Pdg::kOmega2012Minus) { + continue; + } + stack.assign(1, {static_cast(parent.globalIndex()), 0}); + while (!stack.empty()) { + const auto [row, depth] = stack.back(); + stack.pop_back(); + const auto local = row - offset; + if (local < 0 || local >= nRows || !treeRows.insert(row).second) { + continue; + } + const auto particle = mcParticles.rawIteratorAt(local); + writeMcParticle(particle, McSelectionTree); + const auto daughters = particle.daughtersIds(); + if (depth >= McTreeDepth || daughters[0] < 0 || daughters[1] < daughters[0]) { + continue; + } + for (int64_t daughter = daughters[0]; daughter <= daughters[1]; ++daughter) { + stack.emplace_back(daughter, depth + 1); + } + } + } + } + PROCESS_SWITCH(Omega2012TrainingTable, processMcTruth, "Write the pruned McParticles rows of the Omega(2012) decay trees; required with the reconstructed-MC process functions", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)