diff --git a/PWGHF/D2H/Tasks/taskDplus.cxx b/PWGHF/D2H/Tasks/taskDplus.cxx index ca28eacd505..625e12581cc 100644 --- a/PWGHF/D2H/Tasks/taskDplus.cxx +++ b/PWGHF/D2H/Tasks/taskDplus.cxx @@ -17,11 +17,14 @@ /// \author Vít Kučera , CERN /// \author Luca Aglietta , University and INFN Torino +#include + #include "CommonConstants/PhysicsConstants.h" #include "Framework/AnalysisTask.h" #include "Framework/HistogramRegistry.h" #include "Framework/runDataProcessing.h" +#include "PWGHF/Core/CentralityEstimation.h" #include "PWGHF/Core/HfHelper.h" #include "PWGHF/DataModel/CandidateReconstructionTables.h" #include "PWGHF/DataModel/CandidateSelectionTables.h" @@ -30,17 +33,30 @@ using namespace o2; using namespace o2::analysis; using namespace o2::framework; using namespace o2::framework::expressions; +using namespace o2::hf_centrality; + +enum OccupancyEstimator { None = 0, + ITS, + FT0C }; + +enum BHadMothers { NotMatched = 0, + BPlus, + BZero, + Bs, + LambdaBZero }; /// D± analysis task struct HfTaskDplus { Configurable selectionFlagDplus{"selectionFlagDplus", 7, "Selection Flag for DPlus"}; // 7 corresponds to topo+PID cuts Configurable yCandGenMax{"yCandGenMax", 0.5, "max. gen particle rapidity"}; Configurable yCandRecoMax{"yCandRecoMax", 0.8, "max. cand. rapidity"}; + Configurable yGenNBins{"yGenNBins", 100, "number of bins for y axis in sparse for gen candidates"}; + Configurable centEstimator{"centEstimator", 0, "Centrality estimation (None: 0, FT0C: 2, FT0M: 3)"}; + Configurable occEstimator{"occEstimator", 0, "Occupancy estimation (None: 0, ITS: 1, FT0C: 2)"}; Configurable> binsPt{"binsPt", std::vector{hf_cuts_dplus_to_pi_k_pi::vecBinsPt}, "pT bin limits"}; Configurable> classMl{"classMl", {0, 1, 2}, "Indexes of ML scores to be stored. Three indexes max."}; - ConfigurableAxis axisMlScore0{"axisMlScore0", {100, 0., 1.}, "axis for ML output score 0"}; - ConfigurableAxis axisMlScore1{"axisMlScore1", {100, 0., 1.}, "axis for ML output score 1"}; - ConfigurableAxis axisMlScore2{"axisMlScore2", {100, 0., 1.}, "axis for ML output score 2"}; + Configurable storeCentrality{"storeCentrality", false, "Flag to store centrality information"}; + Configurable storeOccupancy{"storeOccupancy", false, "Flag to store occupancy information"}; HfHelper hfHelper; @@ -48,10 +64,16 @@ struct HfTaskDplus { using CandDplusDataWithMl = soa::Filtered>; using CandDplusMcReco = soa::Filtered>; using CandDplusMcRecoWithMl = soa::Filtered>; - using McParticles = soa::Join; + using CandDplusMcGen = soa::Join; + + using CollisionsCent = soa::Join; + using McRecoCollisionsCent = soa::Join; Filter filterDplusFlag = (o2::aod::hf_track_index::hfflag & static_cast(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi))) != static_cast(0); + Preslice mcParticlesPerMcCollision = aod::mcparticle::mcCollisionId; + PresliceUnsorted recoColPerMcCollision = aod::mccollisionlabel::mcCollisionId; + // data Partition selectedDPlusCandidates = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus; Partition selectedDPlusCandidatesWithMl = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus; @@ -64,7 +86,14 @@ struct HfTaskDplus { Partition recoBkgCandidates = nabs(aod::hf_cand_3prong::flagMcMatchRec) != static_cast(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi)) && aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus; Partition recoBkgCandidatesWithMl = nabs(aod::hf_cand_3prong::flagMcMatchRec) != static_cast(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi)) && aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus; - // Generated particles + ConfigurableAxis thnConfigAxisY{"thnConfigAxisY", {40, -1, 1}, "Cand. rapidity bins"}; + ConfigurableAxis thnConfigAxisCent{"thnConfigAxisCent", {110, 0., 110.}, ""}; + ConfigurableAxis thnConfigAxisOccupancy{"thnConfigAxisOccupancy", {14, 0, 14000}, "axis for centrality"}; + ConfigurableAxis thnConfigAxisPtBHad{"thnConfigAxisPtBHad", {25, 0., 50}, "axis for pt of B hadron decayed into D candidate"}; + ConfigurableAxis thnConfigAxisFlagBHad{"thnConfigAxisFlagBHad", {5, 0., 5}, "axis for PDG of B hadron"}; + ConfigurableAxis thnConfigAxisMlScore0{"thnConfigAxisMlScore0", {100, 0., 1.}, "axis for ML output score 0"}; + ConfigurableAxis thnConfigAxisMlScore1{"thnConfigAxisMlScore1", {100, 0., 1.}, "axis for ML output score 1"}; + ConfigurableAxis thnConfigAxisMlScore2{"thnConfigAxisMlScore2", {100, 0., 1.}, "axis for ML output score 2"}; HistogramRegistry registry{ "registry", @@ -84,9 +113,18 @@ struct HfTaskDplus { if ((std::accumulate(doprocess.begin(), doprocess.end(), 0)) != 1) { LOGP(fatal, "Only one process function should be enabled! Please check your configuration!"); } - auto vbins = (std::vector)binsPt; - AxisSpec ptbins = {vbins, "#it{p}_{T} (GeV/#it{c})"}; - AxisSpec massbins = {600, 1.67, 2.27, "inv. mass (K#pi#pi) (GeV/#it{c}^{2})"}; + auto vbins = static_cast>(binsPt); + AxisSpec thnAxisPt = {vbins, "#it{p}_{T} (GeV/#it{c})"}; + AxisSpec thnAxisMass = {600, 1.67, 2.27, "inv. mass (K#pi#pi) (GeV/#it{c}^{2})"}; + AxisSpec thnAxisY = {thnConfigAxisY, "y"}; + AxisSpec thnAxisMlScore0 = {thnConfigAxisMlScore0, "Score 0"}; + AxisSpec thnAxisMlScore1 = {thnConfigAxisMlScore1, "Score 1"}; + AxisSpec thnAxisMlScore2 = {thnConfigAxisMlScore2, "Score 2"}; + AxisSpec thnAxisPtBHad{thnConfigAxisPtBHad, "#it{p}_{T,B} (GeV/#it{c})"}; + AxisSpec thnAxisFlagBHad{thnConfigAxisFlagBHad, "B Hadron flag"}; + AxisSpec thnAxisCent{thnConfigAxisCent, "Centrality"}; + AxisSpec thnAxisOccupancy{thnConfigAxisOccupancy, "Occupancy"}; + registry.add("hMass", "3-prong candidates;inv. mass (#pi K #pi) (GeV/#it{c}^{2});entries", {HistType::kTH2F, {{350, 1.7, 2.05}, {vbins, "#it{p}_{T} (GeV/#it{c})"}}}); registry.add("hEta", "3-prong candidates;candidate #it{#eta};entries", {HistType::kTH2F, {{100, -2., 2.}, {vbins, "#it{p}_{T} (GeV/#it{c})"}}}); registry.add("hCt", "3-prong candidates;proper lifetime (D^{#pm}) * #it{c} (cm);entries", {HistType::kTH2F, {{120, -20., 100.}, {vbins, "#it{p}_{T} (GeV/#it{c})"}}}); @@ -124,13 +162,55 @@ struct HfTaskDplus { registry.add("hPtVsYGen", "MC particles (matched);#it{p}_{T}^{gen.}; #it{y}", {HistType::kTH2F, {{vbins, "#it{p}_{T} (GeV/#it{c})"}, {100, -5., 5.}}}); registry.add("hPtVsYGenPrompt", "MC particles (matched, prompt);#it{p}_{T}^{gen.}; #it{y}", {HistType::kTH2F, {{vbins, "#it{p}_{T} (GeV/#it{c})"}, {100, -5., 5.}}}); registry.add("hPtVsYGenNonPrompt", "MC particles (matched, non-prompt);#it{p}_{T}^{gen.}; #it{y}", {HistType::kTH2F, {{vbins, "#it{p}_{T} (GeV/#it{c})"}, {100, -5., 5.}}}); - if (doprocessDataWithMl) { - registry.add("hSparseMass", "THn for Dplus", HistType::kTHnSparseF, {massbins, ptbins, axisMlScore0, axisMlScore1, axisMlScore2}); + + if (doprocessDataWithMl || doprocessData) { + std::vector axes = {thnAxisMass, thnAxisPt}; + + if (doprocessDataWithMl) { + axes.insert(axes.end(), {thnAxisMlScore0, thnAxisMlScore1, thnAxisMlScore2}); + } + if (storeCentrality) { + axes.insert(axes.end(), {thnAxisCent}); + } + if (storeOccupancy) { + axes.insert(axes.end(), {thnAxisOccupancy}); + } + + registry.add("hSparseMass", "THn for Dplus", HistType::kTHnSparseF, axes); } - if (doprocessMcWithMl) { - registry.add("hSparseMassPrompt", "THn for Dplus Prompt", HistType::kTHnSparseF, {massbins, ptbins, axisMlScore0, axisMlScore1, axisMlScore2}); - registry.add("hSparseMassFD", "THn for Dplus FD", HistType::kTHnSparseF, {massbins, ptbins, axisMlScore0, axisMlScore1, axisMlScore2}); - registry.add("hSparseMassBkg", "THn for Dplus Bkg", HistType::kTHnSparseF, {massbins, ptbins, axisMlScore0, axisMlScore1, axisMlScore2}); + if (doprocessMcWithMl || doprocessMc) { + std::vector axes = {thnAxisMass, thnAxisPt}; + std::vector axesFD = {thnAxisMass, thnAxisPt}; + std::vector axesGenPrompt = {thnAxisPt, thnAxisY}; + std::vector axesGenFD = {thnAxisPt, thnAxisY}; + + axesFD.insert(axesFD.end(), {thnAxisPtBHad}); + axesFD.insert(axesFD.end(), {thnAxisFlagBHad}); + axesGenFD.insert(axesGenFD.end(), {thnAxisPtBHad}); + axesGenFD.insert(axesGenFD.end(), {thnAxisFlagBHad}); + + if (doprocessMcWithMl) { + axes.insert(axes.end(), {thnAxisMlScore0, thnAxisMlScore1, thnAxisMlScore2}); + axesFD.insert(axesFD.end(), {thnAxisMlScore0, thnAxisMlScore1, thnAxisMlScore2}); + } + if (storeCentrality) { + axes.insert(axes.end(), {thnAxisCent}); + axesFD.insert(axesFD.end(), {thnAxisCent}); + axesGenPrompt.insert(axesGenPrompt.end(), {thnAxisCent}); + axesGenFD.insert(axesGenFD.end(), {thnAxisCent}); + } + if (storeOccupancy) { + axes.insert(axes.end(), {thnAxisOccupancy}); + axesFD.insert(axesFD.end(), {thnAxisOccupancy}); + axesGenPrompt.insert(axesGenPrompt.end(), {thnAxisOccupancy}); + axesGenFD.insert(axesGenFD.end(), {thnAxisOccupancy}); + } + registry.add("hSparseMassPrompt", "THn for Dplus Prompt", HistType::kTHnSparseF, axes); + registry.add("hSparseMassFD", "THn for Dplus FD", HistType::kTHnSparseF, axesFD); + registry.add("hSparseMassBkg", "THn for Dplus Bkg", HistType::kTHnSparseF, axes); + registry.add("hSparseMassNotMatched", "THn for Dplus not matched", HistType::kTHnSparseF, axes); + registry.add("hSparseMassGenPrompt", "THn for gen Prompt Dplus", HistType::kTHnSparseF, axesGenPrompt); + registry.add("hSparseMassGenFD", "THn for gen FD Dplus", HistType::kTHnSparseF, axesGenFD); } } @@ -167,25 +247,80 @@ struct HfTaskDplus { // Fill THnSparses for the ML analysis /// \param candidate is a particle candidate + /// \param ptbhad transverse momentum of beauty mother for nonprompt candidates + /// \param flagBHad transverse momentum of beauty mother for nonprompt candidates + /// \param centrality collision centrality + /// \param occupancy collision occupancy template - void fillSparseML(const T1& candidate) + void fillSparseML(const T1& candidate, + float ptbhad, + int flagBHad, + float centrality, + float occupancy) { std::vector outputMl = {-999., -999., -999.}; for (unsigned int iclass = 0; iclass < classMl->size(); iclass++) { outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMl->at(iclass)]; } - if constexpr (isMc) { - if constexpr (isMatched) { - if (candidate.originMcRec() == RecoDecay::OriginType::Prompt) { - registry.fill(HIST("hSparseMassPrompt"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); - } else if (candidate.originMcRec() == RecoDecay::OriginType::NonPrompt) { - registry.fill(HIST("hSparseMassFD"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); + if constexpr (isMc) { // MC + if constexpr (isMatched) { // Matched + if (candidate.originMcRec() == RecoDecay::OriginType::Prompt) { // Prompt + + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassPrompt"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassPrompt"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassPrompt"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy); + } else { + registry.fill(HIST("hSparseMassPrompt"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); + } + + } else if (candidate.originMcRec() == RecoDecay::OriginType::NonPrompt) { // FD + + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassFD"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), ptbhad, flagBHad, outputMl[0], outputMl[1], outputMl[2], centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassFD"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), ptbhad, flagBHad, outputMl[0], outputMl[1], outputMl[2], centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassFD"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), ptbhad, flagBHad, outputMl[0], outputMl[1], outputMl[2], occupancy); + } else { + registry.fill(HIST("hSparseMassFD"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), ptbhad, flagBHad, outputMl[0], outputMl[1], outputMl[2]); + } + + } else { // Bkg + + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassBkg"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassBkg"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassBkg"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy); + } else { + registry.fill(HIST("hSparseMassBkg"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); + } } } else { - registry.fill(HIST("hSparseMassBkg"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassNotMatched"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassNotMatched"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassNotMatched"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy); + } else { + registry.fill(HIST("hSparseMassNotMatched"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); + } + } + } else { // Data + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy); + } else { + registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); } - } else { - registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]); } } @@ -258,11 +393,52 @@ struct HfTaskDplus { registry.fill(HIST("hEtaGen"), particle.eta()); } + // Fill THnSparse of quantities for generated Dplus particles + /// \param particle is a particle with MC information + /// \param ptGenB transverse momentum of beauty mother for nonprompt candidates + /// \param flagGenB transverse momentum of beauty mother for nonprompt candidates + /// \param centrality collision centrality + /// \param occupancy collision occupancy + template + void fillSparseMcGen(const T1& particle, + float ptGenB, + int flagGenB, + float centrality, + float occupancy) + { + auto yGen = RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus); + if (particle.originMcGen() == RecoDecay::OriginType::Prompt) { + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassGenPrompt"), particle.pt(), yGen, centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassGenPrompt"), particle.pt(), yGen, centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassGenPrompt"), particle.pt(), yGen, occupancy); + } else { + registry.fill(HIST("hSparseMassGenPrompt"), particle.pt(), yGen); + } + } else { + if (storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassGenFD"), particle.pt(), yGen, ptGenB, flagGenB, centrality, occupancy); + } else if (storeCentrality && !storeOccupancy) { + registry.fill(HIST("hSparseMassGenFD"), particle.pt(), yGen, ptGenB, flagGenB, centrality); + } else if (!storeCentrality && storeOccupancy) { + registry.fill(HIST("hSparseMassGenFD"), particle.pt(), yGen, ptGenB, flagGenB, occupancy); + } else { + registry.fill(HIST("hSparseMassGenFD"), particle.pt(), yGen, ptGenB, flagGenB); + } + } + } + // Run analysis for the reconstructed Dplus candidates from data /// \param candidates are reconstructed candidates template - void runDataAnalysis(const T1& /*candidates*/) + void runDataAnalysis(const T1& /*candidates*/, CollisionsCent const& /*colls*/) { + float cent{-1.f}; + float occ{-1.f}; + float ptBhad{-1.f}; + int flagBHad{-1}; if constexpr (!fillMl) { for (const auto& candidate : selectedDPlusCandidates) { if ((yCandRecoMax >= 0. && std::abs(hfHelper.yDplus(candidate)) > yCandRecoMax)) { @@ -275,17 +451,34 @@ struct HfTaskDplus { if ((yCandRecoMax >= 0. && std::abs(hfHelper.yDplus(candidate)) > yCandRecoMax)) { continue; } + + if (storeCentrality || storeOccupancy) { + auto collision = candidate.template collision_as(); + if (storeCentrality && centEstimator != CentralityEstimator::None) { + cent = getCentrality(collision); + } + if (storeOccupancy && occEstimator != OccupancyEstimator::None) { + occ = getOccupancy(collision); + } + } + fillHisto(candidate); - fillSparseML(candidate); + fillSparseML(candidate, ptBhad, flagBHad, cent, occ); } } } + // Run analysis for the reconstructed Dplus candidates with MC matching - /// \param candidates are reconstructed candidates - /// \param mcParticles are particles with MC information - template - void runMCAnalysis(const T1& /*recoCandidates*/, const T2& mcParticles) + /// \param recoCandidates are reconstructed candidates + /// \param recoColls are reconstructed collisions + template + void runAnalysisMcRec(McRecoCollisionsCent const& /*recoColls*/) { + float cent{-1}; + float occ{-1}; + float ptBhad{-1}; + int flagBHad{-1}; + // MC rec. w/o Ml if constexpr (!fillMl) { for (const auto& candidate : recoDPlusCandidates) { @@ -307,53 +500,220 @@ struct HfTaskDplus { if ((yCandRecoMax >= 0. && std::abs(hfHelper.yDplus(candidate)) > yCandRecoMax)) { continue; } + ptBhad = candidate.ptBhadMotherPart(); + flagBHad = getBHadMotherFlag(candidate.pdgBhadMotherPart()); + + if (storeCentrality || storeOccupancy) { + auto collision = candidate.template collision_as(); + if (storeCentrality && centEstimator != CentralityEstimator::None) { + cent = getCentrality(collision); + } + if (storeOccupancy && occEstimator != OccupancyEstimator::None) { + occ = getOccupancy(collision); + } + } + fillHisto(candidate); fillHistoMCRec(candidate); - fillSparseML(candidate); + fillSparseML(candidate, ptBhad, flagBHad, cent, occ); } // Bkg + ptBhad = -1; + flagBHad = -1; for (const auto& candidate : recoBkgCandidatesWithMl) { if ((yCandRecoMax >= 0. && std::abs(hfHelper.yDplus(candidate)) > yCandRecoMax)) { continue; } + auto collision = candidate.template collision_as(); + if (storeCentrality && centEstimator != CentralityEstimator::None) { + cent = getCentrality(collision); + } + if (storeOccupancy && occEstimator != OccupancyEstimator::None) { + occ = getOccupancy(collision); + } fillHistoMCRec(candidate); - fillSparseML(candidate); + fillSparseML(candidate, ptBhad, flagBHad, cent, occ); } } + } + + // Run analysis for the generated Dplus candidates + /// \param mcGenCollisions are the generated MC collisions + /// \param mcRecoCollisions are the reconstructed MC collisions + /// \param mcGenParticles are the generated MC particle candidates + template + void runAnalysisMcGen(aod::McCollisions const& mcGenCollisions, + McRecoCollisionsCent const& mcRecoCollisions, + Cand const& mcGenParticles) + { // MC gen. - for (const auto& particle : mcParticles) { - auto yGen = RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus); - if ((yCandGenMax >= 0. && std::abs(yGen) > yCandGenMax) || (std::abs(particle.flagMcMatchGen()) != 1 << aod::hf_cand_3prong::DecayType::DplusToPiKPi)) { - continue; + float cent{-1.}; + float occ{-1.}; + float ptGenB{-1.}; + int flagGenB{-1}; + + for (const auto& mcGenCollision : mcGenCollisions) { + const auto recoCollsPerGenMcColl = mcRecoCollisions.sliceBy(recoColPerMcCollision, mcGenCollision.globalIndex()); + const auto mcParticlesPerGenMcColl = mcGenParticles.sliceBy(mcParticlesPerMcCollision, mcGenCollision.globalIndex()); + if (storeCentrality && centEstimator != CentralityEstimator::None) { + cent = getMcGenCollCentrality(recoCollsPerGenMcColl); } - fillHistoMCGen(particle); + if (storeOccupancy && occEstimator != OccupancyEstimator::None) { + occ = getMcGenCollOccupancy(recoCollsPerGenMcColl); + } + + for (const auto& particle : mcParticlesPerGenMcColl) { + ptGenB = -1; + flagGenB = -1; + auto yGen = RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus); + if ((yCandGenMax >= 0. && std::abs(yGen) > yCandGenMax) || (std::abs(particle.flagMcMatchGen()) != 1 << aod::hf_cand_3prong::DecayType::DplusToPiKPi)) { + continue; + } + if (particle.originMcGen() == RecoDecay::OriginType::NonPrompt) { + auto bHadMother = mcGenParticles.rawIteratorAt(particle.idxBhadMotherPart() - mcGenParticles.offset()); + flagGenB = getBHadMotherFlag(bHadMother.pdgCode()); + ptGenB = bHadMother.pt(); + } + fillHistoMCGen(particle); + if constexpr (fillMl) { + fillSparseMcGen(particle, ptGenB, flagGenB, cent, occ); + } + } + } + } + + /// Get the occupancy + /// \param collision is the collision with the occupancy information + /// \return collision occupancy + template + float getOccupancy(Coll const& collision) + { + float occupancy = -999.; + switch (occEstimator) { + case OccupancyEstimator::ITS: + occupancy = collision.trackOccupancyInTimeRange(); + break; + case OccupancyEstimator::FT0C: + occupancy = collision.ft0cOccupancyInTimeRange(); + break; + default: + LOG(warning) << "Occupancy estimator not valid. Possible values are ITS or FT0C. Fallback to ITS"; + occupancy = collision.trackOccupancyInTimeRange(); + break; + } + return occupancy; + } + + /// \brief Function to get MC collision occupancy + /// \param collSlice collection of reconstructed collisions associated to a generated one + /// \return generated MC collision occupancy + template + int getMcGenCollOccupancy(CCs const& collSlice) + { + float multiplicity{0.f}; + int occupancy = 0; + for (const auto& collision : collSlice) { + float collMult{0.f}; + collMult = collision.numContrib(); + if (collMult > multiplicity) { + occupancy = getOccupancy(collision); + multiplicity = collMult; + } + } // end loop over collisions + + return occupancy; + } + + /// Get the centrality + /// \param collision is the collision with the centrality information + /// \return collision centrality + template + float getCentrality(Coll const& collision) + { + float cent = -999.; + switch (centEstimator) { + case CentralityEstimator::FT0C: + cent = collision.centFT0C(); + break; + case CentralityEstimator::FT0M: + cent = collision.centFT0M(); + break; + default: + LOG(warning) << "Centrality estimator not valid. Possible values are FT0C, FT0M. Fallback to FT0C"; + cent = collision.centFT0C(); + break; + } + return cent; + } + + /// \brief Function to get MC collision centrality + /// \param collSlice collection of reconstructed collisions associated to a generated one + /// \return generated MC collision centrality + template + float getMcGenCollCentrality(CCs const& collSlice) + { + float centrality{-1}; + float multiplicity{0.f}; + for (const auto& collision : collSlice) { + float collMult = collision.numContrib(); + if (collMult > multiplicity) { + centrality = getCentrality(collision); + multiplicity = collMult; + } + } + return centrality; + } + + /// Convert the B hadron mother PDG for non prompt candidates to a flag + /// \param pdg of the b hadron mother + /// \return integer map to specific mothers' PDG codes + int getBHadMotherFlag(const int& flagBHad) + { + if (std::abs(flagBHad) == o2::constants::physics::kBPlus) { + return BHadMothers::BPlus; + } + if (std::abs(flagBHad) == o2::constants::physics::kB0) { + return BHadMothers::BZero; + } + if (std::abs(flagBHad) == o2::constants::physics::kBS) { + return BHadMothers::Bs; + } + if (std::abs(flagBHad) == o2::constants::physics::kLambdaB0) { + return BHadMothers::LambdaBZero; } + return BHadMothers::NotMatched; } // process functions - void processData(CandDplusData const& candidates) + void processData(CandDplusData const& candidates, CollisionsCent const& collisions) { - runDataAnalysis(candidates); + runDataAnalysis(candidates, collisions); } PROCESS_SWITCH(HfTaskDplus, processData, "Process data w/o ML", true); - void processDataWithMl(CandDplusDataWithMl const& candidates) + void processDataWithMl(CandDplusDataWithMl const& candidates, CollisionsCent const& collisions) { - runDataAnalysis(candidates); + runDataAnalysis(candidates, collisions); } PROCESS_SWITCH(HfTaskDplus, processDataWithMl, "Process data with ML", false); - void processMc(CandDplusMcReco const& candidates, - McParticles const& mcParticles) + void processMc(CandDplusMcReco const&, + CandDplusMcGen const& mcGenParticles, + McRecoCollisionsCent const& mcRecoCollisions, + aod::McCollisions const& mcGenCollisions) { - runMCAnalysis(candidates, mcParticles); + runAnalysisMcRec(mcRecoCollisions); + runAnalysisMcGen(mcGenCollisions, mcRecoCollisions, mcGenParticles); } PROCESS_SWITCH(HfTaskDplus, processMc, "Process MC w/o ML", false); - void processMcWithMl(CandDplusMcRecoWithMl const& candidates, - McParticles const& mcParticles) + void processMcWithMl(CandDplusMcRecoWithMl const&, + CandDplusMcGen const& mcGenParticles, + McRecoCollisionsCent const& mcRecoCollisions, + aod::McCollisions const& mcGenCollisions) { - runMCAnalysis(candidates, mcParticles); + runAnalysisMcRec(mcRecoCollisions); + runAnalysisMcGen(mcGenCollisions, mcRecoCollisions, mcGenParticles); } PROCESS_SWITCH(HfTaskDplus, processMcWithMl, "Process MC with ML", false); };