diff --git a/PWGHF/Core/HfMlResponseD0ToKPi.h b/PWGHF/Core/HfMlResponseD0ToKPi.h new file mode 100644 index 00000000000..62e8c201e6b --- /dev/null +++ b/PWGHF/Core/HfMlResponseD0ToKPi.h @@ -0,0 +1,216 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file HfMlResponseD0ToKPi.h +/// \brief Class to compute the ML response for D0 → K∓ π± analysis selections +/// \author Alexandre Bigot , IPHC Strasbourg +/// \author Andrea Tavira García , IJCLab Orsay + +#ifndef PWGHF_CORE_HFMLRESPONSED0TOKPI_H_ +#define PWGHF_CORE_HFMLRESPONSED0TOKPI_H_ + +#include +#include +#include + +#include "PWGHF/Core/HfHelper.h" +#include "PWGHF/Core/HfMlResponse.h" +#include "PWGHF/Core/PDG.h" + +// 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_D0(FEATURE) \ + { \ +#FEATURE, static_cast < uint8_t>(InputFeaturesD0ToKPi::FEATURE) \ + } + +// Check if the index of mCachedIndices (index associated to a FEATURE) +// matches the entry in EnumInputFeatures associated to this FEATURE +// if so, the inputFeatures vector is filled with the FEATURE's value +// by calling the corresponding GETTER from OBJECT +#define CHECK_AND_FILL_VEC_D0_FULL(OBJECT, FEATURE, GETTER) \ + case static_cast(InputFeaturesD0ToKPi::FEATURE): { \ + inputFeatures.emplace_back(OBJECT.GETTER()); \ + break; \ + } + +// Specific case of CHECK_AND_FILL_VEC_D0_FULL(OBJECT, FEATURE, GETTER) +// where OBJECT is named candidate and FEATURE = GETTER +#define CHECK_AND_FILL_VEC_D0(GETTER) \ + case static_cast(InputFeaturesD0ToKPi::GETTER): { \ + inputFeatures.emplace_back(candidate.GETTER()); \ + break; \ + } + +// Variation of CHECK_AND_FILL_VEC_D0_FULL(OBJECT, FEATURE, GETTER) +// where GETTER is a method of hfHelper +#define CHECK_AND_FILL_VEC_D0_HFHELPER(OBJECT, FEATURE, GETTER) \ + case static_cast(InputFeaturesD0ToKPi::FEATURE): { \ + inputFeatures.emplace_back(hfHelper.GETTER(OBJECT)); \ + break; \ + } + +// Variation of CHECK_AND_FILL_VEC_D0_HFHELPER(OBJECT, FEATURE, GETTER) +// where GETTER1 and GETTER2 are methods of hfHelper, and the variable +// is filled depending on whether it is a D0 or a D0bar +#define CHECK_AND_FILL_VEC_D0_HFHELPER_SIGNED(OBJECT, FEATURE, GETTER1, GETTER2) \ + case static_cast(InputFeaturesD0ToKPi::FEATURE): { \ + if (pdgCode == o2::analysis::pdg::kD0) { \ + inputFeatures.emplace_back(hfHelper.GETTER1(OBJECT)); \ + } else { \ + inputFeatures.emplace_back(hfHelper.GETTER2(OBJECT)); \ + } \ + break; \ + } + +namespace o2::analysis +{ +enum class InputFeaturesD0ToKPi : uint8_t { + chi2PCA = 0, + decayLength, + decayLengthXY, + decayLengthNormalised, + decayLengthXYNormalised, + impactParameterNormalised0, + ptProng0, + ptProng1, + impactParameter0, + impactParameter1, + nSigTpcPi0, + nSigTpcKa0, + nSigTofPi0, + nSigTofKa0, + nSigTpcTofPi0, + nSigTpcTofKa0, + nSigTpcPi1, + nSigTpcKa1, + nSigTofPi1, + nSigTofKa1, + nSigTpcTofPi1, + nSigTpcTofKa1, + maxNormalisedDeltaIP, + impactParameterProduct, + cosThetaStar, + cpa, + cpaXY, + ct +}; + +template +class HfMlResponseD0ToKPi : public HfMlResponse +{ + public: + /// Default constructor + HfMlResponseD0ToKPi() = default; + /// Default destructor + virtual ~HfMlResponseD0ToKPi() = default; + + HfHelper hfHelper; + + /// Method to get the input features vector needed for ML inference + /// \param candidate is the D0 candidate + /// \param prong0 is the candidate's prong0 + /// \param prong1 is the candidate's prong1 + /// \return inputFeatures vector + template + std::vector getInputFeatures(T1 const& candidate, + T2 const& prong0, T2 const& prong1, int const& pdgCode) + { + std::vector inputFeatures; + + for (const auto& idx : MlResponse::mCachedIndices) { + switch (idx) { + CHECK_AND_FILL_VEC_D0(chi2PCA); + CHECK_AND_FILL_VEC_D0(decayLength); + CHECK_AND_FILL_VEC_D0(decayLengthXY); + CHECK_AND_FILL_VEC_D0(decayLengthNormalised); + CHECK_AND_FILL_VEC_D0(decayLengthXYNormalised); + CHECK_AND_FILL_VEC_D0(ptProng0); + CHECK_AND_FILL_VEC_D0(ptProng1); + CHECK_AND_FILL_VEC_D0_FULL(candidate, impactParameter0, impactParameter0); + CHECK_AND_FILL_VEC_D0_FULL(candidate, impactParameter1, impactParameter1); + // TPC PID variables + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTpcPi0, tpcNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTpcKa0, tpcNSigmaKa); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTpcPi1, tpcNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTpcKa1, tpcNSigmaKa); + // TOF PID variables + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTofPi0, tofNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTofKa0, tofNSigmaKa); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTofPi1, tofNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTofKa1, tofNSigmaKa); + // Combined PID variables + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTpcTofPi0, tpcTofNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong0, nSigTpcTofKa0, tpcTofNSigmaKa); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTpcTofPi1, tpcTofNSigmaPi); + CHECK_AND_FILL_VEC_D0_FULL(prong1, nSigTpcTofKa1, tpcTofNSigmaKa); + + CHECK_AND_FILL_VEC_D0(maxNormalisedDeltaIP); + CHECK_AND_FILL_VEC_D0_FULL(candidate, impactParameterProduct, impactParameterProduct); + CHECK_AND_FILL_VEC_D0_HFHELPER_SIGNED(candidate, cosThetaStar, cosThetaStarD0, cosThetaStarD0bar); + CHECK_AND_FILL_VEC_D0(cpa); + CHECK_AND_FILL_VEC_D0(cpaXY); + CHECK_AND_FILL_VEC_D0_HFHELPER(candidate, ct, ctD0); + } + } + + return inputFeatures; + } + + protected: + /// Method to fill the map of available input features + void setAvailableInputFeatures() + { + MlResponse::mAvailableInputFeatures = { + FILL_MAP_D0(chi2PCA), + FILL_MAP_D0(decayLength), + FILL_MAP_D0(decayLengthXY), + FILL_MAP_D0(decayLengthNormalised), + FILL_MAP_D0(decayLengthXYNormalised), + FILL_MAP_D0(ptProng0), + FILL_MAP_D0(ptProng1), + FILL_MAP_D0(impactParameter0), + FILL_MAP_D0(impactParameter1), + // TPC PID variables + FILL_MAP_D0(nSigTpcPi0), + FILL_MAP_D0(nSigTpcKa0), + FILL_MAP_D0(nSigTpcPi1), + FILL_MAP_D0(nSigTpcKa1), + // TOF PID variables + FILL_MAP_D0(nSigTofPi0), + FILL_MAP_D0(nSigTofKa0), + FILL_MAP_D0(nSigTofPi1), + FILL_MAP_D0(nSigTofKa1), + // Combined PID variables + FILL_MAP_D0(nSigTpcTofPi0), + FILL_MAP_D0(nSigTpcTofKa0), + FILL_MAP_D0(nSigTpcTofPi1), + FILL_MAP_D0(nSigTpcTofKa1), + + FILL_MAP_D0(maxNormalisedDeltaIP), + FILL_MAP_D0(impactParameterProduct), + FILL_MAP_D0(cosThetaStar), + FILL_MAP_D0(cpa), + FILL_MAP_D0(cpaXY), + FILL_MAP_D0(ct)}; + } +}; + +} // namespace o2::analysis + +#undef FILL_MAP_D0 +#undef CHECK_AND_FILL_VEC_D0_FULL +#undef CHECK_AND_FILL_VEC_D0 +#undef CHECK_AND_FILL_VEC_D0_HFHELPER +#undef CHECK_AND_FILL_VEC_D0_HFHELPER_SIGNED + +#endif // PWGHF_CORE_HFMLRESPONSED0TOKPI_H_ diff --git a/PWGHF/DataModel/CandidateSelectionTables.h b/PWGHF/DataModel/CandidateSelectionTables.h index b27dacb5113..5d8436f47e4 100644 --- a/PWGHF/DataModel/CandidateSelectionTables.h +++ b/PWGHF/DataModel/CandidateSelectionTables.h @@ -43,6 +43,7 @@ DECLARE_SOA_COLUMN(IsRecoTopol, isRecoTopol, int); //! DECLARE_SOA_COLUMN(IsRecoCand, isRecoCand, int); //! DECLARE_SOA_COLUMN(IsRecoPid, isRecoPid, int); DECLARE_SOA_COLUMN(MlProbD0, mlProbD0, std::vector); //! +DECLARE_SOA_COLUMN(MlProbD0bar, mlProbD0bar, std::vector); //! } // namespace hf_sel_candidate_d0 DECLARE_SOA_TABLE(HfSelD0, "AOD", "HFSELD0", //! @@ -54,7 +55,8 @@ DECLARE_SOA_TABLE(HfSelD0, "AOD", "HFSELD0", //! hf_sel_candidate_d0::IsRecoPid); DECLARE_SOA_TABLE(HfMlD0, "AOD", "HFMLD0", //! - hf_sel_candidate_d0::MlProbD0); + hf_sel_candidate_d0::MlProbD0, + hf_sel_candidate_d0::MlProbD0bar); namespace hf_sel_candidate_d0_parametrized_pid { diff --git a/PWGHF/TableProducer/candidateSelectorD0.cxx b/PWGHF/TableProducer/candidateSelectorD0.cxx index 497704bc7e5..aaf1f68b51c 100644 --- a/PWGHF/TableProducer/candidateSelectorD0.cxx +++ b/PWGHF/TableProducer/candidateSelectorD0.cxx @@ -22,6 +22,7 @@ #include "PWGHF/Core/HfHelper.h" #include "PWGHF/Core/HfMlResponse.h" +#include "PWGHF/Core/HfMlResponseD0ToKPi.h" #include "PWGHF/DataModel/CandidateReconstructionTables.h" #include "PWGHF/DataModel/CandidateSelectionTables.h" @@ -61,6 +62,7 @@ struct HfCandidateSelectorD0 { Configurable> 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 nClassesMl{"nClassesMl", (int8_t)hf_cuts_ml::nCutScores, "Number of classes in ML model"}; Configurable enableDebugMl{"enableDebugMl", false, "Flag to enable histograms to monitor BDT application"}; + Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature1", "feature2"}, "Names of ML model input features"}; // CCDB configuration Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; Configurable modelPathsCCDB{"modelPathsCCDB", "EventFiltering/PWGHF/BDTD0", "Path on CCDB"}; @@ -68,14 +70,15 @@ struct HfCandidateSelectorD0 { Configurable timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"}; Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; - o2::analysis::HfMlResponse hfMlResponse; - std::vector outputMl = {}; + o2::analysis::HfMlResponseD0ToKPi hfMlResponse; + std::vector outputMlD0 = {}; + std::vector outputMlD0bar = {}; o2::ccdb::CcdbApi ccdbApi; TrackSelectorPi selectorPion; TrackSelectorKa selectorKaon; HfHelper hfHelper; - using TracksSel = soa::Join; + using TracksSel = soa::Join; // Define histograms AxisSpec axisMassDmeson{200, 1.7f, 2.1f}; @@ -113,8 +116,8 @@ struct HfCandidateSelectorD0 { } else { hfMlResponse.setModelPathsLocal(onnxFileNames); } + hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures); hfMlResponse.init(); - outputMl.assign(((std::vector)cutDirMl).size(), -1.f); // dummy value for ML output } } @@ -260,10 +263,13 @@ struct HfCandidateSelectorD0 { int statusCand = 0; int statusPID = 0; + outputMlD0.clear(); + outputMlD0bar.clear(); + if (!(candidate.hfflag() & 1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) { hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID); if (applyMl) { - hfMlD0Candidate(outputMl); + hfMlD0Candidate(outputMlD0, outputMlD0bar); } continue; } @@ -277,7 +283,7 @@ struct HfCandidateSelectorD0 { if (!selectionTopol(candidate)) { hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID); if (applyMl) { - hfMlD0Candidate(outputMl); + hfMlD0Candidate(outputMlD0, outputMlD0bar); } continue; } @@ -294,7 +300,7 @@ struct HfCandidateSelectorD0 { if (!topolD0 && !topolD0bar) { hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID); if (applyMl) { - hfMlD0Candidate(outputMl); + hfMlD0Candidate(outputMlD0, outputMlD0bar); } continue; } @@ -342,7 +348,7 @@ struct HfCandidateSelectorD0 { if (pidD0 == 0 && pidD0bar == 0) { hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID); if (applyMl) { - hfMlD0Candidate(outputMl); + hfMlD0Candidate(outputMlD0, outputMlD0bar); } continue; } @@ -357,33 +363,38 @@ struct HfCandidateSelectorD0 { if (applyMl) { // ML selections + bool isSelectedMlD0 = false; + bool isSelectedMlD0bar = false; - std::vector inputFeatures{candidate.cpa(), - candidate.cpaXY(), - candidate.decayLength(), - candidate.decayLengthXY(), - candidate.impactParameter0(), - candidate.impactParameter1(), - candidate.impactParameterProduct()}; - - bool isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, ptCand, outputMl); - hfMlD0Candidate(outputMl); + if (statusD0 > 0) { + std::vector inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate, trackPos, trackNeg, o2::analysis::pdg::kD0); + isSelectedMlD0 = hfMlResponse.isSelectedMl(inputFeaturesD0, ptCand, outputMlD0); + } + if (statusD0bar > 0) { + std::vector inputFeaturesD0bar = hfMlResponse.getInputFeatures(candidate, trackPos, trackNeg, o2::analysis::pdg::kD0Bar); + isSelectedMlD0bar = hfMlResponse.isSelectedMl(inputFeaturesD0bar, ptCand, outputMlD0bar); + } - if (!isSelectedMl) { + if (!isSelectedMlD0) { statusD0 = 0; + } + if (!isSelectedMlD0bar) { statusD0bar = 0; } + + hfMlD0Candidate(outputMlD0, outputMlD0bar); + if (enableDebugMl) { - registry.fill(HIST("DebugBdt/hBdtScore1VsStatus"), outputMl[0], statusD0); - registry.fill(HIST("DebugBdt/hBdtScore1VsStatus"), outputMl[0], statusD0bar); - registry.fill(HIST("DebugBdt/hBdtScore2VsStatus"), outputMl[1], statusD0); - registry.fill(HIST("DebugBdt/hBdtScore2VsStatus"), outputMl[1], statusD0bar); - registry.fill(HIST("DebugBdt/hBdtScore3VsStatus"), outputMl[2], statusD0); - registry.fill(HIST("DebugBdt/hBdtScore3VsStatus"), outputMl[2], statusD0bar); - if (statusD0 > 0) { + if (isSelectedMlD0) { + registry.fill(HIST("DebugBdt/hBdtScore1VsStatus"), outputMlD0[0], statusD0); + registry.fill(HIST("DebugBdt/hBdtScore2VsStatus"), outputMlD0[1], statusD0); + registry.fill(HIST("DebugBdt/hBdtScore3VsStatus"), outputMlD0[2], statusD0); registry.fill(HIST("DebugBdt/hMassDmesonSel"), hfHelper.invMassD0ToPiK(candidate)); } - if (statusD0bar > 0) { + if (isSelectedMlD0bar) { + registry.fill(HIST("DebugBdt/hBdtScore1VsStatus"), outputMlD0bar[0], statusD0bar); + registry.fill(HIST("DebugBdt/hBdtScore2VsStatus"), outputMlD0bar[1], statusD0bar); + registry.fill(HIST("DebugBdt/hBdtScore3VsStatus"), outputMlD0bar[2], statusD0bar); registry.fill(HIST("DebugBdt/hMassDmesonSel"), hfHelper.invMassD0barToKPi(candidate)); } }