From a61d0432c2a8b7547755a83315d69666c7af7b8a Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Thu, 21 Aug 2025 20:15:49 +0200 Subject: [PATCH 1/3] PWGEM/Dilepton: 1st version of PID ML --- PWGEM/Dilepton/Core/DielectronCut.h | 57 ++++++---- PWGEM/Dilepton/Core/Dilepton.h | 78 +++++++------ PWGEM/Dilepton/Core/DileptonHadronMPC.h | 104 ++++++++++-------- PWGEM/Dilepton/Core/DileptonMC.h | 86 ++++++++------- PWGEM/Dilepton/Core/PhotonHBT.h | 72 +++++++----- PWGEM/Dilepton/Core/SingleTrackQC.h | 71 +++++++----- PWGEM/Dilepton/Core/SingleTrackQCMC.h | 62 ++++++----- .../TableProducer/skimmerPrimaryElectron.cxx | 89 +++++++++++---- .../skimmerPrimaryElectronQC.cxx | 39 ++----- PWGEM/Dilepton/Tasks/prefilterDielectron.cxx | 56 ++++++---- PWGEM/Dilepton/Utils/MlResponseO2Track.h | 10 +- 11 files changed, 423 insertions(+), 301 deletions(-) diff --git a/PWGEM/Dilepton/Core/DielectronCut.h b/PWGEM/Dilepton/Core/DielectronCut.h index 0dc108b05c5..4c0a6e9d86e 100644 --- a/PWGEM/Dilepton/Core/DielectronCut.h +++ b/PWGEM/Dilepton/Core/DielectronCut.h @@ -162,8 +162,8 @@ class DielectronCut : public TNamed return true; } - template - bool IsSelectedTrack(TTrack const& track, TCollision const& collision = 0) const + template + bool IsSelectedTrack(TTrack const& track) const { if (!track.hasITS()) { return false; @@ -252,36 +252,31 @@ class DielectronCut : public TNamed } // PID cuts - if constexpr (isML) { - if (!PassPIDML(track, collision)) { + if (track.hasITS() && !track.hasTPC() && !track.hasTRD() && !track.hasTOF()) { // ITSsa + float meanClusterSizeITS = track.meanClusterSizeITS() * std::cos(std::atan(track.tgl())); + if (meanClusterSizeITS < mMinMeanClusterSizeITS || mMaxMeanClusterSizeITS < meanClusterSizeITS) { return false; } - } else { - if (track.hasITS() && !track.hasTPC() && !track.hasTRD() && !track.hasTOF()) { // ITSsa - float meanClusterSizeITS = track.meanClusterSizeITS() * std::cos(std::atan(track.tgl())); - if (meanClusterSizeITS < mMinMeanClusterSizeITS || mMaxMeanClusterSizeITS < meanClusterSizeITS) { - return false; - } - } else { // not ITSsa - if (!PassPID(track)) { - return false; - } + } else { // not ITSsa + if (!PassPID(track)) { + return false; } } return true; } - template - bool PassPIDML(TTrack const&, TCollision const&) const + template + bool PassPIDML(TTrack const& track) const { - return false; - /*if (!PassTOFif(track)) { // Allows for pre-selection. But potentially dangerous if analyzers are not aware of it - return false; - }*/ - // std::vector inputFeatures = mPIDMlResponse->getInputFeatures(track, collision); - // float binningFeature = mPIDMlResponse->getBinningFeature(track, collision); - // return mPIDMlResponse->isSelectedMl(inputFeatures, binningFeature); + int pbin = lower_bound(mMLBins.begin(), mMLBins.end(), track.tpcInnerParam()) - mMLBins.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(mMLBins.size()) - 2 < pbin) { + pbin = static_cast(mMLBins.size()) - 2; + } + // LOGF(info, "track.tpcInnerParam() = %f, pbin = %d, track.probElBDT() = %f, mMLCuts[pbin] = %f", track.tpcInnerParam(), pbin, track.probElBDT(), mMLCuts[pbin]); + return track.probElBDT() > mMLCuts[pbin]; } template @@ -307,7 +302,7 @@ class DielectronCut : public TNamed return PassTOFif(track); case static_cast(PIDSchemes::kPIDML): - return true; // don't use kPIDML here. + return PassPIDML(track); case static_cast(PIDSchemes::kTPChadrejORTOFreq_woTOFif): return PassTPConlyhadrej(track) || PassTOFreq(track); @@ -517,6 +512,18 @@ class DielectronCut : public TNamed mPIDMlResponse = mlResponse; } + void SetMLThresholds(const std::vector bins, const std::vector cuts) + { + if (bins.size() != cuts.size() + 1) { + LOG(fatal) << "cuts.size() + 1 mutst be exactly the same as bins.size(). Check your bins and thresholds."; + } + mMLBins = bins; + mMLCuts = cuts; + for (int i = 0; i < static_cast(mMLBins.size()) - 1; i++) { + printf("Dielectron cut: mMLBins[%d] = %3.2f, mMLBins[%d] = %3.2f, mMLCuts[%d] = %3.2f\n", i, mMLBins[i], i + 1, mMLBins[i + 1], i, mMLCuts[i]); + } + } + // Getters bool IsPhotonConversionSelected() const { return mSelectPC; } @@ -597,6 +604,8 @@ class DielectronCut : public TNamed // float mMinP_ITSNsigmaPr{0.0}, mMaxP_ITSNsigmaPr{0.0}; o2::analysis::MlResponseDielectronSingleTrack* mPIDMlResponse{nullptr}; + std::vector mMLBins{}; // binning for a feature variable. e.g. tpcInnerParam + std::vector mMLCuts{}; // threshold for each bin. mMLCuts.size() must be mMLBins.size()-1. ClassDef(DielectronCut, 1); }; diff --git a/PWGEM/Dilepton/Core/Dilepton.h b/PWGEM/Dilepton/Core/Dilepton.h index c928d1cc71f..fc06f13f92b 100644 --- a/PWGEM/Dilepton/Core/Dilepton.h +++ b/PWGEM/Dilepton/Core/Dilepton.h @@ -713,30 +713,42 @@ struct Dilepton { // fDielectronCut.SetPRangeForITSNsigmaPr(dielectroncuts.cfg_min_p_ITSNsigmaPr, dielectroncuts.cfg_max_p_ITSNsigmaPr); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; - } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); - } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); + } + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); + } + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -836,11 +848,11 @@ struct Dilepton { if constexpr (ev_id == 0) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -1377,16 +1389,16 @@ struct Dilepton { } // end of DF - template - bool isPairOK(TCollision const& collision, TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) + template + bool isPairOK(TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -1470,17 +1482,17 @@ struct Dilepton { auto negTracks_per_coll = negTracks.sliceByCached(perCollision, collision.globalIndex(), cache); for (const auto& [pos, neg] : combinations(CombinationsFullIndexPolicy(posTracks_per_coll, negTracks_per_coll))) { // ULS - if (isPairOK(collision, pos, neg, cut, tracks)) { + if (isPairOK(pos, neg, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos.globalIndex(), neg.globalIndex())); } } for (const auto& [pos1, pos2] : combinations(CombinationsStrictlyUpperIndexPolicy(posTracks_per_coll, posTracks_per_coll))) { // LS++ - if (isPairOK(collision, pos1, pos2, cut, tracks)) { + if (isPairOK(pos1, pos2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos1.globalIndex(), pos2.globalIndex())); } } for (const auto& [neg1, neg2] : combinations(CombinationsStrictlyUpperIndexPolicy(negTracks_per_coll, negTracks_per_coll))) { // LS-- - if (isPairOK(collision, neg1, neg2, cut, tracks)) { + if (isPairOK(neg1, neg2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(neg1.globalIndex(), neg2.globalIndex())); } } diff --git a/PWGEM/Dilepton/Core/DileptonHadronMPC.h b/PWGEM/Dilepton/Core/DileptonHadronMPC.h index 11fedd03212..1e0e59dfd2c 100644 --- a/PWGEM/Dilepton/Core/DileptonHadronMPC.h +++ b/PWGEM/Dilepton/Core/DileptonHadronMPC.h @@ -611,30 +611,42 @@ struct DileptonHadronMPC { fDielectronCut.SetPinRangeForPionRejectionTPC(dielectroncuts.cfg_min_pin_pirejTPC, dielectroncuts.cfg_max_pin_pirejTPC); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -701,11 +713,11 @@ struct DileptonHadronMPC { if constexpr (ev_id == 0) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -831,8 +843,8 @@ struct DileptonHadronMPC { return true; } - template - bool fillDileptonHadron(TCollision const& collision, TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks, TRefTrack const& t3) + template + bool fillDileptonHadron(TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks, TRefTrack const& t3) { // this function must be called, if dilepton passes the cut. if constexpr (ev_id == 1) { @@ -860,11 +872,11 @@ struct DileptonHadronMPC { if constexpr (ev_id == 0) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -951,8 +963,8 @@ struct DileptonHadronMPC { return true; } - template - bool fillHadronHadron(TCollision const& collision, TRefTrack const& t1, TRefTrack const& t2, TLeptons const& posLeptons, TLeptons const& negLeptons, TLeptonCut const& cut) + template + bool fillHadronHadron(TRefTrack const& t1, TRefTrack const& t2, TLeptons const& posLeptons, TLeptons const& negLeptons, TLeptonCut const& cut) { if constexpr (ev_id == 0) { if (!fEMTrackCut.IsSelected(t1) || !fEMTrackCut.IsSelected(t2)) { // for charged track @@ -963,11 +975,11 @@ struct DileptonHadronMPC { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { for (const auto& pos : posLeptons) { // leptons per collision if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(pos, collision)) { + if (!cut.template IsSelectedTrack(pos)) { continue; } } else { // cut based - if (!cut.template IsSelectedTrack(pos)) { + if (!cut.template IsSelectedTrack(pos)) { continue; } } @@ -978,11 +990,11 @@ struct DileptonHadronMPC { for (const auto& neg : negLeptons) { // leptons per collision if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(neg, collision)) { + if (!cut.template IsSelectedTrack(neg)) { continue; } } else { // cut based - if (!cut.template IsSelectedTrack(neg)) { + if (!cut.template IsSelectedTrack(neg)) { continue; } } @@ -1105,7 +1117,7 @@ struct DileptonHadronMPC { if (is_pair_ok) { nuls++; for (const auto& refTrack : refTracks_per_coll) { - fillDileptonHadron<0>(collision, pos, neg, cut, tracks, refTrack); + fillDileptonHadron<0>(pos, neg, cut, tracks, refTrack); } } } @@ -1114,7 +1126,7 @@ struct DileptonHadronMPC { if (is_pair_ok) { nlspp++; for (const auto& refTrack : refTracks_per_coll) { - fillDileptonHadron<0>(collision, pos1, pos2, cut, tracks, refTrack); + fillDileptonHadron<0>(pos1, pos2, cut, tracks, refTrack); } } } @@ -1123,7 +1135,7 @@ struct DileptonHadronMPC { if (is_pair_ok) { nlsmm++; for (const auto& refTrack : refTracks_per_coll) { - fillDileptonHadron<0>(collision, neg1, neg2, cut, tracks, refTrack); + fillDileptonHadron<0>(neg1, neg2, cut, tracks, refTrack); } } } @@ -1144,7 +1156,7 @@ struct DileptonHadronMPC { } } for (const auto& [ref1, ref2] : combinations(CombinationsStrictlyUpperIndexPolicy(refTracks_per_coll, refTracks_per_coll))) { - fillHadronHadron<0>(collision, ref1, ref2, posTracks_per_coll, negTracks_per_coll, cut); + fillHadronHadron<0>(ref1, ref2, posTracks_per_coll, negTracks_per_coll, cut); } } @@ -1262,7 +1274,7 @@ struct DileptonHadronMPC { for (const auto& ref1 : selected_refTracks_in_this_event) { // ref-ref mix for (const auto& ref2 : refTracks_from_event_pool) { // LOGF(info, "ref1.pt() = %f, ref2.pt() = %f", ref1.pt(), ref2.pt()); - fillHadronHadron<1>(collision, ref1, ref2, nullptr, nullptr, nullptr); + fillHadronHadron<1>(ref1, ref2, nullptr, nullptr, nullptr); } } } // end of loop over mixed event pool for hadron-hadron @@ -1279,16 +1291,16 @@ struct DileptonHadronMPC { } // end of DF - template - bool isPairOK(TCollision const& collision, TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) + template + bool isPairOK(TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -1348,17 +1360,17 @@ struct DileptonHadronMPC { auto negTracks_per_coll = negTracks.sliceByCached(perCollision, collision.globalIndex(), cache); for (const auto& [pos, neg] : combinations(CombinationsFullIndexPolicy(posTracks_per_coll, negTracks_per_coll))) { // ULS - if (isPairOK(collision, pos, neg, cut, tracks)) { + if (isPairOK(pos, neg, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos.globalIndex(), neg.globalIndex())); } } for (const auto& [pos1, pos2] : combinations(CombinationsStrictlyUpperIndexPolicy(posTracks_per_coll, posTracks_per_coll))) { // LS++ - if (isPairOK(collision, pos1, pos2, cut, tracks)) { + if (isPairOK(pos1, pos2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos1.globalIndex(), pos2.globalIndex())); } } for (const auto& [neg1, neg2] : combinations(CombinationsStrictlyUpperIndexPolicy(negTracks_per_coll, negTracks_per_coll))) { // LS-- - if (isPairOK(collision, neg1, neg2, cut, tracks)) { + if (isPairOK(neg1, neg2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(neg1.globalIndex(), neg2.globalIndex())); } } diff --git a/PWGEM/Dilepton/Core/DileptonMC.h b/PWGEM/Dilepton/Core/DileptonMC.h index 57beff3f8f1..01d4bfabfd9 100644 --- a/PWGEM/Dilepton/Core/DileptonMC.h +++ b/PWGEM/Dilepton/Core/DileptonMC.h @@ -692,30 +692,42 @@ struct DileptonMC { // fDielectronCut.SetPRangeForITSNsigmaPr(dielectroncuts.cfg_min_p_ITSNsigmaPr, dielectroncuts.cfg_max_p_ITSNsigmaPr); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; - } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); - } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); + } + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); + } + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -822,11 +834,11 @@ struct DileptonMC { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -1871,16 +1883,16 @@ struct DileptonMC { } // end of collision loop } - template - bool isPairOK(TCollision const& collision, TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) + template + bool isPairOK(TTrack1 const& t1, TTrack2 const& t2, TCut const& cut, TAllTracks const& tracks) { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(t1, collision) || !cut.template IsSelectedTrack(t2, collision)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } else { // cut-based - if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { + if (!cut.template IsSelectedTrack(t1) || !cut.template IsSelectedTrack(t2)) { return false; } } @@ -1949,7 +1961,7 @@ struct DileptonMC { continue; } - if (isPairOK(collision, pos, neg, cut, tracks)) { + if (isPairOK(pos, neg, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos.globalIndex(), neg.globalIndex())); } } @@ -1965,7 +1977,7 @@ struct DileptonMC { continue; } - if (isPairOK(collision, pos1, pos2, cut, tracks)) { + if (isPairOK(pos1, pos2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(pos1.globalIndex(), pos2.globalIndex())); } } @@ -1980,7 +1992,7 @@ struct DileptonMC { if (cfgEventGeneratorType >= 0 && mccollision_from_neg2.getSubGeneratorId() != cfgEventGeneratorType) { continue; } - if (isPairOK(collision, neg1, neg2, cut, tracks)) { + if (isPairOK(neg1, neg2, cut, tracks)) { passed_pairIds.emplace_back(std::make_pair(neg1.globalIndex(), neg2.globalIndex())); } } @@ -2104,7 +2116,7 @@ struct DileptonMC { continue; } - if (!isPairOK(collision, pos, neg, cut, tracks)) { // without acceptance + if (!isPairOK(pos, neg, cut, tracks)) { // without acceptance continue; } @@ -2242,7 +2254,7 @@ struct DileptonMC { continue; } - if (!isPairOK(collision, pos1, pos2, cut, tracks)) { // without acceptance + if (!isPairOK(pos1, pos2, cut, tracks)) { // without acceptance continue; } @@ -2299,7 +2311,7 @@ struct DileptonMC { if (cfgEventGeneratorType >= 0 && mccollision_from_neg2.getSubGeneratorId() != cfgEventGeneratorType) { continue; } - if (!isPairOK(collision, neg1, neg2, cut, tracks)) { // without acceptance + if (!isPairOK(neg1, neg2, cut, tracks)) { // without acceptance continue; } if ((std::abs(mcneg1.pdgCode()) != pdg_lepton || std::abs(mcneg2.pdgCode()) != pdg_lepton) || (mcneg1.emmceventId() != mcneg2.emmceventId())) { @@ -2316,7 +2328,7 @@ struct DileptonMC { continue; } - if (!isPairOK(collision, neg1, neg2, cut, tracks)) { // without acceptance + if (!isPairOK(neg1, neg2, cut, tracks)) { // without acceptance continue; } diff --git a/PWGEM/Dilepton/Core/PhotonHBT.h b/PWGEM/Dilepton/Core/PhotonHBT.h index 180360e44f3..796f0c5da2b 100644 --- a/PWGEM/Dilepton/Core/PhotonHBT.h +++ b/PWGEM/Dilepton/Core/PhotonHBT.h @@ -589,30 +589,42 @@ struct PhotonHBT { // fDielectronCut.SetPRangeForITSNsigmaPr(dielectroncuts.cfg_min_p_ITSNsigmaPr, dielectroncuts.cfg_max_p_ITSNsigmaPr); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -840,11 +852,11 @@ struct PhotonHBT { continue; } if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut1.template IsSelectedTrack(pos1, collision) || !cut1.template IsSelectedTrack(ele1, collision)) { + if (!cut1.template IsSelectedTrack(pos1) || !cut1.template IsSelectedTrack(ele1)) { continue; } } else { // cut-based - if (!cut1.template IsSelectedTrack(pos1, collision) || !cut1.template IsSelectedTrack(ele1, collision)) { + if (!cut1.template IsSelectedTrack(pos1) || !cut1.template IsSelectedTrack(ele1)) { continue; } } @@ -868,11 +880,11 @@ struct PhotonHBT { continue; } if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut2.template IsSelectedTrack(pos2, collision) || !cut2.template IsSelectedTrack(ele2, collision)) { + if (!cut2.template IsSelectedTrack(pos2) || !cut2.template IsSelectedTrack(ele2)) { continue; } } else { // cut-based - if (!cut2.template IsSelectedTrack(pos2, collision) || !cut2.template IsSelectedTrack(ele2, collision)) { + if (!cut2.template IsSelectedTrack(pos2) || !cut2.template IsSelectedTrack(ele2)) { continue; } } @@ -986,11 +998,11 @@ struct PhotonHBT { continue; } if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut2.template IsSelectedTrack(pos2, collision) || !cut2.template IsSelectedTrack(ele2, collision)) { + if (!cut2.template IsSelectedTrack(pos2) || !cut2.template IsSelectedTrack(ele2)) { continue; } } else { // cut-based - if (!cut2.template IsSelectedTrack(pos2, collision) || !cut2.template IsSelectedTrack(ele2, collision)) { + if (!cut2.template IsSelectedTrack(pos2) || !cut2.template IsSelectedTrack(ele2)) { continue; } } @@ -1354,11 +1366,11 @@ struct PhotonHBT { continue; } if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(pos, collision) || !cut.template IsSelectedTrack(ele, collision)) { + if (!cut.template IsSelectedTrack(pos) || !cut.template IsSelectedTrack(ele)) { continue; } } else { // cut-based - if (!cut.template IsSelectedTrack(pos, collision) || !cut.template IsSelectedTrack(ele, collision)) { + if (!cut.template IsSelectedTrack(pos) || !cut.template IsSelectedTrack(ele)) { continue; } } diff --git a/PWGEM/Dilepton/Core/SingleTrackQC.h b/PWGEM/Dilepton/Core/SingleTrackQC.h index 4f9744ea769..e073b6e2907 100644 --- a/PWGEM/Dilepton/Core/SingleTrackQC.h +++ b/PWGEM/Dilepton/Core/SingleTrackQC.h @@ -269,6 +269,7 @@ struct SingleTrackQC { // fRegistry.add("Track/positive/hTOFNsigmaKa", "TOF n sigma ka;p_{pv} (GeV/c);n #sigma_{K}^{TOF}", kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); // fRegistry.add("Track/positive/hTOFNsigmaPr", "TOF n sigma pr;p_{pv} (GeV/c);n #sigma_{p}^{TOF}", kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); + fRegistry.add("Track/positive/hProbElBDT", "probability to be e from BDT;p_{in} (GeV/c);BDT score;", kTH2F, {{1000, 0, 10}, {100, 0, 1}}, false); fRegistry.add("Track/positive/hMeanClusterSizeITS", "mean cluster size ITS;p_{pv} (GeV/c); on ITS #times cos(#lambda);", kTH2F, {{1000, 0.f, 10.f}, {150, 0, 15}}, false); fRegistry.add("Track/positive/hMeanClusterSizeITSib", "mean cluster size ITS inner barrel;p_{pv} (GeV/c); on ITS #times cos(#lambda);", kTH2F, {{1000, 0.f, 10.f}, {150, 0, 15}}, false); fRegistry.add("Track/positive/hMeanClusterSizeITSob", "mean cluster size ITS outer barrel;p_{pv} (GeV/c); on ITS #times cos(#lambda);", kTH2F, {{1000, 0.f, 10.f}, {150, 0, 15}}, false); @@ -417,30 +418,42 @@ struct SingleTrackQC { // fDielectronCut.SetPRangeForITSNsigmaPr(dielectroncuts.cfg_min_p_ITSNsigmaPr, dielectroncuts.cfg_max_p_ITSNsigmaPr); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; - } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); - } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); + } + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); + } + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -499,6 +512,7 @@ struct SingleTrackQC { fRegistry.fill(HIST("Track/positive/hTPCdEdx"), track.tpcInnerParam(), track.tpcSignal()); fRegistry.fill(HIST("Track/positive/hTOFbeta"), track.p(), track.beta()); + fRegistry.fill(HIST("Track/positive/hProbElBDT"), track.tpcInnerParam(), track.probElBDT()); fRegistry.fill(HIST("Track/positive/hMeanClusterSizeITS"), track.p(), track.meanClusterSizeITS() * std::cos(std::atan(track.tgl()))); fRegistry.fill(HIST("Track/positive/hMeanClusterSizeITSib"), track.p(), track.meanClusterSizeITSib() * std::cos(std::atan(track.tgl()))); fRegistry.fill(HIST("Track/positive/hMeanClusterSizeITSob"), track.p(), track.meanClusterSizeITSob() * std::cos(std::atan(track.tgl()))); @@ -539,6 +553,7 @@ struct SingleTrackQC { fRegistry.fill(HIST("Track/negative/hTPCdEdx"), track.tpcInnerParam(), track.tpcSignal()); fRegistry.fill(HIST("Track/negative/hTOFbeta"), track.p(), track.beta()); + fRegistry.fill(HIST("Track/negative/hProbElBDT"), track.tpcInnerParam(), track.probElBDT()); fRegistry.fill(HIST("Track/negative/hMeanClusterSizeITS"), track.p(), track.meanClusterSizeITS() * std::cos(std::atan(track.tgl()))); fRegistry.fill(HIST("Track/negative/hMeanClusterSizeITSib"), track.p(), track.meanClusterSizeITSib() * std::cos(std::atan(track.tgl()))); fRegistry.fill(HIST("Track/negative/hMeanClusterSizeITSob"), track.p(), track.meanClusterSizeITSob() * std::cos(std::atan(track.tgl()))); @@ -643,11 +658,11 @@ struct SingleTrackQC { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { for (const auto& track : tracks_per_coll) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(track, collision)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } else { // cut-based - if (!cut.template IsSelectedTrack(track)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } @@ -698,11 +713,11 @@ struct SingleTrackQC { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { for (const auto& track : tracks_per_coll) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(track, collision)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } else { // cut-based - if (!cut.template IsSelectedTrack(track)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } diff --git a/PWGEM/Dilepton/Core/SingleTrackQCMC.h b/PWGEM/Dilepton/Core/SingleTrackQCMC.h index 95e4bc97507..16025e245af 100644 --- a/PWGEM/Dilepton/Core/SingleTrackQCMC.h +++ b/PWGEM/Dilepton/Core/SingleTrackQCMC.h @@ -487,30 +487,42 @@ struct SingleTrackQCMC { // fDielectronCut.SetPRangeForITSNsigmaPr(dielectroncuts.cfg_min_p_ITSNsigmaPr, dielectroncuts.cfg_max_p_ITSNsigmaPr); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; - const std::vector labelsClasses = {"Background", "Signal"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = 0.; - cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutNot, o2::cuts_ml::CutSmaller}; + // const std::vector labelsClasses = {"Background", "Signal"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = 0.; + // cutsMlArr[i][1] = dielectroncuts.cutsMl.value[i]; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } @@ -806,11 +818,11 @@ struct SingleTrackQCMC { if constexpr (pairtype == o2::aod::pwgem::dilepton::utils::pairutil::DileptonPairType::kDielectron) { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(track, collision)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } else { // cut-based - if (!cut.template IsSelectedTrack(track)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } @@ -1001,11 +1013,11 @@ struct SingleTrackQCMC { } if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { - if (!cut.template IsSelectedTrack(track, collision)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } else { // cut-based - if (!cut.template IsSelectedTrack(track)) { + if (!cut.template IsSelectedTrack(track)) { continue; } } diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx index bf1e2b61505..0f9cb82aa55 100644 --- a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx @@ -174,6 +174,7 @@ struct skimmerPrimaryElectron { fRegistry.add("Track/hMeanClusterSizeITS", "mean cluster size ITS;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); fRegistry.add("Track/hMeanClusterSizeITSib", "mean cluster size ITSib;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); fRegistry.add("Track/hMeanClusterSizeITSob", "mean cluster size ITSob;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); + fRegistry.add("Track/hProbElBDT", "probability to be e from BDT;p_{in} (GeV/c);BDT score;", kTH2F, {{1000, 0, 10}, {100, 0, 1}}, false); } if (usePIDML) { @@ -370,20 +371,62 @@ struct skimmerPrimaryElectron { return true; } - template - bool isElectron(TTrack const& track) + template + bool isElectron(TCollision const& collision, TTrack const& track, float& probaEl) { + probaEl = 1.f; if (includeITSsa && (track.hasITS() && !track.hasTPC() && !track.hasTRD() && !track.hasTOF())) { return true; } if (usePIDML) { - return true; + if (!isElectron_TOFif(track)) { + return false; + } + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(track); + trackParCov.setPID(o2::track::PID::Electron); + mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); + mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + + std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); + float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); + + // std::vector outputs = {}; + // bool isSelected = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature, outputs); // 0: hadron, 1:electron + // probaEl = outputs[1]; + // outputs.clear(); + // outputs.shrink_to_fit(); + + // std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); + // float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); + + int pbin = lower_bound(binsMl.value.begin(), binsMl.value.end(), binningFeature) - binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(binsMl.value.size()) - 2; + } + // LOGF(info, "track.tpcInnerParam() = %f (GeV/c), pbin = %d", track.tpcInnerParam(), pbin); + + probaEl = mlResponseSingleTrack.getModelOutput(inputFeatures, pbin)[1]; // 0: hadron, 1:electron + return probaEl > cutsMl.value[pbin]; + // return isSelected; } else { return isElectron_TPChadrej(track) || isElectron_TOFreq(track); } } + template + bool isElectron_TOFif(TTrack const& track) + { + bool is_EL_TPC = minTPCNsigmaEl < track.tpcNSigmaEl() && track.tpcNSigmaEl() < maxTPCNsigmaEl; + bool is_EL_TOF = track.hasTOF() ? (std::fabs(track.tofNSigmaEl()) < maxTOFNsigmaEl) : true; // TOFif + return is_EL_TPC && is_EL_TOF; + } + template bool isElectron_TPChadrej(TTrack const& track) { @@ -415,7 +458,7 @@ struct skimmerPrimaryElectron { } template - void fillTrackTable(TCollision const& collision, TTrack const& track) + void fillTrackTable(TCollision const& collision, TTrack const& track, const float probaEl) { if (std::find(stored_trackIds.begin(), stored_trackIds.end(), std::pair{collision.globalIndex(), track.globalIndex()}) == stored_trackIds.end()) { o2::dataformats::DCA mDcaInfoCov; @@ -439,13 +482,6 @@ struct skimmerPrimaryElectron { mcTunedTPCSignal = track.mcTunedTPCSignal(); } - float probaEl = 1.0; - if (usePIDML) { - std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); - float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); - probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); - } - emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(), pt_recalc, eta_recalc, phi_recalc, dcaXY, dcaZ, trackParCov.getSigmaY2(), trackParCov.getSigmaZY(), trackParCov.getSigmaZ2(), @@ -545,6 +581,7 @@ struct skimmerPrimaryElectron { fRegistry.fill(HIST("Track/hMeanClusterSizeITS"), trackParCov.getP(), static_cast(total_cluster_size) / static_cast(nl) * std::cos(std::atan(trackParCov.getTgl()))); fRegistry.fill(HIST("Track/hMeanClusterSizeITSib"), trackParCov.getP(), static_cast(total_cluster_size_ib) / static_cast(nl_ib) * std::cos(std::atan(trackParCov.getTgl()))); fRegistry.fill(HIST("Track/hMeanClusterSizeITSob"), trackParCov.getP(), static_cast(total_cluster_size_ob) / static_cast(nl_ob) * std::cos(std::atan(trackParCov.getTgl()))); + fRegistry.fill(HIST("Track/hProbElBDT"), track.tpcInnerParam(), probaEl); } } } @@ -573,10 +610,11 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop @@ -602,10 +640,11 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop @@ -632,10 +671,11 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop @@ -664,10 +704,11 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop @@ -698,10 +739,11 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracks.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop @@ -729,10 +771,11 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { auto track = trackId.template track_as(); - if (!checkTrack(collision, track) || !isElectron(track)) { + float probaEl = 1.0; + if (!checkTrack(collision, track) || !isElectron(collision, track, probaEl)) { continue; } - fillTrackTable(collision, track); + fillTrackTable(collision, track, probaEl); } } // end of collision loop diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx index 68b0fad8593..89fc6cd09c3 100644 --- a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectronQC.cxx @@ -187,6 +187,7 @@ struct skimmerPrimaryElectronQC { fRegistry.add("Track/hMeanClusterSizeITS", "mean cluster size ITS;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); fRegistry.add("Track/hMeanClusterSizeITSib", "mean cluster size ITSib;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); fRegistry.add("Track/hMeanClusterSizeITSob", "mean cluster size ITSob;p_{pv} (GeV/c); #times cos(#lambda)", kTH2F, {{1000, 0, 10}, {150, 0, 15}}, false); + fRegistry.add("Track/hProbElBDT", "probability to be e from BDT;p_{in} (GeV/c);BDT score;", kTH2F, {{1000, 0, 10}, {100, 0, 1}}, false); fRegistry.add("Pair/hMvsPhiV", "m_{ee} vs. #varphi_{V} ULS;#varphi_{V} (rad.);m_{ee} (GeV/c^{2})", kTH2F, {{180, 0.f, M_PI}, {100, 0, 0.1}}); } @@ -467,7 +468,14 @@ struct skimmerPrimaryElectronQC { if (usePIDML) { std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); - probaEl = mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); + + int pbin = lower_bound(binsMl.value.begin(), binsMl.value.end(), binningFeature) - binsMl.value.begin() - 1; + if (pbin < 0) { + pbin = 0; + } else if (static_cast(binsMl.value.size()) - 2 < pbin) { + pbin = static_cast(binsMl.value.size()) - 2; + } + probaEl = mlResponseSingleTrack.getModelOutput(inputFeatures, pbin)[1]; // 0: hadron, 1:electron } emprimaryelectrons(collision.globalIndex(), track.globalIndex(), track.sign(), @@ -566,6 +574,7 @@ struct skimmerPrimaryElectronQC { fRegistry.fill(HIST("Track/hMeanClusterSizeITS"), trackParCov.getP(), static_cast(total_cluster_size) / static_cast(nl) * std::cos(std::atan(trackParCov.getTgl()))); fRegistry.fill(HIST("Track/hMeanClusterSizeITSib"), trackParCov.getP(), static_cast(total_cluster_size_ib) / static_cast(nl_ib) * std::cos(std::atan(trackParCov.getTgl()))); fRegistry.fill(HIST("Track/hMeanClusterSizeITSob"), trackParCov.getP(), static_cast(total_cluster_size_ob) / static_cast(nl_ob) * std::cos(std::atan(trackParCov.getTgl()))); + fRegistry.fill(HIST("Track/hProbElBDT"), track.tpcInnerParam(), probaEl); } } } @@ -607,7 +616,6 @@ struct skimmerPrimaryElectronQC { std::vector> stored_trackIds; Filter trackFilter = trackcut.minpt < o2::aod::track::pt && nabs(o2::aod::track::eta) < trackcut.maxeta && o2::aod::track::itsChi2NCl < trackcut.maxchi2its && ncheckbit(aod::track::v001::detectorMap, (uint8_t)o2::aod::track::ITS) == true; - Filter pidFilter = trackcut.minTPCNsigmaEl < o2::aod::pidtpc::tpcNSigmaEl && o2::aod::pidtpc::tpcNSigmaEl < trackcut.maxTPCNsigmaEl; using MyFilteredTracks = soa::Filtered; Partition posTracks = o2::aod::track::signed1Pt > 0.f; @@ -642,15 +650,6 @@ struct skimmerPrimaryElectronQC { } } - // if (isDielectronFromPi0(collision, pos, ele)) { - // if ((checkTrackTight(collision, pos) && isElectronTight(pos)) && (checkTrack(collision, ele) && isElectron(ele)) ) { - // fillTrackTable(collision, ele); - // } - // if ((checkTrackTight(collision, ele) && isElectronTight(ele)) && (checkTrack(collision, pos) && isElectron(pos)) ) { - // fillTrackTable(collision, pos); - // } - // } - } // end of ULS pairing } // end of collision loop @@ -690,15 +689,6 @@ struct skimmerPrimaryElectronQC { } } - // if (isDielectronFromPi0(collision, pos, ele)) { - // if ((checkTrackTight(collision, pos) && isElectronTight(pos)) && (checkTrack(collision, ele) && isElectron(ele)) ) { - // fillTrackTable(collision, ele); - // } - // if ((checkTrackTight(collision, ele) && isElectronTight(ele)) && (checkTrack(collision, pos) && isElectron(pos)) ) { - // fillTrackTable(collision, pos); - // } - // } - } // end of ULS pairing } // end of collision loop @@ -741,15 +731,6 @@ struct skimmerPrimaryElectronQC { } } - // if (isDielectronFromPi0(collision, pos, ele)) { - // if ((checkTrackTight(collision, pos) && isElectronTight(pos)) && (checkTrack(collision, ele) && isElectron(ele)) ) { - // fillTrackTable(collision, ele); - // } - // if ((checkTrackTight(collision, ele) && isElectronTight(ele)) && (checkTrack(collision, pos) && isElectron(pos)) ) { - // fillTrackTable(collision, pos); - // } - // } - } // end of ULS pairing } // end of collision loop diff --git a/PWGEM/Dilepton/Tasks/prefilterDielectron.cxx b/PWGEM/Dilepton/Tasks/prefilterDielectron.cxx index d3a12ec6481..711bcc72c48 100644 --- a/PWGEM/Dilepton/Tasks/prefilterDielectron.cxx +++ b/PWGEM/Dilepton/Tasks/prefilterDielectron.cxx @@ -317,30 +317,42 @@ struct prefilterDielectron { fDielectronCut.SetTOFNsigmaElRange(dielectroncuts.cfg_min_TOFNsigmaEl, dielectroncuts.cfg_max_TOFNsigmaEl); if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut - static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; - const std::vector labelsClasses = {"Signal", "Background"}; - const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; - const std::vector labelsBins(nBinsMl, "bin"); - double cutsMlArr[nBinsMl][nClassesMl]; - for (uint32_t i = 0; i < nBinsMl; i++) { - cutsMlArr[i][0] = dielectroncuts.cutsMl.value[i]; - cutsMlArr[i][1] = 0.; + std::vector binsML{}; + binsML.reserve(dielectroncuts.binsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.binsMl.value.size(); i++) { + binsML.emplace_back(dielectroncuts.binsMl.value[i]); } - o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; - - mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); - if (dielectroncuts.loadModelsFromCCDB) { - ccdbApi.init(ccdburl); - mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); - } else { - mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + std::vector thresholdsML{}; + thresholdsML.reserve(dielectroncuts.cutsMl.value.size()); + for (size_t i = 0; i < dielectroncuts.cutsMl.value.size(); i++) { + thresholdsML.emplace_back(dielectroncuts.cutsMl.value[i]); } - mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); - mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); - mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); - - fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); + fDielectronCut.SetMLThresholds(binsML, thresholdsML); + + // static constexpr int nClassesMl = 2; + // const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + // const std::vector labelsClasses = {"Signal", "Background"}; + // const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; + // const std::vector labelsBins(nBinsMl, "bin"); + // double cutsMlArr[nBinsMl][nClassesMl]; + // for (uint32_t i = 0; i < nBinsMl; i++) { + // cutsMlArr[i][0] = dielectroncuts.cutsMl.value[i]; + // cutsMlArr[i][1] = 0.; + // } + // o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + // mlResponseSingleTrack.configure(dielectroncuts.binsMl.value, cutsMl, cutDirMl, nClassesMl); + // if (dielectroncuts.loadModelsFromCCDB) { + // ccdbApi.init(ccdburl); + // mlResponseSingleTrack.setModelPathsCCDB(dielectroncuts.onnxFileNames.value, ccdbApi, dielectroncuts.onnxPathsCCDB.value, dielectroncuts.timestampCCDB.value); + // } else { + // mlResponseSingleTrack.setModelPathsLocal(dielectroncuts.onnxFileNames.value); + // } + // mlResponseSingleTrack.cacheInputFeaturesIndices(dielectroncuts.namesInputFeatures); + // mlResponseSingleTrack.cacheBinningIndex(dielectroncuts.nameBinningFeature); + // mlResponseSingleTrack.init(dielectroncuts.enableOptimizations.value); + + // fDielectronCut.SetPIDMlResponse(&mlResponseSingleTrack); } // end of PID ML } diff --git a/PWGEM/Dilepton/Utils/MlResponseO2Track.h b/PWGEM/Dilepton/Utils/MlResponseO2Track.h index ae0c17096fb..77c05cd21b1 100644 --- a/PWGEM/Dilepton/Utils/MlResponseO2Track.h +++ b/PWGEM/Dilepton/Utils/MlResponseO2Track.h @@ -26,10 +26,9 @@ // Fill the map of available input features // the key is the feature's name (std::string) // the value is the corresponding value in EnumInputFeatures -#define FILL_MAP_O2_TRACK(FEATURE) \ - { \ - #FEATURE, static_cast(InputFeaturesO2Track::FEATURE) \ - } +#define FILL_MAP_O2_TRACK(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesO2Track::FEATURE)} // Check if the index of mCachedIndices (index associated to a FEATURE) // matches the entry in EnumInputFeatures associated to this FEATURE @@ -161,6 +160,7 @@ enum class InputFeaturesO2Track : uint8_t { tpctofNSigmaPr, tpcNClsFound, tpcNClsCrossedRows, + tpcChi2NCl, hasITS, hasTPC, hasTRD, @@ -215,6 +215,7 @@ class MlResponseO2Track : public MlResponse CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaPr, tpcNSigmaPr, tofNSigmaPr, hasTOF); CHECK_AND_FILL_O2_TRACK(tpcNClsFound); CHECK_AND_FILL_O2_TRACK(tpcNClsCrossedRows); + CHECK_AND_FILL_O2_TRACK(tpcChi2NCl); CHECK_AND_FILL_O2_TRACK(hasITS); CHECK_AND_FILL_O2_TRACK(hasTPC); CHECK_AND_FILL_O2_TRACK(hasTRD); @@ -295,6 +296,7 @@ class MlResponseO2Track : public MlResponse FILL_MAP_O2_TRACK(tpctofNSigmaPr), FILL_MAP_O2_TRACK(tpcNClsFound), FILL_MAP_O2_TRACK(tpcNClsCrossedRows), + FILL_MAP_O2_TRACK(tpcChi2NCl), FILL_MAP_O2_TRACK(hasITS), FILL_MAP_O2_TRACK(hasTPC), FILL_MAP_O2_TRACK(hasTRD), From 9dc1db51539d48909e440888bb56ed409d272c6c Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Thu, 21 Aug 2025 18:16:55 +0000 Subject: [PATCH 2/3] Please consider the following formatting changes --- PWGEM/Dilepton/Utils/MlResponseO2Track.h | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/PWGEM/Dilepton/Utils/MlResponseO2Track.h b/PWGEM/Dilepton/Utils/MlResponseO2Track.h index 77c05cd21b1..5e1f2fc4226 100644 --- a/PWGEM/Dilepton/Utils/MlResponseO2Track.h +++ b/PWGEM/Dilepton/Utils/MlResponseO2Track.h @@ -26,9 +26,10 @@ // Fill the map of available input features // the key is the feature's name (std::string) // the value is the corresponding value in EnumInputFeatures -#define FILL_MAP_O2_TRACK(FEATURE) \ - { \ - #FEATURE, static_cast(InputFeaturesO2Track::FEATURE)} +#define FILL_MAP_O2_TRACK(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesO2Track::FEATURE) \ + } // Check if the index of mCachedIndices (index associated to a FEATURE) // matches the entry in EnumInputFeatures associated to this FEATURE From a360a6e64756dea2785652de9aaa5d3275b893a1 Mon Sep 17 00:00:00 2001 From: Daiki Sekihata Date: Thu, 21 Aug 2025 20:24:15 +0200 Subject: [PATCH 3/3] Update DielectronCut.h --- PWGEM/Dilepton/Core/DielectronCut.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/PWGEM/Dilepton/Core/DielectronCut.h b/PWGEM/Dilepton/Core/DielectronCut.h index 4c0a6e9d86e..20fc738b146 100644 --- a/PWGEM/Dilepton/Core/DielectronCut.h +++ b/PWGEM/Dilepton/Core/DielectronCut.h @@ -519,9 +519,9 @@ class DielectronCut : public TNamed } mMLBins = bins; mMLCuts = cuts; - for (int i = 0; i < static_cast(mMLBins.size()) - 1; i++) { - printf("Dielectron cut: mMLBins[%d] = %3.2f, mMLBins[%d] = %3.2f, mMLCuts[%d] = %3.2f\n", i, mMLBins[i], i + 1, mMLBins[i + 1], i, mMLCuts[i]); - } + // for (int i = 0; i < static_cast(mMLBins.size()) - 1; i++) { + // printf("Dielectron cut: mMLBins[%d] = %3.2f, mMLBins[%d] = %3.2f, mMLCuts[%d] = %3.2f\n", i, mMLBins[i], i + 1, mMLBins[i + 1], i, mMLCuts[i]); + // } } // Getters