From 137e1b845f301f56ef7dc75ff6c54a5877f47682 Mon Sep 17 00:00:00 2001 From: KaiCui Date: Thu, 30 Apr 2026 18:29:14 +0200 Subject: [PATCH] Add Model Prediction Add model prediction process two option : cent and mult --- .../Tasks/Strangeness/hStrangeCorrelation.cxx | 213 +++++++++++++++++- 1 file changed, 207 insertions(+), 6 deletions(-) diff --git a/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx b/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx index 96920535e21..f47b0c26eca 100644 --- a/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx +++ b/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx @@ -23,6 +23,8 @@ #include "PWGLF/DataModel/LFHStrangeCorrelationTables.h" #include "PWGLF/DataModel/LFStrangenessTables.h" +#include "PWGLF/DataModel/mcCentrality.h" +#include "PWGLF/Utils/inelGt.h" #include "Common/CCDB/EventSelectionParams.h" #include "Common/Core/RecoDecay.h" @@ -87,6 +89,7 @@ struct HStrangeCorrelation { Service ccdb; Service pdgDB; + o2::pwglf::ParticleCounter mCounter; HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject}; @@ -121,6 +124,8 @@ struct HStrangeCorrelation { Configurable doMCassociation{"doMCassociation", false, "fill everything only for MC associated"}; Configurable doTriggPhysicalPrimary{"doTriggPhysicalPrimary", false, "require physical primary for trigger particles"}; Configurable applyNewMCSelection{"applyNewMCSelection", false, "apply new MC Generated selection"}; + Configurable doSeparateFT0Prediction{"doSeparateFT0Prediction", false, "separate FT0M to FT0A and FT0C in prediction process"}; + Configurable useCentralityinPrediction{"useCentralityinPrediction", false, "if true, use centrality instead of multiplisity"}; } masterConfigurations; // master analysis switches @@ -154,8 +159,9 @@ struct HStrangeCorrelation { ConfigurableAxis axisPtAssoc{"axisPtAssoc", {VARIABLE_WIDTH, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 10.0}, "pt associated axis for histograms"}; ConfigurableAxis axisPtTrigger{"axisPtTrigger", {VARIABLE_WIDTH, 0.0, 1.0, 2.0, 3.0, 100}, "pt associated axis for histograms"}; ConfigurableAxis axisPtQA{"axisPtQA", {VARIABLE_WIDTH, 0.0f, 0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f, 1.3f, 1.4f, 1.5f, 1.6f, 1.7f, 1.8f, 1.9f, 2.0f, 2.2f, 2.4f, 2.6f, 2.8f, 3.0f, 3.2f, 3.4f, 3.6f, 3.8f, 4.0f, 4.4f, 4.8f, 5.2f, 5.6f, 6.0f, 6.5f, 7.0f, 7.5f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 17.0f, 19.0f, 21.0f, 23.0f, 25.0f, 30.0f, 35.0f, 40.0f, 50.0f}, "pt axis for QA histograms"}; - ConfigurableAxis axisMultCount{"axisMultCount", {VARIABLE_WIDTH, 0, 200, 400, 600, 800, 1000, 1400, 1800, 2300, 2800, 3300, 4000, 5000, 6000}, "Mixing bins - multiplicity"}; ConfigurableAxis axisMassNSigma{"axisMassNSigma", {40, -2, 2}, "Axis for mass Nsigma"}; + ConfigurableAxis axisMultiplicity{"axisMultiplicity", {VARIABLE_WIDTH, 0, 20, 40, 60, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300}, "Binning of the Multiplicity axis in model prediction process"}; + } axesConfigurations; // for topo var QA @@ -253,7 +259,7 @@ struct HStrangeCorrelation { } cascadeSelections; struct : ConfigurableGroup { - std::string prefix = "`"; + std::string prefix = "checks"; // cascade selections // more cascade selections in PbPb // Configurable bachBaryonCosPA{"bachBaryonCosPA", 0.9999, "Bachelor baryon CosPA"}; @@ -1568,6 +1574,7 @@ struct HStrangeCorrelation { const AxisSpec preAxisMult{axesConfigurations.axisMult, "mult percentile"}; const AxisSpec axisPtLambda{axesConfigurations.axisPtAssoc, "#it{p}_{T}^{#Lambda} (GeV/c)"}; const AxisSpec axisPtCascade{axesConfigurations.axisPtAssoc, "#it{p}_{T}^{Mother} (GeV/c)"}; + const AxisSpec preAxisMultiplicity{axesConfigurations.axisMultiplicity, "multiplicity"}; // store the original axes in specific TH1Cs for completeness histos.add("axes/hDeltaPhiAxis", "", kTH1C, {preAxisDeltaPhi}); @@ -1576,6 +1583,7 @@ struct HStrangeCorrelation { histos.add("axes/hPtTriggerAxis", "", kTH1C, {preAxisPtTrigger}); histos.add("axes/hVertexZAxis", "", kTH1C, {preAxisVtxZ}); histos.add("axes/hMultAxis", "", kTH1C, {preAxisMult}); + histos.add("axes/hMultiplicityAxis", "", kTH1C, {preAxisMultiplicity}); std::vector edgesDeltaPhiOrig = preAxisDeltaPhi.binEdges; std::vector edgesDeltaEtaOrig = preAxisDeltaEta.binEdges; @@ -1583,6 +1591,7 @@ struct HStrangeCorrelation { std::vector edgesPtTriggerOrig = preAxisPtTrigger.binEdges; std::vector edgesVtxZOrig = preAxisVtxZ.binEdges; std::vector edgesMultOrig = preAxisMult.binEdges; + std::vector edgesMultiplicityOrig = preAxisMultiplicity.binEdges; std::vector rangesDeltaPhi = {static_cast(edgesDeltaPhiOrig[0]), static_cast(edgesDeltaPhiOrig[edgesDeltaPhiOrig.size() - 1])}; std::vector rangesDeltaEta = {static_cast(edgesDeltaEtaOrig[0]), static_cast(edgesDeltaEtaOrig[edgesDeltaEtaOrig.size() - 1])}; @@ -1590,6 +1599,7 @@ struct HStrangeCorrelation { std::vector rangesPtTrigger = {static_cast(edgesPtTriggerOrig[0]), static_cast(edgesPtTriggerOrig[edgesPtTriggerOrig.size() - 1])}; std::vector rangesVtxZ = {static_cast(edgesVtxZOrig[0]), static_cast(edgesVtxZOrig[edgesVtxZOrig.size() - 1])}; std::vector rangesMult = {static_cast(edgesMultOrig[0]), static_cast(edgesMultOrig[edgesMultOrig.size() - 1])}; + std::vector rangesMultiplicity = {static_cast(edgesMultiplicityOrig[0]), static_cast(edgesMultiplicityOrig[edgesMultiplicityOrig.size() - 1])}; axisRanges.emplace_back(rangesDeltaPhi); axisRanges.emplace_back(rangesDeltaEta); @@ -1597,6 +1607,7 @@ struct HStrangeCorrelation { axisRanges.emplace_back(rangesPtTrigger); axisRanges.emplace_back(rangesVtxZ); axisRanges.emplace_back(rangesMult); + axisRanges.emplace_back(rangesMultiplicity); std::vector edgesDeltaPhi; std::vector edgesDeltaEta; @@ -1604,6 +1615,7 @@ struct HStrangeCorrelation { std::vector edgesPtTrigger; std::vector edgesVtxZ; std::vector edgesMult; + std::vector edgesMultiplicity; // v--- skipUnderOverflowInTHn ---v // @@ -1688,13 +1700,26 @@ struct HStrangeCorrelation { for (int i = offset; i < preAxisMult.nBins.value() + 1 - offset; i++) edgesMult.emplace_back(min + static_cast(i) * delta); } + // ===] multiplicity count [=== + if (!preAxisMultiplicity.nBins.has_value()) { + // variable binning, use bins provided + for (int i = offset; i < static_cast(edgesMultiplicityOrig.size()) - offset; i++) + edgesMultiplicity.emplace_back(edgesMultiplicityOrig[i]); + } else { + // fixed binning, generate the bin edges on-the-spot + double min = edgesMultiplicityOrig[0]; + double delta = (edgesMultiplicityOrig[1] - edgesMultiplicityOrig[0]) / preAxisMultiplicity.nBins.value(); + for (int i = offset; i < preAxisMultiplicity.nBins.value() + 1 - offset; i++) + edgesMultiplicity.emplace_back(min + static_cast(i) * delta); + } LOGF(info, "Initialized THnF axis delta-phi with %i bins.", edgesDeltaPhi.size() - 1); LOGF(info, "Initialized THnF axis delta-eta with %i bins.", edgesDeltaEta.size() - 1); LOGF(info, "Initialized THnF axis pTassoc with %i bins.", edgesPtAssoc.size() - 1); LOGF(info, "Initialized THnF axis pTtrigger with %i bins.", edgesPtTrigger.size() - 1); LOGF(info, "Initialized THnF axis vertex-Z with %i bins.", edgesVtxZ.size() - 1); - LOGF(info, "Initialized THnF axis multiplicity with %i bins.", edgesMult.size() - 1); + LOGF(info, "Initialized THnF axis mult cent with %i bins.", edgesMult.size() - 1); + LOGF(info, "Initialized THnF axis multiplicity with %i bins.", edgesMultiplicity.size() - 1); const AxisSpec axisDeltaPhiNDim{edgesDeltaPhi, "#Delta#varphi"}; const AxisSpec axisDeltaEtaNDim{edgesDeltaEta, "#Delta#eta"}; @@ -1702,11 +1727,12 @@ struct HStrangeCorrelation { const AxisSpec axisPtTriggerNDim{edgesPtTrigger, "#it{p}_{T}^{trigger} (GeV/c)"}; const AxisSpec axisVtxZNDim{edgesVtxZ, "vertex Z (cm)"}; const AxisSpec axisMultNDim{edgesMult, "mult percentile"}; + const AxisSpec axisMultiplicityNDim{edgesMultiplicity, "Multiplicity"}; if (doprocessMixedEventHV0sInBuffer || doprocessMixedEventHCascadesInBuffer) { validCollisions.resize(histos.get(HIST("axes/hMultAxis"))->GetNbinsX() * histos.get(HIST("axes/hVertexZAxis"))->GetNbinsX()); - for (std::vector& inner_vec : validCollisions) { - inner_vec.reserve(masterConfigurations.mixingParameter); + for (size_t i = 0; i < validCollisions.size(); ++i) { + validCollisions[i].reserve(masterConfigurations.mixingParameter); } } if (!masterConfigurations.doPPAnalysis) { @@ -1882,7 +1908,31 @@ struct HStrangeCorrelation { histos.add("GeneratedWithPV/hAntiLambdaFromXiZero", "", kTH2F, {axesConfigurations.axisPtQA, axesConfigurations.axisEta}); histos.add("GeneratedWithPV/hAntiLambdaFromXiPlus", "", kTH2F, {axesConfigurations.axisPtQA, axesConfigurations.axisEta}); } - + if (doprocessPrediction) { + mCounter.mPdgDatabase = pdgDB.service; + mCounter.mSelectPrimaries = doAssocPhysicalPrimary.value; + histos.add("Prediction/hEventSelection", "hEventSelection", kTH1F, {{3, 0, 3}}); + TString eventSelLabel[] = {"Read", "INELgt0", "|Z|<10"}; + for (int i = 1; i <= histos.get(HIST("Prediction/hEventSelection"))->GetNbinsX(); i++) { + histos.get(HIST("Prediction/hEventSelection"))->GetXaxis()->SetBinLabel(i, eventSelLabel[i - 1]); + } + histos.add("Prediction/hTrigger", "Trigger Tracks", kTH3F, {axesConfigurations.axisPtQA, axesConfigurations.axisEta, axesConfigurations.axisPhi}); + for (int i = 0; i < AssocParticleTypes; i++) { + if (TESTBIT(doCorrelation, i)) + histos.add(fmt::format("Prediction/h{}", Particlenames[i]).c_str(), "", kTH3F, {axesConfigurations.axisPtQA, axesConfigurations.axisEta, axesConfigurations.axisPhi}); + if (masterConfigurations.useCentralityinPrediction) { + if (TESTBIT(doCorrelation, i)) + histos.add(fmt::format("Prediction/sameEvent/{}", Particlenames[i]).c_str(), "", kTHnF, {axisDeltaPhiNDim, axisDeltaEtaNDim, axisPtAssocNDim, axisPtTriggerNDim, axisVtxZNDim, axisMultNDim}); + } else { + if (TESTBIT(doCorrelation, i)) + histos.add(fmt::format("Prediction/sameEvent/{}", Particlenames[i]).c_str(), "", kTHnF, {axisDeltaPhiNDim, axisDeltaEtaNDim, axisPtAssocNDim, axisPtTriggerNDim, axisVtxZNDim, axisMultiplicityNDim}); + } + } + if (masterConfigurations.doSeparateFT0Prediction) { + histos.addClone("Prediction/sameEvent/", "Prediction/sameEventFT0A/"); + histos.addClone("Prediction/sameEvent/", "Prediction/sameEventFT0C/"); + } + } // visual inspection of sizes histos.print(); @@ -3226,6 +3276,156 @@ struct HStrangeCorrelation { if (masterConfigurations.doFullCorrelationStudy) fillCorrelationsCascade(triggerTracks, associatedCascades, true, true, collision.posX(), collision.posY(), collision.posZ(), cent, bField); } + void processPrediction(soa::Join::iterator const& mcCollision, aod::McParticles const& mcParticles) + { + std::vector triggerIndices; + std::vector> associatedIndices; + std::vector assocHadronIndices; + std::vector piIndices; + std::vector k0ShortIndices; + std::vector lambdaIndices; + std::vector antiLambdaIndices; + std::vector xiMinusIndices; + std::vector xiPlusIndices; + std::vector omegaMinusIndices; + std::vector omegaPlusIndices; + float centMultFT0M = -1; + float centMultFT0A = -1; + float centMultFT0C = -1; + float multFT0M = -1; + float multFT0A = -1; + float multFT0C = -1; + histos.fill(HIST("Prediction/hEventSelection"), 0.5); + if (masterConfigurations.selectINELgtZERO && !o2::pwglf::isINELgt0mc(mcParticles, pdgDB)) { + return; + } + histos.fill(HIST("Prediction/hEventSelection"), 1.5); + if (std::abs(mcCollision.posZ()) > masterConfigurations.zVertexCut) { + return; + } + histos.fill(HIST("Prediction/hEventSelection"), 2.5); + if (masterConfigurations.useCentralityinPrediction) { + centMultFT0M = mcCollision.centFT0M(); + centMultFT0A = mcCollision.centFT0A(); + centMultFT0C = mcCollision.centFT0C(); + } else { + multFT0M = mCounter.countFT0A(mcParticles) + mCounter.countFT0C(mcParticles); + multFT0A = mCounter.countFT0A(mcParticles); + multFT0C = mCounter.countFT0C(mcParticles); + } + int iteratorNum = -1; + for (auto const& mcParticle : mcParticles) { + iteratorNum = iteratorNum + 1; + double geta = mcParticle.eta(); + double gpt = mcParticle.pt(); + double gphi = mcParticle.phi(); + if (std::abs(geta) > etaSel) { + continue; + } + if (std::abs(mcParticle.pdgCode()) == PDG_t::kPiPlus || std::abs(mcParticle.pdgCode()) == PDG_t::kKPlus || std::abs(mcParticle.pdgCode()) == PDG_t::kProton || std::abs(mcParticle.pdgCode()) == PDG_t::kElectron || std::abs(mcParticle.pdgCode()) == PDG_t::kMuonMinus) { + if (!masterConfigurations.doTriggPhysicalPrimary || mcParticle.isPhysicalPrimary()) { + triggerIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hTrigger"), gpt, geta, gphi); + } + if (masterConfigurations.doCorrelationHadron) { + if (!doAssocPhysicalPrimary || mcParticle.isPhysicalPrimary()) { + assocHadronIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hHadron"), gpt, geta, gphi); + } + } + } + if (!doAssocPhysicalPrimary || mcParticle.isPhysicalPrimary()) { + if (std::abs(mcParticle.pdgCode()) == PDG_t::kPiPlus && masterConfigurations.doCorrelationPion) { + piIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hPion"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kK0Short && masterConfigurations.doCorrelationK0Short) { + k0ShortIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hK0Short"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kLambda0 && masterConfigurations.doCorrelationLambda) { + lambdaIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hLambda"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kLambda0Bar && masterConfigurations.doCorrelationAntiLambda) { + antiLambdaIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hAntiLambda"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kXiMinus && masterConfigurations.doCorrelationXiMinus) { + xiMinusIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hXiMinus"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kXiPlusBar && masterConfigurations.doCorrelationXiPlus) { + xiPlusIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hXiPlus"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kOmegaMinus && masterConfigurations.doCorrelationOmegaMinus) { + omegaMinusIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hOmegaMinus"), gpt, geta, gphi); + } + if (mcParticle.pdgCode() == PDG_t::kOmegaPlusBar && masterConfigurations.doCorrelationOmegaPlus) { + omegaPlusIndices.emplace_back(iteratorNum); + histos.fill(HIST("Prediction/hOmegaPlus"), gpt, geta, gphi); + } + } + } + + associatedIndices.emplace_back(k0ShortIndices); + associatedIndices.emplace_back(lambdaIndices); + associatedIndices.emplace_back(antiLambdaIndices); + associatedIndices.emplace_back(xiMinusIndices); + associatedIndices.emplace_back(xiPlusIndices); + associatedIndices.emplace_back(omegaMinusIndices); + associatedIndices.emplace_back(omegaPlusIndices); + associatedIndices.emplace_back(piIndices); + associatedIndices.emplace_back(assocHadronIndices); + for (std::size_t iTrigger = 0; iTrigger < triggerIndices.size(); iTrigger++) { + auto triggerParticle = mcParticles.iteratorAt(triggerIndices[iTrigger]); + // check range of trigger particle + if (triggerParticle.pt() > axisRanges[3][1] || triggerParticle.pt() < axisRanges[3][0]) { + continue; + } + double getatrigger = triggerParticle.eta(); + double gphitrigger = triggerParticle.phi(); + double pttrigger = triggerParticle.pt(); + auto const& mother = triggerParticle.mothers_first_as(); + auto globalIndex = mother.globalIndex(); + static_for<0, 8>([&](auto i) { // associated loop + constexpr int Index = i.value; + for (std::size_t iassoc = 0; iassoc < associatedIndices[Index].size(); iassoc++) { + auto assocParticle = mcParticles.iteratorAt(associatedIndices[Index][iassoc]); + if (triggerIndices[iTrigger] != associatedIndices[Index][iassoc] && globalIndex != assocParticle.globalIndex()) { // avoid self + double getaassoc = assocParticle.eta(); + double gphiassoc = assocParticle.phi(); + double ptassoc = assocParticle.pt(); + double deltaphi = computeDeltaPhi(gphitrigger, gphiassoc); + double deltaeta = getatrigger - getaassoc; + + // skip if basic ranges not met + if (deltaphi < axisRanges[0][0] || deltaphi > axisRanges[0][1]) + continue; + if (deltaeta < axisRanges[1][0] || deltaeta > axisRanges[1][1]) + continue; + if (ptassoc < axisRanges[2][0] || ptassoc > axisRanges[2][1]) + continue; + if (TESTBIT(doCorrelation, i)) { + if (masterConfigurations.useCentralityinPrediction) { + histos.fill(HIST("Prediction/sameEvent/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), centMultFT0M); + if (masterConfigurations.doSeparateFT0Prediction) + histos.fill(HIST("Prediction/sameEventFT0A/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), centMultFT0A); + histos.fill(HIST("Prediction/sameEventFT0C/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), centMultFT0C); + } else { + histos.fill(HIST("Prediction/sameEvent/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), multFT0M); + if (masterConfigurations.doSeparateFT0Prediction) + histos.fill(HIST("Prediction/sameEventFT0A/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), multFT0A); + histos.fill(HIST("Prediction/sameEventFT0C/") + HIST(Particlenames[Index]), computeDeltaPhi(gphitrigger, gphiassoc), deltaeta, ptassoc, pttrigger, mcCollision.posZ(), multFT0C); + } + } + } + } + }); + } + } PROCESS_SWITCH(HStrangeCorrelation, processSelectEventWithTrigger, "Select events with trigger only", true); PROCESS_SWITCH(HStrangeCorrelation, processSameEventHV0s, "Process same events, h-V0s", true); PROCESS_SWITCH(HStrangeCorrelation, processSameEventHCascades, "Process same events, h-Cascades", true); @@ -3242,6 +3442,7 @@ struct HStrangeCorrelation { PROCESS_SWITCH(HStrangeCorrelation, processMCGenerated, "Process MC generated", false); PROCESS_SWITCH(HStrangeCorrelation, processClosureTest, "Process Closure Test", false); PROCESS_SWITCH(HStrangeCorrelation, processFeedDown, "process Feed Down", false); + PROCESS_SWITCH(HStrangeCorrelation, processPrediction, "process model prediction", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)