Skip to content

[PWGHF] ML selections for D0 in D* selector #12176

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mfaggin merged 3 commits into from
Jul 22, 2025
Merged

[PWGHF] ML selections for D0 in D* selector #12176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c9460bf
ML selections for D0
Marcellocosti Jul 15, 2025
50ec9ce
Please consider the following formatting changes
alibuild Jul 21, 2025
930a92e
Merge pull request #29 from alibuild/alibot-cleanup-12176
Marcellocosti Jul 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions PWGHF/Core/HfMlResponseDstarToD0Pi.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,8 @@ enum class InputFeaturesDstarToD0Pi : uint8_t {
ptSoftPi,
impactParameter0,
impactParameter1,
impactParameterXY0,
impactParameterXY1,
impactParameterZ0,
impactParameterZ1,
impParamSoftPi,
Expand Down Expand Up @@ -219,6 +221,28 @@ class HfMlResponseDstarToD0Pi : public HfMlResponse<TypeOutputScore>
return inputFeatures;
}

/// Method to get the input features used for D0 in HF triggers
/// \param candidate is the D* candidate
/// \return inputFeatures vector
template <typename T1>
std::vector<float> getInputFeaturesTrigger(T1 const& candidate)
{
std::vector<float> inputFeatures;

for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
switch (idx) {
CHECK_AND_FILL_VEC_DSTAR(ptProng0);
CHECK_AND_FILL_VEC_DSTAR_GETTER(impactParameterXY0, impactParameter0);
CHECK_AND_FILL_VEC_DSTAR(impactParameterZ0);
CHECK_AND_FILL_VEC_DSTAR(ptProng1);
CHECK_AND_FILL_VEC_DSTAR_GETTER(impactParameterXY1, impactParameter1);
CHECK_AND_FILL_VEC_DSTAR(impactParameterZ1);
}
}

return inputFeatures;
}

protected:
/// Method to fill the map of available input features
void setAvailableInputFeatures()
Expand All @@ -238,6 +262,8 @@ class HfMlResponseDstarToD0Pi : public HfMlResponse<TypeOutputScore>
FILL_MAP_DSTAR(ptSoftPi),
FILL_MAP_DSTAR(impactParameter0),
FILL_MAP_DSTAR(impactParameter1),
FILL_MAP_DSTAR(impactParameterXY0),
FILL_MAP_DSTAR(impactParameterXY1),
FILL_MAP_DSTAR(impactParameterZ0),
FILL_MAP_DSTAR(impactParameterZ1),
FILL_MAP_DSTAR(impParamSoftPi),
Expand Down
35 changes: 34 additions & 1 deletion PWGHF/TableProducer/candidateSelectorDstarToD0Pi.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
/// \author Fabrizio Grosa <fabrizio.grosa@cern.ch>, CERN

#include "PWGHF/Core/HfHelper.h"
#include "PWGHF/Core/HfMlResponseD0ToKPi.h"
#include "PWGHF/Core/HfMlResponseDstarToD0Pi.h"
#include "PWGHF/Core/SelectorCuts.h"
#include "PWGHF/DataModel/CandidateReconstructionTables.h"
Expand Down Expand Up @@ -90,18 +91,27 @@ struct HfCandidateSelectorDstarToD0Pi {

// QA switch
Configurable<bool> activateQA{"activateQA", false, "Flag to enable QA histogram"};
// ML inference
// ML inference D*
Configurable<bool> applyMl{"applyMl", false, "Flag to apply ML selections"};
Configurable<std::vector<double>> binsPtMl{"binsPtMl", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
Configurable<LabeledArray<double>> cutsMl{"cutsMl", {hf_cuts_ml::Cuts[0], hf_cuts_ml::NBinsPt, hf_cuts_ml::NCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsCutScore}, "ML selections per pT bin"};
Configurable<int> nClassesMl{"nClassesMl", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model"};
Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
// ML inference D0
Configurable<bool> applyMlD0Daug{"applyMlD0Daug", false, "Flag to apply ML selections on D0 daughter"};
Configurable<std::vector<double>> binsPtMlD0Daug{"binsPtMlD0Daug", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application on D0 daughter"};
Configurable<std::vector<int>> cutDirMlD0Daug{"cutDirMlD0Daug", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold on D0 daughter"};
Configurable<LabeledArray<double>> cutsMlD0Daug{"cutsMlD0Daug", {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 on D0 daughter"};
Configurable<int> nClassesMlD0Daug{"nClassesMlD0Daug", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model on D0 daughter"};
Configurable<std::vector<std::string>> namesInputFeaturesD0Daug{"namesInputFeaturesD0Daug", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features on D0 daughter"};

// CCDB configuration
Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{""}, "Paths of models on CCDB"};
Configurable<std::vector<std::string>> modelPathsCCDBD0Daug{"modelPathsCCDBD0Daug", std::vector<std::string>{""}, "Paths of models on CCDB for D0 daughter"};
Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"Model.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
Configurable<std::vector<std::string>> onnxFileNamesD0Daug{"onnxFileNamesD0Daug", std::vector<std::string>{"Model.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path) for D0 daughter"};
Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};

Expand All @@ -110,7 +120,9 @@ struct HfCandidateSelectorDstarToD0Pi {

HfHelper hfHelper;
o2::analysis::HfMlResponseDstarToD0Pi<float> hfMlResponse;
o2::analysis::HfMlResponseDstarToD0Pi<float> hfMlResponseD0Daughter;
std::vector<float> outputMlDstarToD0Pi = {};
std::vector<float> outputMlD0ToKPi = {};
o2::ccdb::CcdbApi ccdbApi;

TrackSelectorPi selectorPion;
Expand Down Expand Up @@ -172,6 +184,18 @@ struct HfCandidateSelectorDstarToD0Pi {
hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
hfMlResponse.init();
}

if (applyMlD0Daug) {
hfMlResponseD0Daughter.configure(binsPtMlD0Daug, cutsMlD0Daug, cutDirMlD0Daug, nClassesMlD0Daug);
if (loadModelsFromCCDB) {
ccdbApi.init(ccdbUrl);
hfMlResponseD0Daughter.setModelPathsCCDB(onnxFileNamesD0Daug, ccdbApi, modelPathsCCDBD0Daug, timestampCCDB);
} else {
hfMlResponseD0Daughter.setModelPathsLocal(onnxFileNamesD0Daug);
}
hfMlResponseD0Daughter.cacheInputFeaturesIndices(namesInputFeaturesD0Daug);
hfMlResponseD0Daughter.init();
}
}

/// Conjugate-independent topological cuts on D0
Expand Down Expand Up @@ -234,6 +258,15 @@ struct HfCandidateSelectorDstarToD0Pi {
if (candidate.decayLengthXYD0() > cutsD0->get(binPt, "max decay length XY")) {
return false;
}

if (applyMlD0Daug) {
outputMlD0ToKPi.clear();
std::vector<float> inputFeaturesD0 = hfMlResponseD0Daughter.getInputFeaturesTrigger(candidate);
bool isSelectedMlD0 = hfMlResponseD0Daughter.isSelectedMl(inputFeaturesD0, candpT, outputMlD0ToKPi);
if (!isSelectedMlD0) {
return false;
}
}
return true;
}

Expand Down
Loading