Skip to content

[PWGHF] Adding : ML and skimming in correlator + other minor change #18125

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

Draft
VDiBella wants to merge 1 commit into
base: master
Choose a base branch
from
+285 −52
Draft

[PWGHF] Adding : ML and skimming in correlator + other minor change #18125

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
50 changes: 26 additions & 24 deletions PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
#include <Framework/ASoA.h>
#include <Framework/AnalysisDataModel.h>

#include <cstdint>

namespace o2::aod
{
DECLARE_SOA_TABLE(HfCandDpFullEvs, "AOD", "HFCANDDPFULLEV",
Expand All @@ -38,6 +40,8 @@ DECLARE_SOA_TABLE(HfCandDpMcEvs, "AOD", "HFCANDDPMCEV",

namespace full
{
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpFullEv, hfCandDpFullEv, int32_t, HfCandDpFullEvs, "");
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpMcEv, hfCandDpFullEvMC, int32_t, HfCandDpMcEvs, "");
DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float); //! Radius of secondary vertex (cm)
DECLARE_SOA_COLUMN(PtProng0, ptProng0, float); //! Transverse momentum of prong0 (GeV/c)
DECLARE_SOA_COLUMN(PProng0, pProng0, float); //! Momentum of prong0 (GeV/c)
Expand All @@ -57,8 +61,6 @@ DECLARE_SOA_COLUMN(Eta, eta, float);
DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuth angle of candidate
DECLARE_SOA_COLUMN(E, e, float); //! Energy of candidate (GeV)
DECLARE_SOA_COLUMN(Centrality, centrality, float); //! Collision centrality
DECLARE_SOA_INDEX_COLUMN(HfCandDpMcEv, hfCandDpMcEv); //! The Mc collision index this MC particles belongs to
DECLARE_SOA_INDEX_COLUMN(HfCandDpFullEv, hfCandDpFullEv); //! The collision index this candidate belongs to
DECLARE_SOA_COLUMN(NSigTpcPi0, nSigTpcPi0, float); //! TPC Nsigma separation for prong0 with pion mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcKa0, nSigTpcKa0, float); //! TPC Nsigma separation for prong0 with kaon mass hypothesis
DECLARE_SOA_COLUMN(NSigTofPi0, nSigTofPi0, float); //! TOF Nsigma separation for prong0 with pion mass hypothesis
Expand All @@ -71,12 +73,12 @@ DECLARE_SOA_COLUMN(NSigTpcPi2, nSigTpcPi2, float);
DECLARE_SOA_COLUMN(NSigTpcKa2, nSigTpcKa2, float); //! TPC Nsigma separation for prong2 with kaon mass hypothesis
DECLARE_SOA_COLUMN(NSigTofPi2, nSigTofPi2, float); //! TOF Nsigma separation for prong2 with pion mass hypothesis
DECLARE_SOA_COLUMN(NSigTofKa2, nSigTofKa2, float); //! TOF Nsigma separation for prong2 with kaon mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofPi0, nSigTpcTofPi0, float); //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofKa0, nSigTpcTofKa0, float); //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofPi1, nSigTpcTofPi1, float); //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofKa1, nSigTpcTofKa1, float); //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofPi2, nSigTpcTofPi2, float); //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
DECLARE_SOA_COLUMN(NSigTpcTofKa2, nSigTpcTofKa2, float); //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaPi0, tpcTofNSigmaPi0, float); //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaKa0, tpcTofNSigmaKa0, float); //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaPi1, tpcTofNSigmaPi1, float); //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaKa1, tpcTofNSigmaKa1, float); //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaPi2, tpcTofNSigmaPi2, float); //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
DECLARE_SOA_COLUMN(TPCTofNSigmaKa2, tpcTofNSigmaKa2, float); //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
DECLARE_SOA_COLUMN(DecayLength, decayLength, float); //! Decay length of candidate (cm)
DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float); //! Transverse decay length of candidate (cm)
DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float); //! Normalised decay length of candidate
Expand All @@ -96,7 +98,7 @@ DECLARE_SOA_TABLE(HfCandDpMls, "AOD", "HFCANDDPML",
full::MlScore0,
full::MlScore1)

DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY", o2::soa::Index<>,
full::CandidateSelFlag,
full::M,
full::Pt,
Expand All @@ -107,7 +109,7 @@ DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
hf_cand_mc_flag::OriginMcRec,
hf_cand_mc_flag::FlagMcDecayChanRec)

DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE", o2::soa::Index<>,
hf_cand::Chi2PCA,
full::DecayLength,
full::DecayLengthXY,
Expand All @@ -126,20 +128,20 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
full::NSigTpcKa0,
full::NSigTofPi0,
full::NSigTofKa0,
full::NSigTpcTofPi0,
full::NSigTpcTofKa0,
full::TPCTofNSigmaPi0,
full::TPCTofNSigmaKa0,
full::NSigTpcPi1,
full::NSigTpcKa1,
full::NSigTofPi1,
full::NSigTofKa1,
full::NSigTpcTofPi1,
full::NSigTpcTofKa1,
full::TPCTofNSigmaPi1,
full::TPCTofNSigmaKa1,
full::NSigTpcPi2,
full::NSigTpcKa2,
full::NSigTofPi2,
full::NSigTofKa2,
full::NSigTpcTofPi2,
full::NSigTpcTofKa2,
full::TPCTofNSigmaPi2,
full::TPCTofNSigmaKa2,
full::CandidateSelFlag,
full::M,
full::Pt,
Expand All @@ -155,7 +157,7 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
hf_cand_mc_flag::OriginMcRec,
hf_cand_mc_flag::FlagMcDecayChanRec)

DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL", o2::soa::Index<>,
hf_cand::XSecondaryVertex,
hf_cand::YSecondaryVertex,
hf_cand::ZSecondaryVertex,
Expand Down Expand Up @@ -201,20 +203,20 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
full::NSigTpcKa0,
full::NSigTofPi0,
full::NSigTofKa0,
full::NSigTpcTofPi0,
full::NSigTpcTofKa0,
full::TPCTofNSigmaPi0,
full::TPCTofNSigmaKa0,
full::NSigTpcPi1,
full::NSigTpcKa1,
full::NSigTofPi1,
full::NSigTofKa1,
full::NSigTpcTofPi1,
full::NSigTpcTofKa1,
full::TPCTofNSigmaPi1,
full::TPCTofNSigmaKa1,
full::NSigTpcPi2,
full::NSigTpcKa2,
full::NSigTofPi2,
full::NSigTofKa2,
full::NSigTpcTofPi2,
full::NSigTpcTofKa2,
full::TPCTofNSigmaPi2,
full::TPCTofNSigmaKa2,
full::CandidateSelFlag,
full::M,
full::Pt,
Expand All @@ -233,7 +235,7 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
hf_cand_mc_flag::OriginMcRec,
hf_cand_mc_flag::FlagMcDecayChanRec);

DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP",
DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP", o2::soa::Index<>,
full::Pt,
full::Eta,
full::Phi,
Expand Down
2 changes: 1 addition & 1 deletion PWGHF/HFC/TableProducer/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ o2physics_add_dpl_workflow(correlator-dplus-dminus

o2physics_add_dpl_workflow(correlator-dplus-dplus-reduced
SOURCES correlatorDplusDplusReduced.cxx
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils O2Physics::MLCore
COMPONENT_NAME Analysis)

o2physics_add_dpl_workflow(correlator-dplus-hadrons
Expand Down
130 changes: 124 additions & 6 deletions PWGHF/HFC/TableProducer/correlatorDplusDplusReduced.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
#include "PWGHF/Core/CentralityEstimation.h"
#include "PWGHF/Core/DecayChannels.h"
#include "PWGHF/Core/HfHelper.h"
#include "PWGHF/Core/HfMlResponseDplusToPiKPi.h"
#include "PWGHF/Core/SelectorCuts.h"
#include "PWGHF/DataModel/CandidateReconstructionTables.h"
#include "PWGHF/DataModel/CandidateSelectionTables.h"
#include "PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h"
Expand All @@ -30,20 +32,25 @@
#include "Common/DataModel/Centrality.h"

#include <CCDB/BasicCCDBManager.h>
#include <CCDB/CcdbApi.h>
#include <Framework/ASoA.h>
#include <Framework/AnalysisDataModel.h>
#include <Framework/AnalysisHelpers.h>
#include <Framework/AnalysisTask.h>
#include <Framework/Array2D.h>
#include <Framework/Configurable.h>
#include <Framework/Expressions.h>
#include <Framework/HistogramRegistry.h>
#include <Framework/InitContext.h>
#include <Framework/runDataProcessing.h>

#include <cstdint>
#include <cstdlib>
#include <string>
#include <vector>

using namespace o2;
using namespace o2::analysis;
using namespace o2::framework;
using namespace o2::framework::expressions;
using namespace o2::hf_centrality;
Expand All @@ -69,6 +76,34 @@ struct HfCorrelatorDplusDplusReduced {
Configurable<bool> cfgSkimmedProcessing{"cfgSkimmedProcessing", true, "Enables processing of skimmed datasets"};
Configurable<bool> skipSingleD{"skipSingleD", true, "Skip collisions with one or less D candidates"};

Configurable<bool> applyMl{"applyMl", false, "Flag to apply ML selections"};
Configurable<bool> applySkimming{"applySkimming", false, "Flag to apply Skimming selections"};
Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
Configurable<std::vector<double>> binsPtMl{"binsPtMl", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
Configurable<LabeledArray<double>> cutsMl{"cutsMl", {hf_cuts_ml::Cuts[0], hf_cuts_ml::NBinsPt, hf_cuts_ml::NCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsCutScore}, "ML selections per pT bin"};
Configurable<int> nClassesMl{"nClassesMl", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model"};
Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{"EventFiltering/PWGHF/BDTDPlus"}, "Paths of models on CCDB"};
Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_DPlusToKPiPi.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};

Configurable<std::vector<double>> cutPtSkimming{"cutPtSkimming", {1, 5, 1000}, "pT bin limits for Skimming application"};
Configurable<std::vector<double>> minM{"minM", {0.7, 0.7}, "Mass minimal for the cut for each pt bin"};
Configurable<std::vector<double>> maxM{"maxM", {2.0, 2.1}, "Mass maximal for the cut for each pt bin"};
Configurable<std::vector<double>> minCosTheta{"minCosTheta", {0.96, 0.98}, "CosTheta minimal for the cut for each pt bin"};
Configurable<std::vector<double>> minDecayLength{"minDecayLength", {0.02, 0.03}, "DecayLength minimal for the cut for each pt bin"};
Configurable<std::vector<double>> maxNsigmaTPC{"maxNsigmaTPC", {3, 3}, "NsigmaTPC maximal for the cut for each pt bin"};
Configurable<std::vector<double>> maxNsigmaTOF{"maxNsigmaTOF", {3, 3}, "NsigmaTOF maximal for the cut for each pt bin"};

Configurable<std::vector<double>> binsPtSkimming{"binsPtSkimming", {0}, "pT bin limits for Skimming application"};

HfMlResponseDplusToPiKPi<float> hfMlResponse;

std::vector<float> outputML;
o2::ccdb::CcdbApi ccdbApi;

HfHelper hfHelper;

Service<o2::ccdb::BasicCCDBManager> ccdb;
Expand Down Expand Up @@ -102,6 +137,46 @@ struct HfCorrelatorDplusDplusReduced {
if (cfgSkimmedProcessing) {
zorroSummary.setObject(zorro.getZorroSummary());
}

if (applyMl) {
hfMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
if (loadModelsFromCCDB) {
ccdbApi.init(ccdbUrl);
hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
} else {
hfMlResponse.setModelPathsLocal(onnxFileNames);
}
hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
hfMlResponse.init();
}
}

bool Skimming(auto candidate,
auto PtcutSkimming,
auto Mmin,
auto Mmax,
auto CosThetamin,
auto DecayLengthmin,
auto NsigmaTPCmax,
auto NsigmaTOFmax)
{
if (candidate.pt() < PtcutSkimming[0] || candidate.pt() > PtcutSkimming[PtcutSkimming.size() - 1]) {
return false;
}
for (auto i = 1; i < PtcutSkimming.size(); i++) {
if (candidate.pt() <= PtcutSkimming[i]) {
if (hfHelper.invMassDplusToPiKPi(candidate) < Mmin[i - 1] ||
hfHelper.invMassDplusToPiKPi(candidate) > Mmax[i - 1] ||
candidate.cpa() < CosThetamin[i - 1] ||
candidate.decayLength() < DecayLengthmin[i - 1] ||
candidate.nSigTofKa1() > NsigmaTOFmax[i - 1] ||
candidate.nSigTpcKa1() > NsigmaTPCmax[i - 1]) {
return false;
}
return true;
}
}
return false;
}

template <typename T>
Expand All @@ -127,14 +202,14 @@ struct HfCorrelatorDplusDplusReduced {
channelMc = candidate.flagMcDecayChanRec();
}

std::vector<float> outputMl = {-999., -999.};
std::vector<float> outML = {-999., -999.};
if constexpr (DoMl) {
for (unsigned int iclass = 0; iclass < classMlIndexes->size(); iclass++) {
outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
outML[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
}
rowCandidateMl(
outputMl[0],
outputMl[1]);
outML[0],
outML[1]);
}

float cent{-1.};
Expand Down Expand Up @@ -291,7 +366,9 @@ struct HfCorrelatorDplusDplusReduced {
static int lastRunNumber = -1;
// reserve memory
rowCandidateFullEvents.reserve(collisions.size());
if (fillCandidateLiteTable) {
if (fillCandidateTinyTable) {
rowCandidateTiny.reserve(candidates.size());
} else if (fillCandidateLiteTable) {
rowCandidateLite.reserve(candidates.size());
} else {
rowCandidateFull.reserve(candidates.size());
Expand Down Expand Up @@ -320,6 +397,26 @@ struct HfCorrelatorDplusDplusReduced {
for (const auto& candidate : candidatesInThisCollision) {
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
auto candidateSign = -prongCandidate.sign();

if (applySkimming &&
!Skimming(candidate,
cutPtSkimming,
minM,
maxM,
minCosTheta,
minDecayLength,
maxNsigmaTPC,
maxNsigmaTOF)) {
continue;
}

if (applyMl) {
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outputML);
if (!isSelectedMl) {
continue;
}
}
fillCandidateTable<aod::Collisions>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
}
}
Expand All @@ -332,7 +429,9 @@ struct HfCorrelatorDplusDplusReduced {
{
// reserve memory
rowCandidateFullEvents.reserve(collisions.size());
if (fillCandidateLiteTable) {
if (fillCandidateTinyTable) {
rowCandidateTiny.reserve(candidates.size());
} else if (fillCandidateLiteTable) {
rowCandidateLite.reserve(candidates.size());
} else {
rowCandidateFull.reserve(candidates.size());
Expand All @@ -348,6 +447,25 @@ struct HfCorrelatorDplusDplusReduced {
for (const auto& candidate : candidatesInThisCollision) {
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
auto candidateSign = -prongCandidate.sign();

if (applySkimming &&
!Skimming(candidate,
cutPtSkimming,
minM,
maxM,
minCosTheta,
minDecayLength,
maxNsigmaTPC,
maxNsigmaTOF)) {
continue;
}
if (applyMl) {
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outputML);
if (!isSelectedMl) {
continue;
}
}
fillCandidateTable<aod::Collisions, true>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
}
}
Expand Down
Loading
Loading