From 02434ac3da75bc5185696d781b9b547e0deb1efc Mon Sep 17 00:00:00 2001 From: marcellocosti Date: Tue, 6 Oct 2026 18:51:54 +0200 Subject: [PATCH 1/6] [WIP] ML skimming for D0 and D* --- .../TableProducer/mlBasedTrackSelector.cxx | 620 +++++++++++++++--- 1 file changed, 524 insertions(+), 96 deletions(-) diff --git a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx index 5961630d3d2..cb0174153f9 100644 --- a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx +++ b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx @@ -79,7 +79,7 @@ using namespace o2::framework; using namespace o2::framework::expressions; using namespace o2::constants::physics; -/// Bit positions of aod::HfSelTrack::isSelProng (only Cand3Prong is ever set here). +/// Bit positions of aod::HfSelTrack::isSelProng (only Cand2Prong, Cand3Prong and CandDstar are ever set here). enum CandidateType { Cand2Prong = 0, Cand3Prong, @@ -89,8 +89,14 @@ enum CandidateType { NCandidateTypes }; -/// The two 3-prong channels handled by this task. -enum Channel3Prong { +/// The 2-prong channels handled by this task. +enum Channels2Prong { + ChannelD0ToPiK = 0, + NChannels2Prong +}; + +/// The 3-prong channels handled by this task. +enum Channels3Prong { ChannelDplusToPiKPi = 0, ChannelDsToKKPi, NChannels3Prong @@ -104,35 +110,57 @@ constexpr std::array DecayTypeOfChannel{ /// track-level model, one bit per (channel, role). A track can be pion-like for one channel and /// kaon-like for the other, so the four bits are independent. enum TrackMlRole { - RolePionDplus = 0, + RolePionD0 = 0, + RoleKaonD0, + RolePionDplus, RoleKaonDplus, RolePionDs, RoleKaonDs, + RoleSoftPiDstar, NTrackMlRoles }; constexpr std::array RolePionOfChannel{RolePionDplus, RolePionDs}; constexpr std::array RoleKaonOfChannel{RoleKaonDplus, RoleKaonDs}; -constexpr int NProngs = 3; // prongs of the candidates built here +constexpr int N2Prongs = 2; // Number of prongs for 2-prong candidates +constexpr int N3Prongs = 3; // Number of prongs for 3-prong candidates constexpr int NMassHypos = 2; // mass hypotheses per channel, i.e. the two orderings of the same-sign pair -/// Whether each prong is a kaon, per channel and mass hypothesis. -constexpr std::array, NMassHypos>, NChannels3Prong> IsKaonProng{{ - {{{{false, true, false}}, {{false, true, false}}}}, - {{{{true, true, false}}, {{false, true, true}}}}, -}}; +/// Role of each track in the candidate, per channel and mass hypothesis. +/// The ordering of the channels should strictly follow Channels2Prong +constexpr TrackMlRole trackRoles2Prongs[NChannels2Prong][NMassHypos][N2Prongs] = { + { + {RolePionD0, RoleKaonD0}, // D0 → π± K∓ + {RoleKaonD0, RolePionD0} // D0 → π± K∓ + } +}; + +/// Role of each track in the candidate, per channel and mass hypothesis. +/// The ordering of the channels should strictly follow Channels3Prong +constexpr TrackMlRole trackRoles3Prongs[NChannels3Prong][NMassHypos][N3Prongs] = { + { + {RolePionDplus, RoleKaonDplus, RolePionDplus}, // D± → π± K∓ π± + {RolePionDplus, RoleKaonDplus, RolePionDplus} // D± → π± K∓ π± (swapped, no effect for D+) + }, + { + {RoleKaonDs, RoleKaonDs, RolePionDs}, // D_s± → K± K∓ π± + {RolePionDs, RoleKaonDs, RoleKaonDs} // D_s± → K± K∓ π± (swapped) + } +}; constexpr double PtMaxModel = 1.e10; -constexpr uint32_t MaskSameSignAny = (1u << RolePionDplus) | (1u << RolePionDs) | (1u << RoleKaonDs); -constexpr uint32_t MaskOppSignAny = (1u << RoleKaonDplus) | (1u << RoleKaonDs); +constexpr uint32_t MaskSameSignAny = (1u << RolePionD0) | (1u << RolePionDplus) | (1u << RolePionDs) | (1u << RoleKaonDs) | (1u << RoleSoftPiDstar); +constexpr uint32_t MaskOppSignAny = (1u << RoleKaonD0) | (1u << RoleKaonDplus) | (1u << RoleKaonDs); /// Where the combinatorics spends its time, accumulated in seconds into hTiming. enum TimingStep { TimeLoopTotal = 0, ///< everything inside the triple loop TimeFit2Prong, ///< the 2-track vertex used to skip pairs early + TimeProcessD0, ///< D0 candidate processing time TimeFit3Prong, ///< the 3-track vertex of the surviving triplets + TimeProcessDstar, ///< D* candidate processing time TimeCacheFill, ///< building the per-collision prong cache, propagation included TimeCachePropagate, ///< only the re-propagation inside that build NTimingSteps @@ -143,9 +171,12 @@ enum LoopCounter { CountPairsSeen = 0, ///< (same-sign, opposite-sign) pairs reached CountPairsCandidateOk, ///< ... of those, surviving the ML pair test CountFit2Prong, ///< ... of those, handed to the 2-prong fitter + CountPairsWritten, ///< ... of those, accepted as 2-prong candidates CountPairsRejected2P, ///< ... of those, dropped by the 2-track vertex CountTripletsEnumerated, ///< triplets built from the surviving pairs CountTripletsWritten, ///< ... of those, written to the skim + CountDstarsEnumerated, ///< triplets built from the surviving pairs + CountDstarsWritten, ///< ... of those, written to the skim CountTracksCached, ///< associations put into the prong cache CountTracksPropagated, ///< ... of those, needing a re-propagation to this collision NLoopCounters @@ -155,8 +186,10 @@ enum LoopCounter { enum TrackTimingStep { TrackTimeTotal = 0, ///< the whole per-collision loop: slicing, features, models, table filling TrackTimeFeatures, ///< building the input features, re-propagation included + TrackTimeModelD0, ///< evaluating the D0 model: input vector, ONNX call, thresholds TrackTimeModelDplus, ///< evaluating the D+ model: input vector, ONNX call, thresholds TrackTimeModelDs, ///< evaluating the Ds model + TrackTimeModelDstar, ///< evaluating the D* model NTrackTimingSteps }; @@ -313,7 +346,14 @@ struct HfTrackSelectorTagSelTracks { Configurable fillHistograms{"fillHistograms", true, "fill histograms"}; Configurable ptMinTrack{"ptMinTrack", 0.3f, "min. track pT entering the charm combinatorics"}; Configurable etaMaxTrack{"etaMaxTrack", 0.8f, "max. track eta entering the charm combinatorics"}; + Configurable ptMinSoftPi{"ptMinSoftPi", 0.1f, "min. soft pion track pT entering the charm combinatorics"}; Configurable enableTiming{"enableTiming", false, "fill hTiming with the CPU of the feature building and of each model evaluation (adds two clock reads per call)"}; + // D0 model + Configurable applyMlD0{"applyMlD0", true, "evaluate the D0 track model"}; + Configurable onnxFileNameD0{"onnxFileNameD0", "ModelHandler_D0Tracks.onnx", "ONNX file name of the D0 track model"}; + Configurable> inputFeaturesD0{"inputFeaturesD0", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D0 model, in the order it expects"}; + Configurable thresholdScorePionD0{"thresholdScorePionD0", 0., "min. D0 pion-class score"}; + Configurable thresholdScoreKaonD0{"thresholdScoreKaonD0", 0., "min. D0 kaon-class score"}; // D+ model Configurable applyMlDplus{"applyMlDplus", true, "evaluate the D+ track model"}; Configurable onnxFileNameDplus{"onnxFileNameDplus", "ModelHandler_DplusTracks.onnx", "ONNX file name of the D+ track model"}; @@ -326,10 +366,17 @@ struct HfTrackSelectorTagSelTracks { Configurable> inputFeaturesDs{"inputFeaturesDs", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the Ds model, in the order it expects"}; Configurable thresholdScorePionDs{"thresholdScorePionDs", 0., "min. Ds pion-class score"}; Configurable thresholdScoreKaonDs{"thresholdScoreKaonDs", 0., "min. Ds kaon-class score"}; + // D* model + Configurable applyMlDstar{"applyMlDstar", true, "evaluate the D* track model"}; + Configurable onnxFileNameDstar{"onnxFileNameDstar", "ModelHandler_DstarTracks.onnx", "ONNX file name of the D* track model"}; + Configurable> inputFeaturesDstar{"inputFeaturesDstar", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D* model, in the order it expects"}; + Configurable thresholdScorePionDstar{"thresholdScorePionDstar", 0., "min. D* pion-class score"}; // ONNX runtime Configurable loadModelsFromCcdb{"loadModelsFromCcdb", false, "load the ONNX models from CCDB instead of a local path"}; + Configurable mlModelPathCcdbD0{"mlModelPathCcdbD0", "path/to/ml/models/D0", "CCDB path of the D0 track model"}; Configurable mlModelPathCcdbDplus{"mlModelPathCcdbDplus", "path/to/ml/models/Dplus", "CCDB path of the D+ track model"}; Configurable mlModelPathCcdbDs{"mlModelPathCcdbDs", "path/to/ml/models/Ds", "CCDB path of the Ds track model"}; + Configurable mlModelPathCcdbDstar{"mlModelPathCcdbDstar", "path/to/ml/models/Dstar", "CCDB path of the D* track model"}; Configurable timestampCcdbForMlModels{"timestampCcdbForMlModels", -1, "timestamp of the ONNX files to be queried in CCDB"}; Configurable enableOnnxOptimizations{"enableOnnxOptimizations", true, "enable the ONNX graph optimisations"}; Configurable onnxThreads{"onnxThreads", 1, "number of threads used by the ONNX runtime (0 = let onnxruntime decide)"}; @@ -347,7 +394,9 @@ struct HfTrackSelectorTagSelTracks { double thresholdKaon{0.}; bool enabled{false}; }; - std::array models; + std::array models2Prong; + std::array models3Prong; + HfTrackModel modelSoftPiDstar; Service ccdb{}; o2::ccdb::CcdbApi ccdbApi; @@ -433,16 +482,28 @@ struct HfTrackSelectorTagSelTracks { LOGP(fatal, "One and only one process function of HfTrackSelectorTagSelTracks can be enabled at a time!"); } + if (config.applyMlD0) { + configureModel(models2Prong[ChannelD0ToPiK], config.onnxFileNameD0, config.mlModelPathCcdbD0, config.inputFeaturesD0, + config.thresholdScorePionD0, config.thresholdScoreKaonD0, "D0 track model"); + } if (config.applyMlDplus) { - configureModel(models[ChannelDplusToPiKPi], config.onnxFileNameDplus, config.mlModelPathCcdbDplus, config.inputFeaturesDplus, + configureModel(models3Prong[ChannelDplusToPiKPi], config.onnxFileNameDplus, config.mlModelPathCcdbDplus, config.inputFeaturesDplus, config.thresholdScorePionDplus, config.thresholdScoreKaonDplus, "D+ track model"); } if (config.applyMlDs) { - configureModel(models[ChannelDsToKKPi], config.onnxFileNameDs, config.mlModelPathCcdbDs, config.inputFeaturesDs, + configureModel(models3Prong[ChannelDsToKKPi], config.onnxFileNameDs, config.mlModelPathCcdbDs, config.inputFeaturesDs, config.thresholdScorePionDs, config.thresholdScoreKaonDs, "Ds track model"); } - if (!models[ChannelDplusToPiKPi].enabled && !models[ChannelDsToKKPi].enabled) { - LOGP(fatal, "At least one of the D+ and Ds track models must be enabled!"); + if (config.applyMlDstar) { + double thresholdScoreKaonDstar = 0; // D* model has no kaon score, but the configureModel() function expects a value + configureModel(modelSoftPiDstar, config.onnxFileNameDstar, config.mlModelPathCcdbDstar, config.inputFeaturesDstar, + config.thresholdScorePionDstar, thresholdScoreKaonDstar, "D* track model"); + } + if (!models2Prong[ChannelD0ToPiK].enabled && + !models3Prong[ChannelDplusToPiKPi].enabled && + !models3Prong[ChannelDsToKKPi].enabled && + !modelSoftPiDstar.enabled) { + LOGP(fatal, "At least one of the D0, D+ and Ds track models must be enabled!"); } ccdb->setURL(config.ccdbUrl); @@ -458,19 +519,30 @@ struct HfTrackSelectorTagSelTracks { registry.add("hPtNoCuts", "all track associations;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); registry.add("hPtQuality", "track associations passing the pT floor and the quality cut, scored;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); + registry.add("hPtSoftPiQuality", "track associations passing the pT floor and the quality cut, scored;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); registry.add("hPtQualityRejColl", "track associations passing the pT floor and the quality cut, not scored: collision fails the HF event selection;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); - registry.add("hPtSelected", "track associations selected by at least one model;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); - registry.add("hEtaSelected", "track associations selected by at least one model;#it{#eta};entries", {HistType::kTH1D, {axisEta}}); + registry.add("hPtQualitySoftPiRejColl", "track associations passing the pT floor and the quality cut, not scored: collision fails the HF event selection;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); + registry.add("hPtSelected2Prong", "track associations selected by at least one model;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); + registry.add("hEtaSelected2Prong", "track associations selected by at least one model;#it{#eta};entries", {HistType::kTH1D, {axisEta}}); + registry.add("hPtSelected3Prong", "track associations selected by at least one model;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); + registry.add("hEtaSelected3Prong", "track associations selected by at least one model;#it{#eta};entries", {HistType::kTH1D, {axisEta}}); + registry.add("hPtSelectedSoftPi", "track associations selected as D* soft pions;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPtProng}}); + registry.add("hEtaSelectedSoftPi", "track associations selected as D* soft pions;#it{#eta};entries", {HistType::kTH1D, {axisEta}}); + registry.add("hScorePionD0", "D^{0} pion-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); + registry.add("hScoreKaonD0", "D^{0} kaon-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); registry.add("hScorePionDplus", "D^{#plus} pion-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); registry.add("hScoreKaonDplus", "D^{#plus} kaon-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); registry.add("hScorePionDs", "D_{s}^{#plus} pion-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); registry.add("hScoreKaonDs", "D_{s}^{#plus} kaon-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); + registry.add("hScoreSoftPiDstar", "D^{*} soft pion-class score;#it{p}_{T}^{track} (GeV/#it{c});score;entries", {HistType::kTH2D, {axisPtProng, axisScore}}); // seconds, accumulated over the run via Fill(bin, weight); per-call cost = bin / hPtQuality entries auto hTiming = registry.add("hTiming", "CPU of the ML track selection;;seconds", {HistType::kTH1D, {{NTrackTimingSteps, -0.5f, static_cast(NTrackTimingSteps) - 0.5f}}}); hTiming->GetXaxis()->SetBinLabel(TrackTimeTotal + 1, "track loop total"); hTiming->GetXaxis()->SetBinLabel(TrackTimeFeatures + 1, "features"); + hTiming->GetXaxis()->SetBinLabel(TrackTimeModelD0 + 1, "D0 model"); hTiming->GetXaxis()->SetBinLabel(TrackTimeModelDplus + 1, "D+ model"); hTiming->GetXaxis()->SetBinLabel(TrackTimeModelDs + 1, "Ds model"); + hTiming->GetXaxis()->SetBinLabel(TrackTimeModelDstar + 1, "D* model"); } doTiming = config.enableTiming && config.fillHistograms; } @@ -554,7 +626,7 @@ struct HfTrackSelectorTagSelTracks { // one model for the whole pT range, so the bin index is always 0 const auto output = model.response.getModelOutput(inputFeatures, 0); if (output.size() != scores.size()) { - LOGP(fatal, "the model emitted {} scores, expected {}: only 3-class (bkg/pi/K) models are supported", output.size(), scores.size()); + LOGP(fatal, "the model emitted {} scores, expected {}: only 3-class (bkg/pi/K) models are supported for D0, D+, Ds, only 2-class models are supported for D* candidates", output.size(), scores.size()); } std::copy(output.begin(), output.end(), scores.begin()); @@ -579,6 +651,10 @@ struct HfTrackSelectorTagSelTracks { for (const auto& trackId : trackIndicesCollision) { const auto track = trackId.template track_as(); + const bool isPositive = track.sign() > 0; + if (config.fillHistograms) { + registry.fill(HIST("hPtNoCuts"), track.pt()); + } // One row per track association, always: aod::HfSelTrack is joined with aod::TrackAssoc // downstream, so the two tables have to stay index-aligned. Rejected associations are @@ -586,42 +662,89 @@ struct HfTrackSelectorTagSelTracks { uint32_t statusProng{0}; uint32_t isIdentifiedPid{0}; - if (config.fillHistograms) { - registry.fill(HIST("hPtNoCuts"), track.pt()); + const bool passesQuality = track.pt() >= config.ptMinTrack && std::abs(track.eta()) <= config.etaMaxTrack && track.isGlobalTrackWoDCA(); + const bool passesQualitySoftPi = track.pt() >= config.ptMinSoftPi && std::abs(track.eta()) <= config.etaMaxTrack && track.isGlobalTrackWoDCA(); + if (!scoreCollision) { + if (config.fillHistograms && passesQuality) { + registry.fill(HIST("hPtQualityRejColl"), track.pt()); + } + if (config.fillHistograms && passesQualitySoftPi) { + registry.fill(HIST("hPtQualitySoftPiRejColl"), track.pt()); + } + rowSelectedTrack(statusProng, isIdentifiedPid, isPositive); + continue; } - const bool passesQuality = track.pt() >= config.ptMinTrack && std::abs(track.eta()) <= config.etaMaxTrack && track.isGlobalTrackWoDCA(); - if (passesQuality && !scoreCollision && config.fillHistograms) { - registry.fill(HIST("hPtQualityRejColl"), track.pt()); + HfTrackMlFeatures features{}; + auto tStart = tickIf(doTiming); + fillFeatures(collision, track, centrality, features); + if (doTiming) { + timing[TrackTimeFeatures] += elapsedSeconds(tStart); } - if (passesQuality && scoreCollision) { + std::array scores{}; + bool isSelPion{false}, isSelKaon{false}; + bool from2Prong{false}, from3Prong{false}, fromDstar{false}; + + if (passesQualitySoftPi && modelSoftPiDstar.enabled) { if (config.fillHistograms) { - registry.fill(HIST("hPtQuality"), track.pt()); + registry.fill(HIST("hPtSoftPiQuality"), track.pt()); } - HfTrackMlFeatures features{}; - auto tStart = tickIf(doTiming); - fillFeatures(collision, track, centrality, features); + tStart = tickIf(doTiming); + const bool evaluatedDstar = evaluateModel(modelSoftPiDstar, features, scores, isSelPion, isSelKaon); if (doTiming) { - timing[TrackTimeFeatures] += elapsedSeconds(tStart); + timing[TrackTimeModelDstar] += elapsedSeconds(tStart); } + if (evaluatedDstar) { + if (isSelPion) { + SETBIT(isIdentifiedPid, RoleSoftPiDstar); + fromDstar = true; + } + if (config.fillHistograms) { + registry.fill(HIST("hScoreSoftPiDstar"), features.pt, scores[1]); + } + } + } - std::array scores{}; - bool isSelPion{false}; - bool isSelKaon{false}; + if (passesQuality) { + if (config.fillHistograms) { + registry.fill(HIST("hPtQuality"), track.pt()); + } + + tStart = tickIf(doTiming); + const bool evaluatedD0 = evaluateModel(models2Prong[ChannelD0ToPiK], features, scores, isSelPion, isSelKaon); + if (doTiming) { + timing[TrackTimeModelD0] += elapsedSeconds(tStart); + } + if (evaluatedD0) { + if (isSelPion) { + SETBIT(isIdentifiedPid, RolePionD0); + from2Prong = true; + } + if (isSelKaon) { + SETBIT(isIdentifiedPid, RoleKaonD0); + from2Prong = true; + } + if (config.fillHistograms) { + registry.fill(HIST("hScorePionD0"), features.pt, scores[1]); + registry.fill(HIST("hScoreKaonD0"), features.pt, scores[2]); + } + } tStart = tickIf(doTiming); - const bool evaluatedDplus = evaluateModel(models[ChannelDplusToPiKPi], features, scores, isSelPion, isSelKaon); + const bool evaluatedDplus = evaluateModel(models3Prong[ChannelDplusToPiKPi], features, scores, isSelPion, isSelKaon); if (doTiming) { timing[TrackTimeModelDplus] += elapsedSeconds(tStart); } if (evaluatedDplus) { if (isSelPion) { SETBIT(isIdentifiedPid, RolePionDplus); + from3Prong = true; } if (isSelKaon) { SETBIT(isIdentifiedPid, RoleKaonDplus); + from3Prong = true; } if (config.fillHistograms) { registry.fill(HIST("hScorePionDplus"), features.pt, scores[1]); @@ -630,34 +753,48 @@ struct HfTrackSelectorTagSelTracks { } tStart = tickIf(doTiming); - const bool evaluatedDs = evaluateModel(models[ChannelDsToKKPi], features, scores, isSelPion, isSelKaon); + const bool evaluatedDs = evaluateModel(models3Prong[ChannelDsToKKPi], features, scores, isSelPion, isSelKaon); if (doTiming) { timing[TrackTimeModelDs] += elapsedSeconds(tStart); } if (evaluatedDs) { if (isSelPion) { SETBIT(isIdentifiedPid, RolePionDs); + from3Prong = true; } if (isSelKaon) { SETBIT(isIdentifiedPid, RoleKaonDs); + from3Prong = true; } if (config.fillHistograms) { registry.fill(HIST("hScorePionDs"), features.pt, scores[1]); registry.fill(HIST("hScoreKaonDs"), features.pt, scores[2]); } } + } - // the track enters the combinatorics if it can play any role in any of the two channels - if (isIdentifiedPid != 0u) { - SETBIT(statusProng, CandidateType::Cand3Prong); - if (config.fillHistograms) { - registry.fill(HIST("hPtSelected"), features.pt); - registry.fill(HIST("hEtaSelected"), features.eta); - } + // the track enters the combinatorics if it can play any role in any of the channels + if (from2Prong) { + SETBIT(statusProng, CandidateType::Cand2Prong); + if (config.fillHistograms) { + registry.fill(HIST("hPtSelected2Prong"), features.pt); + registry.fill(HIST("hEtaSelected2Prong"), features.eta); + } + } + if (from3Prong) { + SETBIT(statusProng, CandidateType::Cand3Prong); + if (config.fillHistograms) { + registry.fill(HIST("hPtSelected3Prong"), features.pt); + registry.fill(HIST("hEtaSelected3Prong"), features.eta); + } + } + if (fromDstar) { + SETBIT(statusProng, CandidateType::CandDstar); + if (config.fillHistograms) { + registry.fill(HIST("hPtSelectedSoftPi"), features.pt); + registry.fill(HIST("hEtaSelectedSoftPi"), features.eta); } } - - const bool isPositive = track.sign() > 0; rowSelectedTrack(statusProng, isIdentifiedPid, isPositive); } } @@ -712,10 +849,14 @@ struct HfTrackSelectorTagSelTracks { /// Pre-selection of 3-prong secondary vertices struct HfMlBasedTrackSelector { + Produces rowTrackIndexProng2; Produces rowTrackIndexProng3; + Produces rowTrackIndexDstar; struct : ConfigurableGroup { Configurable fillHistograms{"fillHistograms", true, "fill histograms"}; + Configurable do2Prongs{"do2Prongs", true, "store 2-prong candidates"}; + Configurable doDstar{"doDstar", true, "store D* candidates"}; // preselection Configurable ptTolerance{"ptTolerance", 0.1, "pT tolerance in GeV/c for applying preselections before vertex reconstruction"}; // preselection of 3-prongs using the decay length computed only with the first two tracks @@ -736,12 +877,18 @@ struct HfMlBasedTrackSelector { Configurable ccdbPathGrp{"ccdbPathGrp", "GLO/GRP/GRP", "Path of the grp file (Run 2)"}; Configurable ccdbPathGrpMag{"ccdbPathGrpMag", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object (Run 3)"}; + // D0 cuts + Configurable> binsPtD0ToPiK{"binsPtD0ToPiK", std::vector{hf_cuts_presel_2prong::vecBinsPt}, "pT bin limits for D0->piK pT-dependent cuts"}; + Configurable> cutsD0ToPiK{"cutsD0ToPiK", {hf_cuts_presel_2prong::Cuts[0], hf_cuts_presel_2prong::NBinsPt, hf_cuts_presel_2prong::NCutVars, hf_cuts_presel_2prong::labelsPt, hf_cuts_presel_2prong::labelsCutVar}, "D0->piK selections per pT bin"}; // D+ cuts Configurable> binsPtDplusToPiKPi{"binsPtDplusToPiKPi", std::vector{hf_cuts_presel_3prong::vecBinsPt}, "pT bin limits for D+->piKpi pT-dependent cuts"}; Configurable> cutsDplusToPiKPi{"cutsDplusToPiKPi", {hf_cuts_presel_3prong::Cuts[0], hf_cuts_presel_3prong::NBinsPt, hf_cuts_presel_3prong::NCutVars, hf_cuts_presel_3prong::labelsPt, hf_cuts_presel_3prong::labelsCutVar}, "D+->piKpi selections per pT bin"}; // Ds+ cuts Configurable> binsPtDsToKKPi{"binsPtDsToKKPi", std::vector{hf_cuts_presel_ds::vecBinsPt}, "pT bin limits for Ds+->KKPi pT-dependent cuts"}; Configurable> cutsDsToKKPi{"cutsDsToKKPi", {hf_cuts_presel_ds::Cuts[0], hf_cuts_presel_ds::NBinsPt, hf_cuts_presel_ds::NCutVars, hf_cuts_presel_ds::labelsPt, hf_cuts_presel_ds::labelsCutVar}, "Ds+->KKPi selections per pT bin"}; + // D*+ cuts + Configurable> binsPtDstarToD0Pi{"binsPtDstarToD0Pi", std::vector{hf_cuts_presel_dstar::vecBinsPt}, "pT bin limits for D*+->D0pi pT-dependent cuts"}; + Configurable> cutsDstarToD0Pi{"cutsDstarToD0Pi", {hf_cuts_presel_dstar::Cuts[0], hf_cuts_presel_dstar::NBinsPt, hf_cuts_presel_dstar::NCutVars, hf_cuts_presel_dstar::labelsPt, hf_cuts_presel_dstar::labelsCutVar}, "D*+->D0pi selections per pT bin"}; } config; SliceCache cache; @@ -754,8 +901,13 @@ struct HfMlBasedTrackSelector { // masses of the two mass hypotheses of each channel, in the (same-sign, opposite-sign, // same-sign) prong ordering used by the combinatorics - std::array, NMassHypos>, NChannels3Prong> arrMass3Prong{}; + std::array, NMassHypos>, NChannels3Prong> arrMass3Prong{}; + std::array, NMassHypos>, NChannels2Prong> arrMass2Prong{}; // cuts and pT binning, one entry per channel + std::array, NChannels2Prong> cut2Prong{}; + std::array, NChannels2Prong> binsPt2Prong{}; + LabeledArray cutDstar{}; + std::vector binsPtDstar{}; std::array, NChannels3Prong> cut3Prong{}; std::array, NChannels3Prong> binsPt3Prong{}; @@ -763,6 +915,7 @@ struct HfMlBasedTrackSelector { struct HfProngCandidate { o2::track::TrackParCov parCov; std::array pVec{}; + std::array dcaInfo{}; ///< DCA (xy, z) to the primary vertex uint32_t mask{}; ///< aod::HfSelTrack::isIdentifiedPid int64_t globalIndex{}; ///< index in the track table, written to the skim }; @@ -770,8 +923,8 @@ struct HfMlBasedTrackSelector { /// Cache of propagated tracks for one collision, split by charge. struct HfProngCache { std::vector prongs; - std::vector sameSign; ///< can sit in a same-sign slot: pi(D+), pi(Ds) or K(Ds) - std::vector oppSign; ///< can sit in the opposite-sign slot: K(D+) or K(Ds) + std::vector sameSign; ///< can sit in a same-sign slot: pi(D0), pi(D+), pi(Ds) or K(Ds), softpi(D*) + std::vector oppSign; ///< can sit in the opposite-sign slot: K(D0) K(D+) or K(Ds) void clear() { @@ -794,28 +947,39 @@ struct HfMlBasedTrackSelector { // filter collisions Filter filterSelectCollisions = (aod::hf_sel_collision::whyRejectColl == static_cast(0)); // filter track indices - Filter filterSelectTrackIds = ((aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u); + Filter filterSelectTrackIds = ( (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand2Prong))) != 0u || + (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u || + (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::CandDstar))) != 0u ); Preslice trackIndicesPerCollision = aod::track_association::collisionId; - Partition positiveFor3Prongs = aod::hf_sel_track::isPositive == true && ((aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u); - Partition negativeFor3Prongs = aod::hf_sel_track::isPositive == false && ((aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u); + Partition positiveHfTracks = (aod::hf_sel_track::isPositive == true); + Partition negativeHfTracks = (aod::hf_sel_track::isPositive == false); HistogramRegistry registry{"registry"}; void init(InitContext const&) { - if (!doprocess3Prongs) { + if (!doprocessCandidates) { return; } doTiming = config.enableTiming && config.fillHistograms; + arrMass2Prong[ChannelD0ToPiK] = std::array{std::array{MassPiPlus, MassKPlus}, + std::array{MassKPlus, MassPiPlus}}; + arrMass3Prong[ChannelDplusToPiKPi] = std::array{std::array{MassPiPlus, MassKPlus, MassPiPlus}, std::array{MassPiPlus, MassKPlus, MassPiPlus}}; arrMass3Prong[ChannelDsToKKPi] = std::array{std::array{MassKPlus, MassKPlus, MassPiPlus}, std::array{MassPiPlus, MassKPlus, MassKPlus}}; + // cuts retrieved by json, in the order of Channel3Prong + cut2Prong = {config.cutsD0ToPiK}; + binsPt2Prong = {config.binsPtD0ToPiK}; + cutDstar = {config.cutsDstarToD0Pi}; + binsPtDstar = {config.binsPtDstarToD0Pi}; + // cuts retrieved by json, in the order of Channel3Prong cut3Prong = {config.cutsDplusToPiKPi, config.cutsDsToKKPi}; binsPt3Prong = {config.binsPtDplusToPiKPi, config.binsPtDsToKKPi}; @@ -850,36 +1014,76 @@ struct HfMlBasedTrackSelector { auto hTiming = registry.add("hTiming", "CPU inside the triple loop;;seconds", {HistType::kTH1D, {{NTimingSteps, -0.5f, static_cast(NTimingSteps) - 0.5f}}}); hTiming->GetXaxis()->SetBinLabel(TimeLoopTotal + 1, "triple loop total"); hTiming->GetXaxis()->SetBinLabel(TimeFit2Prong + 1, "2-prong vertex fit"); + hTiming->GetXaxis()->SetBinLabel(TimeProcessD0 + 1, "D0 processing"); hTiming->GetXaxis()->SetBinLabel(TimeFit3Prong + 1, "3-prong vertex fit"); + hTiming->GetXaxis()->SetBinLabel(TimeProcessDstar + 1, "D* processing"); hTiming->GetXaxis()->SetBinLabel(TimeCacheFill + 1, "prong cache fill"); hTiming->GetXaxis()->SetBinLabel(TimeCachePropagate + 1, "cache re-propagation"); auto hLoopCounters = registry.add("hLoopCounters", "Triple loop stages;;entries", {HistType::kTH1D, {{NLoopCounters, -0.5f, static_cast(NLoopCounters) - 0.5f}}}); hLoopCounters->GetXaxis()->SetBinLabel(CountPairsSeen + 1, "pairs seen"); hLoopCounters->GetXaxis()->SetBinLabel(CountPairsCandidateOk + 1, "pairs passing ML roles"); hLoopCounters->GetXaxis()->SetBinLabel(CountFit2Prong + 1, "2-prong fits"); + hLoopCounters->GetXaxis()->SetBinLabel(CountPairsWritten + 1, "pairs written"); hLoopCounters->GetXaxis()->SetBinLabel(CountPairsRejected2P + 1, "pairs rejected by 2-prong"); hLoopCounters->GetXaxis()->SetBinLabel(CountTripletsEnumerated + 1, "triplets enumerated"); hLoopCounters->GetXaxis()->SetBinLabel(CountTripletsWritten + 1, "triplets written"); + hLoopCounters->GetXaxis()->SetBinLabel(CountDstarsEnumerated + 1, "D* enumerated"); + hLoopCounters->GetXaxis()->SetBinLabel(CountDstarsWritten + 1, "D* written"); hLoopCounters->GetXaxis()->SetBinLabel(CountTracksCached + 1, "tracks cached"); hLoopCounters->GetXaxis()->SetBinLabel(CountTracksPropagated + 1, "tracks re-propagated"); + registry.add("hVtx2ProngX", "2-prong candidates;#it{x}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); + registry.add("hVtx2ProngY", "2-prong candidates;#it{y}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); + registry.add("hVtx2ProngZ", "2-prong candidates;#it{z}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -20., 20.}}}); + registry.add("hNCand2Prong", "2-prong candidates preselected;# of candidates;entries", {HistType::kTH1D, {axisNumCands}}); + registry.add("hNCand2ProngVsNTracks", "2-prong candidates preselected;# of selected tracks;# of candidates;entries", {HistType::kTH2D, {axisNumTracks, axisNumCands}}); + registry.add("hMassD0ToPiK", "D^{0} candidates;inv. mass (#pi K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 0., 5.}}}); registry.add("hVtx3ProngX", "3-prong candidates;#it{x}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); registry.add("hVtx3ProngY", "3-prong candidates;#it{y}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); registry.add("hVtx3ProngZ", "3-prong candidates;#it{z}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -20., 20.}}}); registry.add("hNCand3Prong", "3-prong candidates preselected;# of candidates;entries", {HistType::kTH1D, {axisNumCands}}); registry.add("hNCand3ProngVsNTracks", "3-prong candidates preselected;# of selected tracks;# of candidates;entries", {HistType::kTH2D, {axisNumTracks, axisNumCands}}); + registry.add("hNCandDstar", "D* candidates preselected;# of candidates;entries", {HistType::kTH1D, {axisNumCands}}); registry.add("hMassDPlusToPiKPi", "D^{#plus} candidates;inv. mass (#pi K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 0., 5.}}}); registry.add("hMassDsToKKPi", "D_{s}^{#plus} candidates;inv. mass (K K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 0., 5.}}}); + registry.add("hMassDstarToD0Pi", "D*^{#plus} candidates;inv. mass (D0 #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{130, 0.135, 0.2}}}); } } /// Check whether the two prongs can satisfy any mass hypothesis of either channel, based on their ML role masks. /// \param maskOpp is the isIdentifiedPid mask of the opposite-sign prong /// \param maskSame is the isIdentifiedPid mask of the first same-sign prong - static bool isCandidateAllowed(const uint32_t maskOpp, const uint32_t maskSame) + static void checkCandidateAllowed(const uint32_t maskOpp, const uint32_t maskSame, bool& canBe2Prong, bool& canBe3Prong) { + canBe2Prong = TESTBIT(maskOpp, RoleKaonD0) && TESTBIT(maskSame, RolePionD0); + // For D* we check that we can have a D0 and later apply the selection for the soft pion + const bool okDplus = TESTBIT(maskOpp, RoleKaonDplus) && TESTBIT(maskSame, RolePionDplus); const bool okDs = TESTBIT(maskOpp, RoleKaonDs) && (TESTBIT(maskSame, RoleKaonDs) || TESTBIT(maskSame, RolePionDs)); - return okDplus || okDs; + canBe3Prong = okDplus || okDs; + } + + /// Check whether the three ML masks can satisfy any mass hypothesis of either channel + /// \param maskOpp is the mask of the opposite-sign prong (prong 1) + /// \param mask1, mask2 are the masks of the same-sign prongs (prongs 0 and 2) + static bool pairRoleOk(const uint32_t maskOpp, const uint32_t maskSame) + { + const std::array masks{maskSame, maskOpp}; + for (int iChannel = 0; iChannel < NChannels2Prong; iChannel++) { + for (int iHypo = 0; iHypo < NMassHypos; iHypo++) { + bool isHypoAlive = true; + for (int iProng = 0; iProng < N2Prongs; iProng++) { + const int role = trackRoles2Prongs[iChannel][iHypo][iProng]; + if (!TESTBIT(masks[iProng], role)) { + isHypoAlive = false; + break; + } + } + if (isHypoAlive) { + return true; + } + } + } + return false; } /// Check whether the three ML masks can satisfy any mass hypothesis of either channel @@ -887,12 +1091,12 @@ struct HfMlBasedTrackSelector { /// \param mask1, mask2 are the masks of the same-sign prongs (prongs 0 and 2) static bool tripletRoleOk(const uint32_t maskOpp, const uint32_t mask1, const uint32_t mask2) { - const std::array masks{mask1, maskOpp, mask2}; + const std::array masks{mask1, maskOpp, mask2}; for (int iChannel = 0; iChannel < NChannels3Prong; iChannel++) { for (int iHypo = 0; iHypo < NMassHypos; iHypo++) { bool isHypoAlive = true; - for (int iProng = 0; iProng < NProngs; iProng++) { - const int role = IsKaonProng[iChannel][iHypo][iProng] ? RoleKaonOfChannel[iChannel] : RolePionOfChannel[iChannel]; + for (int iProng = 0; iProng < N3Prongs; iProng++) { + const int role = trackRoles3Prongs[iChannel][iHypo][iProng]; if (!TESTBIT(masks[iProng], role)) { isHypoAlive = false; break; @@ -909,19 +1113,24 @@ struct HfMlBasedTrackSelector { /// Require the ML role of each prong to match the mass hypothesis under test. /// The opposite-sign prong is the kaon in both channels; the same-sign pair is (pi, pi) for the /// D+ and (K, pi) or (pi, K) for the Ds, which is exactly what the two mass hypotheses encode. - /// \param isIdentifiedPid0,1,2 are the aod::HfSelTrack::isIdentifiedPid masks of the prongs, in + /// \param isIdentifiedPid is the aod::HfSelTrack::isIdentifiedPid mask of the prong, in /// the (same-sign, opposite-sign, same-sign) ordering /// \param whichHypo information of the mass hypotheses that are still alive /// \param isSelected is a bitmap with selection outcome - void applyMlRoleSelection(const uint32_t isIdentifiedPid0, const uint32_t isIdentifiedPid1, const uint32_t isIdentifiedPid2, - std::array& whichHypo, auto& isSelected) + template + void applyMlRoleSelection(const std::array& isIdentifiedPid, + std::array& whichHypo, + auto& isSelected) { - const std::array isIdentifiedPid{isIdentifiedPid0, isIdentifiedPid1, isIdentifiedPid2}; - - for (int iChannel = 0; iChannel < NChannels3Prong; iChannel++) { - for (int iHypo = 0; iHypo < NMassHypos; iHypo++) { - for (int iProng = 0; iProng < NProngs; iProng++) { - const int role = IsKaonProng[iChannel][iHypo][iProng] ? RoleKaonOfChannel[iChannel] : RolePionOfChannel[iChannel]; + for (size_t iChannel = 0; iChannel < NCandChannels; iChannel++) { + for (size_t iHypo = 0; iHypo < NMassHypos; iHypo++) { + for (size_t iProng = 0; iProng < NProngs; iProng++) { + int role{-1}; + if constexpr (NProngs == 2) { + role = trackRoles2Prongs[iChannel][iHypo][iProng]; + } else if constexpr (NProngs == 3) { + role = trackRoles3Prongs[iChannel][iHypo][iProng]; + } if (!TESTBIT(isIdentifiedPid[iProng], role)) { CLRBIT(whichHypo[iChannel], iHypo); break; @@ -1053,7 +1262,11 @@ struct HfMlBasedTrackSelector { return; } + // Precompute, no dependence on channel or mass hypothesis const auto pt = RecoDecay::pt(pVecCand); + const auto cpa = RecoDecay::cpa(primVtx, secVtx, pVecCand); + const auto decayLength = RecoDecay::distance(primVtx, secVtx); + for (int iChannel = 0; iChannel < NChannels3Prong; iChannel++) { if (!TESTBIT(isSelected, iChannel)) { continue; @@ -1067,14 +1280,12 @@ struct HfMlBasedTrackSelector { } // cos of pointing angle - const auto cpa = RecoDecay::cpa(primVtx, secVtx, pVecCand); if (cpa < cut3Prong[iChannel].get(binPt, 2u)) { // 2u == "cosp" CLRBIT(isSelected, iChannel); continue; } // decay length - const auto decayLength = RecoDecay::distance(primVtx, secVtx); if (decayLength < cut3Prong[iChannel].get(binPt, 3u)) { // 3u == "decL" CLRBIT(isSelected, iChannel); } @@ -1083,10 +1294,11 @@ struct HfMlBasedTrackSelector { /// Translate the compact per-channel selection bitmap into the hfflag written to aod::Hf3Prongs, /// which uses the hf_cand_3prong::DecayType bit positions. - uint8_t hfFlagOfSelection(const uint32_t isSelected) const + uint8_t hfFlagOfSelection(const uint32_t isSelected, bool is3Prong) const { uint8_t hfFlag{0}; - for (int iChannel = 0; iChannel < NChannels3Prong; iChannel++) { + size_t nChannels = is3Prong ? static_cast(NChannels3Prong) : static_cast(NChannels2Prong); + for (size_t iChannel = 0; iChannel < nChannels; iChannel++) { if (TESTBIT(isSelected, iChannel)) { SETBIT(hfFlag, DecayTypeOfChannel[iChannel]); } @@ -1112,7 +1324,8 @@ struct HfMlBasedTrackSelector { std::array whichHypo3Prong{}; whichHypo3Prong.fill(BIT(NMassHypos) - 1); // all mass hypotheses alive - applyMlRoleSelection(isIdentifiedPid0, isIdentifiedPid1, isIdentifiedPid2, whichHypo3Prong, isSelected3ProngCand); + std::array isIdentifiedPid{isIdentifiedPid0, isIdentifiedPid1, isIdentifiedPid2}; + applyMlRoleSelection(isIdentifiedPid, whichHypo3Prong, isSelected3ProngCand); if (isSelected3ProngCand == 0) { return; } @@ -1156,7 +1369,7 @@ struct HfMlBasedTrackSelector { } // fill table row - rowTrackIndexProng3(collision.globalIndex(), globalIndex0, globalIndex1, globalIndex2, hfFlagOfSelection(isSelected3ProngCand)); + rowTrackIndexProng3(collision.globalIndex(), globalIndex0, globalIndex1, globalIndex2, hfFlagOfSelection(isSelected3ProngCand, true)); // fill histograms if (config.fillHistograms) { @@ -1184,6 +1397,167 @@ struct HfMlBasedTrackSelector { } } + /// Method to perform selections for 2-prong candidates before vertex reconstruction + /// \param pVecTrack0 is the momentum array of the first daughter track + /// \param pVecTrack1 is the momentum array of the second daughter track + /// \param dcaTrack0 is the dcaXY of the first daughter track + /// \param dcaTrack1 is the dcaXY of the second daughter track + /// \param cutStatus is a 2D array with outcome of each selection (filled only in debug mode) + /// \param whichHypo information of the mass hypoteses that were selected + /// \param isSelected is a bitmap with selection outcome + /// \param pt2Prong is the pt of the 2-prong candidate + // void applySelection2Prong(const T1& pVecCand, const T2& secVtx, const T3& primVtx, T4& cutStatus, auto& isSelected) + template + void applySelection2Prong(const T1& pVecCand, T2 const& prong0, T2 const& prong1, const T3& secVtx, const T4& primVtx, T5& whichHypo, auto& isSelected) + { + + const auto pt2Prong = RecoDecay::pt(pVecCand[0], pVecCand[1]); + const auto impParProduct = prong0.dcaInfo[0] * prong1.dcaInfo[0]; + const auto cpa = RecoDecay::cpa(primVtx, secVtx, pVecCand); + + for (int iDecay2P = 0; iDecay2P < NChannels2Prong; iDecay2P++) { + + // return immediately if it is outside the defined pT bins + const auto binPt = findBin(&binsPt2Prong[iDecay2P], pt2Prong); + if (binPt == -1) { + CLRBIT(isSelected, iDecay2P); + continue; + } + + // invariant mass + double massHypos[2] = {0., 0.}; + + if (TESTBIT(isSelected, iDecay2P)) { + const double minMass = cut2Prong[iDecay2P].get(binPt, 0u); + const double maxMass = cut2Prong[iDecay2P].get(binPt, 1u); + if (minMass >= 0. && maxMass > 0.) { + const std::array, N2Prongs> arr2Mom{prong0.pVec, prong1.pVec}; + massHypos[0] = RecoDecay::m2(arr2Mom, arrMass2Prong[iDecay2P][0]); + massHypos[1] = RecoDecay::m2(arr2Mom, arrMass2Prong[iDecay2P][1]); + const double min2 = minMass * minMass; + const double max2 = maxMass * maxMass; + if (massHypos[0] < min2 || massHypos[0] >= max2) { + CLRBIT(whichHypo[iDecay2P], 0); + } + if (massHypos[1] < min2 || massHypos[1] >= max2) { + CLRBIT(whichHypo[iDecay2P], 1); + } + if (whichHypo[iDecay2P] == 0) { + CLRBIT(isSelected, iDecay2P); + } + } + } + + // imp. par. product cut + if (TESTBIT(isSelected, iDecay2P)) { + if (impParProduct > cut2Prong[iDecay2P].get(binPt, 3u)) { + CLRBIT(isSelected, iDecay2P); + } + } + + // cos of pointing angle + if (TESTBIT(isSelected, iDecay2P)) { + if (cpa < cut2Prong[iDecay2P].get(binPt, 2u)) { // 2u == "cospIndex[iDecay2P]" + CLRBIT(isSelected, iDecay2P); + } + } + } + } + + /// Reconstruct and preselect one 2-prong candidate, and fill its table row if it survives. + /// \param collision is the collision the candidate belongs to + /// \param prong0,1 are the track informations + template + void processPair(SelectedCollisions::iterator const& collision, + TProngInfo const& prong0, TProngInfo const& prong1) + { + uint32_t isSelected2ProngCand = BIT(NChannels2Prong) - 1; + std::array whichHypo2Prong{}; + whichHypo2Prong.fill(BIT(NMassHypos) - 1); // 2 bits on, all mass hypotheses alive + + std::array isIdentifiedPid{prong0.mask, prong1.mask}; + applyMlRoleSelection(isIdentifiedPid, whichHypo2Prong, isSelected2ProngCand); + if (isSelected2ProngCand == 0) { + return; + } + + // get secondary vertex + const auto& secondaryVertex2 = df2.getPCACandidate(); + // get track momenta + std::array pvec0{}, pvec1{}; + df2.getTrack(0).getPxPyPzGlo(pvec0); + df2.getTrack(1).getPxPyPzGlo(pvec1); + const auto pVecCandProng2 = RecoDecay::pVec(pvec0, pvec1); + + // 2-prong selections + const std::array pvCoord{collision.posX(), collision.posY(), collision.posZ()}; + applySelection2Prong(pVecCandProng2, prong0, prong1, secondaryVertex2, pvCoord, whichHypo2Prong, isSelected2ProngCand); + + // Only D0 for now + if (!TESTBIT(isSelected2ProngCand, ChannelD0ToPiK)) { + return; + } + + // fill table row + rowTrackIndexProng2(collision.globalIndex(), prong0.globalIndex, prong1.globalIndex, hfFlagOfSelection(isSelected2ProngCand, false)); + ++counters[CountPairsWritten]; + + // fill histograms + if (config.fillHistograms) { + registry.fill(HIST("hVtx2ProngX"), secondaryVertex2[0]); + registry.fill(HIST("hVtx2ProngY"), secondaryVertex2[1]); + registry.fill(HIST("hVtx2ProngZ"), secondaryVertex2[2]); + const std::array arr2Mom{pvec0, pvec1}; + registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][0])); + registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][1])); + } + } + + /// Method to perform selections for D* candidates before vertex reconstruction + /// \param collision is the collision the candidate belongs to + /// \param lastFilledD0 is the index of the last filled D0 candidate in the table + /// \param softPiGlobalIndex is the global index of the soft pion track + /// \param pVecTrack0 is the momentum array of the first daughter track (same charge) + /// \param pVecTrack1 is the momentum array of the second daughter track (opposite charge) + /// \param pVecTrack2 is the momentum array of the third daughter track (same charge) + void processDstar(SelectedCollisions::iterator const& collision, int lastFilledD0, int64_t softPiGlobalIndex, + std::array const& pVecTrack0, std::array const& pVecTrack1, std::array const& pVecTrack2) + { + const std::array arrMom{pVecTrack0, pVecTrack1, pVecTrack2}; + const std::array arrMomD0{pVecTrack0, pVecTrack1}; + const auto pt = RecoDecay::pt(pVecTrack0, pVecTrack1, pVecTrack2) + config.ptTolerance; // add tolerance because of no reco decay vertex + + // pT + const auto binPt = findBin(config.binsPtDstarToD0Pi, pt); + // return immediately if it is outside the defined pT bins + if (binPt == -1) { + return; + } + + // D0 mass + const double deltaMassD0 = config.cutsDstarToD0Pi->get(binPt, 1u); // 1u == deltaMassD0Index + const double invMassD0 = RecoDecay::m(arrMomD0, std::array{MassPiPlus, MassKPlus}); + if (std::abs(invMassD0 - MassD0) > deltaMassD0) { + return; + } + + // D*+ mass + const double maxDeltaMass = config.cutsDstarToD0Pi->get(binPt, 0u); // 0u == deltaMassIndex + const double invMassDstar = RecoDecay::m(arrMom, std::array{MassPiPlus, MassKPlus, MassPiPlus}); + const double deltaMass = invMassDstar - invMassD0; + if (deltaMass > maxDeltaMass) { + return; + } + + // D* candidate is accepted + + // fill table row and fill histograms + rowTrackIndexDstar(collision.globalIndex(), softPiGlobalIndex, lastFilledD0); + if (config.fillHistograms) { + registry.fill(HIST("hMassDstarToD0Pi"), deltaMass); + } + } + /// Build the per-collision cache of prongs propagated to the collision's primary vertex. /// \param collision is the collision under study /// \param trackIndices are the track associations of this collision @@ -1196,13 +1570,12 @@ struct HfMlBasedTrackSelector { const auto thisCollId = collision.globalIndex(); for (const auto& trackIndex : trackIndices) { const auto track = trackIndex.template track_as(); - HfProngCandidate prong{getTrackParCov(track), track.pVector(), trackIndex.isIdentifiedPid(), track.globalIndex()}; + HfProngCandidate prong{getTrackParCov(track), track.pVector(), {track.dcaXY(), track.dcaZ()}, trackIndex.isIdentifiedPid(), track.globalIndex()}; ++counters[CountTracksCached]; if (thisCollId != track.collisionId()) { // this is not the "default" collision for this track, we have to re-propagate it const auto tStartPropagate = tickIf(doTiming); ++counters[CountTracksPropagated]; - std::array dcaInfo{track.dcaXY(), track.dcaZ()}; - o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, prong.parCov, 2.f, noMatCorr, &dcaInfo); + o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, prong.parCov, 2.f, noMatCorr, &prong.dcaInfo); getPxPyPz(prong.parCov, prong.pVec); if (doTiming) { timing[TimeCachePropagate] += elapsedSeconds(tStartPropagate); @@ -1232,8 +1605,8 @@ struct HfMlBasedTrackSelector { { const std::vector& sameSignList = sameSignCache.sameSign; const std::vector& oppSignList = oppSignCache.oppSign; - constexpr std::size_t NSameSignProngs = NProngs - 1; // two of the three prongs carry the candidate charge - if (sameSignList.size() < NSameSignProngs || oppSignList.empty()) { + constexpr std::size_t NSameSign3Prongs = N3Prongs - 1; // two of the three prongs carry the candidate charge + if ((!config.do2Prongs && !config.doDstar && sameSignList.size() < NSameSign3Prongs) || oppSignList.empty()) { return; } const auto tStartLoop = tickIf(doTiming); @@ -1243,16 +1616,23 @@ struct HfMlBasedTrackSelector { for (const int iOpp : oppSignList) { // o2-linter: disable=const-ref-in-for-loop (int elements) const auto& prongOpp = oppSignCache.prongs[iOpp]; - for (std::size_t i1 = 0; i1 + 1 < sameSignList.size(); ++i1) { + for (std::size_t i1 = 0; i1 < sameSignList.size(); ++i1) { const auto& prong1 = sameSignCache.prongs[sameSignList[i1]]; ++counters[CountPairsSeen]; - if (!isCandidateAllowed(prongOpp.mask, prong1.mask)) { + bool canBe2Prong{false}, canBe3Prong{false}; + bool isTwoProngVtxGoodFor3Prongs{true}; + checkCandidateAllowed(prongOpp.mask, prong1.mask, canBe2Prong, canBe3Prong); + if (!canBe2Prong && !canBe3Prong) { + continue; + } + if (!canBe3Prong && !config.do2Prongs && !config.doDstar) { continue; } ++counters[CountPairsCandidateOk]; + int lastFilledD0 = rowTrackIndexProng2.lastIndex(); // optional 2-track vertex to skip the pair early - if (useTwoTrackVertex) { + if (useTwoTrackVertex || config.do2Prongs || config.doDstar) { const auto tStart2Prong = tickIf(doTiming); ++counters[CountFit2Prong]; int nVtxFrom2ProngFitter = 0; @@ -1260,28 +1640,67 @@ struct HfMlBasedTrackSelector { nVtxFrom2ProngFitter = df2.process(prong1.parCov, prongOpp.parCov); } catch (...) { } - const bool pairOk = (nVtxFrom2ProngFitter != 0) && - isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); if (doTiming) { timing[TimeFit2Prong] += elapsedSeconds(tStart2Prong); } - if (!pairOk) { + const bool hasVertices = (nVtxFrom2ProngFitter != 0); + if (!hasVertices) { ++counters[CountPairsRejected2P]; continue; } + const auto tStartD0 = tickIf(doTiming); + if (config.do2Prongs || config.doDstar) { + processPair(collision, prong1, prongOpp); + } + if (doTiming) { + timing[TimeProcessD0] += elapsedSeconds(tStartD0); + } + isTwoProngVtxGoodFor3Prongs = isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); + } + + // D* -> D0 pi+ candidates + const auto tStartDstar = tickIf(doTiming); + if (config.doDstar && + lastFilledD0 != rowTrackIndexProng2.lastIndex()) { // we have a new D0 candidate, so we can try to form a D* candidate + + lastFilledD0 = rowTrackIndexProng2.lastIndex(); + // Loop over all tracks to search for soft pions + for (std::size_t iSoftPi = 0; iSoftPi < sameSignList.size(); ++iSoftPi) { + const auto& prongSoftPi = sameSignCache.prongs[sameSignList[iSoftPi]]; + + if (iSoftPi == i1) { + continue; // skip the same track + } + if (!TESTBIT(prongSoftPi.mask, RoleSoftPiDstar)) { + continue; // skip if the second same-sign prong is not a soft pion candidate + } + ++counters[CountDstarsEnumerated]; + + // Update the index of the last filled D0 candidate + processDstar(collision, lastFilledD0, prongSoftPi.globalIndex, + prong1.pVec, prongOpp.pVec, prongSoftPi.pVec); + } + } + if (doTiming) { + timing[TimeProcessDstar] += elapsedSeconds(tStartDstar); } + if (!canBe3Prong) { continue; } + if (!isTwoProngVtxGoodFor3Prongs) { + ++counters[CountPairsRejected2P]; + continue; + } + // 3-prong candidates for (std::size_t i2 = i1 + 1; i2 < sameSignList.size(); ++i2) { const auto& prong2 = sameSignCache.prongs[sameSignList[i2]]; if (!tripletRoleOk(prongOpp.mask, prong1.mask, prong2.mask)) { continue; } ++counters[CountTripletsEnumerated]; - processTriplet(collision, prong1.parCov, prongOpp.parCov, prong2.parCov, - prong1.pVec, prongOpp.pVec, prong2.pVec, - prong1.mask, prongOpp.mask, prong2.mask, - prong1.globalIndex, prongOpp.globalIndex, prong2.globalIndex); + prong1.pVec, prongOpp.pVec, prong2.pVec, + prong1.mask, prongOpp.mask, prong2.mask, + prong1.globalIndex, prongOpp.globalIndex, prong2.globalIndex); } } } @@ -1290,10 +1709,10 @@ struct HfMlBasedTrackSelector { } } - void process3Prongs(SelectedCollisions const& collisions, - aod::BCsWithTimestamps const&, - FilteredTrackAssocSel const&, - aod::TracksWCovDca const& /*tracks*/) + void processCandidates(SelectedCollisions const& collisions, + aod::BCsWithTimestamps const&, + FilteredTrackAssocSel const&, + aod::TracksWCovDca const& /*tracks*/) { for (const auto& collision : collisions) { // set the magnetic field from CCDB @@ -1303,15 +1722,17 @@ struct HfMlBasedTrackSelector { df3.setBz(o2::base::Propagator::Instance()->getNominalBz()); // used to calculate number of candidates per event + auto nCand2 = rowTrackIndexProng2.lastIndex(); auto nCand3 = rowTrackIndexProng3.lastIndex(); + auto nCandDstars = rowTrackIndexDstar.lastIndex(); timing.fill(0.); counters.fill(0); const auto thisCollId = collision.globalIndex(); - const auto allPos = positiveFor3Prongs->sliceByCached(aod::track::collisionId, thisCollId, cache); - const auto allNeg = negativeFor3Prongs->sliceByCached(aod::track::collisionId, thisCollId, cache); + const auto allPos = positiveHfTracks->sliceByCached(aod::track::collisionId, thisCollId, cache); + const auto allNeg = negativeHfTracks->sliceByCached(aod::track::collisionId, thisCollId, cache); const int nTracks = allPos.size() + allNeg.size(); // We fill the prong cache for each collision, so that the prongs are propagated to the PV only once. @@ -1325,8 +1746,12 @@ struct HfMlBasedTrackSelector { runCombinatorics(collision, cachePos, cacheNeg); runCombinatorics(collision, cacheNeg, cachePos); + nCand2 = rowTrackIndexProng2.lastIndex() - nCand2; // number of 2-prong candidates in this collision + counters[CountPairsWritten] = nCand2; nCand3 = rowTrackIndexProng3.lastIndex() - nCand3; // number of 3-prong candidates in this collision counters[CountTripletsWritten] = nCand3; + nCandDstars = rowTrackIndexDstar.lastIndex() - nCandDstars; // number of D* candidates in this collision + counters[CountDstarsWritten] = nCandDstars; if (config.fillHistograms) { // Fill(bin, weight) so the bins accumulate seconds and counts over the whole run @@ -1337,18 +1762,21 @@ struct HfMlBasedTrackSelector { registry.fill(HIST("hLoopCounters"), iStep, static_cast(counters[iStep])); } registry.fill(HIST("hNTracks"), nTracks); + registry.fill(HIST("hNCand2Prong"), nCand2); registry.fill(HIST("hNCand3Prong"), nCand3); + registry.fill(HIST("hNCandDstar"), nCandDstars); registry.fill(HIST("hNCand3ProngVsNTracks"), nTracks, nCand3); } } + LOG(info) << "Processed " << collisions.size() << " collisions for 3-prong candidates"; } - PROCESS_SWITCH(HfMlBasedTrackSelector, process3Prongs, "Process 3-prong skim", true); + PROCESS_SWITCH(HfMlBasedTrackSelector, processCandidates, "Process 3-prong skim", true); - void processNo3Prongs(SelectedCollisions const&) + void processDummy(SelectedCollisions const&) { // dummy } - PROCESS_SWITCH(HfMlBasedTrackSelector, processNo3Prongs, "Do not process 3-prongs", false); + PROCESS_SWITCH(HfMlBasedTrackSelector, processDummy, "Do not process 3-prongs", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) From 763c0b3f1ac8347ec9705cc48dcfcec081a56335 Mon Sep 17 00:00:00 2001 From: marcellocosti Date: Thu, 8 Oct 2026 17:20:20 +0200 Subject: [PATCH 2/6] [WIP] Implement Fabrizio comments --- PWGHF/D2H/TableProducer/CMakeLists.txt | 5 + .../D2H/TableProducer/dplusCombinatorics.cxx | 2644 +++++++++++++++++ .../TableProducer/mlBasedTrackSelector.cxx | 173 +- PWGHF/DataModel/CharmTrackScoreScanTables.h | 64 + PWGHF/TableProducer/candidateCreatorDstar.cxx | 31 +- PWGLF/TableProducer/Resonances/CMakeLists.txt | 8 +- PWGLF/Tasks/Resonances/CMakeLists.txt | 28 +- 7 files changed, 2879 insertions(+), 74 deletions(-) create mode 100644 PWGHF/D2H/TableProducer/dplusCombinatorics.cxx create mode 100644 PWGHF/DataModel/CharmTrackScoreScanTables.h diff --git a/PWGHF/D2H/TableProducer/CMakeLists.txt b/PWGHF/D2H/TableProducer/CMakeLists.txt index ebc99c83782..a7e8a7bf781 100644 --- a/PWGHF/D2H/TableProducer/CMakeLists.txt +++ b/PWGHF/D2H/TableProducer/CMakeLists.txt @@ -102,6 +102,11 @@ o2physics_add_dpl_workflow(data-creator-hidden-charm-reduced PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(dplus-combinatorics + SOURCES dplusCombinatorics.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils O2Physics::AnalysisCore O2Physics::MLCore + COMPONENT_NAME Analysis) + # Converters o2physics_add_dpl_workflow(converter-reduced-3-prongs-ml diff --git a/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx new file mode 100644 index 00000000000..c0be83ee581 --- /dev/null +++ b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx @@ -0,0 +1,2644 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file dplusCombinatorics.cxx +/// \brief D± → π± K∓ π± and D±s → K± K∓ π± secondary-vertex finding driven by track-level ML +/// +/// Standalone alternative to the 3-prong branch of trackIndexSkimCreator.cxx. The linear, +/// pT-binned single-track selections of the skim creator are replaced by *two* multi-class BDTs, +/// one per species, each scoring every track as {not-a-daughter, pion-from-X, kaon-from-X}. +/// Inference runs once per (track, collision) association, so a tighter track selection +/// translates directly into fewer triplets entering the vertex fitter. +/// +/// Two tasks live here: +/// - HfDplusTrackSelMl : builds the track features and either dumps them for training (MC) +/// or evaluates both models and flags pion-/kaon-like tracks per species +/// - HfDplusCombinatorics: the triple loop over the flagged tracks, filling aod::Hf3Prongs +/// +/// The combinatorics offers three strategies, so that the cost of each can be measured: +/// - processDplus : D+ only, the original implementation, kept so that the earlier +/// efficiency/rate comparison against the skim creator stays reproducible +/// - processSeparateLoops: one triple loop per species, each over its own track sets. A triplet +/// that satisfies both hypotheses is fitted -- and written -- twice. +/// - processMergedLoop : a single triple loop over the union of the selected tracks, with a +/// per-triplet hypothesis mask. A triplet satisfying both hypotheses is +/// fitted once and written once, with both flags set. This is what the +/// standard skim creator does, and whether it pays off depends entirely +/// on how much the two species' track sets overlap. +/// +/// Both new strategies cache the propagated track parameters once per collision instead of +/// re-propagating inside the innermost loop, so that they differ only in the loop structure. +/// +/// The efficiency/rate QA used to compare this workflow with the standard track-index skim +/// creator lives in a separate executable, dplusCombinatoricsQa.cxx, so that it can be piped +/// after either skimmer without two devices declaring aod::Hf3Prongs as output. +/// +/// \author Fabrizio Chinu , Universita and INFN Torino + +#include "PWGHF/Core/HfMlResponseHfTracks.h" +#include "PWGHF/Core/SelectorCuts.h" +#include "PWGHF/DataModel/AliasTables.h" +#include "PWGHF/DataModel/CharmTrackScoreScanTables.h" +#include "PWGHF/DataModel/TrackIndexSkimmingTables.h" +#include "PWGHF/Utils/utilsAnalysis.h" +#include "PWGHF/Utils/utilsBfieldCCDB.h" +#include "PWGHF/Utils/utilsEvSelHf.h" + +#include "Common/Core/RecoDecay.h" +#include "Common/Core/ZorroSummary.h" +#include "Common/Core/trackUtilities.h" +#include "Common/DataModel/Centrality.h" +#include "Common/DataModel/CollisionAssociationTables.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/Multiplicity.h" +#include "Common/DataModel/PIDResponseTPC.h" +#include "Common/DataModel/TrackSelectionTables.h" + +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::analysis; +using namespace o2::aod; +using namespace o2::framework; +using namespace o2::framework::expressions; +using namespace o2::hf_centrality; +using namespace o2::hf_evsel; +using namespace o2::constants::physics; + +namespace o2::aod +{ +namespace hf_skim_sel_collision +{ +DECLARE_SOA_COLUMN(WhyRejectColl, whyRejectColl, o2::hf_evsel::HfCollisionRejectionMask); //! reason for rejection (0 = selected) +DECLARE_SOA_COLUMN(Centrality, centrality, float); //! centrality percentile of the collision +} // namespace hf_skim_sel_collision + +DECLARE_SOA_TABLE(HfSkimSelColls, "AOD", "HFSKIMSELCOLL", //! Collision selection for the charm combinatorics + hf_skim_sel_collision::WhyRejectColl, + hf_skim_sel_collision::Centrality); + +namespace hf_skim_track_sel +{ +// D0 -> K pi model +DECLARE_SOA_COLUMN(ScoreBkgDzero, scoreBkgDzero, float); //! ML score, background class +DECLARE_SOA_COLUMN(ScorePionDzero, scorePionDzero, float); //! ML score, pion-from-D0 class +DECLARE_SOA_COLUMN(ScoreKaonDzero, scoreKaonDzero, float); //! ML score, kaon-from-D0 class +DECLARE_SOA_COLUMN(IsSelPionDzero, isSelPionDzero, bool); //! passes the D0 pion-score threshold +DECLARE_SOA_COLUMN(IsSelKaonDzero, isSelKaonDzero, bool); //! passes the D0 kaon-score threshold + // D+ -> K pi pi model +DECLARE_SOA_COLUMN(ScoreBkgDplus, scoreBkgDplus, float); //! ML score, background class +DECLARE_SOA_COLUMN(ScorePionDplus, scorePionDplus, float); //! ML score, pion-from-D+ class +DECLARE_SOA_COLUMN(ScoreKaonDplus, scoreKaonDplus, float); //! ML score, kaon-from-D+ class +DECLARE_SOA_COLUMN(IsSelPionDplus, isSelPionDplus, bool); //! passes the D+ pion-score threshold +DECLARE_SOA_COLUMN(IsSelKaonDplus, isSelKaonDplus, bool); //! passes the D+ kaon-score threshold +// Ds -> K K pi model +DECLARE_SOA_COLUMN(ScoreBkgDs, scoreBkgDs, float); //! ML score, background class +DECLARE_SOA_COLUMN(ScorePionDs, scorePionDs, float); //! ML score, pion-from-Ds class +DECLARE_SOA_COLUMN(ScoreKaonDs, scoreKaonDs, float); //! ML score, kaon-from-Ds class +DECLARE_SOA_COLUMN(IsSelPionDs, isSelPionDs, bool); //! passes the Ds pion-score threshold +DECLARE_SOA_COLUMN(IsSelKaonDs, isSelKaonDs, bool); //! passes the Ds kaon-score threshold +// D* -> D0 pi -> pi K pi model +DECLARE_SOA_COLUMN(ScoreBkgDstar, scoreBkgDstar, float); //! ML score, background class +DECLARE_SOA_COLUMN(ScorePionDstar, scorePionDstar, float); //! ML score, pion-from-D* class +DECLARE_SOA_COLUMN(ScoreKaonDstar, scoreKaonDstar, float); //! ML score, kaon-from-D* class +DECLARE_SOA_COLUMN(IsSelPionDstar, isSelPionDstar, bool); //! passes the D* pion-score threshold +DECLARE_SOA_COLUMN(IsSelKaonDstar, isSelKaonDstar, bool); //! passes the D* kaon-score threshold +// common +DECLARE_SOA_COLUMN(IsPositive, isPositive, bool); //! charge sign of the track +DECLARE_SOA_COLUMN(IsSelAny, isSelAny, bool); //! passes at least one of the four selections +} // namespace hf_skim_track_sel + +DECLARE_SOA_TABLE(HfSkimTrackSels, "AOD", "HFSKIMTRKSEL", //! Track-level ML selection, aligned with aod::TrackAssoc + hf_skim_track_sel::ScoreBkgDzero, + hf_skim_track_sel::ScorePionDzero, + hf_skim_track_sel::ScoreKaonDzero, + hf_skim_track_sel::IsSelPionDzero, + hf_skim_track_sel::IsSelKaonDzero, + hf_skim_track_sel::ScoreBkgDplus, + hf_skim_track_sel::ScorePionDplus, + hf_skim_track_sel::ScoreKaonDplus, + hf_skim_track_sel::IsSelPionDplus, + hf_skim_track_sel::IsSelKaonDplus, + hf_skim_track_sel::ScoreBkgDs, + hf_skim_track_sel::ScorePionDs, + hf_skim_track_sel::ScoreKaonDs, + hf_skim_track_sel::IsSelPionDs, + hf_skim_track_sel::IsSelKaonDs, + hf_skim_track_sel::ScoreBkgDstar, + hf_skim_track_sel::ScorePionDstar, + hf_skim_track_sel::ScoreKaonDstar, + hf_skim_track_sel::IsSelPionDstar, + hf_skim_track_sel::IsSelKaonDstar, + hf_skim_track_sel::IsPositive, + hf_skim_track_sel::IsSelAny); + +namespace hf_skim_track_sel_mask +{ +DECLARE_SOA_COLUMN(CollIdx, collIdx, int); //! global index of the collision the association refers to +DECLARE_SOA_COLUMN(SelMask, selMask, uint16_t); //! bit0 -> pi(D0) + //! bit1 -> K(D0) + //! bit2 -> pi(D+) + //! bit3 -> K(D+) + //! bit4 -> pi(Ds) + //! bit5 -> K(Ds) + //! bit6 -> pi(D*) + //! bit7 -> K(D*) + //! bit8 -> positive charge +} // namespace hf_skim_track_sel_mask + +/// One row per *selected* (track, collision) association. This is the smallest object from which +/// the overlap between the four role selections, and the number of triplets each combinatorics +/// strategy would enumerate, can be reconstructed exactly offline: the 16 mask multiplicities per +/// (collision, charge) are a sufficient statistic for both. +DECLARE_SOA_TABLE(HfSkimTrackMasks, "AOD", "HFSKIMTRKMASK", //! Compact dump of the per-track role selections + hf_skim_track_sel_mask::CollIdx, + hf_skim_track_sel_mask::SelMask); + +namespace hf_skim_track_train +{ +// features +DECLARE_SOA_COLUMN(Pt, pt, float); //! +DECLARE_SOA_COLUMN(Eta, eta, float); //! +DECLARE_SOA_COLUMN(DcaXY, dcaXY, float); //! +DECLARE_SOA_COLUMN(DcaZ, dcaZ, float); //! +DECLARE_SOA_COLUMN(SigmaDcaXY, sigmaDcaXY, float); //! +DECLARE_SOA_COLUMN(SigmaDcaZ, sigmaDcaZ, float); //! +DECLARE_SOA_COLUMN(NormDcaXY, normDcaXY, float); //! +DECLARE_SOA_COLUMN(NormDcaZ, normDcaZ, float); //! +DECLARE_SOA_COLUMN(Signed1Pt, signed1Pt, float); //! +DECLARE_SOA_COLUMN(Tgl, tgl, float); //! +DECLARE_SOA_COLUMN(Sign, sign, float); //! +DECLARE_SOA_COLUMN(IsPvContributor, isPvContributor, float); //! +DECLARE_SOA_COLUMN(ItsNCls, itsNCls, float); //! +DECLARE_SOA_COLUMN(ItsNClsInnerBarrel, itsNClsInnerBarrel, float); //! +DECLARE_SOA_COLUMN(ItsChi2NCl, itsChi2NCl, float); //! +DECLARE_SOA_COLUMN(TpcNClsFound, tpcNClsFound, float); //! +DECLARE_SOA_COLUMN(TpcNClsCrossedRows, tpcNClsCrossedRows, float); //! +DECLARE_SOA_COLUMN(TpcCrossedRowsOverFindableCls, tpcCrossedRowsOverFindableCls, float); //! +DECLARE_SOA_COLUMN(TpcChi2NCl, tpcChi2NCl, float); //! +DECLARE_SOA_COLUMN(TpcFractionSharedCls, tpcFractionSharedCls, float); //! +DECLARE_SOA_COLUMN(TpcNSigmaPi, tpcNSigmaPi, float); //! +DECLARE_SOA_COLUMN(TpcNSigmaKa, tpcNSigmaKa, float); //! +DECLARE_SOA_COLUMN(Centrality, centrality, float); //! +// labels and auxiliary information, not meant to be used as features +DECLARE_SOA_COLUMN(TrackIndex, trackIndex, int64_t); //! global index of the track: the same track can be associated to several collisions, and therefore appear more than once, so the train/test split has to group by it +DECLARE_SOA_COLUMN(McMask, mcMask, uint16_t); //! MC label quality mask (bit 15 = fake track) +DECLARE_SOA_COLUMN(HasCharmBeautyAncestor, hasCharmBeautyAncestor, bool); //! any charm/beauty hadron in the ancestry +// Selection outcomes +DECLARE_SOA_COLUMN(Label, label, int8_t); //! D+ model: -1 = not usable, 0 = bkg, 1 = pion from D+, 2 = kaon from D+ +DECLARE_SOA_COLUMN(PtCand, ptCand, float); //! pT of the mother D+ (-1 if none) +DECLARE_SOA_COLUMN(IsPrompt, isPrompt, int8_t); //! 1 prompt, 0 non-prompt, -1 not applicable +DECLARE_SOA_COLUMN(Channel, channel, int8_t); //! -1 none, 0 D+ -> K pi pi, 1 D+ -> anti-K*0 pi +DECLARE_SOA_COLUMN(PdgMother, pdgMother, int); //! PDG code of the mother particle +} // namespace hf_skim_track_train + +DECLARE_SOA_TABLE(HfSkimTrackTrains, "AOD", "HFSKIMTRKTRAIN", //! Track-level training sample, one row per usable association + hf_skim_track_train::Pt, + hf_skim_track_train::Eta, + hf_skim_track_train::DcaXY, + hf_skim_track_train::DcaZ, + hf_skim_track_train::SigmaDcaXY, + hf_skim_track_train::SigmaDcaZ, + hf_skim_track_train::NormDcaXY, + hf_skim_track_train::NormDcaZ, + hf_skim_track_train::Signed1Pt, + hf_skim_track_train::Tgl, + hf_skim_track_train::Sign, + hf_skim_track_train::IsPvContributor, + hf_skim_track_train::ItsNCls, + hf_skim_track_train::ItsNClsInnerBarrel, + hf_skim_track_train::ItsChi2NCl, + hf_skim_track_train::TpcNClsFound, + hf_skim_track_train::TpcNClsCrossedRows, + hf_skim_track_train::TpcCrossedRowsOverFindableCls, + hf_skim_track_train::TpcChi2NCl, + hf_skim_track_train::TpcFractionSharedCls, + hf_skim_track_train::TpcNSigmaPi, + hf_skim_track_train::TpcNSigmaKa, + hf_skim_track_train::Centrality, + hf_skim_track_train::TrackIndex, + hf_skim_track_train::Label, + hf_skim_track_train::PtCand, + hf_skim_track_train::IsPrompt, + hf_skim_track_train::Channel, + hf_skim_track_train::PdgMother, + hf_skim_track_train::McMask, + hf_skim_track_train::HasCharmBeautyAncestor); +} // namespace o2::aod + +/// Bit set in aod::McTrackLabels::mcMask when the MC label of the track is not reliable. +constexpr uint16_t McMaskFakeTrackBit = 15; + +// HF mothers to investigate +constexpr std::array PdgMotherList = {Pdg::kD0, Pdg::kDPlus, Pdg::kDS, Pdg::kDStar}; + +/// Decay topologies used to identify the channels in the MC kinematic tree. +constexpr int NDaughtersTwoBody = 2; +constexpr int NDaughtersDzeroToPiK = 2; +constexpr int NPionsDzeroToPiK = 1; +constexpr int NKaonsDzeroToPiK = 1; +constexpr int NDaughtersDplusToPiKPi = 3; +constexpr int NPionsDplusToPiKPi = 2; +constexpr int NKaonsDplusToPiKPi = 1; +constexpr int NDaughtersDsToKKPi = 3; +constexpr int NKaonsDsToKKPi = 2; +constexpr int NPionsDsToKKPi = 1; +constexpr int NDaughtersDstarToD0Pi = 2; +constexpr int ND0DstarToD0Pi = 1; +constexpr int NPionsDstarToD0Pi = 1; + +/// Resonant sub-channels recorded next to the training label, so that the signal definition can be +/// narrowed offline without re-running the workflow. +enum Channels : int { + kNone = -1, ///< no mother found, or mother not in the list of interest + kDzeroDirect, + kDplusDirect, + kDplusKstar, + kDsPhiPi, + kDsKstarK, + kDsDirect, + kDstarToD0Pi, +}; +// constexpr int8_t ChannelNone = -1; +// constexpr int8_t ChannelDzeroDirect = 0; ///< D0 -> K- pi+ +// constexpr int8_t ChannelDplusDirect = 0; ///< D+ -> K- pi+ pi+ +// constexpr int8_t ChannelDplusKstar = 1; ///< D+ -> anti-K*(892)0 pi+, anti-K*0 -> K- pi+ +// constexpr int8_t ChannelDsPhiPi = 0; ///< Ds+ -> phi(1020) pi+, phi -> K+ K- +// constexpr int8_t ChannelDsKstarK = 1; ///< Ds+ -> anti-K*(892)0 K+, anti-K*0 -> K- pi+ +// constexpr int8_t ChannelDsDirect = 2; ///< Ds+ -> K+ K- pi+ non resonant + +/// Above this value the chi2 cut on the two-track vertex is considered disabled. +constexpr double MaxTwoTrackChi2PcaDisabled = 1.e9; + +/// Margin applied to the kinematic cuts of the cheap pre-filter that runs before the track is +/// re-propagated. Propagation without material corrections preserves q/pT and tgl, so pT and eta +/// are in practice unchanged; the margin only absorbs numerical drift. The exact cut is applied +/// afterwards, on the propagated values. +constexpr float PreselKinematicMargin = 0.1f; + +/// Bits of the compact per-association role mask, both in aod::HfSkimTrackMasks and in the +/// per-collision prong cache of the combinatorics. +/// Output-class count of the origin-aware (5-class) track model. +constexpr int NClassesWithOrigin = 5; + +enum TrackRoleBit : uint16_t { + BitPionDzero = 0, ///< pion-like for D0 -> K pi + BitKaonDzero, ///< kaon-like for D0 -> K pi + BitPionDplus, ///< pion-like for D+ -> K pi pi + BitKaonDplus, ///< kaon-like for D+ -> K pi pi + BitPionDs, ///< pion-like for Ds -> K K pi + BitKaonDs, ///< kaon-like for Ds -> K K pi + BitPionDstar, ///< pion-like for D* -> D0 pi + BitKaonDstar, ///< kaon-like for D* -> D0 pi + BitPositive ///< charge sign, so that one row fully describes the association +}; + +enum TrackType : uint8_t { + Pion, + Kaon +}; + +/// CPU accounting. steady_clock::now() costs ~20 ns through the vDSO, so the timers are placed +/// only around blocks that themselves cost microseconds -- a model evaluation, a vertex fit, one +/// collision's worth of looping -- and never inside the per-triplet enumeration, whose own cost +/// is comparable to a clock read. Everything not attributed to a measured block is reported as a +/// remainder, so the bins always add up to the total. +using StopWatch = std::chrono::steady_clock; + +inline double elapsedSeconds(StopWatch::time_point const& since) +{ + return std::chrono::duration(StopWatch::now() - since).count(); +} + +/// Timing bins of hf-dplus-track-sel-ml. +enum TimingTrackSel : int { + TimeTrackSelTotal = 0, + TimeTrackSelEvSel, + TimeTrackSelFeatures, + TimeTrackSelInferenceDzero, + TimeTrackSelInferenceDplus, + TimeTrackSelInferenceDs, + TimeTrackSelInferenceDstar, + TimeTrackSelRemainder, + NTimingTrackSel +}; + +/// Timing bins of hf-dplus-combinatorics. +enum TimingCombinatorics : int { + TimeCombTotal = 0, + TimeCombCache, + TimeCombFit2Prong, + TimeCombFit3Prong, + TimeCombRemainder, + NTimingCombinatorics +}; + +/// Counter bins of hf-dplus-track-sel-ml. +enum CounterTrackSel : int { + CountTrackSelCollisions = 0, + CountAssociations, + CountCheapPresel, + CountPresel, + CountInferenceDzero, + CountInferenceDplus, + CountInferenceDs, + CountInferenceDstar, + CountSelPionDzero, + CountSelKaonDzero, + CountSelPionDplus, + CountSelKaonDplus, + CountSelPionDs, + CountSelKaonDs, + CountSelPionDstar, + CountSelKaonDstar, + CountSelAny, + NCounterTrackSel +}; + +/// Outcomes of the heavy-flavour ancestry walk of the written training rows (hHfAncestryWalk). +enum HfAncestryWalkOutcome : int { + HfWalkNoMcParticle = 0, + HfWalkNotFound, + HfWalkFound, + HfWalkDepthLimit, + NHfAncestryWalkOutcomes +}; + +/// Counter bins of hf-dplus-combinatorics. +enum CounterCombinatorics : int { + CountCollisions = 0, + CountTriplets, ///< triplets whose hypothesis mask is non-empty, i.e. reaching the preselection + CountTripletsPresel,///< triplets passing the mass/pT preselection, i.e. reaching the 3-prong fitter + CountFit2Prong, + CountFit3Prong, + CountRows, ///< rows written to aod::Hf3Prongs + CountRowsDplus, + CountRowsDs, + CountRowsBoth, + NCounterCombinatorics +}; + +/// Deterministic pseudo-random number in [0, 1) from a 64-bit index. +/// Used to downsample the background class reproducibly, without any RNG state. +inline float deterministicUniform(uint64_t index, uint64_t seed) +{ + // splitmix64 + uint64_t z = index + seed + 0x9e3779b97f4a7c15ULL; + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL; + z = z ^ (z >> 31); + return static_cast(z >> 11) / static_cast(1ULL << 53); +} + +/// Track-level feature building, training-sample production and ML inference for both species. +struct HfDplusTrackSelMl { + Produces rowSelectedCollision; + Produces rowSelectedTrack; + Produces rowSelMask; + Produces rowTrackTrain; + + struct : ConfigurableGroup { + Configurable fillHistograms{"fillHistograms", true, "fill QA histograms"}; + Configurable enableTiming{"enableTiming", false, "account the CPU time of the individual stages into hTiming"}; + Configurable fillSelMaskTable{"fillSelMaskTable", false, "write one aod::HfSkimTrackMasks row per selected association (for the offline overlap study)"}; + Configurable rejBkgCollisions{"rejBkgCollisions", true, "Reject tracks from background events"}; + // loose pre-filter applied before inference (kept minimal on purpose) + Configurable ptMinTrack{"ptMinTrack", 0.3f, "min. track pT for the charm combinatorics"}; + Configurable etaMaxTrack{"etaMaxTrack", 0.8f, "max. |eta| for the charm combinatorics"}; + // Run 3 equivalent of the old ITS-refit requirement. It is NOT spelled + // TESTBIT(track.flags(), o2::aod::track::ITSrefit) + // the way the track-index skim creator does it: ITSrefit belongs to TrackFlagsRun2Enum and + // equals 0x1, so TESTBIT(flags, 0x1) tests bit *index* 1, i.e. mask 0x2 -- which in Run 3 is + // PVContributor, not ITSrefit. That silently required charm daughters to contribute to the + // primary vertex: verified on this sample, every selected track had isPvContributor == 1, + // and it removed ~40% of the genuine signal, preferentially the displaced tracks the + // selection is supposed to keep. + Configurable requireIts{"requireIts", true, "require the track to have ITS information"}; + Configurable requireItsIbHit{"requireItsIbHit", true, "require a hit in one of the two innermost ITS layers"}; + // D0 model + Configurable applyMlDzero{"applyMlDzero", true, "evaluate the D+ track model"}; + Configurable nClassesDzero{"nClassesDzero", 3, "output classes of the D+ model: 3 (bkg/pi/K) or 5 (bkg/prompt pi/prompt K/non-prompt pi/non-prompt K)"}; + Configurable> binsPtMlDzero{"binsPtMlDzero", std::vector{0., 1.e10}, "pT bin limits of the D0 track models"}; + Configurable> onnxFileNamesDzero{"onnxFileNamesDzero", std::vector{"ModelHandler_DzeroTracks.onnx"}, "ONNX file names of the D0 models, one per pT bin"}; + Configurable> inputFeaturesDzero{"inputFeaturesDzero", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D0 models, in the order expected by them"}; + Configurable> thresholdScorePionDzero{"thresholdScorePionDzero", std::vector{0.}, "min. D0 pion-class score, one per pT bin"}; + Configurable> thresholdScoreKaonDzero{"thresholdScoreKaonDzero", std::vector{0.}, "min. D0 kaon-class score, one per pT bin"}; + // D+ model + Configurable applyMlDplus{"applyMlDplus", true, "evaluate the D+ track model"}; + Configurable nClassesDplus{"nClassesDplus", 3, "output classes of the D+ model: 3 (bkg/pi/K) or 5 (bkg/prompt pi/prompt K/non-prompt pi/non-prompt K)"}; + Configurable> binsPtMlDplus{"binsPtMlDplus", std::vector{0., 1.e10}, "pT bin limits of the D+ track models"}; + Configurable> onnxFileNamesDplus{"onnxFileNamesDplus", std::vector{"ModelHandler_DplusTracks.onnx"}, "ONNX file names of the D+ models, one per pT bin"}; + Configurable> inputFeaturesDplus{"inputFeaturesDplus", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D+ models, in the order expected by them"}; + Configurable> thresholdScorePionDplus{"thresholdScorePionDplus", std::vector{0.}, "min. D+ pion-class score, one per pT bin"}; + Configurable> thresholdScoreKaonDplus{"thresholdScoreKaonDplus", std::vector{0.}, "min. D+ kaon-class score, one per pT bin"}; + // Ds model + Configurable applyMlDs{"applyMlDs", true, "evaluate the Ds track model"}; + Configurable nClassesDs{"nClassesDs", 3, "output classes of the Ds model: 3 or 5, as for the D+"}; + Configurable> binsPtMlDs{"binsPtMlDs", std::vector{0., 1.e10}, "pT bin limits of the Ds track models"}; + Configurable> onnxFileNamesDs{"onnxFileNamesDs", std::vector{"ModelHandler_DsTracks.onnx"}, "ONNX file names of the Ds models, one per pT bin"}; + Configurable> inputFeaturesDs{"inputFeaturesDs", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the Ds models, in the order expected by them"}; + Configurable> thresholdScorePionDs{"thresholdScorePionDs", std::vector{0.}, "min. Ds pion-class score, one per pT bin"}; + Configurable> thresholdScoreKaonDs{"thresholdScoreKaonDs", std::vector{0.}, "min. Ds kaon-class score, one per pT bin"}; + // Ds model + Configurable applyMlDstar{"applyMlDstar", true, "evaluate the Dstar track model"}; + Configurable nClassesDstar{"nClassesDstar", 2, "output classes of the Dstar model: 2, bkg or pion"}; + Configurable> binsPtMlDstar{"binsPtMlDstar", std::vector{0., 1.e10}, "pT bin limits of the Dstar track models"}; + Configurable> onnxFileNamesDstar{"onnxFileNamesDstar", std::vector{"ModelHandler_DstarTracks.onnx"}, "ONNX file names of the Dstar models, one per pT bin"}; + Configurable> inputFeaturesDstar{"inputFeaturesDstar", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the Dstar models, in the order expected by them"}; + Configurable> thresholdScorePionDstar{"thresholdScorePionDstar", std::vector{0.}, "min. Dstar pion-class score, one per pT bin"}; + // ONNX runtime + Configurable loadModelsFromCcdb{"loadModelsFromCcdb", false, "load the ONNX models from CCDB instead of a local path"}; + Configurable mlModelPathCcdb{"mlModelPathCcdb", "Users/f/fchinu/DplusTracks", "CCDB path of the ML models"}; + Configurable timestampCcdbForMlModels{"timestampCcdbForMlModels", -1, "timestamp of the ONNX files to be queried in CCDB"}; + Configurable enableOnnxOptimizations{"enableOnnxOptimizations", true, "enable the ONNX graph optimisations"}; + Configurable onnxThreads{"onnxThreads", 1, "number of threads used by the ONNX runtime (0 = let onnxruntime decide)"}; + // training-sample production + Configurable bkgSamplingFraction{"bkgSamplingFraction", 0.05f, "fraction of background tracks written to the training tree"}; + Configurable bkgSamplingSeed{"bkgSamplingSeed", 12345, "seed of the deterministic background downsampling"}; + Configurable keepFakeTracksInBkg{"keepFakeTracksInBkg", true, "keep tracks with an unreliable MC label in the background class"}; + // RecoDecay::getMother walks the ancestry breadth-first, keeping every stage in a vector of + // vectors and deduplicating each stage with a linear std::find. With an unbounded depth that + // walk runs back to the beam particles: in Pb-Pb the fan-out makes it quadratic and enormous + // (tens of GB of retained heap, minutes per timeframe). It is physically wrong too -- through + // the fragmentation record almost any track in an event containing a D+ acquires a D+ + // "ancestor". A real charm decay product is a direct daughter, or sits behind one resonance + // or one strange decay, so depth 3 covers every genuine chain. + Configurable maxAncestryDepth{"maxAncestryDepth", 3, "max. ancestry depth for the same-species veto and the prompt/non-prompt flag"}; + // The charm/beauty ancestor flag walks the decay chain only (it stops at partons, see + // continueAncestryWalk), so it does not suffer from the fan-out above and can go much deeper. + Configurable maxAncestryDepthHfFlag{"maxAncestryDepthHfFlag", 30, "max. ancestry depth of the charm/beauty ancestor flag (reaching it sets the flag)"}; + // CCDB + Configurable isRun2{"isRun2", false, "enable Run 2 or Run 3 GRP objects for the magnetic field"}; + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable ccdbPathLut{"ccdbPathLut", "GLO/Param/MatLUT", "path for LUT parametrization"}; + Configurable ccdbPathGrp{"ccdbPathGrp", "GLO/GRP/GRP", "path of the grp file (Run 2)"}; + Configurable ccdbPathGrpMag{"ccdbPathGrpMag", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object (Run 3)"}; + } config; + + HfEventSelection hfEvSel; // event selection and monitoring, shared with the skim creator + Service ccdb{}; // needed for the event selection and the magnetic field + o2::ccdb::CcdbApi ccdbApi; + o2::base::MatLayerCylSet* lut{}; + o2::base::Propagator::MatCorrType noMatCorr = o2::base::Propagator::MatCorrType::USEMatCorrNONE; + int runNumber{}; + + /// One track-level model: the response, its per-pT-bin score thresholds and whether it runs. + struct TrackModel { + HfMlResponseHfTracks response{}; + std::vector thresholdPion{}; + std::vector thresholdKaon{}; + bool enabled{false}; + /// 3 = bkg / pion / kaon, 5 = bkg / prompt pi / prompt K / non-prompt pi / non-prompt K. + /// With 5 the two origin scores of a role are summed into that role's score, so the working + /// point stays one threshold per role and the two model kinds are directly comparable. + int nClasses{3}; + }; + TrackModel modelDzero; + TrackModel modelDplus; + TrackModel modelDs; + TrackModel modelDstar; + + // running variables for the re-propagation of tracks to a non-default collision + o2::dataformats::DCA dcaInfoCov; + o2::dataformats::VertexBase vtx; + + // per-collision timing accumulators, in seconds + std::array timing{}; + + using CollisionsWithCent = soa::Join; // PVMults: INEL>0 flag of the HF event selection + using TracksWithSelAndPid = soa::Join; + using TracksWithSelAndPidMc = soa::Join; + + Preslice trackIndicesPerCollision = aod::track_association::collisionId; + + HistogramRegistry registry{"registry"}; + // OutputObj zorroSummary{"zorroSummary"}; + + /// Configure one model from its configurables. Returns false if the model is disabled. + /// \param model is the model to configure + /// \param binsPt is the pT binning + /// \param onnxFiles are the ONNX file names, one per pT bin + /// \param features are the input feature names, in the order the model expects + /// \param thrPion, thrKaon are the per-bin score thresholds + /// \param name is used in the error messages + bool configureModel(TrackModel& model, + std::vector const& binsPt, + std::vector const& onnxFiles, + std::vector const& features, + std::vector const& thrPion, + std::vector const& thrKaon, + int nClasses, + const char* name) + { + if (nClasses != static_cast(HfTrackMlClass::NClasses) && nClasses != NClassesWithOrigin) { + LOGP(fatal, "{}: nClasses must be {} or {}, got {}!", name, + static_cast(HfTrackMlClass::NClasses), NClassesWithOrigin, nClasses); + } + model.nClasses = nClasses; + const std::size_t nModels = binsPt.size() - 1; + if (onnxFiles.size() != nModels) { + LOGP(fatal, "{}: number of ONNX files ({}) does not match the number of pT bins ({})!", name, onnxFiles.size(), nModels); + } + if (thrPion.size() != nModels || thrKaon.size() != nModels) { + LOGP(fatal, "{}: number of score thresholds (pion {}, kaon {}) does not match the number of pT bins ({})!", name, thrPion.size(), thrKaon.size(), nModels); + } + model.thresholdPion = thrPion; + model.thresholdKaon = thrKaon; + + // the per-class thresholds are applied by the task, so no cut is delegated to MlResponse + const auto nOut = static_cast(nClasses); + const std::vector cutDir(nOut, o2::cuts_ml::CutDirection::CutNot); + const std::vector dummyCutValues(nModels * nOut, 0.); + const o2::framework::LabeledArray dummyCuts{dummyCutValues.data(), static_cast(nModels), static_cast(nOut), {}, {}}; + + model.response.configure(binsPt, dummyCuts, cutDir, static_cast(nOut)); + if (config.loadModelsFromCcdb) { + ccdbApi.init(config.ccdbUrl); + model.response.setModelPathsCCDB(onnxFiles, ccdbApi, std::vector(nModels, config.mlModelPathCcdb.value), config.timestampCcdbForMlModels); + } else { + model.response.setModelPathsLocal(onnxFiles); + } + model.response.cacheInputFeaturesIndices(features); + model.response.init(config.enableOnnxOptimizations, config.onnxThreads); + model.enabled = true; + LOGP(info, "{}: {} model(s) configured with {} input features and {} output classes", name, nModels, features.size(), nClasses); + return true; + } + + void init(InitContext const&) + { + const std::array doProcess = {doprocessTraining, doprocessTrainingData, doprocessApplication}; + if (std::accumulate(doProcess.begin(), doProcess.end(), 0) != 1) { + LOGP(fatal, "One and only one process function of HfDplusTrackSelMl must be enabled at a time!"); + } + + // hfEvSel.init(registry, &zorroSummary); + hfEvSel.init(registry); + + ccdb->setURL(config.ccdbUrl); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->get(config.ccdbPathLut)); + + if (config.fillHistograms) { + const AxisSpec axisPt{200, 0., 20., "#it{p}_{T}^{track} (GeV/#it{c})"}; + const AxisSpec axisScore{100, 0., 1., "ML score"}; + const AxisSpec axisNTracks{500, 0., 5000., "n. tracks"}; + + registry.add("hPtAllTracks", "all associated tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hPtPreselected", "tracks passing the loose pre-filter;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + + if (doprocessApplication) { + registry.add("hScoreBkgDzero", "D^{0} background score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScorePionDzero", "D^{0} pion score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScoreKaonDzero", "D^{0} kaon score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hPtSelPionDzero", "D^{0} pion-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hPtSelKaonDzero", "D^{0} kaon-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hNTracksSelPionDzeroPerColl", "D^{0} pion-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksSelKaonDzeroPerColl", "D^{0} kaon-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hScoreBkgDplus", "D^{+} background score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScorePionDplus", "D^{+} pion score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScoreKaonDplus", "D^{+} kaon score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hPtSelPionDplus", "D^{+} pion-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hPtSelKaonDplus", "D^{+} kaon-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hNTracksSelPionDplusPerColl", "D^{+} pion-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksSelKaonDplusPerColl", "D^{+} kaon-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hScoreBkgDs", "D_{s} background score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScorePionDs", "D_{s} pion score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScoreKaonDs", "D_{s} kaon score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hPtSelPionDs", "D_{s} pion-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hPtSelKaonDs", "D_{s} kaon-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hNTracksSelPionDsPerColl", "D_{s} pion-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksSelKaonDsPerColl", "D_{s} kaon-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hScoreBkgDstar", "D^{*} background score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScorePionDstar", "D^{*} pion score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hScoreKaonDstar", "D^{*} kaon score;score;entries", {HistType::kTH1D, {axisScore}}); + registry.add("hPtSelPionDstar", "D^{*} pion-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hPtSelKaonDstar", "D^{*} kaon-like tracks;#it{p}_{T}^{track} (GeV/#it{c});entries", {HistType::kTH1D, {axisPt}}); + registry.add("hNTracksSelPionDstarPerColl", "D^{*} pion-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksSelKaonDstarPerColl", "D^{*} kaon-like tracks per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksPreselPerColl", "tracks passing the pre-filter per collision;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + registry.add("hNTracksSelAnyPerColl", "tracks passing at least one role selection;n. tracks;entries", {HistType::kTH1D, {axisNTracks}}); + // Role overlap, summed over the whole sample. The 16 bins are the multiplicities of the + // four-bit role mask: this is what tells whether the two species should share one triple + // loop or run two of their own. + registry.add("hRoleMask", "role mask of the selected associations;mask (bit 0 #pi D^{+}, 1 K D^{+}, 2 #pi D_{s}, 3 K D_{s});entries", {HistType::kTH1D, {{16, -0.5, 15.5}}}); + } + if (doprocessTraining) { + registry.add("hLabelDzero", "D^{0} training sample;label;entries", {HistType::kTH1D, {{3, -0.5, 2.5}}}); + registry.get(HIST("hLabelDzero"))->GetXaxis()->SetBinLabel(1, "background"); + registry.get(HIST("hLabelDzero"))->GetXaxis()->SetBinLabel(2, "#pi from D^{0}"); + registry.get(HIST("hLabelDzero"))->GetXaxis()->SetBinLabel(3, "K from D^{0}"); + registry.add("hLabelDplus", "D^{+} training sample;label;entries", {HistType::kTH1D, {{3, -0.5, 2.5}}}); + registry.get(HIST("hLabelDplus"))->GetXaxis()->SetBinLabel(1, "background"); + registry.get(HIST("hLabelDplus"))->GetXaxis()->SetBinLabel(2, "#pi from D^{+}"); + registry.get(HIST("hLabelDplus"))->GetXaxis()->SetBinLabel(3, "K from D^{+}"); + registry.add("hLabelDs", "D_{s} training sample;label;entries", {HistType::kTH1D, {{3, -0.5, 2.5}}}); + registry.get(HIST("hLabelDs"))->GetXaxis()->SetBinLabel(1, "background"); + registry.get(HIST("hLabelDs"))->GetXaxis()->SetBinLabel(2, "#pi from D_{s}"); + registry.get(HIST("hLabelDs"))->GetXaxis()->SetBinLabel(3, "K from D_{s}"); + registry.add("hLabelDstar", "D^{*} training sample;label;entries", {HistType::kTH1D, {{3, -0.5, 2.5}}}); + registry.get(HIST("hLabelDstar"))->GetXaxis()->SetBinLabel(1, "background"); + registry.get(HIST("hLabelDstar"))->GetXaxis()->SetBinLabel(2, "#pi from D^{*}"); + registry.get(HIST("hLabelDstar"))->GetXaxis()->SetBinLabel(3, "K from D^{*}"); + registry.add("hChannelDzero", "resonant sub-channel of the labelled D^{0};channel;entries", {HistType::kTH1D, {{2, -0.5, 1.5}}}); + registry.add("hChannelDplus", "resonant sub-channel of the labelled D^{+};channel;entries", {HistType::kTH1D, {{2, -0.5, 1.5}}}); + registry.add("hChannelDs", "resonant sub-channel of the labelled D_{s};channel;entries", {HistType::kTH1D, {{3, -0.5, 2.5}}}); + registry.add("hChannelDstar", "resonant sub-channel of the labelled D^{*};channel;entries", {HistType::kTH1D, {{2, -0.5, 1.5}}}); + float hfAncestryWalkOutcomes = static_cast(NHfAncestryWalkOutcomes); + registry.add("hHfAncestryWalk", "charm/beauty ancestry of the written rows;;entries", {HistType::kTH1D, {{hfAncestryWalkOutcomes, -0.5, hfAncestryWalkOutcomes - 0.5}}}); + const char* walkLabels[NHfAncestryWalkOutcomes] = {"no MC particle", "no HF ancestor", "HF ancestor", "depth limit"}; + for (int i = 0; i < NHfAncestryWalkOutcomes; ++i) { + registry.get(HIST("hHfAncestryWalk"))->GetXaxis()->SetBinLabel(i + 1, walkLabels[i]); + } + } + } + + // CPU accounting. Always booked, so that a run with the timers off is recognisable by an + // empty histogram rather than by a missing one. + registry.add("hTiming", "CPU time by stage;;CPU time (s)", {HistType::kTH1D, {{NTimingTrackSel, -0.5, NTimingTrackSel - 0.5}}}); + const char* timingLabels[NTimingTrackSel] = {"total", "event selection", "features", "inference D0", "inference D+", "inference Ds", "inference Dstar", "remainder"}; + for (int i = 0; i < NTimingTrackSel; ++i) { + registry.get(HIST("hTiming"))->GetXaxis()->SetBinLabel(i + 1, timingLabels[i]); + } + registry.add("hCounters", "counters;;entries", {HistType::kTH1D, {{NCounterTrackSel, -0.5, NCounterTrackSel - 0.5}}}); + const char* counterLabels[NCounterTrackSel] = {"collisions", "associations", "cheap presel", "presel", "inference D0", "inference D+", "inference Ds", "inference Dstar", + "#pi D^{0}", "K D^{0}", "#pi D^{+}", "K D^{+}", "#pi D_{s}", "K D_{s}", "#pi D^{*}", "K D^{*}", "any"}; + for (int i = 0; i < NCounterTrackSel; ++i) { + registry.get(HIST("hCounters"))->GetXaxis()->SetBinLabel(i + 1, counterLabels[i]); + } + + if (doprocessApplication) { + if (config.applyMlDzero) { + configureModel(modelDzero, config.binsPtMlDzero, config.onnxFileNamesDzero, config.inputFeaturesDzero, + config.thresholdScorePionDzero, config.thresholdScoreKaonDzero, + config.nClassesDzero, "D0 track model"); + } + if (config.applyMlDplus) { + configureModel(modelDplus, config.binsPtMlDplus, config.onnxFileNamesDplus, config.inputFeaturesDplus, + config.thresholdScorePionDplus, config.thresholdScoreKaonDplus, + config.nClassesDplus, "D+ track model"); + } + if (config.applyMlDs) { + configureModel(modelDs, config.binsPtMlDs, config.onnxFileNamesDs, config.inputFeaturesDs, + config.thresholdScorePionDs, config.thresholdScoreKaonDs, + config.nClassesDs, "Ds track model"); + } + if (config.applyMlDstar) { + configureModel(modelDstar, config.binsPtMlDstar, config.onnxFileNamesDstar, config.inputFeaturesDstar, + config.thresholdScorePionDstar, std::vector{0.}, + config.nClassesDstar, "Dstar track model"); + } + if (!modelDzero.enabled && !modelDplus.enabled && !modelDs.enabled && !modelDstar.enabled) { + LOGP(fatal, "processApplication is enabled but neither model is: nothing would be selected."); + } + } + } + + /// Propagation-independent part of the preselection, evaluated on the un-propagated track. + /// Re-propagating a track to a collision that is not its own costs a full covariance + /// propagation, so it must not be paid for tracks that no charm analysis would ever use. The + /// ITS requirements here are exact (propagation cannot change them); the kinematic ones carry a + /// margin and are re-applied exactly, on the propagated values, by isPreselected. + /// \param track is the track + /// \return true if the track is worth propagating + template + bool passesCheapPreselection(TTrack const& track) + { + if (config.requireIts && !track.hasITS()) { + return false; + } + if (config.requireItsIbHit) { + const auto clusterMap = track.itsClusterMap(); + if (!(TESTBIT(clusterMap, 0) || TESTBIT(clusterMap, 1))) { + return false; + } + } + if (track.pt() < config.ptMinTrack * (1.f - PreselKinematicMargin)) { + return false; + } + if (std::abs(track.eta()) > config.etaMaxTrack + PreselKinematicMargin) { + return false; + } + return true; + } + + template + bool isPreselected(TTrack const& track, float trackPt, float trackEta) + { + if (trackPt < config.ptMinTrack) { + return false; + } + if (std::abs(trackEta) > config.etaMaxTrack) { + return false; + } + if (config.requireIts && !track.hasITS()) { + return false; + } + if (config.requireItsIbHit) { + const auto clusterMap = track.itsClusterMap(); + if (!(TESTBIT(clusterMap, 0) || TESTBIT(clusterMap, 1))) { + return false; + } + } + return true; + } + + /// Fill the feature container for one (track, collision) association. + /// Tracks that do not belong to the collision under study are re-propagated to it, so that the + /// impact parameter and its resolution refer to the correct primary vertex. + /// \param collision is the collision the track is associated to + /// \param track is the track + /// \param centrality is the centrality percentile of the collision + /// \param features is the container to be filled + template + void fillFeatures(TCollision const& collision, TTrack const& track, float centrality, HfTrackMlFeatures& features) + { + float trackPt = track.pt(); + float trackEta = track.eta(); + float dcaXY = track.dcaXY(); + float dcaZ = track.dcaZ(); + float sigmaDcaXY = std::sqrt(std::abs(track.sigmaDcaXY2())); + float sigmaDcaZ = std::sqrt(std::abs(track.sigmaDcaZ2())); + + if (track.collisionId() != collision.globalIndex()) { + // not the default collision of this track: re-propagate to get DCA and its covariance + auto trackParCov = getTrackParCov(track); + dcaInfoCov.set(999.f, 999.f, 999.f, 999.f, 999.f); + vtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); + vtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); + if (o2::base::Propagator::Instance()->propagateToDCABxByBz(vtx, trackParCov, 2.f, noMatCorr, &dcaInfoCov)) { + trackPt = trackParCov.getPt(); + trackEta = trackParCov.getEta(); + dcaXY = dcaInfoCov.getY(); + dcaZ = dcaInfoCov.getZ(); + sigmaDcaXY = std::sqrt(std::abs(dcaInfoCov.getSigmaY2())); + sigmaDcaZ = std::sqrt(std::abs(dcaInfoCov.getSigmaZ2())); + } + } + + constexpr float MinSigmaDca = 1.e-6f; // guard against a null resolution before dividing + + features.pt = trackPt; + features.eta = trackEta; + features.dcaXY = dcaXY; + features.dcaZ = dcaZ; + features.sigmaDcaXY = sigmaDcaXY; + features.sigmaDcaZ = sigmaDcaZ; + features.normDcaXY = dcaXY / std::max(sigmaDcaXY, MinSigmaDca); + features.normDcaZ = dcaZ / std::max(sigmaDcaZ, MinSigmaDca); + features.signed1Pt = track.signed1Pt(); + features.tgl = track.tgl(); + features.sign = static_cast(track.sign()); + features.isPvContributor = static_cast(track.isPVContributor()); + features.itsNCls = static_cast(track.itsNCls()); + features.itsNClsInnerBarrel = static_cast(track.itsNClsInnerBarrel()); + features.itsChi2NCl = track.itsChi2NCl(); + features.tpcNClsFound = static_cast(track.tpcNClsFound()); + features.tpcNClsCrossedRows = static_cast(track.tpcNClsCrossedRows()); + features.tpcCrossedRowsOverFindableCls = track.tpcCrossedRowsOverFindableCls(); + features.tpcChi2NCl = track.tpcChi2NCl(); + features.tpcFractionSharedCls = track.tpcFractionSharedCls(); + features.tpcNSigmaPi = track.tpcNSigmaPi(); + features.tpcNSigmaKa = track.tpcNSigmaKa(); + features.centrality = centrality; + } + + /// Index of the closest ancestor with the requested |PDG| code, searched breadth-first up to + /// maxDepth generations, or -1 if there is none. + /// + /// Deliberately NOT RecoDecay::getMother. That helper iterates the mother indices as a *range* + /// for (i = particleMother.mothersIds().front(); i <= particleMother.mothersIds().back(); ++i) + /// but mothersIds() is an array of the actual mother indices, not a contiguous slice. A hadron + /// produced in string fragmentation has two mothers -- the distant endpoints of the string -- so + /// that range spans hundreds of unrelated particles (median 650, up to 1880 in this sample) and + /// any D+ sitting inside the window is returned as an "ancestor". Measured on AO2D_002, this + /// mislabelled 2869 fragmentation pions and kaons as D+ daughters against 1594 genuine ones, + /// pushing the pion:kaon ratio to 4.5 where the K pi pi topology demands 2.0. + /// The same pattern appears in getCharmHadronOrigin, so the prompt/non-prompt flag is derived + /// here too rather than through that helper. + /// \param particle is the MC particle to start from + /// \param mcParticles is the MC particle table + /// \param pdgWanted is the |PDG| code looked for (particles and antiparticles both match) + /// \param maxDepth is the maximum number of generations to walk + /// \return the global index of the matched ancestor, or -1 + template + int findAncestorWithPdg(TMcParticle const& particle, TMcParticles const& mcParticles, int pdgWanted, int maxDepth) + { + std::vector current{particle.globalIndex()}; + std::vector next{}; + for (int depth = 0; depth < maxDepth && !current.empty(); ++depth) { + next.clear(); + for (const auto idx : current) { // o2-linter: disable=const-ref-in-for-loop (int elements) + const auto thisParticle = mcParticles.rawIteratorAt(idx - mcParticles.offset()); + if (!thisParticle.has_mothers()) { + continue; + } + for (const auto motherId : thisParticle.mothersIds()) { // o2-linter: disable=const-ref-in-for-loop (int elements) + if (motherId < 0) { + continue; + } + const auto mother = mcParticles.rawIteratorAt(motherId - mcParticles.offset()); + if (std::abs(mother.pdgCode()) == pdgWanted) { + return motherId; + } + next.push_back(motherId); + } + } + current.swap(next); + } + return -1; + } + + /// Whether a particle decays into exactly the two requested species. + /// \param particle is the mother + /// \param pdgA, pdgB are the |PDG| codes of the two daughters (they may be equal) + /// \return true if the decay is a two-body one into pdgA pdgB + template + bool decaysToTwoBody(TMcParticle const& particle, TMcParticles const&, int pdgA, int pdgB) + { + int nDaughters{0}, nA{0}, nB{0}; + for (const auto& daughter : particle.template daughters_as()) { + ++nDaughters; + const int pdgDaughter = std::abs(daughter.pdgCode()); + if (pdgDaughter == pdgA) { + ++nA; + } + if (pdgDaughter == pdgB) { + ++nB; + } + } + if (nDaughters != NDaughtersTwoBody) { + return false; + } + return (pdgA == pdgB) ? (nA == NDaughtersTwoBody) : (nA == 1 && nB == 1); + } + + /// Sub-channel of a D0 decay into the K pi final state, or ChannelNone if it is not one. + /// \param mother is the D0 MC particle + /// \param mcParticles is the MC particle table + /// \return ChannelDzeroDirect or ChannelNone + template + int8_t getDzeroToKPiChannel(TMcParticle const& mother, TMcParticles const&) // mcParticles) + { + LOG(info) << "Running getDzeroToKPiChannel"; + int nDaughters{0}, nPions{0}, nKaons{0}, nOther{0}; + for (const auto& daughter : mother.template daughters_as()) { + ++nDaughters; + LOG(info) << "D0 daughter PDG code: " << daughter.pdgCode(); + switch (std::abs(daughter.pdgCode())) { + case kPiPlus: + ++nPions; + break; + case kKPlus: + ++nKaons; + break; + default: + break; + } + } + LOG(info) << "D0 daughters: " << nDaughters << ", pions: " << nPions << ", kaons: " << nKaons << ", other: " << nOther; + if (nDaughters == NDaughtersDzeroToPiK && nPions == NPionsDzeroToPiK && nKaons == NKaonsDzeroToPiK) { + LOG(info) << "D0 decay channel identified as Channels::kDzeroDirect"; + return Channels::kDzeroDirect; + } + LOG(info) << "D0 decay channel not identified, returning Channels::kNone"; + return Channels::kNone; + } + + /// Sub-channel of a D0 decay into the K pi final state, or ChannelNone if it is not one. + /// \param mother is the D* MC particle + /// \param mcParticles is the MC particle table + /// \return ChannelDstarToD0Pi or ChannelNone + template + int8_t getDstarToD0PiChannel(TMcParticle const& mother, TMcParticles const&) // mcParticles) + { + LOG(info) << "Running getDstarToD0PiChannel"; + int nDaughters{0}, nPions{0}, nD0{0}, nOther{0}; + for (const auto& daughter : mother.template daughters_as()) { + ++nDaughters; + LOG(info) << "D* daughter PDG code: " << daughter.pdgCode(); + switch (std::abs(daughter.pdgCode())) { + case kPiPlus: + ++nPions; + break; + case kD0: + ++nD0; + break; + default: + ++nOther; + break; + } + } + LOG(info) << "D* daughters: " << nDaughters << ", pions: " << nPions << ", D0: " << nD0 << ", other: " << nOther; + if (nDaughters == NDaughtersDstarToD0Pi && nPions == NPionsDstarToD0Pi && nD0 == ND0DstarToD0Pi) { + LOG(info) << "D* decay channel identified as Channels::kDstarToD0Pi"; + return Channels::kDstarToD0Pi; + } + LOG(info) << "D* decay channel not identified, returning Channels::kNone"; + return Channels::kNone; + } + + /// Sub-channel of a D+ decay into the K pi pi final state, or Channels::kNone if it is not one. + /// \param mother is the D+ MC particle + /// \param mcParticles is the MC particle table + /// \return Channels::kDplusDirect, Channels::kDplusKstar or Channels::kNone + template + int8_t getDplusToKPiPiChannel(TMcParticle const& mother, TMcParticles const& mcParticles) + { + int nDaughters{0}, nPions{0}, nKaons{0}, nKstars{0}; + int64_t indexKstar{-1}; + for (const auto& daughter : mother.template daughters_as()) { + ++nDaughters; + switch (std::abs(daughter.pdgCode())) { + case kPiPlus: + ++nPions; + break; + case kKPlus: + ++nKaons; + break; + case Pdg::kK0Star892: + ++nKstars; + indexKstar = daughter.globalIndex(); + break; + default: + break; + } + } + if (nDaughters == NDaughtersDplusToPiKPi && nPions == NPionsDplusToPiKPi && nKaons == NKaonsDplusToPiKPi) { + return Channels::kDplusDirect; + } + if (nDaughters == NDaughtersTwoBody && nPions == 1 && nKstars == 1) { + const auto kstar = mcParticles.rawIteratorAt(indexKstar - mcParticles.offset()); + if (decaysToTwoBody(kstar, mcParticles, kKPlus, kPiPlus)) { + return Channels::kDplusKstar; + } + } + return Channels::kNone; + } + + /// Sub-channel of a Ds decay into the K K pi final state, or Channels::kNone if it is not one. + /// The three sub-channels share the same final state -- two kaons and one pion -- so the role of + /// a given track follows from its own PDG code, whatever the intermediate resonance. + /// \param mother is the Ds MC particle + /// \param mcParticles is the MC particle table + /// \return Channels::kDsPhiPi, Channels::kDsKstarK, Channels::kDsDirect or Channels::kNone + template + int8_t getDsToKKPiChannel(TMcParticle const& mother, TMcParticles const& mcParticles) + { + int nDaughters{0}, nPions{0}, nKaons{0}, nPhis{0}, nKstars{0}; + int64_t indexPhi{-1}, indexKstar{-1}; + for (const auto& daughter : mother.template daughters_as()) { + ++nDaughters; + switch (std::abs(daughter.pdgCode())) { + case kPiPlus: + ++nPions; + break; + case kKPlus: + ++nKaons; + break; + case Pdg::kPhi: + ++nPhis; + indexPhi = daughter.globalIndex(); + break; + case Pdg::kK0Star892: + ++nKstars; + indexKstar = daughter.globalIndex(); + break; + default: + break; + } + } + if (nDaughters == NDaughtersTwoBody && nPions == 1 && nPhis == 1) { + const auto phi = mcParticles.rawIteratorAt(indexPhi - mcParticles.offset()); + if (decaysToTwoBody(phi, mcParticles, kKPlus, kKPlus)) { + return Channels::kDsPhiPi; + } + } + if (nDaughters == NDaughtersTwoBody && nKaons == 1 && nKstars == 1) { + const auto kstar = mcParticles.rawIteratorAt(indexKstar - mcParticles.offset()); + if (decaysToTwoBody(kstar, mcParticles, kKPlus, kPiPlus)) { + return Channels::kDsKstarK; + } + } + if (nDaughters == NDaughtersDsToKKPi && nKaons == NKaonsDsToKKPi && nPions == NPionsDsToKKPi) { + return Channels::kDsDirect; + } + return Channels::kNone; + } + + /// Whether a PDG code is a hadron containing a charm or beauty (anti)quark, quarkonia included. + /// Decided from the quark digits of the PDG numbering scheme rather than from a list of species, + /// so that the excited and rarer states a list misses (D*0, Ds*, Sigma_c, Xi_c0, Omega_c, B*, + /// Sigma_b, charmonia, ...) are covered too. Diquarks (e.g. 4101) and nuclei are not hadrons here. + static bool isHeavyFlavourHadron(int pdgCode) + { + const int pdgAbs = std::abs(pdgCode); + if (pdgAbs < 100 || pdgAbs >= 1000000000) { // partons, leptons, bosons, generator codes; nuclei + return false; + } + const int nq1 = (pdgAbs / 1000) % 10; + const int nq2 = (pdgAbs / 100) % 10; + const int nq3 = (pdgAbs / 10) % 10; + if (pdgAbs < 10000 && nq1 != 0 && nq3 == 0) { // diquark + return false; + } + const auto isHeavyQuark = [](int nq) { return nq == 4 || nq == 5; }; + return isHeavyQuark(nq1) || isHeavyQuark(nq2) || isHeavyQuark(nq3); + } + + /// Whether the ancestry walk continues through a particle. It stops at quarks, gluons, diquarks + /// and generator-internal objects (strings, clusters, codes 81-100): above a hadron those are the + /// parton shower, whose history fans out to the beams (see maxAncestryDepth). A decay product of + /// a heavy-flavour hadron reaches it through hadrons, leptons and photons only -- unless the + /// generator decays it into partons first (e.g. PYTHIA B decays), which this walk does not follow. + static bool continueAncestryWalk(int pdgCode) + { + const int pdgAbs = std::abs(pdgCode); + if (pdgAbs <= 8 || pdgAbs == 21 || (pdgAbs >= 81 && pdgAbs <= 100)) { + return false; + } + return !(pdgAbs >= 1000 && pdgAbs < 10000 && (pdgAbs / 10) % 10 == 0); + } + + /// Whether the track's MC particle is a heavy-flavour hadron or descends from one, at any depth + /// of the decay chain (feed-down through strange hadrons, conversions and interactions in the + /// detector included), up to maxAncestryDepthHfFlag generations. + /// Stored as an auxiliary column, so that the background veto is applied offline. Deliberately + /// called only for the tracks that are actually written out. + /// \param track is the track, joined with the MC labels + /// \param mcParticles is the MC particle table + /// \return true if a charm or beauty hadron is found among the ancestors; if the depth limit is + /// reached first, also true (and counted in hHfAncestryWalk), so that the veto errs on + /// the side of dropping a track + template + bool hasHeavyFlavourAncestor(TTrack const& track, TMcParticles const& mcParticles) + { + if (!track.has_mcParticle()) { + fillHfAncestryWalk(HfWalkNoMcParticle); + return false; + } + const auto mcParticle = track.template mcParticle_as(); + if (isHeavyFlavourHadron(mcParticle.pdgCode())) { + fillHfAncestryWalk(HfWalkFound); + return true; + } + std::vector current{mcParticle.globalIndex()}; + std::vector next{}; + std::vector visited{mcParticle.globalIndex()}; + for (int depth = 0; depth < config.maxAncestryDepthHfFlag; ++depth) { + next.clear(); + for (const auto idx : current) { // o2-linter: disable=const-ref-in-for-loop (int elements) + const auto particle = mcParticles.rawIteratorAt(idx - mcParticles.offset()); + if (!particle.has_mothers()) { + continue; + } + for (const auto motherId : particle.mothersIds()) { // o2-linter: disable=const-ref-in-for-loop (int elements) + if (motherId < 0 || std::find(visited.begin(), visited.end(), motherId) != visited.end()) { + continue; + } + visited.push_back(motherId); + const auto mother = mcParticles.rawIteratorAt(motherId - mcParticles.offset()); + if (isHeavyFlavourHadron(mother.pdgCode())) { + fillHfAncestryWalk(HfWalkFound); + return true; + } + if (continueAncestryWalk(mother.pdgCode())) { + next.push_back(motherId); + } + } + } + if (next.empty()) { + fillHfAncestryWalk(HfWalkNotFound); + return false; + } + current.swap(next); + } + fillHfAncestryWalk(HfWalkDepthLimit); + return true; + } + + void fillHfAncestryWalk(int outcome) + { + if (config.fillHistograms) { + registry.fill(HIST("hHfAncestryWalk"), outcome); + } + } + + /// Whether the charm hadron descends from a beauty hadron. + /// getCharmHadronOrigin walks the mother indices the same broken way as getMother, so the + /// search is done here instead. + template + bool isFromBeauty(TMcParticle const& mother, TMcParticles const& mcParticles) + { + for (const int pdgBeauty : {Pdg::kBPlus, Pdg::kB0, Pdg::kBS, Pdg::kLambdaB0}) { // o2-linter: disable=const-ref-in-for-loop (int elements) + if (findAncestorWithPdg(mother, mcParticles, pdgBeauty, config.maxAncestryDepth) > -1) { + return true; + } + } + return false; + } + + /// Charm hadron the track's MC particle is a *prong* of, or -1 if it is not one. + /// + /// A prong is a direct daughter of the charm hadron, or a direct daughter of a phi(1020) or a + /// K*(892)0 that is itself a direct daughter of the charm hadron. Anything deeper is a decay or + /// interaction product and must not be labelled as a prong, even though a plain depth-2 + /// ancestry search would happily find the charm hadron: a kaon from D+ -> K pi pi decaying in + /// flight gives a pion whose grandmother is the D+, and calling that pion a D+ prong would teach + /// the model to accept tracks that are not there. + /// \param particle is the MC particle of the track + /// \param mcParticles is the MC particle table + /// \param pdgMother is the |PDG| code of the charm hadron + /// \return the global index of the charm hadron, or -1 + template + int64_t findCharmMotherOfProng(TMcParticle const& particle, TMcParticles const& mcParticles, int& pdgMother) + { + if (!particle.has_mothers()) { + return -1; + } + for (const auto motherId : particle.mothersIds()) { // o2-linter: disable=const-ref-in-for-loop (int elements) + if (motherId < 0 || motherId >= particle.globalIndex()) { + continue; + } + const auto mother = mcParticles.rawIteratorAt(motherId - mcParticles.offset()); + const int pdgMotherAbs = std::abs(mother.pdgCode()); + if (std::find(PdgMotherList.begin(), PdgMotherList.end(), pdgMotherAbs) != PdgMotherList.end()) { + LOG(info) << "Found charm mother of prong: PDG " << pdgMotherAbs << ", global index " << mother.globalIndex(); + pdgMother = pdgMotherAbs; + return motherId; + } + if (pdgMotherAbs != Pdg::kPhi && pdgMotherAbs != Pdg::kK0Star892) { + continue; + } + if (!mother.has_mothers()) { + continue; + } + for (const auto grandMotherId : mother.mothersIds()) { // o2-linter: disable=const-ref-in-for-loop (int elements) + if (grandMotherId < 0 || grandMotherId >= mother.globalIndex()) { + continue; + } + const auto grandMother = mcParticles.rawIteratorAt(grandMotherId - mcParticles.offset()); + if (std::find(PdgMotherList.begin(), PdgMotherList.end(), std::abs(grandMother.pdgCode())) != PdgMotherList.end()) { + LOG(info) << "Found charm grandmother of prong: PDG " << std::abs(grandMother.pdgCode()) << ", global index " << grandMother.globalIndex(); + pdgMother = std::abs(grandMother.pdgCode()); + return grandMotherId; + } + } + } + return -1; + } + + /// Sub-channel of a charm hadron decay, or Channels::kNone if it is not one of the wanted final + /// states. The sub-channel is used to narrow the signal definition offline, but is not part of + /// the training label itself. + /// \param mother is the charm hadron MC particle + /// \param mcParticles is the MC particle table + /// \return the sub-channel, or Channels::kNone + template + int8_t flagSubChannel(TMcParticle const& mother, TMcParticles const& mcParticles) + { + const int pdgMotherAbs = std::abs(mother.pdgCode()); + LOG(info) << "Flagging sub-channel for charm mother PDG " << pdgMotherAbs; + if (pdgMotherAbs == Pdg::kD0) { + return getDzeroToKPiChannel(mother, mcParticles); + } + if (pdgMotherAbs == Pdg::kDPlus) { + return getDplusToKPiPiChannel(mother, mcParticles); + } + if (pdgMotherAbs == Pdg::kDS) { + return getDsToKKPiChannel(mother, mcParticles); + } + if (pdgMotherAbs == Pdg::kDStar) { + LOG(info) << "D* found, checking daughters ... "; + return getDstarToD0PiChannel(mother, mcParticles); + } + return Channels::kNone; + } + + /// Training label of a track for one charm species. + /// + /// Signal: the track is MC-matched to a pion or a kaon that is a prong (see + /// findCharmMotherOfProng) of a charm hadron of the requested species, produced in the + /// injected (non-background) event and decaying into the wanted final state. The resonant + /// sub-channels (D+ -> K*0 pi, Ds -> phi pi, Ds -> K*0 K) are included and the sub-channel is + /// recorded, so that the signal definition can still be narrowed offline. + /// Background: the track comes from the background event and has no hadron of that species + /// anywhere in its ancestry -- background events do contain real charm, so the veto matters. + /// Anything else -- in particular an underlying-event track of the injected event, and a track + /// from a charm hadron of the *other* species -- is not labelled: it is neither a daughter nor + /// a clean negative example. + /// + /// \param track is the track, joined with the MC labels + /// \param mcParticles is the MC particle table + /// \param pdgMother is the |PDG| code of the charm hadron + /// \param keepIfBackground is the outcome of the background downsampling, decided by the caller + /// beforehand so that the ancestry veto -- the expensive part of the labelling -- is not + /// paid for the background tracks that are discarded anyway + /// \param ptMother, isPrompt, channel are filled with the properties of the mother, if any + /// \return the label, or -1 if the track must not enter this species' training sample + template + int8_t getTrainingLabel(TTrack const& track, TMcParticles const& mcParticles, + int& pdgMother, bool keepIfBackground, + float& ptMother, int8_t& isPrompt, int8_t& channel) + { + ptMother = -1.f; + isPrompt = -1; + channel = Channels::kNone; + + if (!track.has_mcParticle()) { + // a track without an MC particle cannot be a charm daughter; treat it as background + return (config.keepFakeTracksInBkg && keepIfBackground) ? static_cast(HfTrackMlClass::Background) : -1; + } + const auto mcParticle = track.template mcParticle_as(); + + const int pdgAbs = std::abs(mcParticle.pdgCode()); + const bool isPionOrKaon = (pdgAbs == kPiPlus || pdgAbs == kKPlus); + + if (!config.rejBkgCollisions || (config.rejBkgCollisions && !mcParticle.fromBackgroundEvent())) { + if (!isPionOrKaon) { + return -1; + } + LOG(info) << "\n\nFlagging track with PDG code " << mcParticle.pdgCode() << " ... "; + for (const auto motherId : mcParticle.mothersIds()) { + const auto m = mcParticles.rawIteratorAt(motherId - mcParticles.offset()); + LOG(info) << "kaon " << mcParticle.globalIndex() << " mother id " << motherId + << " PDG " << m.pdgCode() << " nDaughters " << m.daughtersIds().size(); + } + const int64_t indexMother = findCharmMotherOfProng(mcParticle, mcParticles, pdgMother); + if (indexMother < 0) { + return -1; + } + const auto mother = mcParticles.rawIteratorAt(indexMother - mcParticles.offset()); + channel = flagSubChannel(mother, mcParticles); + if (channel == Channels::kNone) { + return -1; + } + ptMother = mother.pt(); + isPrompt = isFromBeauty(mother, mcParticles) ? 0 : 1; + if (pdgAbs == kKPlus && channel == Channels::kDstarToD0Pi) { + LOG(info) << "Track is a kaon from D* -> D0 pi, which is not a prong of the D* itself"; + } + return (pdgAbs == kKPlus) ? static_cast(HfTrackMlClass::Kaon) : static_cast(HfTrackMlClass::Pion); + } + + if (!keepIfBackground) { // downsampled away: skip the ancestry walk entirely + return -1; + } + if (findAncestorWithPdg(mcParticle, mcParticles, pdgMother, config.maxAncestryDepth) > -1) { + return -1; + } + return static_cast(HfTrackMlClass::Background); + } + + /// Evaluate one model on a set of features. + /// \param model is the model to evaluate + /// \param features are the track features + /// \param scores is filled with the three class scores (-1 if the track is outside the pT range) + /// \param isSelPion, isSelKaon are filled with the threshold decisions + /// \return true if the model was actually evaluated + bool evaluateModel(TrackModel& model, HfTrackMlFeatures const& features, + std::array& scores, bool& isSelPion, bool& isSelKaon) + { + scores = {-1.f, -1.f, -1.f}; + isSelPion = false; + isSelKaon = false; + if (!model.enabled) { + return false; + } + const int nModel = model.response.getModelBin(features.pt); + if (nModel < 0) { // outside the pT range covered by the models + return false; + } + auto inputFeatures = model.response.getInputFeatures(features); + const auto output = model.response.getModelOutput(inputFeatures, nModel); + + // Collapse whatever the model emits onto (background, pion-role, kaon-role). For the 5-class + // model the two origin scores of a role are summed: P(pion) = P(prompt pion) + P(non-prompt + // pion), which is the probability the track is that role's daughter irrespective of where the + // charm hadron came from. That keeps one threshold per role and makes the 3- and 5-class + // working points mean the same thing. + scores[0] = output[0]; + if (model.nClasses == NClassesWithOrigin) { + scores[1] = output[1] + output[3]; + scores[2] = output[2] + output[4]; + } else { + scores[1] = output[1]; + scores[2] = output[2]; + } + + isSelPion = scores[1] > model.thresholdPion[nModel]; + isSelKaon = scores[2] > model.thresholdKaon[nModel]; + return true; + } + + /// One row of aod::HfSkimTrackSels for a track that was not scored. + void fillEmptyTrackRow(bool isPositive) + { + rowSelectedTrack(-1.f, -1.f, -1.f, false, false, + -1.f, -1.f, -1.f, false, false, + -1.f, -1.f, -1.f, false, false, + -1.f, -1.f, -1.f, false, false, + isPositive, false); + } + + /// Common loop over the collisions: event selection, then per-track work. + /// With IsTraining and !IsMc the training table is filled on data: no labels (all -1), and every + /// preselected association is downsampled with bkgSamplingFraction, to study the feature + /// distributions the model sees on data. + template + void runTrackSelection(TCollisions const& collisions, + aod::TrackAssoc const& trackIndices, + TTracks const& /*tracks*/, + TMcParticles const& mcParticles, + aod::BcFullInfos const& /*bcs*/) + { + const bool doTiming = config.enableTiming; + rowSelectedCollision.reserve(collisions.size()); + + for (const auto& collision : collisions) { + const auto tStartCollision = StopWatch::now(); + timing.fill(0.); + registry.fill(HIST("hCounters"), CountTrackSelCollisions); + + const auto tStartEvSel = StopWatch::now(); + float centrality{-1.f}; + const auto rejectionMask = hfEvSel.getHfCollisionRejectionMask(collision, centrality, ccdb, registry); + if (config.fillHistograms) { + hfEvSel.fillHistograms(collision, rejectionMask, centrality); + } + rowSelectedCollision(rejectionMask, centrality); + if (doTiming) { + timing[TimeTrackSelEvSel] += elapsedSeconds(tStartEvSel); + } + + const auto thisCollId = collision.globalIndex(); + const auto groupedTrackIndices = trackIndices.sliceBy(trackIndicesPerCollision, thisCollId); + + if (rejectionMask != 0) { + // collision rejected: still fill one row per association to keep the table aligned + if constexpr (!IsTraining) { + for (auto iAssoc{0u}; iAssoc < groupedTrackIndices.size(); ++iAssoc) { + fillEmptyTrackRow(false); + } + } + continue; + } + + // set the magnetic field for this run + const auto bc = collision.template bc_as(); + initCCDB(bc, runNumber, ccdb, config.isRun2 ? config.ccdbPathGrp : config.ccdbPathGrpMag, lut, config.isRun2); + + [[maybe_unused]] int nPresel{0}; + [[maybe_unused]] int nSelPionDzero{0}; + [[maybe_unused]] int nSelKaonDzero{0}; + [[maybe_unused]] int nSelPionDplus{0}; + [[maybe_unused]] int nSelKaonDplus{0}; + [[maybe_unused]] int nSelPionDs{0}; + [[maybe_unused]] int nSelKaonDs{0}; + [[maybe_unused]] int nSelPionDstar{0}; + [[maybe_unused]] int nSelKaonDstar{0}; + [[maybe_unused]] int nSelAny{0}; + + for (const auto& trackId : groupedTrackIndices) { + const auto track = trackId.template track_as(); + registry.fill(HIST("hCounters"), CountAssociations); + + if (config.fillHistograms) { + registry.fill(HIST("hPtAllTracks"), track.pt()); + } + + if (!passesCheapPreselection(track)) { + if constexpr (!IsTraining) { + // one row per association, so that the table stays aligned with aod::TrackAssoc + fillEmptyTrackRow(track.sign() > 0); + } + continue; + } + registry.fill(HIST("hCounters"), CountCheapPresel); + + const auto tStartFeatures = StopWatch::now(); + HfTrackMlFeatures features{}; + fillFeatures(collision, track, centrality, features); + if (doTiming) { + timing[TimeTrackSelFeatures] += elapsedSeconds(tStartFeatures); + } + + const bool preselected = isPreselected(track, features.pt, features.eta); + if (preselected) { + ++nPresel; + registry.fill(HIST("hCounters"), CountPresel); + if (config.fillHistograms) { + registry.fill(HIST("hPtPreselected"), features.pt); + } + } + + if constexpr (IsTraining) { + if (!preselected) { + continue; + } + // The downsampling verdict is drawn first, from a deterministic hash of the track index, + // so that the labelling can skip the charm ancestry veto for the background tracks that + // would be thrown away right afterwards. Signal tracks are never downsampled. + const bool keepIfBackground = + deterministicUniform(static_cast(track.globalIndex()), config.bkgSamplingSeed) < config.bkgSamplingFraction; + + if constexpr (!IsMc) { + // data: no labels, a uniform random subsample of the preselected associations + if (!keepIfBackground) { + continue; + } + rowTrackTrain(features.pt, features.eta, features.dcaXY, features.dcaZ, + features.sigmaDcaXY, features.sigmaDcaZ, features.normDcaXY, features.normDcaZ, + features.signed1Pt, features.tgl, features.sign, features.isPvContributor, + features.itsNCls, features.itsNClsInnerBarrel, features.itsChi2NCl, + features.tpcNClsFound, features.tpcNClsCrossedRows, features.tpcCrossedRowsOverFindableCls, + features.tpcChi2NCl, features.tpcFractionSharedCls, + features.tpcNSigmaPi, features.tpcNSigmaKa, features.centrality, + track.globalIndex(), + -1, -1.f, -1, Channels::kNone, -1, + 0, false); + continue; + } else { + float pt{-1.f}; + int pdgMother{-1}; + int8_t isPrompt{-1}, channel{Channels::kNone}; + const int8_t label = getTrainingLabel( + track, mcParticles, pdgMother, keepIfBackground, + pt, isPrompt, channel); + + if (label < 0) { // usable for neither model + continue; + } + // only now, on the tracks that are actually written out, pay for the ancestry scan + const bool hasCharmBeautyAncestor = hasHeavyFlavourAncestor(track, mcParticles); + if (config.fillHistograms) { + if (channel == Channels::kDzeroDirect) { + registry.fill(HIST("hLabelDzero"), label); + if (channel >= 0) { + registry.fill(HIST("hChannelDzero"), channel); + } + } + if (channel == Channels::kDplusDirect || + channel == Channels::kDplusKstar) { + registry.fill(HIST("hLabelDplus"), label); + if (channel >= 0) { + registry.fill(HIST("hChannelDplus"), channel); + } + } + if (channel == Channels::kDsDirect || + channel == Channels::kDsKstarK || + channel == Channels::kDsPhiPi) { + registry.fill(HIST("hLabelDs"), label); + if (channel >= 0) { + registry.fill(HIST("hChannelDs"), channel); + } + } + if (channel == Channels::kDstarToD0Pi) { + registry.fill(HIST("hLabelDstar"), label); + if (channel >= 0) { + registry.fill(HIST("hChannelDstar"), channel); + } + } + } + rowTrackTrain(features.pt, features.eta, features.dcaXY, features.dcaZ, + features.sigmaDcaXY, features.sigmaDcaZ, features.normDcaXY, features.normDcaZ, + features.signed1Pt, features.tgl, features.sign, features.isPvContributor, + features.itsNCls, features.itsNClsInnerBarrel, features.itsChi2NCl, + features.tpcNClsFound, features.tpcNClsCrossedRows, features.tpcCrossedRowsOverFindableCls, + features.tpcChi2NCl, features.tpcFractionSharedCls, + features.tpcNSigmaPi, features.tpcNSigmaKa, features.centrality, + track.globalIndex(), + label, pt, isPrompt, channel, pdgMother, + track.mcMask(), hasCharmBeautyAncestor); + } + } else { + if (!preselected) { + fillEmptyTrackRow(track.sign() > 0); + continue; + } + + std::array scoresDzero{-1.f, -1.f, -1.f}; + std::array scoresDplus{-1.f, -1.f, -1.f}; + std::array scoresDs{-1.f, -1.f, -1.f}; + std::array scoresDstar{-1.f, -1.f, -1.f}; + bool isSelPionDzero{false}, isSelKaonDzero{false}; + bool isSelPionDplus{false}, isSelKaonDplus{false}; + bool isSelPionDs{false}, isSelKaonDs{false}; + bool isSelPionDstar{false}, isSelKaonDstar{false}; + + const auto tStartDzero = StopWatch::now(); + if (evaluateModel(modelDzero, features, scoresDzero, isSelPionDzero, isSelKaonDzero)) { + registry.fill(HIST("hCounters"), CountInferenceDzero); + } + if (doTiming) { + timing[TimeTrackSelInferenceDzero] += elapsedSeconds(tStartDzero); + } + + const auto tStartDplus = StopWatch::now(); + if (evaluateModel(modelDplus, features, scoresDplus, isSelPionDplus, isSelKaonDplus)) { + registry.fill(HIST("hCounters"), CountInferenceDplus); + } + if (doTiming) { + timing[TimeTrackSelInferenceDplus] += elapsedSeconds(tStartDplus); + } + + const auto tStartDs = StopWatch::now(); + if (evaluateModel(modelDs, features, scoresDs, isSelPionDs, isSelKaonDs)) { + registry.fill(HIST("hCounters"), CountInferenceDs); + } + if (doTiming) { + timing[TimeTrackSelInferenceDs] += elapsedSeconds(tStartDs); + } + + const auto tStartDstar = StopWatch::now(); + if (evaluateModel(modelDstar, features, scoresDstar, isSelPionDstar, isSelKaonDstar)) { + registry.fill(HIST("hCounters"), CountInferenceDstar); + } + if (doTiming) { + timing[TimeTrackSelInferenceDstar] += elapsedSeconds(tStartDstar); + } + + uint16_t mask{0}; + if (isSelPionDzero) { + SETBIT(mask, BitPionDzero); + ++nSelPionDzero; + registry.fill(HIST("hCounters"), CountSelPionDzero); + } + if (isSelKaonDzero) { + SETBIT(mask, BitKaonDzero); + ++nSelKaonDzero; + registry.fill(HIST("hCounters"), CountSelKaonDzero); + } + if (isSelPionDplus) { + SETBIT(mask, BitPionDplus); + ++nSelPionDplus; + registry.fill(HIST("hCounters"), CountSelPionDplus); + } + if (isSelKaonDplus) { + SETBIT(mask, BitKaonDplus); + ++nSelKaonDplus; + registry.fill(HIST("hCounters"), CountSelKaonDplus); + } + if (isSelPionDs) { + SETBIT(mask, BitPionDs); + ++nSelPionDs; + registry.fill(HIST("hCounters"), CountSelPionDs); + } + if (isSelKaonDs) { + SETBIT(mask, BitKaonDs); + ++nSelKaonDs; + registry.fill(HIST("hCounters"), CountSelKaonDs); + } + if (isSelPionDstar) { + SETBIT(mask, BitPionDstar); + ++nSelPionDstar; + registry.fill(HIST("hCounters"), CountSelPionDstar); + } + if (isSelKaonDstar) { + SETBIT(mask, BitKaonDstar); + ++nSelKaonDstar; + registry.fill(HIST("hCounters"), CountSelKaonDstar); + } + const bool isSelAny = mask != 0; + if (isSelAny) { + ++nSelAny; + registry.fill(HIST("hCounters"), CountSelAny); + } + + if (config.fillHistograms) { + registry.fill(HIST("hScoreBkgDzero"), scoresDzero[0]); + registry.fill(HIST("hScorePionDzero"), scoresDzero[1]); + registry.fill(HIST("hScoreKaonDzero"), scoresDzero[2]); + registry.fill(HIST("hScoreBkgDplus"), scoresDplus[0]); + registry.fill(HIST("hScorePionDplus"), scoresDplus[1]); + registry.fill(HIST("hScoreKaonDplus"), scoresDplus[2]); + registry.fill(HIST("hScoreBkgDs"), scoresDs[0]); + registry.fill(HIST("hScorePionDs"), scoresDs[1]); + registry.fill(HIST("hScoreKaonDs"), scoresDs[2]); + registry.fill(HIST("hScoreBkgDstar"), scoresDstar[0]); + registry.fill(HIST("hScorePionDstar"), scoresDstar[1]); + registry.fill(HIST("hScoreKaonDstar"), scoresDstar[2]); + if (isSelPionDzero) { + registry.fill(HIST("hPtSelPionDzero"), features.pt); + } + if (isSelKaonDzero) { + registry.fill(HIST("hPtSelKaonDzero"), features.pt); + } + if (isSelPionDplus) { + registry.fill(HIST("hPtSelPionDplus"), features.pt); + } + if (isSelKaonDplus) { + registry.fill(HIST("hPtSelKaonDplus"), features.pt); + } + if (isSelPionDs) { + registry.fill(HIST("hPtSelPionDs"), features.pt); + } + if (isSelKaonDs) { + registry.fill(HIST("hPtSelKaonDs"), features.pt); + } + if (isSelPionDstar) { + registry.fill(HIST("hPtSelPionDstar"), features.pt); + } + if (isSelKaonDstar) { + registry.fill(HIST("hPtSelKaonDstar"), features.pt); + } + if (isSelAny) { + registry.fill(HIST("hRoleMask"), mask); + } + } + + const bool isPositive = track.sign() > 0; + rowSelectedTrack(scoresDzero[0], scoresDzero[1], scoresDzero[2], isSelPionDzero, isSelKaonDzero, + scoresDplus[0], scoresDplus[1], scoresDplus[2], isSelPionDplus, isSelKaonDplus, + scoresDs[0], scoresDs[1], scoresDs[2], isSelPionDs, isSelKaonDs, + scoresDstar[0], scoresDstar[1], scoresDstar[2], isSelPionDstar, isSelKaonDstar, + isPositive, isSelAny); + if (config.fillSelMaskTable && isSelAny) { + uint16_t maskOut = mask; + if (isPositive) { + SETBIT(maskOut, BitPositive); + } + rowSelMask(static_cast(thisCollId), maskOut); + } + } + } + + if constexpr (!IsTraining) { + if (config.fillHistograms) { + registry.fill(HIST("hNTracksPreselPerColl"), nPresel); + registry.fill(HIST("hNTracksSelPionDzeroPerColl"), nSelPionDzero); + registry.fill(HIST("hNTracksSelKaonDzeroPerColl"), nSelKaonDzero); + registry.fill(HIST("hNTracksSelPionDplusPerColl"), nSelPionDplus); + registry.fill(HIST("hNTracksSelKaonDplusPerColl"), nSelKaonDplus); + registry.fill(HIST("hNTracksSelPionDsPerColl"), nSelPionDs); + registry.fill(HIST("hNTracksSelKaonDsPerColl"), nSelKaonDs); + registry.fill(HIST("hNTracksSelPionDstarPerColl"), nSelPionDstar); + registry.fill(HIST("hNTracksSelKaonDstarPerColl"), nSelKaonDstar); + registry.fill(HIST("hNTracksSelAnyPerColl"), nSelAny); + } + } + + if (doTiming) { + timing[TimeTrackSelTotal] = elapsedSeconds(tStartCollision); + timing[TimeTrackSelRemainder] = timing[TimeTrackSelTotal] - timing[TimeTrackSelEvSel] - timing[TimeTrackSelFeatures] - timing[TimeTrackSelInferenceDplus] - timing[TimeTrackSelInferenceDs]; + for (int iStage = 0; iStage < NTimingTrackSel; ++iStage) { + registry.fill(HIST("hTiming"), iStage, timing[iStage]); + } + } + } + } + + void processTraining(CollisionsWithCent const& collisions, + aod::TrackAssoc const& trackIndices, + TracksWithSelAndPidMc const& tracks, + aod::McParticles const& mcParticles, + aod::BcFullInfos const& bcs) + { + runTrackSelection(collisions, trackIndices, tracks, mcParticles, bcs); + } + PROCESS_SWITCH(HfDplusTrackSelMl, processTraining, "Produce the track-level training sample (MC only)", false); + + void processTrainingData(CollisionsWithCent const& collisions, + aod::TrackAssoc const& trackIndices, + TracksWithSelAndPid const& tracks, + aod::BcFullInfos const& bcs) + { + runTrackSelection(collisions, trackIndices, tracks, tracks, bcs); + } + PROCESS_SWITCH(HfDplusTrackSelMl, processTrainingData, "Write the unlabelled feature table on data, downsampled with bkgSamplingFraction", false); + + void processApplication(CollisionsWithCent const& collisions, + aod::TrackAssoc const& trackIndices, + TracksWithSelAndPid const& tracks, + aod::BcFullInfos const& bcs) + { + runTrackSelection(collisions, trackIndices, tracks, tracks, bcs); + } + PROCESS_SWITCH(HfDplusTrackSelMl, processApplication, "Apply the track-level ML selection", true); +}; + +/// Charm 3-prong combinatorics over the ML-selected tracks. +struct HfDplusCombinatorics { + Produces rowTrackIndexProng3; + Produces rowCandScoreScan; + + struct : ConfigurableGroup { + Configurable fillHistograms{"fillHistograms", true, "fill QA histograms"}; + Configurable fillScoreScanTable{"fillScoreScanTable", false, "write one aod::HfCandScoreScans row per candidate, so the working-point scan can be redone offline"}; + Configurable enableTiming{"enableTiming", false, "account the CPU time of the individual stages into hTiming"}; + Configurable ptTolerance{"ptTolerance", 0.1, "pT tolerance in GeV/c for the preselections before vertex reconstruction"}; + // D+ preselections, same layout and defaults as in the track-index skim creator + Configurable> binsPtDplusToPiKPi{"binsPtDplusToPiKPi", std::vector{hf_cuts_presel_3prong::vecBinsPt}, "pT bin limits for D+->piKpi pT-dependent cuts"}; + Configurable> cutsDplusToPiKPi{"cutsDplusToPiKPi", {hf_cuts_presel_3prong::Cuts[0], hf_cuts_presel_3prong::NBinsPt, hf_cuts_presel_3prong::NCutVars, hf_cuts_presel_3prong::labelsPt, hf_cuts_presel_3prong::labelsCutVar}, "D+->piKpi selections per pT bin"}; + // Ds preselections, same layout and defaults as in the track-index skim creator + Configurable> binsPtDsToKKPi{"binsPtDsToKKPi", std::vector{hf_cuts_presel_ds::vecBinsPt}, "pT bin limits for Ds->KKpi pT-dependent cuts"}; + Configurable> cutsDsToKKPi{"cutsDsToKKPi", {hf_cuts_presel_ds::Cuts[0], hf_cuts_presel_ds::NBinsPt, hf_cuts_presel_ds::NCutVars, hf_cuts_presel_ds::labelsPt, hf_cuts_presel_ds::labelsCutVar}, "Ds->KKpi selections per pT bin"}; + // two-track speedups, carried over from the skim creator so that the CPU comparison is fair + Configurable minTwoTrackDecayLengthFor3Prongs{"minTwoTrackDecayLengthFor3Prongs", 0., "minimum decay length computed with 2 tracks, to speed up the combinatorics"}; + Configurable maxTwoTrackChi2PcaFor3Prongs{"maxTwoTrackChi2PcaFor3Prongs", 1.e10, "maximum chi2 PCA computed with 2 tracks, to speed up the combinatorics"}; + // vertexing + Configurable propagateToPCA{"propagateToPCA", true, "create tracks version propagated to PCA"}; + Configurable useAbsDCA{"useAbsDCA", false, "minimise abs. distance rather than chi2"}; + Configurable useWeightedFinalPCA{"useWeightedFinalPCA", false, "recalculate vertex position using track covariances, effective only if useAbsDCA is true"}; + Configurable maxR{"maxR", 200., "reject PCA's above this radius"}; + Configurable maxDZIni{"maxDZIni", 4., "reject (if>0) PCA candidate if tracks DZ exceeds threshold"}; + Configurable minParamChange{"minParamChange", 1.e-3, "stop iterations if largest change of any X is smaller than this"}; + Configurable minRelChi2Change{"minRelChi2Change", 0.9, "stop iterations if chi2/chi2old > this"}; + // CCDB + Configurable isRun2{"isRun2", false, "enable Run 2 or Run 3 GRP objects for the magnetic field"}; + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable ccdbPathLut{"ccdbPathLut", "GLO/Param/MatLUT", "path for LUT parametrization"}; + Configurable ccdbPathGrp{"ccdbPathGrp", "GLO/GRP/GRP", "path of the grp file (Run 2)"}; + Configurable ccdbPathGrpMag{"ccdbPathGrpMag", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object (Run 3)"}; + } config; + + SliceCache cache; + o2::vertexing::DCAFitterN<2> df2; // 2-prong fitter, only for the two-track speedups + o2::vertexing::DCAFitterN<3> df3; // 3-prong vertex fitter + Service ccdb{}; + o2::base::MatLayerCylSet* lut{}; + o2::base::Propagator::MatCorrType noMatCorr = o2::base::Propagator::MatCorrType::USEMatCorrNONE; + int runNumber{}; + + /// Prong ordering conventions of aod::Hf3Prongs, shared with the track-index skim creator: + /// prong 1 is always the opposite-sign track, prongs 0 and 2 the same-sign pair. + std::array massDplusToPiKPi{MassPiPlus, MassKPlus, MassPiPlus}; + std::array, 2> massDsToKKPi{std::array{MassKPlus, MassKPlus, MassPiPlus}, // hypothesis 0: K K pi + std::array{MassPiPlus, MassKPlus, MassKPlus}}; // hypothesis 1: pi K K + + /// One selected track, propagated to the collision under study once and reused by every triplet + /// it takes part in. The original D+-only implementation re-propagated the innermost prong + /// inside the innermost loop, i.e. once per triplet rather than once per track. + struct ProngCandidate { + o2::track::TrackParCov parCov; + std::array pVec; + int64_t trackIndex; + uint8_t mask; + // kept so that every built candidate can record the per-role minimum score it came from, + // which is what makes the offline working-point scan possible from a single run + float scorePionDplus; + float scoreKaonDplus; + float scorePionDs; + float scoreKaonDs; + }; + + /// Per-collision working area, allocated once and cleared for each collision. + struct ProngCache { + std::vector prongs; + std::vector pionDplus; + std::vector kaonDplus; + std::vector pionDs; + std::vector kaonDs; + std::vector dsSameSign; ///< can sit in a same-sign slot of a Ds candidate: pi(Ds) or K(Ds) + std::vector sameSignAny; ///< can sit in a same-sign slot of any candidate: pi(D+), pi(Ds) or K(Ds) + std::vector oppSignAny; ///< can sit in the opposite-sign slot: K(D+) or K(Ds) + + void clear() + { + prongs.clear(); + pionDplus.clear(); + kaonDplus.clear(); + pionDs.clear(); + kaonDs.clear(); + dsSameSign.clear(); + sameSignAny.clear(); + oppSignAny.clear(); + } + }; + ProngCache cachePos; + ProngCache cacheNeg; + + // per-collision accumulators + std::array timing{}; + std::array counters{}; + + using SelectedCollisions = soa::Filtered>; + using TracksWithCov = aod::TracksWCovDcaExtra; + using FilteredTrackAssocSel = soa::Filtered>; + + Filter filterSelectCollisions = (aod::hf_skim_sel_collision::whyRejectColl == static_cast(0)); + Filter filterSelectTrackIds = (aod::hf_skim_track_sel::isSelAny == true); + + // Registers the slice cache for the filtered association table. Without it the partitions below + // cannot be sliced per collision and the device dies with + // (AOD/TRACKASSOC/0) /fIndexCollisions not found neither in sorted or unsorted cache + Preslice trackIndicesPerCollision = aod::track_association::collisionId; + + // Partitions used by the legacy D+-only process function. + Partition positivePions = aod::hf_skim_track_sel::isPositive == true && aod::hf_skim_track_sel::isSelPionDplus == true; + Partition negativePions = aod::hf_skim_track_sel::isPositive == false && aod::hf_skim_track_sel::isSelPionDplus == true; + Partition positiveKaons = aod::hf_skim_track_sel::isPositive == true && aod::hf_skim_track_sel::isSelKaonDplus == true; + Partition negativeKaons = aod::hf_skim_track_sel::isPositive == false && aod::hf_skim_track_sel::isSelKaonDplus == true; + + HistogramRegistry registry{"registry"}; + + void init(InitContext const&) + { + const std::array doProcess = {doprocessDplus, doprocessDplusCached, doprocessSeparateLoops, doprocessMergedLoop, doprocessDummy}; + if (std::accumulate(doProcess.begin(), doProcess.end(), 0) != 1) { + LOGP(fatal, "One and only one process function of HfDplusCombinatorics must be enabled at a time!"); + } + + // CPU accounting and counters are booked unconditionally: an empty histogram is a clearer + // statement than a missing one. + registry.add("hTiming", "CPU time by stage;;CPU time (s)", {HistType::kTH1D, {{NTimingCombinatorics, -0.5, NTimingCombinatorics - 0.5}}}); + const char* timingLabels[NTimingCombinatorics] = {"total", "prong cache", "2-prong fits", "3-prong fits", "remainder"}; + for (int i = 0; i < NTimingCombinatorics; ++i) { + registry.get(HIST("hTiming"))->GetXaxis()->SetBinLabel(i + 1, timingLabels[i]); + } + registry.add("hCounters", "counters;;entries", {HistType::kTH1D, {{NCounterCombinatorics, -0.5, NCounterCombinatorics - 0.5}}}); + const char* counterLabels[NCounterCombinatorics] = {"collisions", "triplets", "triplets presel", "2-prong fits", "3-prong fits", + "rows", "rows D^{+}", "rows D_{s}", "rows both"}; + for (int i = 0; i < NCounterCombinatorics; ++i) { + registry.get(HIST("hCounters"))->GetXaxis()->SetBinLabel(i + 1, counterLabels[i]); + } + + if (doprocessDummy) { + return; + } + + ccdb->setURL(config.ccdbUrl); + ccdb->setCaching(true); + ccdb->setLocalObjectValidityChecking(); + lut = o2::base::MatLayerCylSet::rectifyPtrFromFile(ccdb->get(config.ccdbPathLut)); + + df2.setPropagateToPCA(config.propagateToPCA); + df2.setMaxR(config.maxR); + df2.setMaxDZIni(config.maxDZIni); + df2.setMinParamChange(config.minParamChange); + df2.setMinRelChi2Change(config.minRelChi2Change); + df2.setUseAbsDCA(config.useAbsDCA); + df2.setWeightedFinalPCA(config.useWeightedFinalPCA); + + df3.setPropagateToPCA(config.propagateToPCA); + df3.setMaxR(config.maxR); + df3.setMaxDZIni(config.maxDZIni); + df3.setMinParamChange(config.minParamChange); + df3.setMinRelChi2Change(config.minRelChi2Change); + df3.setUseAbsDCA(config.useAbsDCA); + df3.setWeightedFinalPCA(config.useWeightedFinalPCA); + + if (config.fillHistograms) { + registry.add("hMassDplusToPiKPi", "D^{+} candidates;inv. mass (#pi K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 1.6, 2.1}}}); + registry.add("hMassDsToKKPi", "D_{s} candidates;inv. mass (K K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 1.6, 2.3}}}); + registry.add("hMassKK", "D_{s} candidates;inv. mass (K K) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{200, 0.98, 1.18}}}); + registry.add("hPtCand", "3-prong candidates;#it{p}_{T} (GeV/#it{c});entries", {HistType::kTH1D, {{200, 0., 20.}}}); + registry.add("hNCandPerColl", "candidates per collision;n. candidates;entries", {HistType::kTH1D, {{500, -0.5, 499.5}}}); + registry.add("hVtxX", "secondary vertex;x (cm);entries", {HistType::kTH1D, {{200, -2., 2.}}}); + } + } + + /// Preselection of a D+ candidate before vertex reconstruction: pT bin, invariant mass and + /// minimum prong pT. Mirrors applyPreselection3Prong of the track-index skim creator, reduced + /// to the D+ hypothesis. + /// \param pVecTrack0, pVecTrack1, pVecTrack2 are the prong momenta (π, K, π ordering) + /// \return true if the candidate survives + template + bool applyPreselectionDplus(TVec const& pVecTrack0, TVec const& pVecTrack1, TVec const& pVecTrack2) + { + const auto pt = RecoDecay::pt(pVecTrack0, pVecTrack1, pVecTrack2) + config.ptTolerance; + const auto binPt = findBin(config.binsPtDplusToPiKPi, pt); + if (binPt == -1) { + return false; + } + + const double minMass = config.cutsDplusToPiKPi->get(binPt, 0u); + const double maxMass = config.cutsDplusToPiKPi->get(binPt, 1u); + if (minMass >= 0. && maxMass > 0.) { + const std::array arrMom{pVecTrack0, pVecTrack1, pVecTrack2}; + const double mass2 = RecoDecay::m2(arrMom, massDplusToPiKPi); + if (mass2 < minMass * minMass || mass2 >= maxMass * maxMass) { + return false; + } + } + + const auto ptProngMin = config.cutsDplusToPiKPi->get(binPt, 4u); + const auto pt2ProngMin = ptProngMin * ptProngMin; + if (RecoDecay::pt2(pVecTrack0) < pt2ProngMin || RecoDecay::pt2(pVecTrack1) < pt2ProngMin || RecoDecay::pt2(pVecTrack2) < pt2ProngMin) { + return false; + } + + return true; + } + + /// Preselection of a Ds candidate before vertex reconstruction, for the mass hypotheses still + /// allowed by the track selection. Mirrors applyPreselection3Prong plus applyPreselectionPhiDecay + /// of the track-index skim creator: the two hypotheses are K K π (bit 0) and π K K (bit 1), the + /// opposite-sign prong being the middle one in both, and each surviving hypothesis must also + /// have its kaon pair compatible with the φ(1020). + /// \param pVecTrack0, pVecTrack1, pVecTrack2 are the prong momenta + /// \param hypoMask is the incoming hypothesis mask, narrowed in place + /// \return true if at least one hypothesis survives + template + bool applyPreselectionDs(TVec const& pVecTrack0, TVec const& pVecTrack1, TVec const& pVecTrack2, uint8_t& hypoMask) + { + const auto pt = RecoDecay::pt(pVecTrack0, pVecTrack1, pVecTrack2) + config.ptTolerance; + const auto binPt = findBin(config.binsPtDsToKKPi, pt); + if (binPt == -1) { + hypoMask = 0; + return false; + } + + const auto ptProngMin = config.cutsDsToKKPi->get(binPt, 4u); + const auto pt2ProngMin = ptProngMin * ptProngMin; + if (RecoDecay::pt2(pVecTrack0) < pt2ProngMin || RecoDecay::pt2(pVecTrack1) < pt2ProngMin || RecoDecay::pt2(pVecTrack2) < pt2ProngMin) { + hypoMask = 0; + return false; + } + + const double minMass = config.cutsDsToKKPi->get(binPt, 0u); + const double maxMass = config.cutsDsToKKPi->get(binPt, 1u); + const std::array arrMom{pVecTrack0, pVecTrack1, pVecTrack2}; + if (minMass >= 0. && maxMass > 0.) { + for (uint8_t iHypo = 0; iHypo < 2; ++iHypo) { + if (!TESTBIT(hypoMask, iHypo)) { + continue; + } + const double mass2 = RecoDecay::m2(arrMom, massDsToKKPi[iHypo]); + if (mass2 < minMass * minMass || mass2 >= maxMass * maxMass) { + CLRBIT(hypoMask, iHypo); + } + } + if (hypoMask == 0) { + return false; + } + } + + // φ(1020) window on the kaon pair: prongs (0, 1) for the K K π hypothesis, (1, 2) for π K K + const double deltaMassKKMax = config.cutsDsToKKPi->get(binPt, 5u); + if (deltaMassKKMax > 0.) { + if (TESTBIT(hypoMask, 0)) { + const double massKK = RecoDecay::m(std::array{pVecTrack0, pVecTrack1}, std::array{MassKPlus, MassKPlus}); + if (std::abs(massKK - MassPhi) > deltaMassKKMax) { + CLRBIT(hypoMask, 0); + } + } + if (TESTBIT(hypoMask, 1)) { + const double massKK = RecoDecay::m(std::array{pVecTrack1, pVecTrack2}, std::array{MassKPlus, MassKPlus}); + if (std::abs(massKK - MassPhi) > deltaMassKKMax) { + CLRBIT(hypoMask, 1); + } + } + } + + return hypoMask != 0; + } + + /// Selection after vertex reconstruction: cos(pointing angle) and decay length. + /// Mirrors applySelection3Prong of the track-index skim creator. + template + bool applySelectionAfterVertexing(TBins const& binsPt, TCuts const& cuts, TVec const& pVecCand, TVtx const& secVtx, TPv const& primVtx) + { + const auto binPt = findBin(binsPt, RecoDecay::pt(pVecCand)); + if (binPt == -1) { + return false; + } + if (RecoDecay::cpa(primVtx, secVtx, pVecCand) < cuts->get(binPt, 2u)) { + return false; + } + if (RecoDecay::distance(primVtx, secVtx) < cuts->get(binPt, 3u)) { + return false; + } + return true; + } + + /// Check whether the vertex built with the first two prongs is compatible with a 3-prong decay. + /// Used only if the corresponding speedup configurables are set. + template + bool isTwoTrackVertexSelected(TVtx const& secVtx, TPv const& primVtx, TFitter const& fitter) + { + if (config.maxTwoTrackChi2PcaFor3Prongs < MaxTwoTrackChi2PcaDisabled && fitter.getChi2AtPCACandidate() > config.maxTwoTrackChi2PcaFor3Prongs) { + return false; + } + if (config.minTwoTrackDecayLengthFor3Prongs > 0.f && RecoDecay::distance(primVtx, secVtx) < config.minTwoTrackDecayLengthFor3Prongs) { + return false; + } + return true; + } + + /// Whether the two-track speedups are configured at all. + bool twoTrackSpeedupsEnabled() + { + return config.minTwoTrackDecayLengthFor3Prongs > 0.f || config.maxTwoTrackChi2PcaFor3Prongs < MaxTwoTrackChi2PcaDisabled; + } + + /// Fill the per-collision prong cache: propagate every selected association to this collision + /// once, and build the per-role index lists the loops iterate over. + /// \param collision is the collision under study + /// \param groupedAssocs are its selected (track, collision) associations + /// \param tracks is the track table + template + void buildProngCache(TCollision const& collision, TAssocs const& groupedAssocs, TTracks const& /*tracks*/) + { + cachePos.clear(); + cacheNeg.clear(); + const auto thisCollId = collision.globalIndex(); + + for (const auto& assocId : groupedAssocs) { + const auto track = assocId.template track_as(); + + uint8_t mask{0}; + if (assocId.isSelPionDplus()) { + SETBIT(mask, BitPionDplus); + } + if (assocId.isSelKaonDplus()) { + SETBIT(mask, BitKaonDplus); + } + if (assocId.isSelPionDs()) { + SETBIT(mask, BitPionDs); + } + if (assocId.isSelKaonDs()) { + SETBIT(mask, BitKaonDs); + } + if (mask == 0) { + continue; + } + + ProngCandidate prong{getTrackParCov(track), track.pVector(), track.globalIndex(), mask, + assocId.scorePionDplus(), assocId.scoreKaonDplus(), + assocId.scorePionDs(), assocId.scoreKaonDs()}; + if (thisCollId != track.collisionId()) { + // not the "default" collision of this track: re-propagate it, once + std::array dcaInfo{-999.f, -999.f}; + o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, prong.parCov, 2.f, noMatCorr, &dcaInfo); + getPxPyPz(prong.parCov, prong.pVec); + } + + auto& cacheThis = assocId.isPositive() ? cachePos : cacheNeg; + const int index = static_cast(cacheThis.prongs.size()); + cacheThis.prongs.push_back(prong); + if (TESTBIT(mask, BitPionDplus)) { + cacheThis.pionDplus.push_back(index); + } + if (TESTBIT(mask, BitKaonDplus)) { + cacheThis.kaonDplus.push_back(index); + } + if (TESTBIT(mask, BitPionDs)) { + cacheThis.pionDs.push_back(index); + } + if (TESTBIT(mask, BitKaonDs)) { + cacheThis.kaonDs.push_back(index); + } + if (TESTBIT(mask, BitPionDs) || TESTBIT(mask, BitKaonDs)) { + cacheThis.dsSameSign.push_back(index); + } + if (TESTBIT(mask, BitPionDplus) || TESTBIT(mask, BitPionDs) || TESTBIT(mask, BitKaonDs)) { + cacheThis.sameSignAny.push_back(index); + } + if (TESTBIT(mask, BitKaonDplus) || TESTBIT(mask, BitKaonDs)) { + cacheThis.oppSignAny.push_back(index); + } + } + } + + /// Run the 2-prong speedup fit on a (same-sign, opposite-sign) pair. + /// \return true if the pair is worth continuing with + template + bool passesTwoTrackSpeedup(ProngCandidate const& prongSameSign, ProngCandidate const& prongOppSign, TPv const& pvCoord, bool doTiming) + { + const auto tStart = StopWatch::now(); + int nVtx{0}; + ++counters[CountFit2Prong]; + try { + nVtx = df2.process(prongSameSign.parCov, prongOppSign.parCov); + } catch (...) { + nVtx = 0; + } + const bool ok = (nVtx != 0) && isTwoTrackVertexSelected(df2.getPCACandidate(), pvCoord, df2); + if (doTiming) { + timing[TimeCombFit2Prong] += elapsedSeconds(tStart); + } + return ok; + } + + /// Fit the 3-prong vertex and, if it converges, hand back the refitted prong momenta. + /// \return true if the fit converged + bool fitThreeProngVertex(ProngCandidate const& p0, ProngCandidate const& p1, ProngCandidate const& p2, + std::array& pvec0, std::array& pvec1, std::array& pvec2, + bool doTiming) + { + const auto tStart = StopWatch::now(); + int nVtx{0}; + ++counters[CountFit3Prong]; + try { + nVtx = df3.process(p0.parCov, p1.parCov, p2.parCov); + } catch (...) { + nVtx = 0; + } + if (nVtx != 0) { + df3.getTrack(0).getPxPyPzGlo(pvec0); + df3.getTrack(1).getPxPyPzGlo(pvec1); + df3.getTrack(2).getPxPyPzGlo(pvec2); + } + if (doTiming) { + timing[TimeCombFit3Prong] += elapsedSeconds(tStart); + } + return nVtx != 0; + } + + /// The triple loop itself, shared by every strategy so that they differ only in the track lists + /// they iterate over and in the hypotheses they are allowed to test. + /// + /// The prong ordering handed to the fitter and written to aod::Hf3Prongs is + /// (same-sign 1, opposite-sign, same-sign 2), the convention of the track-index skim creator. + /// + /// \param collision is the collision under study + /// \param sameSign, oppSign are the per-charge prong caches, in the roles this charge + /// combination gives them + /// \param sameSignList, oppSignList are the index lists to loop over + /// \param pairIsUseful(maskOpp, maskSame1) rejects a (same-sign, opposite-sign) pair before the + /// 2-prong speedup fit is paid for it + /// \param hypothesisOf(maskOpp, mask1, mask2, dplusOk, dsHypoMask) fills the hypotheses this + /// strategy is allowed to test on one triplet + /// \param nCand counts the rows written + template + void runTripleLoop(TCollision const& collision, + ProngCache const& sameSign, ProngCache const& oppSign, + std::vector const& sameSignList, std::vector const& oppSignList, + TPairFilter pairIsUseful, THypothesis hypothesisOf, + bool doTiming, int& nCand) + { + if (sameSignList.size() < 2 || oppSignList.empty()) { + return; + } + const auto thisCollId = collision.globalIndex(); + const std::array pvCoord{collision.posX(), collision.posY(), collision.posZ()}; + const bool useTwoTrackSpeedups = twoTrackSpeedupsEnabled(); + + for (const int iOpp : oppSignList) { // o2-linter: disable=const-ref-in-for-loop (int elements) + const auto& prongOpp = oppSign.prongs[iOpp]; + + for (std::size_t i1 = 0; i1 + 1 < sameSignList.size(); ++i1) { + const auto& prong1 = sameSign.prongs[sameSignList[i1]]; + if (!pairIsUseful(prongOpp.mask, prong1.mask)) { + continue; + } + if (useTwoTrackSpeedups && !passesTwoTrackSpeedup(prong1, prongOpp, pvCoord, doTiming)) { + continue; + } + + for (std::size_t i2 = i1 + 1; i2 < sameSignList.size(); ++i2) { + const auto& prong2 = sameSign.prongs[sameSignList[i2]]; + + bool dplusOk{false}; + uint8_t dsHypo{0}; + hypothesisOf(prongOpp.mask, prong1.mask, prong2.mask, dplusOk, dsHypo); + if (!dplusOk && dsHypo == 0) { + continue; + } + ++counters[CountTriplets]; + + if (dplusOk && !applyPreselectionDplus(prong1.pVec, prongOpp.pVec, prong2.pVec)) { + dplusOk = false; + } + if (dsHypo != 0) { + applyPreselectionDs(prong1.pVec, prongOpp.pVec, prong2.pVec, dsHypo); + } + if (!dplusOk && dsHypo == 0) { + continue; + } + ++counters[CountTripletsPresel]; + + std::array pvec0{}; + std::array pvec1{}; + std::array pvec2{}; + if (!fitThreeProngVertex(prong1, prongOpp, prong2, pvec0, pvec1, pvec2, doTiming)) { + continue; + } + + const auto& secondaryVertex = df3.getPCACandidate(); + const auto pVecCand = RecoDecay::pVec(pvec0, pvec1, pvec2); + + uint8_t hfFlag{0}; + if (dplusOk && applySelectionAfterVertexing(config.binsPtDplusToPiKPi, config.cutsDplusToPiKPi, pVecCand, secondaryVertex, pvCoord)) { + SETBIT(hfFlag, hf_cand_3prong::DecayType::DplusToPiKPi); + } + if (dsHypo != 0 && applySelectionAfterVertexing(config.binsPtDsToKKPi, config.cutsDsToKKPi, pVecCand, secondaryVertex, pvCoord)) { + SETBIT(hfFlag, hf_cand_3prong::DecayType::DsToKKPi); + } + if (hfFlag == 0) { + continue; + } + + rowTrackIndexProng3(thisCollId, prong1.trackIndex, prongOpp.trackIndex, prong2.trackIndex, hfFlag); + ++nCand; + ++counters[CountRows]; + const bool isDplus = TESTBIT(hfFlag, hf_cand_3prong::DecayType::DplusToPiKPi); + const bool isDs = TESTBIT(hfFlag, hf_cand_3prong::DecayType::DsToKKPi); + + if (config.fillScoreScanTable) { + // Prong order is (sameSign1, oppSign, sameSign2). For the D+ the opposite-sign prong + // is the kaon and both same-sign prongs are pions. For the Ds the opposite-sign prong + // is always a kaon, and the two hypotheses differ in which same-sign prong is the + // pion -- dsHypo carries only the hypotheses that survived the mass and phi windows. + const bool hypoKKPi = isDs && TESTBIT(dsHypo, 0); // prong1 = K, prong2 = pi + const bool hypoPiKK = isDs && TESTBIT(dsHypo, 1); // prong1 = pi, prong2 = K + rowCandScoreScan(RecoDecay::pt(pVecCand), + isDplus ? std::min(prong1.scorePionDplus, prong2.scorePionDplus) : kScoreScanNone, + isDplus ? prongOpp.scoreKaonDplus : kScoreScanNone, + hypoKKPi ? std::min(prong1.scoreKaonDs, prongOpp.scoreKaonDs) : kScoreScanNone, + hypoKKPi ? prong2.scorePionDs : kScoreScanNone, + hypoPiKK ? std::min(prong2.scoreKaonDs, prongOpp.scoreKaonDs) : kScoreScanNone, + hypoPiKK ? prong1.scorePionDs : kScoreScanNone); + } + if (isDplus) { + ++counters[CountRowsDplus]; + } + if (isDs) { + ++counters[CountRowsDs]; + } + if (isDplus && isDs) { + ++counters[CountRowsBoth]; + } + + if (config.fillHistograms) { + const std::array arrMom{pvec0, pvec1, pvec2}; + if (isDplus) { + registry.fill(HIST("hMassDplusToPiKPi"), RecoDecay::m(arrMom, massDplusToPiKPi)); + } + if (isDs) { + for (uint8_t iHypo = 0; iHypo < 2; ++iHypo) { + if (!TESTBIT(dsHypo, iHypo)) { + continue; + } + registry.fill(HIST("hMassDsToKKPi"), RecoDecay::m(arrMom, massDsToKKPi[iHypo])); + const auto& pVecKaonA = (iHypo == 0) ? pvec0 : pvec1; + const auto& pVecKaonB = (iHypo == 0) ? pvec1 : pvec2; + registry.fill(HIST("hMassKK"), RecoDecay::m(std::array{pVecKaonA, pVecKaonB}, std::array{MassKPlus, MassKPlus})); + } + } + registry.fill(HIST("hPtCand"), RecoDecay::pt(pVecCand)); + registry.fill(HIST("hVtxX"), secondaryVertex[0]); + } + } + } + } + } + + /// Hypothesis evaluators. The prong ordering is (same-sign 1, opposite-sign, same-sign 2), so + /// D+ -> π K π needs both same-sign prongs pion-like, and Ds -> K K π needs the opposite-sign + /// prong kaon-like plus one kaon and one pion among the same-sign pair -- two orderings, which + /// become the two mass hypotheses of aod::Hf3Prongs. + static void hypothesisDplusOnly(uint8_t maskOpp, uint8_t mask1, uint8_t mask2, bool& dplusOk, uint8_t& dsHypo) + { + dsHypo = 0; + dplusOk = TESTBIT(maskOpp, BitKaonDplus) && TESTBIT(mask1, BitPionDplus) && TESTBIT(mask2, BitPionDplus); + } + + static void hypothesisDsOnly(uint8_t maskOpp, uint8_t mask1, uint8_t mask2, bool& dplusOk, uint8_t& dsHypo) + { + dplusOk = false; + dsHypo = 0; + if (!TESTBIT(maskOpp, BitKaonDs)) { + return; + } + if (TESTBIT(mask1, BitKaonDs) && TESTBIT(mask2, BitPionDs)) { + SETBIT(dsHypo, 0); + } + if (TESTBIT(mask1, BitPionDs) && TESTBIT(mask2, BitKaonDs)) { + SETBIT(dsHypo, 1); + } + } + + static void hypothesisBoth(uint8_t maskOpp, uint8_t mask1, uint8_t mask2, bool& dplusOk, uint8_t& dsHypo) + { + hypothesisDplusOnly(maskOpp, mask1, mask2, dplusOk, dsHypo); + bool dummy{false}; + hypothesisDsOnly(maskOpp, mask1, mask2, dummy, dsHypo); + } + + static bool pairUsefulDplus(uint8_t maskOpp, uint8_t mask1) + { + return TESTBIT(maskOpp, BitKaonDplus) && TESTBIT(mask1, BitPionDplus); + } + + static bool pairUsefulDs(uint8_t maskOpp, uint8_t mask1) + { + return TESTBIT(maskOpp, BitKaonDs) && (TESTBIT(mask1, BitKaonDs) || TESTBIT(mask1, BitPionDs)); + } + + static bool pairUsefulBoth(uint8_t maskOpp, uint8_t mask1) + { + return pairUsefulDplus(maskOpp, mask1) || pairUsefulDs(maskOpp, mask1); + } + + /// Per-collision preamble shared by the two new strategies: magnetic field, prong cache, timers. + template + void prepareCollision(TCollision const& collision, TAssocs const& trackIndices, TTracks const& tracks, bool doTiming) + { + const auto bc = collision.template bc_as(); + initCCDB(bc, runNumber, ccdb, config.isRun2 ? config.ccdbPathGrp : config.ccdbPathGrpMag, lut, config.isRun2); + const auto bz = o2::base::Propagator::Instance()->getNominalBz(); + df2.setBz(bz); + df3.setBz(bz); + + const auto tStart = StopWatch::now(); + const auto groupedAssocs = trackIndices.sliceBy(trackIndicesPerCollision, collision.globalIndex()); + buildProngCache(collision, groupedAssocs, tracks); + if (doTiming) { + timing[TimeCombCache] += elapsedSeconds(tStart); + } + } + + /// Book-keeping at the end of a collision: fill the timing and counter histograms. + void finaliseCollision(StopWatch::time_point const& tStartCollision, bool doTiming, int nCand) + { + if (doTiming) { + timing[TimeCombTotal] = elapsedSeconds(tStartCollision); + timing[TimeCombRemainder] = timing[TimeCombTotal] - timing[TimeCombCache] - timing[TimeCombFit2Prong] - timing[TimeCombFit3Prong]; + for (int iStage = 0; iStage < NTimingCombinatorics; ++iStage) { + registry.fill(HIST("hTiming"), iStage, timing[iStage]); + } + } + for (int iCounter = 0; iCounter < NCounterCombinatorics; ++iCounter) { + if (counters[iCounter] != 0) { + registry.fill(HIST("hCounters"), iCounter, static_cast(counters[iCounter])); + } + } + if (config.fillHistograms) { + registry.fill(HIST("hNCandPerColl"), nCand); + } + } + + /// Legacy D+-only loop, kept verbatim except for the renamed selection columns so that the + /// earlier efficiency/rate/CPU comparison against the track-index skim creator stays + /// reproducible. It re-propagates the innermost prong inside the innermost loop, which is + /// exactly what the prong cache of the newer strategies removes. + template + int buildCandidatesLegacy(TCollision const& collision, + TTracks const& /*tracks*/, + TSameSign const& sameSignTracks, + TOppSign const& oppSignTracks) + { + const auto thisCollId = collision.globalIndex(); + const std::array pvCoord{collision.posX(), collision.posY(), collision.posZ()}; + int nCand{0}; + + for (const auto& oppSignId : oppSignTracks) { + const auto trackKa = oppSignId.template track_as(); + auto trackParVarKa = getTrackParCov(trackKa); + std::array pVecTrackKa{trackKa.pVector()}; + if (thisCollId != trackKa.collisionId()) { + std::array dcaInfo{-999.f, -999.f}; + o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarKa, 2.f, noMatCorr, &dcaInfo); + getPxPyPz(trackParVarKa, pVecTrackKa); + } + + for (auto sameSignId1 = sameSignTracks.begin(); sameSignId1 != sameSignTracks.end(); ++sameSignId1) { + const auto trackPi1 = sameSignId1.template track_as(); + auto trackParVarPi1 = getTrackParCov(trackPi1); + std::array pVecTrackPi1{trackPi1.pVector()}; + if (thisCollId != trackPi1.collisionId()) { + std::array dcaInfo{-999.f, -999.f}; + o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarPi1, 2.f, noMatCorr, &dcaInfo); + getPxPyPz(trackParVarPi1, pVecTrackPi1); + } + + // optional two-track speedups, evaluated once per (π, K) pair + if (twoTrackSpeedupsEnabled()) { + int nVtxFrom2ProngFitter{0}; + ++counters[CountFit2Prong]; + try { + nVtxFrom2ProngFitter = df2.process(trackParVarPi1, trackParVarKa); + } catch (...) { + continue; + } + if (nVtxFrom2ProngFitter == 0) { + continue; + } + if (!isTwoTrackVertexSelected(df2.getPCACandidate(), pvCoord, df2)) { + continue; + } + } + + for (auto sameSignId2 = sameSignId1 + 1; sameSignId2 != sameSignTracks.end(); ++sameSignId2) { + const auto trackPi2 = sameSignId2.template track_as(); + auto trackParVarPi2 = getTrackParCov(trackPi2); + std::array pVecTrackPi2{trackPi2.pVector()}; + if (thisCollId != trackPi2.collisionId()) { + std::array dcaInfo{-999.f, -999.f}; + o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, trackParVarPi2, 2.f, noMatCorr, &dcaInfo); + getPxPyPz(trackParVarPi2, pVecTrackPi2); + } + + ++counters[CountTriplets]; + if (!applyPreselectionDplus(pVecTrackPi1, pVecTrackKa, pVecTrackPi2)) { + continue; + } + ++counters[CountTripletsPresel]; + + int nVtxFrom3ProngFitter{0}; + ++counters[CountFit3Prong]; + try { + nVtxFrom3ProngFitter = df3.process(trackParVarPi1, trackParVarKa, trackParVarPi2); + } catch (...) { + continue; + } + if (nVtxFrom3ProngFitter == 0) { + continue; + } + + const auto& secondaryVertex = df3.getPCACandidate(); + std::array pvec0{}; + std::array pvec1{}; + std::array pvec2{}; + df3.getTrack(0).getPxPyPzGlo(pvec0); + df3.getTrack(1).getPxPyPzGlo(pvec1); + df3.getTrack(2).getPxPyPzGlo(pvec2); + const auto pVecCand = RecoDecay::pVec(pvec0, pvec1, pvec2); + + if (!applySelectionAfterVertexing(config.binsPtDplusToPiKPi, config.cutsDplusToPiKPi, pVecCand, secondaryVertex, pvCoord)) { + continue; + } + + uint8_t hfFlag{0}; + SETBIT(hfFlag, hf_cand_3prong::DecayType::DplusToPiKPi); + rowTrackIndexProng3(thisCollId, trackPi1.globalIndex(), trackKa.globalIndex(), trackPi2.globalIndex(), hfFlag); + if (config.fillScoreScanTable) { + // the legacy path re-propagates per triplet and keeps no prong cache, so it has no + // scores to report; the row still has to exist to keep the two tables in lockstep + rowCandScoreScan(RecoDecay::pt(pVecCand), kScoreScanNone, kScoreScanNone, + kScoreScanNone, kScoreScanNone, kScoreScanNone, kScoreScanNone); + } + ++nCand; + ++counters[CountRows]; + ++counters[CountRowsDplus]; + + if (config.fillHistograms) { + const std::array arrMom{pvec0, pvec1, pvec2}; + registry.fill(HIST("hMassDplusToPiKPi"), RecoDecay::m(arrMom, massDplusToPiKPi)); + registry.fill(HIST("hPtCand"), RecoDecay::pt(pVecCand)); + registry.fill(HIST("hVtxX"), secondaryVertex[0]); + } + } + } + } + + return nCand; + } + + void processDplus(SelectedCollisions const& collisions, + FilteredTrackAssocSel const& /*trackIndices*/, + TracksWithCov const& tracks, + aod::BCsWithTimestamps const& /*bcWithTimeStamps*/) + { + const bool doTiming = config.enableTiming; + for (const auto& collision : collisions) { + const auto tStartCollision = StopWatch::now(); + timing.fill(0.); + counters.fill(0); + counters[CountCollisions] = 1; + const auto thisCollId = collision.globalIndex(); + + const auto bc = collision.bc_as(); + initCCDB(bc, runNumber, ccdb, config.isRun2 ? config.ccdbPathGrp : config.ccdbPathGrpMag, lut, config.isRun2); + const auto bz = o2::base::Propagator::Instance()->getNominalBz(); + df2.setBz(bz); + df3.setBz(bz); + + const auto groupedPosPions = positivePions->sliceByCached(aod::track::collisionId, thisCollId, cache); + const auto groupedNegPions = negativePions->sliceByCached(aod::track::collisionId, thisCollId, cache); + const auto groupedPosKaons = positiveKaons->sliceByCached(aod::track::collisionId, thisCollId, cache); + const auto groupedNegKaons = negativeKaons->sliceByCached(aod::track::collisionId, thisCollId, cache); + + int nCand{0}; + // D+ -> pi+ K- pi+ + nCand += buildCandidatesLegacy(collision, tracks, groupedPosPions, groupedNegKaons); + // D- -> pi- K+ pi- + nCand += buildCandidatesLegacy(collision, tracks, groupedNegPions, groupedPosKaons); + + finaliseCollision(tStartCollision, doTiming, nCand); + } + } + PROCESS_SWITCH(HfDplusCombinatorics, processDplus, "D+ only, original implementation (no prong cache)", true); + + void processDplusCached(SelectedCollisions const& collisions, + FilteredTrackAssocSel const& trackIndices, + TracksWithCov const& tracks, + aod::BCsWithTimestamps const& /*bcWithTimeStamps*/) + { + const bool doTiming = config.enableTiming; + for (const auto& collision : collisions) { + const auto tStartCollision = StopWatch::now(); + timing.fill(0.); + counters.fill(0); + counters[CountCollisions] = 1; + + prepareCollision(collision, trackIndices, tracks, doTiming); + + int nCand{0}; + runTripleLoop(collision, cachePos, cacheNeg, cachePos.pionDplus, cacheNeg.kaonDplus, pairUsefulDplus, hypothesisDplusOnly, doTiming, nCand); + runTripleLoop(collision, cacheNeg, cachePos, cacheNeg.pionDplus, cachePos.kaonDplus, pairUsefulDplus, hypothesisDplusOnly, doTiming, nCand); + + finaliseCollision(tStartCollision, doTiming, nCand); + } + } + PROCESS_SWITCH(HfDplusCombinatorics, processDplusCached, "D+ only, with the per-collision prong cache", false); + + void processSeparateLoops(SelectedCollisions const& collisions, + FilteredTrackAssocSel const& trackIndices, + TracksWithCov const& tracks, + aod::BCsWithTimestamps const& /*bcWithTimeStamps*/) + { + const bool doTiming = config.enableTiming; + for (const auto& collision : collisions) { + const auto tStartCollision = StopWatch::now(); + timing.fill(0.); + counters.fill(0); + counters[CountCollisions] = 1; + + prepareCollision(collision, trackIndices, tracks, doTiming); + + int nCand{0}; + // D+ -> pi+ K- pi+ and charge conjugate + runTripleLoop(collision, cachePos, cacheNeg, cachePos.pionDplus, cacheNeg.kaonDplus, pairUsefulDplus, hypothesisDplusOnly, doTiming, nCand); + runTripleLoop(collision, cacheNeg, cachePos, cacheNeg.pionDplus, cachePos.kaonDplus, pairUsefulDplus, hypothesisDplusOnly, doTiming, nCand); + // Ds -> K+ K- pi+ and charge conjugate, as a second and independent loop: a triplet good for + // both species is fitted twice and written twice + runTripleLoop(collision, cachePos, cacheNeg, cachePos.dsSameSign, cacheNeg.kaonDs, pairUsefulDs, hypothesisDsOnly, doTiming, nCand); + runTripleLoop(collision, cacheNeg, cachePos, cacheNeg.dsSameSign, cachePos.kaonDs, pairUsefulDs, hypothesisDsOnly, doTiming, nCand); + + finaliseCollision(tStartCollision, doTiming, nCand); + } + } + PROCESS_SWITCH(HfDplusCombinatorics, processSeparateLoops, "One independent triple loop per species", false); + + void processMergedLoop(SelectedCollisions const& collisions, + FilteredTrackAssocSel const& trackIndices, + TracksWithCov const& tracks, + aod::BCsWithTimestamps const& /*bcWithTimeStamps*/) + { + const bool doTiming = config.enableTiming; + for (const auto& collision : collisions) { + const auto tStartCollision = StopWatch::now(); + timing.fill(0.); + counters.fill(0); + counters[CountCollisions] = 1; + + prepareCollision(collision, trackIndices, tracks, doTiming); + + int nCand{0}; + // a single loop over the union of the selected tracks, with a per-triplet hypothesis mask + runTripleLoop(collision, cachePos, cacheNeg, cachePos.sameSignAny, cacheNeg.oppSignAny, pairUsefulBoth, hypothesisBoth, doTiming, nCand); + runTripleLoop(collision, cacheNeg, cachePos, cacheNeg.sameSignAny, cachePos.oppSignAny, pairUsefulBoth, hypothesisBoth, doTiming, nCand); + + finaliseCollision(tStartCollision, doTiming, nCand); + } + } + PROCESS_SWITCH(HfDplusCombinatorics, processMergedLoop, "One triple loop over the union of the selected tracks", false); + + void processDummy(aod::Collisions const&) + { + // no-op, to be enabled when only the training sample is produced + } + PROCESS_SWITCH(HfDplusCombinatorics, processDummy, "Do nothing", false); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{ + adaptAnalysisTask(cfgc), + adaptAnalysisTask(cfgc)}; +} diff --git a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx index cb0174153f9..1fa246b96d6 100644 --- a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx +++ b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx @@ -151,8 +151,8 @@ constexpr TrackMlRole trackRoles3Prongs[NChannels3Prong][NMassHypos][N3Prongs] = constexpr double PtMaxModel = 1.e10; -constexpr uint32_t MaskSameSignAny = (1u << RolePionD0) | (1u << RolePionDplus) | (1u << RolePionDs) | (1u << RoleKaonDs) | (1u << RoleSoftPiDstar); -constexpr uint32_t MaskOppSignAny = (1u << RoleKaonD0) | (1u << RoleKaonDplus) | (1u << RoleKaonDs); +constexpr uint32_t MaskSameSignAny = (1u << RolePionD0) | (1u << RoleKaonD0) | (1u << RolePionDplus) | (1u << RolePionDs) | (1u << RoleKaonDs) | (1u << RoleSoftPiDstar); +constexpr uint32_t MaskOppSignAny = (1u << RolePionD0) | (1u << RoleKaonD0) | (1u << RoleKaonDplus) | (1u << RoleKaonDs); /// Where the combinatorics spends its time, accumulated in seconds into hTiming. enum TimingStep { @@ -349,25 +349,25 @@ struct HfTrackSelectorTagSelTracks { Configurable ptMinSoftPi{"ptMinSoftPi", 0.1f, "min. soft pion track pT entering the charm combinatorics"}; Configurable enableTiming{"enableTiming", false, "fill hTiming with the CPU of the feature building and of each model evaluation (adds two clock reads per call)"}; // D0 model - Configurable applyMlD0{"applyMlD0", true, "evaluate the D0 track model"}; + Configurable applyMlD0{"applyMlD0", false, "evaluate the D0 track model"}; Configurable onnxFileNameD0{"onnxFileNameD0", "ModelHandler_D0Tracks.onnx", "ONNX file name of the D0 track model"}; Configurable> inputFeaturesD0{"inputFeaturesD0", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D0 model, in the order it expects"}; Configurable thresholdScorePionD0{"thresholdScorePionD0", 0., "min. D0 pion-class score"}; Configurable thresholdScoreKaonD0{"thresholdScoreKaonD0", 0., "min. D0 kaon-class score"}; // D+ model - Configurable applyMlDplus{"applyMlDplus", true, "evaluate the D+ track model"}; + Configurable applyMlDplus{"applyMlDplus", false, "evaluate the D+ track model"}; Configurable onnxFileNameDplus{"onnxFileNameDplus", "ModelHandler_DplusTracks.onnx", "ONNX file name of the D+ track model"}; Configurable> inputFeaturesDplus{"inputFeaturesDplus", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D+ model, in the order it expects"}; Configurable thresholdScorePionDplus{"thresholdScorePionDplus", 0., "min. D+ pion-class score"}; Configurable thresholdScoreKaonDplus{"thresholdScoreKaonDplus", 0., "min. D+ kaon-class score"}; // Ds model - Configurable applyMlDs{"applyMlDs", true, "evaluate the Ds track model"}; + Configurable applyMlDs{"applyMlDs", false, "evaluate the Ds track model"}; Configurable onnxFileNameDs{"onnxFileNameDs", "ModelHandler_DsTracks.onnx", "ONNX file name of the Ds track model"}; Configurable> inputFeaturesDs{"inputFeaturesDs", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the Ds model, in the order it expects"}; Configurable thresholdScorePionDs{"thresholdScorePionDs", 0., "min. Ds pion-class score"}; Configurable thresholdScoreKaonDs{"thresholdScoreKaonDs", 0., "min. Ds kaon-class score"}; // D* model - Configurable applyMlDstar{"applyMlDstar", true, "evaluate the D* track model"}; + Configurable applyMlDstar{"applyMlDstar", false, "evaluate the D* track model"}; Configurable onnxFileNameDstar{"onnxFileNameDstar", "ModelHandler_DstarTracks.onnx", "ONNX file name of the D* track model"}; Configurable> inputFeaturesDstar{"inputFeaturesDstar", std::vector{"pt", "eta", "dcaXY", "dcaZ", "sigmaDcaXY", "sigmaDcaZ", "normDcaXY", "normDcaZ", "signed1Pt", "tgl", "sign", "isPvContributor", "itsNCls", "itsNClsInnerBarrel", "itsChi2NCl", "tpcNClsFound", "tpcCrossedRowsOverFindableCls", "tpcChi2NCl", "tpcFractionSharedCls", "tpcNSigmaPi", "tpcNSigmaKa"}, "input features of the D* model, in the order it expects"}; Configurable thresholdScorePionDstar{"thresholdScorePionDstar", 0., "min. D* pion-class score"}; @@ -439,6 +439,7 @@ struct HfTrackSelectorTagSelTracks { /// \param features are the input feature names, in the order the model expects /// \param thrPion, thrKaon are the score thresholds /// \param name is used in the log message + template void configureModel(HfTrackModel& model, std::string const& onnxFile, std::string const& ccdbPath, @@ -452,7 +453,7 @@ struct HfTrackSelectorTagSelTracks { const std::vector binsPtSingle{0., PtMaxModel}; const std::vector onnxFiles{onnxFile}; - const auto nOut = static_cast(HfTrackMlClass::NClasses); + const auto nOut = IsBinaryModel ? static_cast(2) : static_cast(HfTrackMlClass::NClasses); const std::vector cutDir(nOut, o2::cuts_ml::CutDirection::CutNot); const std::vector dummyCutValues(nOut, 0.); const o2::framework::LabeledArray dummyCuts{dummyCutValues.data(), 1u, static_cast(nOut), {}, {}}; @@ -483,21 +484,21 @@ struct HfTrackSelectorTagSelTracks { } if (config.applyMlD0) { - configureModel(models2Prong[ChannelD0ToPiK], config.onnxFileNameD0, config.mlModelPathCcdbD0, config.inputFeaturesD0, + configureModel(models2Prong[ChannelD0ToPiK], config.onnxFileNameD0, config.mlModelPathCcdbD0, config.inputFeaturesD0, config.thresholdScorePionD0, config.thresholdScoreKaonD0, "D0 track model"); } if (config.applyMlDplus) { - configureModel(models3Prong[ChannelDplusToPiKPi], config.onnxFileNameDplus, config.mlModelPathCcdbDplus, config.inputFeaturesDplus, + configureModel(models3Prong[ChannelDplusToPiKPi], config.onnxFileNameDplus, config.mlModelPathCcdbDplus, config.inputFeaturesDplus, config.thresholdScorePionDplus, config.thresholdScoreKaonDplus, "D+ track model"); } if (config.applyMlDs) { - configureModel(models3Prong[ChannelDsToKKPi], config.onnxFileNameDs, config.mlModelPathCcdbDs, config.inputFeaturesDs, + configureModel(models3Prong[ChannelDsToKKPi], config.onnxFileNameDs, config.mlModelPathCcdbDs, config.inputFeaturesDs, config.thresholdScorePionDs, config.thresholdScoreKaonDs, "Ds track model"); } if (config.applyMlDstar) { double thresholdScoreKaonDstar = 0; // D* model has no kaon score, but the configureModel() function expects a value - configureModel(modelSoftPiDstar, config.onnxFileNameDstar, config.mlModelPathCcdbDstar, config.inputFeaturesDstar, - config.thresholdScorePionDstar, thresholdScoreKaonDstar, "D* track model"); + configureModel(modelSoftPiDstar, config.onnxFileNameDstar, config.mlModelPathCcdbDstar, config.inputFeaturesDstar, + config.thresholdScorePionDstar, thresholdScoreKaonDstar, "D* track model"); } if (!models2Prong[ChannelD0ToPiK].enabled && !models3Prong[ChannelDplusToPiKPi].enabled && @@ -613,10 +614,11 @@ struct HfTrackSelectorTagSelTracks { /// \param scores is filled with (background, pion-role, kaon-role); -1 if not evaluated /// \param isSelPion, isSelKaon are filled with the threshold decisions /// \return true if the model was actually evaluated + template bool evaluateModel(HfTrackModel& model, HfTrackMlFeatures const& features, - std::array& scores, bool& isSelPion, bool& isSelKaon) + std::array& scores, bool& isSelPion, bool& isSelKaon) { - scores = {-1.f, -1.f, -1.f}; + scores.fill(-1.f); isSelPion = false; isSelKaon = false; if (!model.enabled) { @@ -625,13 +627,15 @@ struct HfTrackSelectorTagSelTracks { auto inputFeatures = model.response.getInputFeatures(features); // one model for the whole pT range, so the bin index is always 0 const auto output = model.response.getModelOutput(inputFeatures, 0); - if (output.size() != scores.size()) { - LOGP(fatal, "the model emitted {} scores, expected {}: only 3-class (bkg/pi/K) models are supported for D0, D+, Ds, only 2-class models are supported for D* candidates", output.size(), scores.size()); + if (output.size() != NClasses) { + LOGP(fatal, "the model emitted {} scores, expected {}: only 3-class (bkg/pi/K) models are supported for D0, D+, Ds, only 2-class models are supported for D* candidates", output.size(), NClasses); } std::copy(output.begin(), output.end(), scores.begin()); isSelPion = scores[1] > model.thresholdPion; - isSelKaon = scores[2] > model.thresholdKaon; + if constexpr (NClasses > 2) { + isSelKaon = scores[2] > model.thresholdKaon; + } return true; } @@ -683,6 +687,7 @@ struct HfTrackSelectorTagSelTracks { } std::array scores{}; + std::array scoresSoftPi{}; bool isSelPion{false}, isSelKaon{false}; bool from2Prong{false}, from3Prong{false}, fromDstar{false}; @@ -692,7 +697,7 @@ struct HfTrackSelectorTagSelTracks { } tStart = tickIf(doTiming); - const bool evaluatedDstar = evaluateModel(modelSoftPiDstar, features, scores, isSelPion, isSelKaon); + const bool evaluatedDstar = evaluateModel(modelSoftPiDstar, features, scoresSoftPi, isSelPion, isSelKaon); if (doTiming) { timing[TrackTimeModelDstar] += elapsedSeconds(tStart); } @@ -702,7 +707,7 @@ struct HfTrackSelectorTagSelTracks { fromDstar = true; } if (config.fillHistograms) { - registry.fill(HIST("hScoreSoftPiDstar"), features.pt, scores[1]); + registry.fill(HIST("hScoreSoftPiDstar"), features.pt, scoresSoftPi[1]); } } } @@ -856,6 +861,7 @@ struct HfMlBasedTrackSelector { struct : ConfigurableGroup { Configurable fillHistograms{"fillHistograms", true, "fill histograms"}; Configurable do2Prongs{"do2Prongs", true, "store 2-prong candidates"}; + Configurable do3Prongs{"do3Prongs", true, "store 3-prong candidates"}; Configurable doDstar{"doDstar", true, "store D* candidates"}; // preselection Configurable ptTolerance{"ptTolerance", 0.1, "pT tolerance in GeV/c for applying preselections before vertex reconstruction"}; @@ -906,7 +912,7 @@ struct HfMlBasedTrackSelector { // cuts and pT binning, one entry per channel std::array, NChannels2Prong> cut2Prong{}; std::array, NChannels2Prong> binsPt2Prong{}; - LabeledArray cutDstar{}; + LabeledArray cutsDstar{}; std::vector binsPtDstar{}; std::array, NChannels3Prong> cut3Prong{}; std::array, NChannels3Prong> binsPt3Prong{}; @@ -963,6 +969,9 @@ struct HfMlBasedTrackSelector { if (!doprocessCandidates) { return; } + if (!config.do2Prongs && !config.do3Prongs && !config.doDstar) { + LOGP(fatal, "At least one of do2Prongs, do3Prongs or doDstar must be true!"); + } doTiming = config.enableTiming && config.fillHistograms; arrMass2Prong[ChannelD0ToPiK] = std::array{std::array{MassPiPlus, MassKPlus}, @@ -977,7 +986,7 @@ struct HfMlBasedTrackSelector { // cuts retrieved by json, in the order of Channel3Prong cut2Prong = {config.cutsD0ToPiK}; binsPt2Prong = {config.binsPtD0ToPiK}; - cutDstar = {config.cutsDstarToD0Pi}; + cutsDstar = {config.cutsDstarToD0Pi}; binsPtDstar = {config.binsPtDstarToD0Pi}; // cuts retrieved by json, in the order of Channel3Prong @@ -1036,7 +1045,7 @@ struct HfMlBasedTrackSelector { registry.add("hVtx2ProngZ", "2-prong candidates;#it{z}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -20., 20.}}}); registry.add("hNCand2Prong", "2-prong candidates preselected;# of candidates;entries", {HistType::kTH1D, {axisNumCands}}); registry.add("hNCand2ProngVsNTracks", "2-prong candidates preselected;# of selected tracks;# of candidates;entries", {HistType::kTH2D, {axisNumTracks, axisNumCands}}); - registry.add("hMassD0ToPiK", "D^{0} candidates;inv. mass (#pi K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 0., 5.}}}); + registry.add("hMassD0ToPiK", "D^{0} candidates;inv. mass (#pi K) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{500, 0., 5.}}}); registry.add("hVtx3ProngX", "3-prong candidates;#it{x}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); registry.add("hVtx3ProngY", "3-prong candidates;#it{y}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -2., 2.}}}); registry.add("hVtx3ProngZ", "3-prong candidates;#it{z}_{sec. vtx.} (cm);entries", {HistType::kTH1D, {{1000, -20., 20.}}}); @@ -1406,10 +1415,12 @@ struct HfMlBasedTrackSelector { /// \param whichHypo information of the mass hypoteses that were selected /// \param isSelected is a bitmap with selection outcome /// \param pt2Prong is the pt of the 2-prong candidate - // void applySelection2Prong(const T1& pVecCand, const T2& secVtx, const T3& primVtx, T4& cutStatus, auto& isSelected) template void applySelection2Prong(const T1& pVecCand, T2 const& prong0, T2 const& prong1, const T3& secVtx, const T4& primVtx, T5& whichHypo, auto& isSelected) { + if (isSelected == 0) { + return; + } const auto pt2Prong = RecoDecay::pt(pVecCand[0], pVecCand[1]); const auto impParProduct = prong0.dcaInfo[0] * prong1.dcaInfo[0]; @@ -1461,6 +1472,22 @@ struct HfMlBasedTrackSelector { CLRBIT(isSelected, iDecay2P); } } + + // additional check for D0 to be used in D* finding + if (iDecay2P == ChannelD0ToPiK && config.doDstar && TESTBIT(isSelected, iDecay2P)) { + const auto binPtDstar = findBin(&binsPtDstar, pt2Prong * 1.2); // assuming the D* pT about 20% higher than the one of the D0 to be safe + if (binPtDstar >= 0) { + whichHypo[NChannels2Prong] = whichHypo[hf_cand_2prong::DecayType::D0ToPiK]; + const double deltaMass = cutsDstar.get(binPtDstar, 1u); + + if (TESTBIT(whichHypo[iDecay2P], 0) && (massHypos[0] > (MassD0 + deltaMass) * (MassD0 + deltaMass) || massHypos[0] < (MassD0 - deltaMass) * (MassD0 - deltaMass))) { + CLRBIT(whichHypo[NChannels2Prong], 0); + } + if (TESTBIT(whichHypo[iDecay2P], 1) && (massHypos[1] > (MassD0 + deltaMass) * (MassD0 + deltaMass) || massHypos[1] < (MassD0 - deltaMass) * (MassD0 - deltaMass))) { + CLRBIT(whichHypo[NChannels2Prong], 1); + } + } + } } } @@ -1469,10 +1496,10 @@ struct HfMlBasedTrackSelector { /// \param prong0,1 are the track informations template void processPair(SelectedCollisions::iterator const& collision, - TProngInfo const& prong0, TProngInfo const& prong1) + TProngInfo const& prong0, TProngInfo const& prong1, + std::array& whichHypo2Prong) { uint32_t isSelected2ProngCand = BIT(NChannels2Prong) - 1; - std::array whichHypo2Prong{}; whichHypo2Prong.fill(BIT(NMassHypos) - 1); // 2 bits on, all mass hypotheses alive std::array isIdentifiedPid{prong0.mask, prong1.mask}; @@ -1508,8 +1535,12 @@ struct HfMlBasedTrackSelector { registry.fill(HIST("hVtx2ProngY"), secondaryVertex2[1]); registry.fill(HIST("hVtx2ProngZ"), secondaryVertex2[2]); const std::array arr2Mom{pvec0, pvec1}; - registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][0])); - registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][1])); + if (TESTBIT(whichHypo2Prong[ChannelD0ToPiK], 0)) { + registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][0])); + } + if (TESTBIT(whichHypo2Prong[ChannelD0ToPiK], 1)) { + registry.fill(HIST("hMassD0ToPiK"), RecoDecay::m(arr2Mom, arrMass2Prong[ChannelD0ToPiK][1])); + } } } @@ -1528,21 +1559,21 @@ struct HfMlBasedTrackSelector { const auto pt = RecoDecay::pt(pVecTrack0, pVecTrack1, pVecTrack2) + config.ptTolerance; // add tolerance because of no reco decay vertex // pT - const auto binPt = findBin(config.binsPtDstarToD0Pi, pt); + const auto binPt = findBin(&binsPtDstar, pt); // return immediately if it is outside the defined pT bins if (binPt == -1) { return; } // D0 mass - const double deltaMassD0 = config.cutsDstarToD0Pi->get(binPt, 1u); // 1u == deltaMassD0Index + const double deltaMassD0 = cutsDstar.get(binPt, 1u); // 1u == deltaMassD0Index const double invMassD0 = RecoDecay::m(arrMomD0, std::array{MassPiPlus, MassKPlus}); if (std::abs(invMassD0 - MassD0) > deltaMassD0) { return; } // D*+ mass - const double maxDeltaMass = config.cutsDstarToD0Pi->get(binPt, 0u); // 0u == deltaMassIndex + const double maxDeltaMass = cutsDstar.get(binPt, 0u); // 0u == deltaMassIndex const double invMassDstar = RecoDecay::m(arrMom, std::array{MassPiPlus, MassKPlus, MassPiPlus}); const double deltaMass = invMassDstar - invMassD0; if (deltaMass > maxDeltaMass) { @@ -1600,13 +1631,21 @@ struct HfMlBasedTrackSelector { /// \param sameSignCache holds the prongs of the charge the candidate carries; its sameSign /// list supplies the two same-sign slots /// \param oppSignCache holds the opposite charge; its oppSign list supplies the kaon slot + template void runCombinatorics(SelectedCollisions::iterator const& collision, HfProngCache const& sameSignCache, HfProngCache const& oppSignCache) { + + bool store2Prongs = Store2Prongs && config.do2Prongs; + bool storeDstar = Store2Prongs && config.doDstar; + const std::vector& sameSignList = sameSignCache.sameSign; const std::vector& oppSignList = oppSignCache.oppSign; constexpr std::size_t NSameSign3Prongs = N3Prongs - 1; // two of the three prongs carry the candidate charge - if ((!config.do2Prongs && !config.doDstar && sameSignList.size() < NSameSign3Prongs) || oppSignList.empty()) { + if (oppSignList.empty()) { + return; + } + if (!store2Prongs && !storeDstar && sameSignList.size() < NSameSign3Prongs) { return; } const auto tStartLoop = tickIf(doTiming); @@ -1625,14 +1664,16 @@ struct HfMlBasedTrackSelector { if (!canBe2Prong && !canBe3Prong) { continue; } - if (!canBe3Prong && !config.do2Prongs && !config.doDstar) { + if (!canBe3Prong && !store2Prongs && !storeDstar) { continue; } ++counters[CountPairsCandidateOk]; int lastFilledD0 = rowTrackIndexProng2.lastIndex(); // optional 2-track vertex to skip the pair early - if (useTwoTrackVertex || config.do2Prongs || config.doDstar) { + bool do2ProngVertex = useTwoTrackVertex || store2Prongs || storeDstar; + std::array whichHypo2Prong{}; // Add one more for D0 from D* selection + if (do2ProngVertex) { const auto tStart2Prong = tickIf(doTiming); ++counters[CountFit2Prong]; int nVtxFrom2ProngFitter = 0; @@ -1648,43 +1689,65 @@ struct HfMlBasedTrackSelector { ++counters[CountPairsRejected2P]; continue; } - const auto tStartD0 = tickIf(doTiming); - if (config.do2Prongs || config.doDstar) { - processPair(collision, prong1, prongOpp); + if (store2Prongs || storeDstar) { + const auto tStartD0 = tickIf(doTiming); + processPair(collision, prong1, prongOpp, whichHypo2Prong); + if (doTiming) { + timing[TimeProcessD0] += elapsedSeconds(tStartD0); + } } - if (doTiming) { - timing[TimeProcessD0] += elapsedSeconds(tStartD0); + if (useTwoTrackVertex) { + isTwoProngVtxGoodFor3Prongs = isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); } - isTwoProngVtxGoodFor3Prongs = isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); } // D* -> D0 pi+ candidates - const auto tStartDstar = tickIf(doTiming); - if (config.doDstar && + const auto tStartDstar = tickIf(doTiming); + if (storeDstar && lastFilledD0 != rowTrackIndexProng2.lastIndex()) { // we have a new D0 candidate, so we can try to form a D* candidate lastFilledD0 = rowTrackIndexProng2.lastIndex(); // Loop over all tracks to search for soft pions - for (std::size_t iSoftPi = 0; iSoftPi < sameSignList.size(); ++iSoftPi) { - const auto& prongSoftPi = sameSignCache.prongs[sameSignList[iSoftPi]]; - - if (iSoftPi == i1) { - continue; // skip the same track + if (TESTBIT(whichHypo2Prong[NChannels2Prong], 0)) { // D0 -> pi+ K-, loop on positive soft pions + for (std::size_t iSoftPi = 0; iSoftPi < sameSignList.size(); ++iSoftPi) { + const auto& prongSoftPi = sameSignCache.prongs[sameSignList[iSoftPi]]; + + if (iSoftPi == static_cast(i1)) { + continue; // skip the same track + } + if (!TESTBIT(prongSoftPi.mask, RoleSoftPiDstar)) { + continue; // skip if the second same-sign prong is not a soft pion candidate + } + ++counters[CountDstarsEnumerated]; + + // Update the index of the last filled D0 candidate + processDstar(collision, lastFilledD0, prongSoftPi.globalIndex, + prong1.pVec, prongOpp.pVec, prongSoftPi.pVec); } - if (!TESTBIT(prongSoftPi.mask, RoleSoftPiDstar)) { - continue; // skip if the second same-sign prong is not a soft pion candidate + } + if (TESTBIT(whichHypo2Prong[NChannels2Prong], 1)) { // D0bar -> K+ pi-, loop on negative soft pions + for (std::size_t iSoftPi = 0; iSoftPi < oppSignList.size(); ++iSoftPi) { + const auto& prongSoftPi = oppSignCache.prongs[oppSignList[iSoftPi]]; + + if (oppSignList[iSoftPi] == iOpp) { + continue; // skip the same track + } + if (!TESTBIT(prongSoftPi.mask, RoleSoftPiDstar)) { + continue; // skip if the second same-sign prong is not a soft pion candidate + } + ++counters[CountDstarsEnumerated]; + + // Update the index of the last filled D0 candidate + processDstar(collision, lastFilledD0, prongSoftPi.globalIndex, + prong1.pVec, prongOpp.pVec, prongSoftPi.pVec); } - ++counters[CountDstarsEnumerated]; - - // Update the index of the last filled D0 candidate - processDstar(collision, lastFilledD0, prongSoftPi.globalIndex, - prong1.pVec, prongOpp.pVec, prongSoftPi.pVec); } } if (doTiming) { timing[TimeProcessDstar] += elapsedSeconds(tStartDstar); } + if (!config.do3Prongs) { continue; } if (!canBe3Prong) { continue; } if (!isTwoProngVtxGoodFor3Prongs) { ++counters[CountPairsRejected2P]; @@ -1743,8 +1806,8 @@ struct HfMlBasedTrackSelector { // D+ -> pi+ K- pi+ / Ds -> K+ K- pi+ and their charge conjugates: the same-sign pair carries // the charge of the candidate, the opposite-sign prong is the kaon in both channels - runCombinatorics(collision, cachePos, cacheNeg); - runCombinatorics(collision, cacheNeg, cachePos); + runCombinatorics(collision, cachePos, cacheNeg); + runCombinatorics(collision, cacheNeg, cachePos); nCand2 = rowTrackIndexProng2.lastIndex() - nCand2; // number of 2-prong candidates in this collision counters[CountPairsWritten] = nCand2; diff --git a/PWGHF/DataModel/CharmTrackScoreScanTables.h b/PWGHF/DataModel/CharmTrackScoreScanTables.h new file mode 100644 index 00000000000..97335ebd823 --- /dev/null +++ b/PWGHF/DataModel/CharmTrackScoreScanTables.h @@ -0,0 +1,64 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file CharmTrackScoreScanTables.h +/// \brief Per-candidate ML scores, so a track-level working-point scan can be redone offline +/// +/// Written by dplusCombinatorics.cxx and consumed by dplusCombinatoricsQa.cxx. It lives in a +/// header rather than being declared twice because the two declarations must not drift: a +/// mismatch would be read back as silently reinterpreted columns, not as an error. +/// +/// The point of the table is that one *loose* run gives the whole scan. A candidate is built iff +/// every prong clears its role's threshold, and thresholds only rise as the working point +/// tightens, so the candidate set at any tighter point is a strict subset of the loose one. Given +/// the per-role minimum score each candidate was built from, membership at a tighter point is a +/// comparison, not another run. +/// +/// Scores are per *role*, not one global minimum: the pion and kaon thresholds sit at different +/// absolute values (0.027 vs 0.032 at the 90% working point), so a single minimum over all three +/// prongs could not be compared against either of them. +/// +/// \author Fabrizio Chinu , Universita and INFN Torino + +#ifndef PWGHF_DATAMODEL_CHARMTRACKSCORESCANTABLES_H_ +#define PWGHF_DATAMODEL_CHARMTRACKSCORESCANTABLES_H_ + +#include +#include + +namespace o2::aod +{ +/// Score written for a hypothesis a candidate does not have. Below every real score (which lie in +/// [0, 1]) and below every threshold, so such a hypothesis never survives a scan point. +constexpr float kScoreScanNone = -1.f; + +namespace hf_cand_score_scan +{ +DECLARE_SOA_COLUMN(PtCand, ptCand, float); //! pT of the candidate +DECLARE_SOA_COLUMN(MinScorePionDplus, minScorePionDplus, float); //! D+: min pion score over the two pion prongs +DECLARE_SOA_COLUMN(ScoreKaonDplus, scoreKaonDplus, float); //! D+: kaon score of the opposite-sign prong +DECLARE_SOA_COLUMN(MinScoreKaonDsHypo0, minScoreKaonDsHypo0, float); //! Ds (K K pi): min kaon score over the two kaon prongs +DECLARE_SOA_COLUMN(ScorePionDsHypo0, scorePionDsHypo0, float); //! Ds (K K pi): pion score of the pion prong +DECLARE_SOA_COLUMN(MinScoreKaonDsHypo1, minScoreKaonDsHypo1, float); //! Ds (pi K K): min kaon score over the two kaon prongs +DECLARE_SOA_COLUMN(ScorePionDsHypo1, scorePionDsHypo1, float); //! Ds (pi K K): pion score of the pion prong +} // namespace hf_cand_score_scan + +DECLARE_SOA_TABLE(HfCandScoreScans, "AOD", "HFCANDSCORE", //! One row per aod::Hf3Prongs row, in the same order + hf_cand_score_scan::PtCand, + hf_cand_score_scan::MinScorePionDplus, + hf_cand_score_scan::ScoreKaonDplus, + hf_cand_score_scan::MinScoreKaonDsHypo0, + hf_cand_score_scan::ScorePionDsHypo0, + hf_cand_score_scan::MinScoreKaonDsHypo1, + hf_cand_score_scan::ScorePionDsHypo1); +} // namespace o2::aod + +#endif // PWGHF_DATAMODEL_CHARMTRACKSCORESCANTABLES_H_ diff --git a/PWGHF/TableProducer/candidateCreatorDstar.cxx b/PWGHF/TableProducer/candidateCreatorDstar.cxx index b6cfbd8da88..4e5b65122d6 100644 --- a/PWGHF/TableProducer/candidateCreatorDstar.cxx +++ b/PWGHF/TableProducer/candidateCreatorDstar.cxx @@ -117,7 +117,8 @@ struct HfCandidateCreatorDstar { using TracksWCovExtraPidPiKa = soa::Join; - AxisSpec ptAxis = {100, 0., 2.0, "#it{p}_{T} (GeV/#it{c}"}; + AxisSpec massAxis = {130, 0.135, 0.2, "Mass (GeV/#it{c}^{2})"}; + AxisSpec ptAxis = {50, 0., 25.0, "#it{p}_{T} (GeV/#it{c}"}; AxisSpec dcaAxis = {200, -500., 500., "#it{d}_{xy,z} (#mum)"}; std::shared_ptr hCandidates; @@ -137,6 +138,7 @@ struct HfCandidateCreatorDstar { {"QA/hDCAXYPi", "DCAxy of Soft Pi;#it{p}_{T} (GeV/#it{c};#it{d}_{xy}) (#mum);entries", {HistType::kTH2F, {ptAxis, dcaAxis}}}, {"QA/hDCAZPi", "DCAz of Soft Pi;#it{p}_{T} (GeV/#it{c};#it{d}_{z}) (#mum);entries", {HistType::kTH2F, {ptAxis, dcaAxis}}}, + {"QA/hDeltaMassDStar", "D* candidates", {HistType::kTH1F, {massAxis}}}, {"QA/hPtPi", "#pi candidates", {HistType::kTH1F, {ptAxis}}}, {"QA/hPtD0Prong0", "D^{0} candidates' prong0", {HistType::kTH1F, {ptAxis}}}, {"QA/hPtD0Prong1", "D^{0} candidates' prong1", {HistType::kTH1F, {ptAxis}}}, @@ -387,6 +389,32 @@ struct HfCandidateCreatorDstar { fillProngPid(trackPi, rowProngSoftPiPidKa); if (fillHistograms) { + float massDstar{}, massD0{}; + if (signSoftPi > 0) { + massDstar = RecoDecay::m(std::array{std::array{pVecSoftPi[0], pVecSoftPi[1], pVecSoftPi[2]}, + std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, + std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, + std::array{constants::physics::MassPiPlus, + constants::physics::MassPiPlus, + constants::physics::MassKPlus}); + massD0 = RecoDecay::m(std::array{std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, + std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, + std::array{constants::physics::MassPiPlus, + constants::physics::MassKPlus}); + } else { + massDstar = RecoDecay::m(std::array{std::array{pVecSoftPi[0], pVecSoftPi[1], pVecSoftPi[2]}, + std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, + std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, + std::array{constants::physics::MassPiPlus, + constants::physics::MassKPlus, + constants::physics::MassPiPlus}); + massD0 = RecoDecay::m(std::array{std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, + std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, + std::array{constants::physics::MassKPlus, + constants::physics::MassPiPlus}); + } + auto deltaMass = std::abs(massDstar - massD0); + registry.fill(HIST("QA/hDeltaMassDStar"), deltaMass); registry.fill(HIST("QA/hPtD0"), ptD0); registry.fill(HIST("QA/hPtPi"), RecoDecay::pt(pVecSoftPi)); registry.fill(HIST("QA/hPtD0Prong0"), RecoDecay::pt(pVecD0Prong0)); @@ -661,6 +689,7 @@ struct HfCandidateCreatorDstarExpressions { } if (indexRecDstar > -1) { + LOG(info) << "Matched D* with index: " << rowCandidateDstar.globalIndex() << " matched with MC index: " << indexRecDstar; flagDstar = signDstar * hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPi; } if (indexRecD0 > -1) { diff --git a/PWGLF/TableProducer/Resonances/CMakeLists.txt b/PWGLF/TableProducer/Resonances/CMakeLists.txt index bafd0b12ea0..0a0f66a0f8c 100644 --- a/PWGLF/TableProducer/Resonances/CMakeLists.txt +++ b/PWGLF/TableProducer/Resonances/CMakeLists.txt @@ -35,10 +35,10 @@ o2physics_add_dpl_workflow(resonance-initializer PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase COMPONENT_NAME Analysis) -o2physics_add_dpl_workflow(resonance-module-initializer - SOURCES resonanceModuleInitializer.cxx - PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase - COMPONENT_NAME Analysis) +# o2physics_add_dpl_workflow(resonance-module-initializer +# SOURCES resonanceModuleInitializer.cxx +# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase +# COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(resonance-merge-df SOURCES resonanceMergeDF.cxx diff --git a/PWGLF/Tasks/Resonances/CMakeLists.txt b/PWGLF/Tasks/Resonances/CMakeLists.txt index 8c289ef8675..b73291d29f0 100644 --- a/PWGLF/Tasks/Resonances/CMakeLists.txt +++ b/PWGLF/Tasks/Resonances/CMakeLists.txt @@ -194,20 +194,20 @@ o2physics_add_dpl_workflow(xi1530analysisqa PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore COMPONENT_NAME Analysis) -o2physics_add_dpl_workflow(omega2012analysis - SOURCES omega2012Analysis.cxx - PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore - COMPONENT_NAME Analysis) - -o2physics_add_dpl_workflow(omega2012-training-table - SOURCES omega2012TrainingTable.cxx - PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore - COMPONENT_NAME Analysis) - -o2physics_add_dpl_workflow(xi1820analysis - SOURCES xi1820Analysis.cxx - PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore - COMPONENT_NAME Analysis) +# o2physics_add_dpl_workflow(omega2012analysis +# SOURCES omega2012Analysis.cxx +# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore +# COMPONENT_NAME Analysis) + +# o2physics_add_dpl_workflow(omega2012-training-table +# SOURCES omega2012TrainingTable.cxx +# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore +# COMPONENT_NAME Analysis) +# +# o2physics_add_dpl_workflow(xi1820analysis +# SOURCES xi1820Analysis.cxx +# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore +# COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(initializereventqa SOURCES initializereventqa.cxx From 002ece4a1b97ea073537866c2c47131e94eb6e43 Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Tue, 27 Oct 2026 09:12:45 +0000 Subject: [PATCH 3/6] Please consider the following formatting changes --- .../D2H/TableProducer/dplusCombinatorics.cxx | 85 +++++++++---------- .../TableProducer/mlBasedTrackSelector.cxx | 48 ++++++----- PWGHF/TableProducer/candidateCreatorDstar.cxx | 20 ++--- 3 files changed, 77 insertions(+), 76 deletions(-) diff --git a/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx index c0be83ee581..349512519c9 100644 --- a/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx +++ b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx @@ -61,7 +61,6 @@ #include "Common/DataModel/Multiplicity.h" #include "Common/DataModel/PIDResponseTPC.h" #include "Common/DataModel/TrackSelectionTables.h" - #include "Tools/ML/MlResponse.h" #include @@ -126,7 +125,7 @@ DECLARE_SOA_COLUMN(ScorePionDzero, scorePionDzero, float); //! ML score, pion-fr DECLARE_SOA_COLUMN(ScoreKaonDzero, scoreKaonDzero, float); //! ML score, kaon-from-D0 class DECLARE_SOA_COLUMN(IsSelPionDzero, isSelPionDzero, bool); //! passes the D0 pion-score threshold DECLARE_SOA_COLUMN(IsSelKaonDzero, isSelKaonDzero, bool); //! passes the D0 kaon-score threshold - // D+ -> K pi pi model + // D+ -> K pi pi model DECLARE_SOA_COLUMN(ScoreBkgDplus, scoreBkgDplus, float); //! ML score, background class DECLARE_SOA_COLUMN(ScorePionDplus, scorePionDplus, float); //! ML score, pion-from-D+ class DECLARE_SOA_COLUMN(ScoreKaonDplus, scoreKaonDplus, float); //! ML score, kaon-from-D+ class @@ -145,8 +144,8 @@ DECLARE_SOA_COLUMN(ScoreKaonDstar, scoreKaonDstar, float); //! ML score, kaon-fr DECLARE_SOA_COLUMN(IsSelPionDstar, isSelPionDstar, bool); //! passes the D* pion-score threshold DECLARE_SOA_COLUMN(IsSelKaonDstar, isSelKaonDstar, bool); //! passes the D* kaon-score threshold // common -DECLARE_SOA_COLUMN(IsPositive, isPositive, bool); //! charge sign of the track -DECLARE_SOA_COLUMN(IsSelAny, isSelAny, bool); //! passes at least one of the four selections +DECLARE_SOA_COLUMN(IsPositive, isPositive, bool); //! charge sign of the track +DECLARE_SOA_COLUMN(IsSelAny, isSelAny, bool); //! passes at least one of the four selections } // namespace hf_skim_track_sel DECLARE_SOA_TABLE(HfSkimTrackSels, "AOD", "HFSKIMTRKSEL", //! Track-level ML selection, aligned with aod::TrackAssoc @@ -176,15 +175,15 @@ DECLARE_SOA_TABLE(HfSkimTrackSels, "AOD", "HFSKIMTRKSEL", //! Track-level ML sel namespace hf_skim_track_sel_mask { DECLARE_SOA_COLUMN(CollIdx, collIdx, int); //! global index of the collision the association refers to -DECLARE_SOA_COLUMN(SelMask, selMask, uint16_t); //! bit0 -> pi(D0) - //! bit1 -> K(D0) - //! bit2 -> pi(D+) - //! bit3 -> K(D+) - //! bit4 -> pi(Ds) - //! bit5 -> K(Ds) - //! bit6 -> pi(D*) - //! bit7 -> K(D*) - //! bit8 -> positive charge +DECLARE_SOA_COLUMN(SelMask, selMask, uint16_t); //! bit0 -> pi(D0) + //! bit1 -> K(D0) + //! bit2 -> pi(D+) + //! bit3 -> K(D+) + //! bit4 -> pi(Ds) + //! bit5 -> K(Ds) + //! bit6 -> pi(D*) + //! bit7 -> K(D*) + //! bit8 -> positive charge } // namespace hf_skim_track_sel_mask /// One row per *selected* (track, collision) association. This is the smallest object from which @@ -198,32 +197,32 @@ DECLARE_SOA_TABLE(HfSkimTrackMasks, "AOD", "HFSKIMTRKMASK", //! Compact dump of namespace hf_skim_track_train { // features -DECLARE_SOA_COLUMN(Pt, pt, float); //! -DECLARE_SOA_COLUMN(Eta, eta, float); //! -DECLARE_SOA_COLUMN(DcaXY, dcaXY, float); //! -DECLARE_SOA_COLUMN(DcaZ, dcaZ, float); //! -DECLARE_SOA_COLUMN(SigmaDcaXY, sigmaDcaXY, float); //! -DECLARE_SOA_COLUMN(SigmaDcaZ, sigmaDcaZ, float); //! -DECLARE_SOA_COLUMN(NormDcaXY, normDcaXY, float); //! -DECLARE_SOA_COLUMN(NormDcaZ, normDcaZ, float); //! -DECLARE_SOA_COLUMN(Signed1Pt, signed1Pt, float); //! -DECLARE_SOA_COLUMN(Tgl, tgl, float); //! -DECLARE_SOA_COLUMN(Sign, sign, float); //! -DECLARE_SOA_COLUMN(IsPvContributor, isPvContributor, float); //! -DECLARE_SOA_COLUMN(ItsNCls, itsNCls, float); //! -DECLARE_SOA_COLUMN(ItsNClsInnerBarrel, itsNClsInnerBarrel, float); //! -DECLARE_SOA_COLUMN(ItsChi2NCl, itsChi2NCl, float); //! -DECLARE_SOA_COLUMN(TpcNClsFound, tpcNClsFound, float); //! -DECLARE_SOA_COLUMN(TpcNClsCrossedRows, tpcNClsCrossedRows, float); //! +DECLARE_SOA_COLUMN(Pt, pt, float); //! +DECLARE_SOA_COLUMN(Eta, eta, float); //! +DECLARE_SOA_COLUMN(DcaXY, dcaXY, float); //! +DECLARE_SOA_COLUMN(DcaZ, dcaZ, float); //! +DECLARE_SOA_COLUMN(SigmaDcaXY, sigmaDcaXY, float); //! +DECLARE_SOA_COLUMN(SigmaDcaZ, sigmaDcaZ, float); //! +DECLARE_SOA_COLUMN(NormDcaXY, normDcaXY, float); //! +DECLARE_SOA_COLUMN(NormDcaZ, normDcaZ, float); //! +DECLARE_SOA_COLUMN(Signed1Pt, signed1Pt, float); //! +DECLARE_SOA_COLUMN(Tgl, tgl, float); //! +DECLARE_SOA_COLUMN(Sign, sign, float); //! +DECLARE_SOA_COLUMN(IsPvContributor, isPvContributor, float); //! +DECLARE_SOA_COLUMN(ItsNCls, itsNCls, float); //! +DECLARE_SOA_COLUMN(ItsNClsInnerBarrel, itsNClsInnerBarrel, float); //! +DECLARE_SOA_COLUMN(ItsChi2NCl, itsChi2NCl, float); //! +DECLARE_SOA_COLUMN(TpcNClsFound, tpcNClsFound, float); //! +DECLARE_SOA_COLUMN(TpcNClsCrossedRows, tpcNClsCrossedRows, float); //! DECLARE_SOA_COLUMN(TpcCrossedRowsOverFindableCls, tpcCrossedRowsOverFindableCls, float); //! -DECLARE_SOA_COLUMN(TpcChi2NCl, tpcChi2NCl, float); //! -DECLARE_SOA_COLUMN(TpcFractionSharedCls, tpcFractionSharedCls, float); //! -DECLARE_SOA_COLUMN(TpcNSigmaPi, tpcNSigmaPi, float); //! -DECLARE_SOA_COLUMN(TpcNSigmaKa, tpcNSigmaKa, float); //! -DECLARE_SOA_COLUMN(Centrality, centrality, float); //! +DECLARE_SOA_COLUMN(TpcChi2NCl, tpcChi2NCl, float); //! +DECLARE_SOA_COLUMN(TpcFractionSharedCls, tpcFractionSharedCls, float); //! +DECLARE_SOA_COLUMN(TpcNSigmaPi, tpcNSigmaPi, float); //! +DECLARE_SOA_COLUMN(TpcNSigmaKa, tpcNSigmaKa, float); //! +DECLARE_SOA_COLUMN(Centrality, centrality, float); //! // labels and auxiliary information, not meant to be used as features -DECLARE_SOA_COLUMN(TrackIndex, trackIndex, int64_t); //! global index of the track: the same track can be associated to several collisions, and therefore appear more than once, so the train/test split has to group by it -DECLARE_SOA_COLUMN(McMask, mcMask, uint16_t); //! MC label quality mask (bit 15 = fake track) +DECLARE_SOA_COLUMN(TrackIndex, trackIndex, int64_t); //! global index of the track: the same track can be associated to several collisions, and therefore appear more than once, so the train/test split has to group by it +DECLARE_SOA_COLUMN(McMask, mcMask, uint16_t); //! MC label quality mask (bit 15 = fake track) DECLARE_SOA_COLUMN(HasCharmBeautyAncestor, hasCharmBeautyAncestor, bool); //! any charm/beauty hadron in the ancestry // Selection outcomes DECLARE_SOA_COLUMN(Label, label, int8_t); //! D+ model: -1 = not usable, 0 = bkg, 1 = pion from D+, 2 = kaon from D+ @@ -291,7 +290,7 @@ constexpr int NPionsDstarToD0Pi = 1; /// Resonant sub-channels recorded next to the training label, so that the signal definition can be /// narrowed offline without re-running the workflow. enum Channels : int { - kNone = -1, ///< no mother found, or mother not in the list of interest + kNone = -1, ///< no mother found, or mother not in the list of interest kDzeroDirect, kDplusDirect, kDplusKstar, @@ -408,11 +407,11 @@ enum HfAncestryWalkOutcome : int { /// Counter bins of hf-dplus-combinatorics. enum CounterCombinatorics : int { CountCollisions = 0, - CountTriplets, ///< triplets whose hypothesis mask is non-empty, i.e. reaching the preselection - CountTripletsPresel,///< triplets passing the mass/pT preselection, i.e. reaching the 3-prong fitter + CountTriplets, ///< triplets whose hypothesis mask is non-empty, i.e. reaching the preselection + CountTripletsPresel, ///< triplets passing the mass/pT preselection, i.e. reaching the 3-prong fitter CountFit2Prong, CountFit3Prong, - CountRows, ///< rows written to aod::Hf3Prongs + CountRows, ///< rows written to aod::Hf3Prongs CountRowsDplus, CountRowsDs, CountRowsBoth, @@ -1505,8 +1504,8 @@ struct HfDplusTrackSelMl { int pdgMother{-1}; int8_t isPrompt{-1}, channel{Channels::kNone}; const int8_t label = getTrainingLabel( - track, mcParticles, pdgMother, keepIfBackground, - pt, isPrompt, channel); + track, mcParticles, pdgMother, keepIfBackground, + pt, isPrompt, channel); if (label < 0) { // usable for neither model continue; diff --git a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx index 1fa246b96d6..9600525084d 100644 --- a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx +++ b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx @@ -131,23 +131,21 @@ constexpr int NMassHypos = 2; // mass hypotheses per channel, i.e. the two order /// The ordering of the channels should strictly follow Channels2Prong constexpr TrackMlRole trackRoles2Prongs[NChannels2Prong][NMassHypos][N2Prongs] = { { - {RolePionD0, RoleKaonD0}, // D0 → π± K∓ - {RoleKaonD0, RolePionD0} // D0 → π± K∓ - } -}; + {RolePionD0, RoleKaonD0}, // D0 → π± K∓ + {RoleKaonD0, RolePionD0} // D0 → π± K∓ + }}; /// Role of each track in the candidate, per channel and mass hypothesis. /// The ordering of the channels should strictly follow Channels3Prong constexpr TrackMlRole trackRoles3Prongs[NChannels3Prong][NMassHypos][N3Prongs] = { { - {RolePionDplus, RoleKaonDplus, RolePionDplus}, // D± → π± K∓ π± - {RolePionDplus, RoleKaonDplus, RolePionDplus} // D± → π± K∓ π± (swapped, no effect for D+) + {RolePionDplus, RoleKaonDplus, RolePionDplus}, // D± → π± K∓ π± + {RolePionDplus, RoleKaonDplus, RolePionDplus} // D± → π± K∓ π± (swapped, no effect for D+) }, { {RoleKaonDs, RoleKaonDs, RolePionDs}, // D_s± → K± K∓ π± {RolePionDs, RoleKaonDs, RoleKaonDs} // D_s± → K± K∓ π± (swapped) - } -}; + }}; constexpr double PtMaxModel = 1.e10; @@ -439,7 +437,7 @@ struct HfTrackSelectorTagSelTracks { /// \param features are the input feature names, in the order the model expects /// \param thrPion, thrKaon are the score thresholds /// \param name is used in the log message - template + template void configureModel(HfTrackModel& model, std::string const& onnxFile, std::string const& ccdbPath, @@ -485,15 +483,15 @@ struct HfTrackSelectorTagSelTracks { if (config.applyMlD0) { configureModel(models2Prong[ChannelD0ToPiK], config.onnxFileNameD0, config.mlModelPathCcdbD0, config.inputFeaturesD0, - config.thresholdScorePionD0, config.thresholdScoreKaonD0, "D0 track model"); + config.thresholdScorePionD0, config.thresholdScoreKaonD0, "D0 track model"); } if (config.applyMlDplus) { configureModel(models3Prong[ChannelDplusToPiKPi], config.onnxFileNameDplus, config.mlModelPathCcdbDplus, config.inputFeaturesDplus, - config.thresholdScorePionDplus, config.thresholdScoreKaonDplus, "D+ track model"); + config.thresholdScorePionDplus, config.thresholdScoreKaonDplus, "D+ track model"); } if (config.applyMlDs) { configureModel(models3Prong[ChannelDsToKKPi], config.onnxFileNameDs, config.mlModelPathCcdbDs, config.inputFeaturesDs, - config.thresholdScorePionDs, config.thresholdScoreKaonDs, "Ds track model"); + config.thresholdScorePionDs, config.thresholdScoreKaonDs, "Ds track model"); } if (config.applyMlDstar) { double thresholdScoreKaonDstar = 0; // D* model has no kaon score, but the configureModel() function expects a value @@ -614,7 +612,7 @@ struct HfTrackSelectorTagSelTracks { /// \param scores is filled with (background, pion-role, kaon-role); -1 if not evaluated /// \param isSelPion, isSelKaon are filled with the threshold decisions /// \return true if the model was actually evaluated - template + template bool evaluateModel(HfTrackModel& model, HfTrackMlFeatures const& features, std::array& scores, bool& isSelPion, bool& isSelKaon) { @@ -921,7 +919,7 @@ struct HfMlBasedTrackSelector { struct HfProngCandidate { o2::track::TrackParCov parCov; std::array pVec{}; - std::array dcaInfo{}; ///< DCA (xy, z) to the primary vertex + std::array dcaInfo{}; ///< DCA (xy, z) to the primary vertex uint32_t mask{}; ///< aod::HfSelTrack::isIdentifiedPid int64_t globalIndex{}; ///< index in the track table, written to the skim }; @@ -953,9 +951,9 @@ struct HfMlBasedTrackSelector { // filter collisions Filter filterSelectCollisions = (aod::hf_sel_collision::whyRejectColl == static_cast(0)); // filter track indices - Filter filterSelectTrackIds = ( (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand2Prong))) != 0u || - (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u || - (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::CandDstar))) != 0u ); + Filter filterSelectTrackIds = ((aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand2Prong))) != 0u || + (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::Cand3Prong))) != 0u || + (aod::hf_sel_track::isSelProng & static_cast(BIT(CandidateType::CandDstar))) != 0u); Preslice trackIndicesPerCollision = aod::track_association::collisionId; @@ -1702,7 +1700,7 @@ struct HfMlBasedTrackSelector { } // D* -> D0 pi+ candidates - const auto tStartDstar = tickIf(doTiming); + const auto tStartDstar = tickIf(doTiming); if (storeDstar && lastFilledD0 != rowTrackIndexProng2.lastIndex()) { // we have a new D0 candidate, so we can try to form a D* candidate @@ -1747,8 +1745,12 @@ struct HfMlBasedTrackSelector { timing[TimeProcessDstar] += elapsedSeconds(tStartDstar); } - if (!config.do3Prongs) { continue; } - if (!canBe3Prong) { continue; } + if (!config.do3Prongs) { + continue; + } + if (!canBe3Prong) { + continue; + } if (!isTwoProngVtxGoodFor3Prongs) { ++counters[CountPairsRejected2P]; continue; @@ -1761,9 +1763,9 @@ struct HfMlBasedTrackSelector { } ++counters[CountTripletsEnumerated]; processTriplet(collision, prong1.parCov, prongOpp.parCov, prong2.parCov, - prong1.pVec, prongOpp.pVec, prong2.pVec, - prong1.mask, prongOpp.mask, prong2.mask, - prong1.globalIndex, prongOpp.globalIndex, prong2.globalIndex); + prong1.pVec, prongOpp.pVec, prong2.pVec, + prong1.mask, prongOpp.mask, prong2.mask, + prong1.globalIndex, prongOpp.globalIndex, prong2.globalIndex); } } } diff --git a/PWGHF/TableProducer/candidateCreatorDstar.cxx b/PWGHF/TableProducer/candidateCreatorDstar.cxx index 4e5b65122d6..f8839d1790a 100644 --- a/PWGHF/TableProducer/candidateCreatorDstar.cxx +++ b/PWGHF/TableProducer/candidateCreatorDstar.cxx @@ -394,24 +394,24 @@ struct HfCandidateCreatorDstar { massDstar = RecoDecay::m(std::array{std::array{pVecSoftPi[0], pVecSoftPi[1], pVecSoftPi[2]}, std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, - std::array{constants::physics::MassPiPlus, - constants::physics::MassPiPlus, - constants::physics::MassKPlus}); + std::array{constants::physics::MassPiPlus, + constants::physics::MassPiPlus, + constants::physics::MassKPlus}); massD0 = RecoDecay::m(std::array{std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, - std::array{constants::physics::MassPiPlus, - constants::physics::MassKPlus}); + std::array{constants::physics::MassPiPlus, + constants::physics::MassKPlus}); } else { massDstar = RecoDecay::m(std::array{std::array{pVecSoftPi[0], pVecSoftPi[1], pVecSoftPi[2]}, std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, - std::array{constants::physics::MassPiPlus, - constants::physics::MassKPlus, - constants::physics::MassPiPlus}); + std::array{constants::physics::MassPiPlus, + constants::physics::MassKPlus, + constants::physics::MassPiPlus}); massD0 = RecoDecay::m(std::array{std::array{pVecD0Prong0[0], pVecD0Prong0[1], pVecD0Prong0[2]}, std::array{pVecD0Prong1[0], pVecD0Prong1[1], pVecD0Prong1[2]}}, - std::array{constants::physics::MassKPlus, - constants::physics::MassPiPlus}); + std::array{constants::physics::MassKPlus, + constants::physics::MassPiPlus}); } auto deltaMass = std::abs(massDstar - massD0); registry.fill(HIST("QA/hDeltaMassDStar"), deltaMass); From bf92d882159e96d6a34820623ab7762363947b1e Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Tue, 27 Oct 2026 09:14:39 +0000 Subject: [PATCH 4/6] Please consider the following formatting changes --- PWGHF/D2H/TableProducer/dplusCombinatorics.cxx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx index 349512519c9..35341b7d543 100644 --- a/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx +++ b/PWGHF/D2H/TableProducer/dplusCombinatorics.cxx @@ -125,7 +125,7 @@ DECLARE_SOA_COLUMN(ScorePionDzero, scorePionDzero, float); //! ML score, pion-fr DECLARE_SOA_COLUMN(ScoreKaonDzero, scoreKaonDzero, float); //! ML score, kaon-from-D0 class DECLARE_SOA_COLUMN(IsSelPionDzero, isSelPionDzero, bool); //! passes the D0 pion-score threshold DECLARE_SOA_COLUMN(IsSelKaonDzero, isSelKaonDzero, bool); //! passes the D0 kaon-score threshold - // D+ -> K pi pi model + // D+ -> K pi pi model DECLARE_SOA_COLUMN(ScoreBkgDplus, scoreBkgDplus, float); //! ML score, background class DECLARE_SOA_COLUMN(ScorePionDplus, scorePionDplus, float); //! ML score, pion-from-D+ class DECLARE_SOA_COLUMN(ScoreKaonDplus, scoreKaonDplus, float); //! ML score, kaon-from-D+ class From 01100902be867da0dff7bee346d664d4205f77cf Mon Sep 17 00:00:00 2001 From: marcellocosti Date: Tue, 27 Oct 2026 10:28:05 +0100 Subject: [PATCH 5/6] Restore CMakeFiles --- .../TableProducer/mlBasedTrackSelector.cxx | 72 +++++++------------ PWGLF/TableProducer/Resonances/CMakeLists.txt | 8 +-- PWGLF/Tasks/Resonances/CMakeLists.txt | 28 ++++---- 3 files changed, 44 insertions(+), 64 deletions(-) diff --git a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx index 9600525084d..9aaf7ec568b 100644 --- a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx +++ b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx @@ -13,6 +13,7 @@ /// \brief D± → π± K∓ π± and D_s± → K± K∓ π± skimming with track-level ML /// TEST ONLY, PLEASE DO NOT USE FOR ANALYSIS /// \author Fabrizio Chinu , Universita and INFN Torino +/// \author Marcello Di Costanzo , Turin Polytechnic University and INFN Turin #include "PWGHF/Core/CentralityEstimation.h" #include "PWGHF/Core/HfMlResponseHfTracks.h" @@ -102,7 +103,9 @@ enum Channels3Prong { NChannels3Prong }; -constexpr std::array DecayTypeOfChannel{ +constexpr std::array DecayTypeOfChannel2Prong{ + hf_cand_2prong::DecayType::D0ToPiK}; +constexpr std::array DecayTypeOfChannel3Prong{ hf_cand_3prong::DecayType::DplusToPiKPi, hf_cand_3prong::DecayType::DsToKKPi}; @@ -120,9 +123,6 @@ enum TrackMlRole { NTrackMlRoles }; -constexpr std::array RolePionOfChannel{RolePionDplus, RolePionDs}; -constexpr std::array RoleKaonOfChannel{RoleKaonDplus, RoleKaonDs}; - constexpr int N2Prongs = 2; // Number of prongs for 2-prong candidates constexpr int N3Prongs = 3; // Number of prongs for 3-prong candidates constexpr int NMassHypos = 2; // mass hypotheses per channel, i.e. the two orderings of the same-sign pair @@ -150,7 +150,7 @@ constexpr TrackMlRole trackRoles3Prongs[NChannels3Prong][NMassHypos][N3Prongs] = constexpr double PtMaxModel = 1.e10; constexpr uint32_t MaskSameSignAny = (1u << RolePionD0) | (1u << RoleKaonD0) | (1u << RolePionDplus) | (1u << RolePionDs) | (1u << RoleKaonDs) | (1u << RoleSoftPiDstar); -constexpr uint32_t MaskOppSignAny = (1u << RolePionD0) | (1u << RoleKaonD0) | (1u << RoleKaonDplus) | (1u << RoleKaonDs); +constexpr uint32_t MaskOppSignAny = (1u << RolePionD0) | (1u << RoleKaonD0) | (1u << RoleKaonDplus) | (1u << RoleKaonDs) | (1u << RoleSoftPiDstar); /// Where the combinatorics spends its time, accumulated in seconds into hTiming. enum TimingStep { @@ -1061,7 +1061,7 @@ struct HfMlBasedTrackSelector { /// \param maskSame is the isIdentifiedPid mask of the first same-sign prong static void checkCandidateAllowed(const uint32_t maskOpp, const uint32_t maskSame, bool& canBe2Prong, bool& canBe3Prong) { - canBe2Prong = TESTBIT(maskOpp, RoleKaonD0) && TESTBIT(maskSame, RolePionD0); + canBe2Prong = (TESTBIT(maskOpp, RoleKaonD0) && TESTBIT(maskSame, RolePionD0)) || (TESTBIT(maskOpp, RolePionD0) && TESTBIT(maskSame, RoleKaonD0)); // For D* we check that we can have a D0 and later apply the selection for the soft pion const bool okDplus = TESTBIT(maskOpp, RoleKaonDplus) && TESTBIT(maskSame, RolePionDplus); @@ -1069,30 +1069,6 @@ struct HfMlBasedTrackSelector { canBe3Prong = okDplus || okDs; } - /// Check whether the three ML masks can satisfy any mass hypothesis of either channel - /// \param maskOpp is the mask of the opposite-sign prong (prong 1) - /// \param mask1, mask2 are the masks of the same-sign prongs (prongs 0 and 2) - static bool pairRoleOk(const uint32_t maskOpp, const uint32_t maskSame) - { - const std::array masks{maskSame, maskOpp}; - for (int iChannel = 0; iChannel < NChannels2Prong; iChannel++) { - for (int iHypo = 0; iHypo < NMassHypos; iHypo++) { - bool isHypoAlive = true; - for (int iProng = 0; iProng < N2Prongs; iProng++) { - const int role = trackRoles2Prongs[iChannel][iHypo][iProng]; - if (!TESTBIT(masks[iProng], role)) { - isHypoAlive = false; - break; - } - } - if (isHypoAlive) { - return true; - } - } - } - return false; - } - /// Check whether the three ML masks can satisfy any mass hypothesis of either channel /// \param maskOpp is the mask of the opposite-sign prong (prong 1) /// \param mask1, mask2 are the masks of the same-sign prongs (prongs 0 and 2) @@ -1307,7 +1283,9 @@ struct HfMlBasedTrackSelector { size_t nChannels = is3Prong ? static_cast(NChannels3Prong) : static_cast(NChannels2Prong); for (size_t iChannel = 0; iChannel < nChannels; iChannel++) { if (TESTBIT(isSelected, iChannel)) { - SETBIT(hfFlag, DecayTypeOfChannel[iChannel]); + const int decayType = is3Prong ? DecayTypeOfChannel3Prong[iChannel] + : DecayTypeOfChannel2Prong[iChannel]; + SETBIT(hfFlag, decayType); } } return hfFlag; @@ -1491,16 +1469,18 @@ struct HfMlBasedTrackSelector { /// Reconstruct and preselect one 2-prong candidate, and fill its table row if it survives. /// \param collision is the collision the candidate belongs to - /// \param prong0,1 are the track informations + /// \param prongPos is the positive track information + /// \param prongNeg are the negative track information template void processPair(SelectedCollisions::iterator const& collision, - TProngInfo const& prong0, TProngInfo const& prong1, - std::array& whichHypo2Prong) + TProngInfo const& prongPos, TProngInfo const& prongNeg, + std::array& whichHypo2Prong) { uint32_t isSelected2ProngCand = BIT(NChannels2Prong) - 1; whichHypo2Prong.fill(BIT(NMassHypos) - 1); // 2 bits on, all mass hypotheses alive + whichHypo2Prong[NChannels2Prong] = 0; // D0 for D* - std::array isIdentifiedPid{prong0.mask, prong1.mask}; + std::array isIdentifiedPid{prongPos.mask, prongNeg.mask}; applyMlRoleSelection(isIdentifiedPid, whichHypo2Prong, isSelected2ProngCand); if (isSelected2ProngCand == 0) { return; @@ -1516,7 +1496,7 @@ struct HfMlBasedTrackSelector { // 2-prong selections const std::array pvCoord{collision.posX(), collision.posY(), collision.posZ()}; - applySelection2Prong(pVecCandProng2, prong0, prong1, secondaryVertex2, pvCoord, whichHypo2Prong, isSelected2ProngCand); + applySelection2Prong(pVecCandProng2, prongPos, prongNeg, secondaryVertex2, pvCoord, whichHypo2Prong, isSelected2ProngCand); // Only D0 for now if (!TESTBIT(isSelected2ProngCand, ChannelD0ToPiK)) { @@ -1524,7 +1504,7 @@ struct HfMlBasedTrackSelector { } // fill table row - rowTrackIndexProng2(collision.globalIndex(), prong0.globalIndex, prong1.globalIndex, hfFlagOfSelection(isSelected2ProngCand, false)); + rowTrackIndexProng2(collision.globalIndex(), prongPos.globalIndex, prongNeg.globalIndex, hfFlagOfSelection(isSelected2ProngCand, false)); ++counters[CountPairsWritten]; // fill histograms @@ -1653,8 +1633,8 @@ struct HfMlBasedTrackSelector { for (const int iOpp : oppSignList) { // o2-linter: disable=const-ref-in-for-loop (int elements) const auto& prongOpp = oppSignCache.prongs[iOpp]; - for (std::size_t i1 = 0; i1 < sameSignList.size(); ++i1) { - const auto& prong1 = sameSignCache.prongs[sameSignList[i1]]; + for (std::size_t iSame1 = 0; iSame1 < sameSignList.size(); ++iSame1) { + const auto& prong1 = sameSignCache.prongs[sameSignList[iSame1]]; ++counters[CountPairsSeen]; bool canBe2Prong{false}, canBe3Prong{false}; bool isTwoProngVtxGoodFor3Prongs{true}; @@ -1687,6 +1667,9 @@ struct HfMlBasedTrackSelector { ++counters[CountPairsRejected2P]; continue; } + if (useTwoTrackVertex) { + isTwoProngVtxGoodFor3Prongs = isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); + } if (store2Prongs || storeDstar) { const auto tStartD0 = tickIf(doTiming); processPair(collision, prong1, prongOpp, whichHypo2Prong); @@ -1694,9 +1677,6 @@ struct HfMlBasedTrackSelector { timing[TimeProcessD0] += elapsedSeconds(tStartD0); } } - if (useTwoTrackVertex) { - isTwoProngVtxGoodFor3Prongs = isTwoTrackVertexSelectedFor3Prongs(df2.getPCACandidate(), pvCoord2Prong, df2); - } } // D* -> D0 pi+ candidates @@ -1710,7 +1690,7 @@ struct HfMlBasedTrackSelector { for (std::size_t iSoftPi = 0; iSoftPi < sameSignList.size(); ++iSoftPi) { const auto& prongSoftPi = sameSignCache.prongs[sameSignList[iSoftPi]]; - if (iSoftPi == static_cast(i1)) { + if (iSoftPi == static_cast(iSame1)) { continue; // skip the same track } if (!TESTBIT(prongSoftPi.mask, RoleSoftPiDstar)) { @@ -1737,7 +1717,7 @@ struct HfMlBasedTrackSelector { // Update the index of the last filled D0 candidate processDstar(collision, lastFilledD0, prongSoftPi.globalIndex, - prong1.pVec, prongOpp.pVec, prongSoftPi.pVec); + prongOpp.pVec, prong1.pVec, prongSoftPi.pVec); } } } @@ -1756,8 +1736,8 @@ struct HfMlBasedTrackSelector { continue; } // 3-prong candidates - for (std::size_t i2 = i1 + 1; i2 < sameSignList.size(); ++i2) { - const auto& prong2 = sameSignCache.prongs[sameSignList[i2]]; + for (std::size_t iSame2 = iSame1 + 1; iSame2 < sameSignList.size(); ++iSame2) { + const auto& prong2 = sameSignCache.prongs[sameSignList[iSame2]]; if (!tripletRoleOk(prongOpp.mask, prong1.mask, prong2.mask)) { continue; } diff --git a/PWGLF/TableProducer/Resonances/CMakeLists.txt b/PWGLF/TableProducer/Resonances/CMakeLists.txt index 0a0f66a0f8c..bafd0b12ea0 100644 --- a/PWGLF/TableProducer/Resonances/CMakeLists.txt +++ b/PWGLF/TableProducer/Resonances/CMakeLists.txt @@ -35,10 +35,10 @@ o2physics_add_dpl_workflow(resonance-initializer PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase COMPONENT_NAME Analysis) -# o2physics_add_dpl_workflow(resonance-module-initializer -# SOURCES resonanceModuleInitializer.cxx -# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase -# COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(resonance-module-initializer + SOURCES resonanceModuleInitializer.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DetectorsBase + COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(resonance-merge-df SOURCES resonanceMergeDF.cxx diff --git a/PWGLF/Tasks/Resonances/CMakeLists.txt b/PWGLF/Tasks/Resonances/CMakeLists.txt index b73291d29f0..8c289ef8675 100644 --- a/PWGLF/Tasks/Resonances/CMakeLists.txt +++ b/PWGLF/Tasks/Resonances/CMakeLists.txt @@ -194,20 +194,20 @@ o2physics_add_dpl_workflow(xi1530analysisqa PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore COMPONENT_NAME Analysis) -# o2physics_add_dpl_workflow(omega2012analysis -# SOURCES omega2012Analysis.cxx -# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore -# COMPONENT_NAME Analysis) - -# o2physics_add_dpl_workflow(omega2012-training-table -# SOURCES omega2012TrainingTable.cxx -# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore -# COMPONENT_NAME Analysis) -# -# o2physics_add_dpl_workflow(xi1820analysis -# SOURCES xi1820Analysis.cxx -# PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore -# COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(omega2012analysis + SOURCES omega2012Analysis.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore + COMPONENT_NAME Analysis) + +o2physics_add_dpl_workflow(omega2012-training-table + SOURCES omega2012TrainingTable.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore + COMPONENT_NAME Analysis) + +o2physics_add_dpl_workflow(xi1820analysis + SOURCES xi1820Analysis.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore + COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(initializereventqa SOURCES initializereventqa.cxx From bb0a35ce9939691350c6a2962635f426745a79bb Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Tue, 27 Oct 2026 09:30:20 +0000 Subject: [PATCH 6/6] Please consider the following formatting changes --- PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx index 9aaf7ec568b..a0c3d02fd41 100644 --- a/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx +++ b/PWGHF/D2H/TableProducer/mlBasedTrackSelector.cxx @@ -1478,7 +1478,7 @@ struct HfMlBasedTrackSelector { { uint32_t isSelected2ProngCand = BIT(NChannels2Prong) - 1; whichHypo2Prong.fill(BIT(NMassHypos) - 1); // 2 bits on, all mass hypotheses alive - whichHypo2Prong[NChannels2Prong] = 0; // D0 for D* + whichHypo2Prong[NChannels2Prong] = 0; // D0 for D* std::array isIdentifiedPid{prongPos.mask, prongNeg.mask}; applyMlRoleSelection(isIdentifiedPid, whichHypo2Prong, isSelected2ProngCand);