Skip to content
216 changes: 216 additions & 0 deletions PWGHF/Core/HfMlResponseD0ToKPi.h
Original file line number Diff line number Diff line change
@@ -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 <alexandre.bigot@cern.ch>, IPHC Strasbourg
/// \author Andrea Tavira García <tavira-garcia@ijclab.in2p3.fr>, IJCLab Orsay

#ifndef PWGHF_CORE_HFMLRESPONSED0TOKPI_H_
#define PWGHF_CORE_HFMLRESPONSED0TOKPI_H_

#include <map>
#include <string>
#include <vector>

#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<uint8_t>(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<uint8_t>(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<uint8_t>(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<uint8_t>(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,
Comment thread
fgrosa marked this conversation as resolved.
nSigTpcTofPi1,
nSigTpcTofKa1,
maxNormalisedDeltaIP,
impactParameterProduct,
cosThetaStar,
cpa,
cpaXY,
ct
};

template <typename TypeOutputScore = float>
class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
{
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 <typename T1, typename T2>
std::vector<float> getInputFeatures(T1 const& candidate,
T2 const& prong0, T2 const& prong1, int const& pdgCode)
{
std::vector<float> inputFeatures;

for (const auto& idx : MlResponse<TypeOutputScore>::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<TypeOutputScore>::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_
4 changes: 3 additions & 1 deletion PWGHF/DataModel/CandidateSelectionTables.h
Original file line number Diff line number Diff line change
Expand Up @@ -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<float>); //!
DECLARE_SOA_COLUMN(MlProbD0bar, mlProbD0bar, std::vector<float>); //!
} // namespace hf_sel_candidate_d0

DECLARE_SOA_TABLE(HfSelD0, "AOD", "HFSELD0", //!
Expand All @@ -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
{
Expand Down
65 changes: 38 additions & 27 deletions PWGHF/TableProducer/candidateSelectorD0.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -61,21 +62,23 @@ struct HfCandidateSelectorD0 {
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<int8_t> nClassesMl{"nClassesMl", (int8_t)hf_cuts_ml::nCutScores, "Number of classes in ML model"};
Configurable<bool> enableDebugMl{"enableDebugMl", false, "Flag to enable histograms to monitor BDT application"};
Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
// CCDB configuration
Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
Configurable<std::string> modelPathsCCDB{"modelPathsCCDB", "EventFiltering/PWGHF/BDTD0", "Path on CCDB"};
Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_D0ToKPi.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
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"};

o2::analysis::HfMlResponse<float> hfMlResponse;
std::vector<float> outputMl = {};
o2::analysis::HfMlResponseD0ToKPi<float> hfMlResponse;
std::vector<float> outputMlD0 = {};
std::vector<float> outputMlD0bar = {};
o2::ccdb::CcdbApi ccdbApi;
TrackSelectorPi selectorPion;
TrackSelectorKa selectorKaon;
HfHelper hfHelper;

using TracksSel = soa::Join<aod::TracksWDcaExtra, aod::TracksPidPi, aod::TracksPidKa>;
using TracksSel = soa::Join<aod::TracksWDcaExtra, aod::TracksPidPiExt, aod::TracksPidKaExt>;

// Define histograms
AxisSpec axisMassDmeson{200, 1.7f, 2.1f};
Expand Down Expand Up @@ -113,8 +116,8 @@ struct HfCandidateSelectorD0 {
} else {
hfMlResponse.setModelPathsLocal(onnxFileNames);
}
hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
hfMlResponse.init();
outputMl.assign(((std::vector<int>)cutDirMl).size(), -1.f); // dummy value for ML output
}
}

Expand Down Expand Up @@ -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;
}
Expand All @@ -277,7 +283,7 @@ struct HfCandidateSelectorD0 {
if (!selectionTopol<reconstructionType>(candidate)) {
hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID);
if (applyMl) {
hfMlD0Candidate(outputMl);
hfMlD0Candidate(outputMlD0, outputMlD0bar);
}
continue;
}
Expand All @@ -294,7 +300,7 @@ struct HfCandidateSelectorD0 {
if (!topolD0 && !topolD0bar) {
hfSelD0Candidate(statusD0, statusD0bar, statusHFFlag, statusTopol, statusCand, statusPID);
if (applyMl) {
hfMlD0Candidate(outputMl);
hfMlD0Candidate(outputMlD0, outputMlD0bar);
}
continue;
}
Expand Down Expand Up @@ -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;
}
Expand All @@ -357,33 +363,38 @@ struct HfCandidateSelectorD0 {

if (applyMl) {
// ML selections
bool isSelectedMlD0 = false;
bool isSelectedMlD0bar = false;

std::vector<float> 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<float> inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate, trackPos, trackNeg, o2::analysis::pdg::kD0);
isSelectedMlD0 = hfMlResponse.isSelectedMl(inputFeaturesD0, ptCand, outputMlD0);
}
if (statusD0bar > 0) {
std::vector<float> 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));
}
}
Expand Down