diff --git a/PWGHF/Core/HfMlResponseXicToXiPiPi.h b/PWGHF/Core/HfMlResponseXicToXiPiPi.h index 4bb4a4d023f..6e8b598a166 100644 --- a/PWGHF/Core/HfMlResponseXicToXiPiPi.h +++ b/PWGHF/Core/HfMlResponseXicToXiPiPi.h @@ -17,19 +17,21 @@ #define PWGHF_CORE_HFMLRESPONSEXICTOXIPIPI_H_ #include "PWGHF/Core/HfMlResponse.h" +#include "PWGHF/DataModel/CandidateReconstructionTables.h" #include "Tools/ML/MlResponse.h" #include +#include +#include #include // 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_XICTOXIPIPI(FEATURE) \ - { \ - #FEATURE, static_cast(InputFeaturesXicToXiPiPi::FEATURE) \ - } +#define FILL_MAP_XICTOXIPIPI(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesXicToXiPiPi::FEATURE)} // Check if the index of mCachedIndices (index associated to a FEATURE) // matches the entry in EnumInputFeatures associated to this FEATURE @@ -88,10 +90,31 @@ enum class InputFeaturesXicToXiPiPi : uint8_t { nSigTofPiFromXicPlus1, nSigTofBachelorPi, nSigTofPiFromLambda, - nSigTofPrFromLambda + nSigTofPrFromLambda, + // kf-specific variables + kfDecayLength, + kfDecayLengthNormalised, + kfDecayLengthXY, + kfDecayLengthXYNormalised, + kfCascadeChi2, + kfV0Chi2, + chi2XicPlusToPV, + chi2TopoXicPlusToPV, + chi2XiToPV, + chi2Pi0ToPV, + chi2Pi1ToPV, + chi2DevPi0Pi1, + chi2DevPi0Xi, + chi2DevPi1Xi, + dcaPi0Pi1, + dcaPi0Xi, + dcaPi1Xi, + dcaXYPi0Pi1, + dcaXYPi0Xi, + dcaXYPi1Xi }; -template +template class HfMlResponseXicToXiPiPi : public HfMlResponse { public: @@ -147,6 +170,30 @@ class HfMlResponseXicToXiPiPi : public HfMlResponse CHECK_AND_FILL_VEC_XICTOXIPIPI(nSigTofPiFromLambda); CHECK_AND_FILL_VEC_XICTOXIPIPI(nSigTofPrFromLambda); } + if constexpr (ReconstructionType == aod::hf_cand::VertexerType::KfParticle) { + switch (idx) { + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfDecayLength); + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfDecayLengthNormalised); + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfDecayLengthXY); + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfDecayLengthXYNormalised); + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfCascadeChi2); + CHECK_AND_FILL_VEC_XICTOXIPIPI(kfV0Chi2); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2XicPlusToPV); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2TopoXicPlusToPV); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2XiToPV); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2Pi0ToPV); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2Pi1ToPV); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2DevPi0Pi1); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2DevPi0Xi); + CHECK_AND_FILL_VEC_XICTOXIPIPI(chi2DevPi1Xi); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaPi0Pi1); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaPi0Xi); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaPi1Xi); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaXYPi0Pi1); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaXYPi0Xi); + CHECK_AND_FILL_VEC_XICTOXIPIPI(dcaXYPi1Xi); + } + } } return inputFeatures; @@ -193,6 +240,30 @@ class HfMlResponseXicToXiPiPi : public HfMlResponse FILL_MAP_XICTOXIPIPI(nSigTofBachelorPi), FILL_MAP_XICTOXIPIPI(nSigTofPiFromLambda), FILL_MAP_XICTOXIPIPI(nSigTofPrFromLambda)}; + if constexpr (ReconstructionType == aod::hf_cand::VertexerType::KfParticle) { + std::map mapKfFeatures{ + FILL_MAP_XICTOXIPIPI(kfDecayLength), + FILL_MAP_XICTOXIPIPI(kfDecayLengthNormalised), + FILL_MAP_XICTOXIPIPI(kfDecayLengthXY), + FILL_MAP_XICTOXIPIPI(kfDecayLengthXYNormalised), + FILL_MAP_XICTOXIPIPI(kfCascadeChi2), + FILL_MAP_XICTOXIPIPI(kfV0Chi2), + FILL_MAP_XICTOXIPIPI(chi2XicPlusToPV), + FILL_MAP_XICTOXIPIPI(chi2TopoXicPlusToPV), + FILL_MAP_XICTOXIPIPI(chi2XiToPV), + FILL_MAP_XICTOXIPIPI(chi2Pi0ToPV), + FILL_MAP_XICTOXIPIPI(chi2Pi1ToPV), + FILL_MAP_XICTOXIPIPI(chi2DevPi0Pi1), + FILL_MAP_XICTOXIPIPI(chi2DevPi0Xi), + FILL_MAP_XICTOXIPIPI(chi2DevPi1Xi), + FILL_MAP_XICTOXIPIPI(dcaPi0Pi1), + FILL_MAP_XICTOXIPIPI(dcaPi0Xi), + FILL_MAP_XICTOXIPIPI(dcaPi1Xi), + FILL_MAP_XICTOXIPIPI(dcaXYPi0Pi1), + FILL_MAP_XICTOXIPIPI(dcaXYPi0Xi), + FILL_MAP_XICTOXIPIPI(dcaXYPi1Xi)}; + MlResponse::mAvailableInputFeatures.insert(mapKfFeatures.begin(), mapKfFeatures.end()); + } } }; diff --git a/PWGHF/TableProducer/candidateSelectorXicToXiPiPi.cxx b/PWGHF/TableProducer/candidateSelectorXicToXiPiPi.cxx index 8972554209a..588dd15a1bb 100644 --- a/PWGHF/TableProducer/candidateSelectorXicToXiPiPi.cxx +++ b/PWGHF/TableProducer/candidateSelectorXicToXiPiPi.cxx @@ -92,7 +92,8 @@ struct HfCandidateSelectorXicToXiPiPi { 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::HfMlResponseXicToXiPiPi hfMlResponse; + o2::analysis::HfMlResponseXicToXiPiPi hfMlResponseDca; + o2::analysis::HfMlResponseXicToXiPiPi hfMlResponseKf; std::vector outputMlXicToXiPiPi; o2::ccdb::CcdbApi ccdbApi; TrackSelectorPi selectorPion; @@ -121,9 +122,13 @@ struct HfCandidateSelectorXicToXiPiPi { void init(InitContext&) { - if ((doprocessData + doprocessMc) != 1) { + if ((doprocessDataWithDCAFitter + doprocessDataWithKFParticle + doprocessMcWithDCAFitter + doprocessMcWithKFParticle) != 1) { LOGP(fatal, "Enable exactly one process function at a time."); } + const bool doprocessData = doprocessDataWithDCAFitter || doprocessDataWithKFParticle; + const bool doprocessMc = doprocessMcWithDCAFitter || doprocessMcWithKFParticle; + const bool doprocessDCAFitter = doprocessDataWithDCAFitter || doprocessMcWithDCAFitter; + [[maybe_unused]] const bool doprocessKFParticle = doprocessDataWithKFParticle || doprocessMcWithKFParticle; if (usePid) { // pion @@ -251,15 +256,22 @@ struct HfCandidateSelectorXicToXiPiPi { } if (applyMl) { - hfMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); - if (loadModelsFromCCDB) { - ccdbApi.init(ccdbUrl); - hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); + auto initMlResponse = [&](auto& hfMlResponseGeneric) { + hfMlResponseGeneric.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelsFromCCDB) { + ccdbApi.init(ccdbUrl); + hfMlResponseGeneric.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); + } else { + hfMlResponseGeneric.setModelPathsLocal(onnxFileNames); + } + hfMlResponseGeneric.cacheInputFeaturesIndices(namesInputFeatures); + hfMlResponseGeneric.init(); + }; + if (doprocessDCAFitter) { + initMlResponse(hfMlResponseDca); } else { - hfMlResponse.setModelPathsLocal(onnxFileNames); + initMlResponse(hfMlResponseKf); } - hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures); - hfMlResponse.init(); } } @@ -669,7 +681,7 @@ struct HfCandidateSelectorXicToXiPiPi { /// \param hfCandXic Xic candidate /// \param statusXicToXiPiPi Flag to store selection status as defined in hf_sel_candidate_xic::XicToXiPiPiSelectionStep /// \param isMatchedSignal Flag to indicate if the candidate is matched to a genereated XiCplus MC particle - template + template void isBdtSelected(XicCandidate const& hfCandXic, int& statusXicToXiPiPi, bool const isMatchedSignal = false) @@ -677,8 +689,13 @@ struct HfCandidateSelectorXicToXiPiPi { bool isSelectedMlXicToXiPiPi = false; float const ptCandXic = hfCandXic.pt(); - std::vector inputFeaturesXicToXiPiPi = hfMlResponse.getInputFeatures(hfCandXic); - isSelectedMlXicToXiPiPi = hfMlResponse.isSelectedMl(inputFeaturesXicToXiPiPi, ptCandXic, outputMlXicToXiPiPi); + if constexpr (ReconstructionType == aod::hf_cand::VertexerType::DCAFitter) { + std::vector inputFeaturesXicToXiPiPi = hfMlResponseDca.getInputFeatures(hfCandXic); + isSelectedMlXicToXiPiPi = hfMlResponseDca.isSelectedMl(inputFeaturesXicToXiPiPi, ptCandXic, outputMlXicToXiPiPi); + } else { + std::vector inputFeaturesXicToXiPiPi = hfMlResponseKf.getInputFeatures(hfCandXic); + isSelectedMlXicToXiPiPi = hfMlResponseKf.isSelectedMl(inputFeaturesXicToXiPiPi, ptCandXic, outputMlXicToXiPiPi); + } hfMlXicToXiPiPiCandidate(outputMlXicToXiPiPi); @@ -699,8 +716,8 @@ struct HfCandidateSelectorXicToXiPiPi { } } - void processData(aod::HfCandXic const& hfCandsXic, - TracksExtraWPid const& tracks) + template + void runData(CandType const& hfCandsXic, TracksExtraWPid const& tracks) { for (const auto& hfCandXic : hfCandsXic) { int statusXicToXiPiPi = 0; @@ -719,7 +736,7 @@ struct HfCandidateSelectorXicToXiPiPi { // ML selection if (applyMl) { - isBdtSelected(hfCandXic, statusXicToXiPiPi); + isBdtSelected(hfCandXic, statusXicToXiPiPi); } else { registry.fill(HIST("hSelCandidates"), BdtSelected); } @@ -727,10 +744,9 @@ struct HfCandidateSelectorXicToXiPiPi { hfSelXicToXiPiPiCandidate(statusXicToXiPiPi); } } - PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processData, "Select candidates without MC matching information", true); - void processMc(soa::Join const& hfCandsXic, - TracksExtraWPid const& tracks) + template + void runMc(CandType const& hfCandsXic, TracksExtraWPid const& tracks) { for (const auto& hfCandXic : hfCandsXic) { int statusXicToXiPiPi = 0; @@ -754,7 +770,7 @@ struct HfCandidateSelectorXicToXiPiPi { // ML selection if (applyMl) { - isBdtSelected(hfCandXic, statusXicToXiPiPi, isMatchedCandidate); + isBdtSelected(hfCandXic, statusXicToXiPiPi, isMatchedCandidate); } else if (isMatchedCandidate) { registry.fill(HIST("hSelCandidatesRecSig"), BdtSelected); } else { @@ -764,7 +780,34 @@ struct HfCandidateSelectorXicToXiPiPi { hfSelXicToXiPiPiCandidate(statusXicToXiPiPi); } } - PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processMc, "Select candidates with MC matching information", false); + + void processDataWithDCAFitter(aod::HfCandXic const& hfCandsXic, + TracksExtraWPid const& tracks) + { + runData(hfCandsXic, tracks); + } + PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processDataWithDCAFitter, "Select candidates without MC matching information using DCAFitter", true); + + void processDataWithKFParticle(soa::Join const& hfCandsXic, + TracksExtraWPid const& tracks) + { + runData(hfCandsXic, tracks); + } + PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processDataWithKFParticle, "Select candidates without MC matching information using KFParticle", false); + + void processMcWithDCAFitter(soa::Join const& hfCandsXic, + TracksExtraWPid const& tracks) + { + runMc(hfCandsXic, tracks); + } + PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processMcWithDCAFitter, "Select candidates with MC matching information using DCAFitter", false); + + void processMcWithKFParticle(soa::Join const& hfCandsXic, + TracksExtraWPid const& tracks) + { + runMc(hfCandsXic, tracks); + } + PROCESS_SWITCH(HfCandidateSelectorXicToXiPiPi, processMcWithKFParticle, "Select candidates with MC matching information using KFParticle", false); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)