Skip to content

[PWGCF] correlations.cxx: prompt + non-prompt separation #18432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jaelpark wants to merge 3 commits into
base: master
Choose a base branch
from
Open

[PWGCF] correlations.cxx: prompt + non-prompt separation #18432

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
369c4ea
Various HF fixed on CF correlations.cxx
jaelpark Oct 23, 2026
5221251
Various HF fixes for CF-correlations.cxx 2.
jaelpark Oct 23, 2026
fbe06e4
Linter fixes
jaelpark Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 74 additions & 11 deletions PWGCF/Tasks/correlations.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
#include <experimental/type_traits>
#include <iterator>
#include <memory>
#include <set>
#include <string>
#include <tuple>
#include <type_traits>
Expand Down Expand Up @@ -105,7 +106,7 @@ struct CorrelationTask {
O2_DEFINE_CONFIGURABLE(cfgTwoTrackCut, float, -1, "Two track cut: -1 = off; >0 otherwise distance value (suggested: 0.02)");
O2_DEFINE_CONFIGURABLE(cfgTwoTrackCutMinRadius, float, 0.8f, "Two track cut: radius in m from which two track cuts are applied");
O2_DEFINE_CONFIGURABLE(cfgLocalEfficiency, int, 0, "0 = OFF and 1 = ON for local efficiency");
O2_DEFINE_CONFIGURABLE(cfgDropStepRECO, bool, false, "choice to drop step RECO if efficiency correction is used")
O2_DEFINE_CONFIGURABLE(cfgDropStepRECO, bool, false, "choice to drop step RECO if efficiency correction is used");
O2_DEFINE_CONFIGURABLE(cfgCentBinsForMC, int, 0, "0 = OFF and 1 = ON for data like multiplicity/centrality bins for MC steps");
O2_DEFINE_CONFIGURABLE(cfgRequireRecoCollision, int, 1, "0 = all generated collisions; 1 = only generated collisions with exactly 1 reconstructed collision; 2 = select reconstructed collision with largest number of contributors per generated collision")
O2_DEFINE_CONFIGURABLE(cfgTrackBitMask, uint16_t, 0, "BitMask for track selection systematics; refer to the enum TrackSelectionCuts in filtering task");
Expand All @@ -118,6 +119,7 @@ struct CorrelationTask {
O2_DEFINE_CONFIGURABLE(cfgEfficiencyTrigger, std::string, "", "CCDB path to efficiency object for trigger particles")
O2_DEFINE_CONFIGURABLE(cfgEfficiencyAssociated, std::string, "", "CCDB path to efficiency object for associated particles")
O2_DEFINE_CONFIGURABLE(cfgUseRun3MagField, bool, false, "false: Run 2 GRPObject; true: Run 3 GRPMagField")
O2_DEFINE_CONFIGURABLE(cfgMultWeight, std::string, "", "Multiplicity/Centrality weight")

O2_DEFINE_CONFIGURABLE(cfgNoMixedEvents, int, 5, "Number of mixed events per event")
O2_DEFINE_CONFIGURABLE(cfgRejectMixedPhiProngEvents, bool, true, "Reject associated hadrons from either mixed-phi prong event")
Expand All @@ -132,6 +134,8 @@ struct CorrelationTask {
O2_DEFINE_CONFIGURABLE(cfgPtDepMLbkg, std::vector<float>, {}, "pT interval for ML training")
O2_DEFINE_CONFIGURABLE(cfgPtCentDepMLbkgSel, std::vector<float>, {}, "Bkg ML selection")
O2_DEFINE_CONFIGURABLE(cfgPtCentDepMLpromptSel, std::vector<float>, {}, "Prompt ML selection")
O2_DEFINE_CONFIGURABLE(cfgPtCentDepMLnonpromptSel, std::vector<float>, {}, "Non-prompt ML selection (maximum cut on the non-prompt score)")
O2_DEFINE_CONFIGURABLE(cfgMcClassSel, int, -1, "MC class for 2-prong efficiency: -1=auto(promptSel empty?all:prompt), 0=all, 1=prompt, 2=non-prompt")

std::unordered_map<int64_t, int64_t> cfTrackToCFCollision;

Expand Down Expand Up @@ -183,6 +187,7 @@ struct CorrelationTask {
bool mPairCuts = false;
THn* mEfficiencyTrigger = nullptr;
THn* mEfficiencyAssociated = nullptr;
TH1* mMultWeight = nullptr;
bool efficiencyLoaded = false;
} cfg;

Expand Down Expand Up @@ -220,6 +225,14 @@ struct CorrelationTask {
if (!cfgPtCentDepMLpromptSel->empty() && cfgPtCentDepMLpromptSel->size() != cfgPtDepMLbkg->size()) {
LOGF(fatal, "cfgPtDepMLbkg and cfgPtCentDepMLpromptSel must be same size.");
}
if (!cfgPtCentDepMLnonpromptSel->empty() && cfgPtCentDepMLnonpromptSel->size() != cfgPtDepMLbkg->size())
LOGF(fatal, "cfgPtDepMLbkg and cfgPtCentDepMLnonpromptSel must be same size.");
if (cfgPtCentDepMLpromptSel->size() > 0)
LOGF(info, "Prompt selection scores in use.");
else
LOGF(info, "Using inclusive selections.");
if (cfgPtCentDepMLnonpromptSel->size() > 0)
LOGF(info, "Non-prompt selection scores in use.");
}
registry.add("yields", "multiplicity/centrality vs pT vs eta", {HistType::kTH3F, {{100, 0, 100, "/multiplicity/centrality"}, {40, 0, 20, "p_{T}"}, {100, -2, 2, "#eta"}}});
registry.add("etaphi", "multiplicity/centrality vs eta vs phi", {HistType::kTH3F, {{100, 0, 100, "multiplicity/centrality"}, {100, -2, 2, "#eta"}, {200, 0, o2::constants::math::TwoPI, "#varphi"}}});
Expand All @@ -229,6 +242,7 @@ struct CorrelationTask {
const AxisSpec& a = AxisSpec(axisInvMass);
AxisSpec axisSpecMass = {1000, a.binEdges[0], a.binEdges[a.getNbins()]};
registry.add("invMass", "2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisPtTrigger, axisMultiplicity, axisVertex}});
registry.add("invMassCorrected", "2-prong invariant mass with efficiency correction (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisPtTrigger, axisMultiplicity, axisVertex}});
if (doprocessSame2Prong2Prong || doprocessSame2Prong2ProngML) {
registry.add("invMassTwoPart", "2D 2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisSpecMass, axisPtTrigger, axisPtAssoc, axisMultiplicity}});
registry.add("invMassTwoPartDPhi", "2D 2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisSpecMass, axisPtTrigger, axisPtAssoc, axisDeltaPhi}});
Expand Down Expand Up @@ -395,12 +409,16 @@ struct CorrelationTask {
return (track.decay() != aod::cf2prongtrack::D0ToPiK ||
(track.mlProbD0()[0] < cfgPtCentDepMLbkgSel->at(idx) && // not background
(cfgPtCentDepMLpromptSel->empty() ||
track.mlProbD0()[1] > cfgPtCentDepMLpromptSel->at(idx)))) && // prompt
track.mlProbD0()[1] > cfgPtCentDepMLpromptSel->at(idx)) && // prompt
(cfgPtCentDepMLnonpromptSel->empty() ||
track.mlProbD0()[2] < cfgPtCentDepMLnonpromptSel->at(idx)))) && // not non-prompt

((track.decay() != aod::cf2prongtrack::D0barToKPi && track.decay() != aod::cf2prongtrack::D0barToKPiExclusive) ||
(track.mlProbD0bar()[0] < cfgPtCentDepMLbkgSel->at(idx) && // not background
(cfgPtCentDepMLpromptSel->empty() ||
track.mlProbD0bar()[1] > cfgPtCentDepMLpromptSel->at(idx)))); // prompt
track.mlProbD0bar()[1] > cfgPtCentDepMLpromptSel->at(idx)) && // prompt
(cfgPtCentDepMLnonpromptSel->empty() ||
track.mlProbD0bar()[2] < cfgPtCentDepMLnonpromptSel->at(idx)))); // not non-prompt
}

template <class T>
Expand Down Expand Up @@ -428,6 +446,8 @@ struct CorrelationTask {
using HasInvMass = decltype(std::declval<T&>().invMass());
template <class T>
using HasPDGCode = decltype(std::declval<T&>().pdgCode());
template <class T>
using HasTrackType = decltype(std::declval<T&>().trackType());

template <typename TCollision, typename TTracks1, typename TTracks2>
void fillQA(const TCollision& collision, float multiplicity, float posZ, const TTracks1& tracks1, const TTracks2& tracks2)
Expand All @@ -442,7 +462,15 @@ struct CorrelationTask {
continue;
}
}
float triggerWeight = 1.0f;
if (cfg.mMultWeight) {
triggerWeight *= cfg.mMultWeight->GetBinContent(cfg.mMultWeight->FindBin(multiplicity));
}
if (cfg.mEfficiencyTrigger) {
triggerWeight *= getEfficiencyCorrection(cfg.mEfficiencyTrigger, track1.eta(), track1.pt(), multiplicity, collision.posZ());
}
registry.fill(HIST("invMass"), track1.invMass(), track1.pt(), multiplicity, posZ);
registry.fill(HIST("invMassCorrected"), track1.invMass(), track1.pt(), multiplicity, posZ, triggerWeight);
for (const auto& track2 : tracks2) {
if constexpr (std::experimental::is_detected<HasInvMass, typename TTracks2::iterator>::value && std::experimental::is_detected<HasDecay, typename TTracks2::iterator>::value) {
if (doprocessSame2Prong2Prong || doprocessMixed2Prong2Prong || doprocessSame2Prong2ProngML || doprocessMixed2Prong2ProngML) {
Expand Down Expand Up @@ -628,6 +656,9 @@ struct CorrelationTask {
}
}

static std::set<std::tuple<int, int>> duplicates; // workaround for old derived datasets involving 2prongs when cfgDecayParticleMask = 0
duplicates.clear();

for (const auto& track1 : tracks1) {
// LOGF(info, "Track %f | %f | %f %d %d", track1.eta(), track1.phi(), track1.pt(), track1.isGlobalTrack(), track1.isGlobalTrackSDD());

Expand All @@ -637,6 +668,15 @@ struct CorrelationTask {
}
}

if constexpr (std::experimental::is_detected<HasProng0Id, typename TTracks1::iterator>::value) {
if (cfgDecayParticleMask == 0) {
auto a = std::make_tuple(track1.cfTrackProng0Id(), track1.cfTrackProng1Id()); // workaround allows for use of D0bar in old derived datasets
if (duplicates.contains(a))
continue;
duplicates.insert(a);
}
}

// sign check and PDG code special cases
if constexpr (std::experimental::is_detected<HasPDGCode, typename TTracks1::iterator>::value) {
// If the MC trigger particle is on the trigger PDG code list, we will accept them regardless of their charge.
Expand Down Expand Up @@ -691,6 +731,9 @@ struct CorrelationTask {

float triggerWeight = eventWeight;
if constexpr (step == CorrelationContainer::kCFStepCorrected) {
if (cfg.mMultWeight) {
triggerWeight *= cfg.mMultWeight->GetBinContent(cfg.mMultWeight->FindBin(multiplicity));
}
if (cfg.mEfficiencyTrigger) {
triggerWeight *= getEfficiencyCorrection(cfg.mEfficiencyTrigger, track1.eta(), track1.pt(), multiplicity, posZ);
}
Expand Down Expand Up @@ -777,6 +820,7 @@ struct CorrelationTask {
// If decay attributes are found for the second track/particle, we assume 2p-2p correlation
if constexpr (std::experimental::is_detected<HasMcDecay, typename TTracks2::iterator>::value) {
if ((((track2.mcDecay()) != aod::cf2prongtrack::D0ToPiK) && ((track2.mcDecay()) != aod::cf2prongtrack::D0barToKPiExclusive)) || (!cfgPtCentDepMLpromptSel->empty() && (track2.decay() & aod::cf2prongmcpart::Prompt) == 0)) {

continue;
}
} else if constexpr (std::experimental::is_detected<HasDecay, typename TTracks2::iterator>::value) {
Expand Down Expand Up @@ -929,6 +973,11 @@ struct CorrelationTask {
}
LOGF(info, "Loaded efficiency histogram for associated particles from %s (%p)", cfgEfficiencyAssociated.value.c_str(), static_cast<void*>(cfg.mEfficiencyAssociated));
}
if (cfgMultWeight.value.empty() == false) {
TFile* pf = TFile::Open(cfgMultWeight.value.c_str(), "READ");
cfg.mMultWeight = reinterpret_cast<TH1*>(pf->Get("hWeight"));
LOGF(info, "Loaded multiplicity weight histogram for associated particles from %s", cfgMultWeight.value.c_str());
}
cfg.efficiencyLoaded = true;
}

Expand Down Expand Up @@ -998,6 +1047,7 @@ struct CorrelationTask {
same->fillEvent(multiplicity, CorrelationContainer::kCFStepReconstructed);
fillCorrelations<CorrelationContainer::kCFStepReconstructed>(same, tracks1, tracks2, multiplicity, collision.posZ(), field, 1.0f);
}

if (hasEfficiency) {
same->fillEvent(multiplicity, CorrelationContainer::kCFStepCorrected);
fillCorrelations<CorrelationContainer::kCFStepCorrected>(same, tracks1, tracks2, multiplicity, collision.posZ(), field, 1.0f);
Expand Down Expand Up @@ -1363,6 +1413,19 @@ struct CorrelationTask {
}
// Primaries
p2indexCache.clear();
std::set<int> signalMC; // mc global index cache
std::set<int> reflectionMC; // mc global index cache
// Resolve the MC class to use for the efficiency denominator/numerator.
// cfgMcClassSel: -1=auto (promptSel empty -> all, else prompt), 0=all, 1=prompt, 2=non-prompt.
int mcClass = cfgMcClassSel.value;
if (mcClass < 0) {
mcClass = cfgPtCentDepMLpromptSel->empty() ? 0 : 1;
}
auto mcClassSel = [&](const auto& mp) -> bool {
bool isP = (mp.decay() & aod::cf2prongmcpart::Prompt) != 0;
bool isNP = (mp.decay() & aod::cf2prongmcpart::NonPrompt) != 0;
return (mcClass == 0) || (mcClass == 1 && isP) || (mcClass == 2 && isNP);
};
for (const auto& mcParticle : mcParticles) {
if (std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), mcParticle.pdgCode()) != cfgMcTriggerPDGs->end()) {
if ((mcParticle.mcDecay() != aod::cf2prongtrack::D0ToPiK) && (mcParticle.mcDecay() != aod::cf2prongtrack::D0barToKPiExclusive)) {
Expand All @@ -1372,9 +1435,8 @@ struct CorrelationTask {
continue; // daughters not found
}
if constexpr (!reflectionSpec) {
if (cfgPtCentDepMLpromptSel->empty() || (mcParticle.decay() & aod::cf2prongmcpart::Prompt) != 0) {
if (mcClassSel(mcParticle))
same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
}
}
p2indexCache.push_back(mcParticle.globalIndex());
}
Expand Down Expand Up @@ -1409,17 +1471,18 @@ struct CorrelationTask {
return false;
}
const auto& mcParticle = mcParticles.iteratorAt(*m - mcParticles.begin().globalIndex());
// check MC class and return true if not the target class (wont be a fake)
if (!mcClassSel(mcParticle))
return true; // a valid candidate but not the target class. Return true so we won't register a fake.
if constexpr (!reflectionSpec) {
if (!cfgPtCentDepMLpromptSel->empty() && (mcParticle.decay() & aod::cf2prongmcpart::Prompt) == 0) {
return true; // a valid candidate but not a prompt
}
same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoPrimaries, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
} else {
if ((mcParticle.mcDecay() == aod::cf2prongtrack::D0barToKPiExclusive && (p2track.decay() == aod::cf2prongtrack::D0barToKPiExclusive || p2track.decay() == aod::cf2prongtrack::D0barToKPi)) ||
(mcParticle.mcDecay() == aod::cf2prongtrack::D0ToPiK && p2track.decay() == aod::cf2prongtrack::D0ToPiK)) {
if (!signalMC.contains(mcParticle.globalIndex()) && ((mcParticle.mcDecay() == aod::cf2prongtrack::D0barToKPiExclusive && (p2track.decay() == aod::cf2prongtrack::D0barToKPiExclusive || p2track.decay() == aod::cf2prongtrack::D0barToKPi)) || (mcParticle.mcDecay() == aod::cf2prongtrack::D0ToPiK && p2track.decay() == aod::cf2prongtrack::D0ToPiK))) {
registry.fill(HIST("invMassSignal"), p2track.invMass(), p2track.pt(), multiplicity);
} else { // one particle may be filled into both histograms through duplicates
signalMC.insert(mcParticle.globalIndex());
} else if (!reflectionMC.contains(mcParticle.globalIndex()) && !signalMC.contains(mcParticle.globalIndex()) && ((mcParticle.mcDecay() == aod::cf2prongtrack::D0barToKPiExclusive && p2track.decay() == aod::cf2prongtrack::D0ToPiK) || (mcParticle.mcDecay() == aod::cf2prongtrack::D0ToPiK && (p2track.decay() == aod::cf2prongtrack::D0barToKPi || p2track.decay() == aod::cf2prongtrack::D0barToKPiExclusive)))) {
registry.fill(HIST("invMassReflected"), p2track.invMass(), p2track.pt(), multiplicity);
reflectionMC.insert(mcParticle.globalIndex());
}
}
return true;
Expand Down
Loading