@@ -1693,6 +1693,9 @@ struct HfTrackIndexSkimCreator {
16931693 const std::vector<double > ptBinsMl{0 ., 1 .e10 };
16941694 const std::vector<int > cutDirMl{o2::cuts_ml::CutDirection::CutGreater, o2::cuts_ml::CutDirection::CutSmaller, o2::cuts_ml::CutDirection::CutSmaller};
16951695 const std::array<LabeledArray<double >, kN3ProngDecaysUsedMlForHfFilters > thresholdMlScore3Prongs{config.thresholdMlScoreDplusToPiKPi , config.thresholdMlScoreLcToPiKP , config.thresholdMlScoreDsToPiKK , config.thresholdMlScoreXicToPiKP };
1696+ const std::vector<std::string> inputFeatures2Prongs = {" ptProng0" , " dcaXyProng0" , " dcaZProng0" , " ptProng1" , " dcaXyProng1" , " dcaZProng1" };
1697+ const std::vector<std::string> inputFeatures3Prongs = {" ptProng0" , " dcaXyProng0" , " dcaZProng0" , " ptProng1" , " dcaXyProng1" , " dcaZProng1" , " ptProng2" , " dcaXyProng2" , " dcaZProng2" };
1698+ const std::vector<std::string> inputFeatures3ProngsWithPid = {" ptProng0" , " dcaXyProng0" , " dcaZProng0" , " ptProng1" , " dcaXyProng1" , " dcaZProng1" , " ptProng2" , " dcaXyProng2" , " dcaZProng2" , " tpcNSigmaPrProng0" , " tpcNSigmaPrProng2" , " tpcNSigmaPiProng0" , " tpcNSigmaPiProng2" , " tpcNSigmaKaProng1" };
16961699
16971700 // initialise 2-prong ML response
16981701 hfMlResponse2Prongs.configure (ptBinsMl, config.thresholdMlScoreD0ToKPi , cutDirMl, 3 );
@@ -1702,6 +1705,7 @@ struct HfTrackIndexSkimCreator {
17021705 } else {
17031706 hfMlResponse2Prongs.setModelPathsLocal (onnxFileNames2Prongs);
17041707 }
1708+ hfMlResponse2Prongs.cacheInputFeaturesIndices (inputFeatures2Prongs);
17051709 hfMlResponse2Prongs.init ();
17061710
17071711 // initialise 3-prong ML responses
@@ -1717,6 +1721,11 @@ struct HfTrackIndexSkimCreator {
17171721 } else {
17181722 hfMlResponse3Prongs[iDecay3P].setModelPathsLocal (onnxFileNames3Prongs[iDecay3P]);
17191723 }
1724+ if ((doprocess2And3ProngsWithPvRefitWithPidForHfFiltersBdt || doprocess2And3ProngsNoPvRefitWithPidForHfFiltersBdt) && iDecay3P == aod::hf_cand_3prong::DecayType::LcToPKPi) {
1725+ hfMlResponse3Prongs[iDecay3P].cacheInputFeaturesIndices (inputFeatures3ProngsWithPid);
1726+ } else {
1727+ hfMlResponse3Prongs[iDecay3P].cacheInputFeaturesIndices (inputFeatures3Prongs);
1728+ }
17201729 hfMlResponse3Prongs[iDecay3P].init ();
17211730 }
17221731 }
0 commit comments