From 666fa6136ab03b2e2b55db50b1c2539241af31f8 Mon Sep 17 00:00:00 2001 From: marcellocosti Date: Tue, 6 Oct 2026 18:51:54 +0200 Subject: [PATCH] [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 32a7ba5ad8c..6f63cc6673c 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)