Skip to content

Commit

Permalink
PWGHF: B+ reduced workflow with ML (AliceO2Group#5964)
Browse files Browse the repository at this point in the history
* updateDraftBplusReduced

* clang

* hfflag removal

* Extending hfflag removal to B0
  • Loading branch information
apalasciano authored May 13, 2024
1 parent 3f872f9 commit f347b6b
Show file tree
Hide file tree
Showing 19 changed files with 1,535 additions and 581 deletions.
221 changes: 221 additions & 0 deletions PWGHF/Core/HfMlResponseBplusToD0Pi.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,221 @@
// 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 HfMlResponseBplusToD0Pi.h
/// \brief Class to compute the ML response for B± → D0(bar) π± analysis selections
/// \author Antonio Palasciano <[email protected]>, INFN Bari

#ifndef PWGHF_CORE_HFMLRESPONSEBPLUSTOD0PI_H_
#define PWGHF_CORE_HFMLRESPONSEBPLUSTOD0PI_H_

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

#include "PWGHF/Core/HfMlResponse.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_BPLUS(FEATURE) \
{ \
#FEATURE, static_cast < uint8_t>(InputFeaturesBplusToD0Pi::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_BPLUS_FULL(OBJECT, FEATURE, GETTER) \
case static_cast<uint8_t>(InputFeaturesBplusToD0Pi::FEATURE): { \
inputFeatures.emplace_back(OBJECT.GETTER()); \
break; \
}

// 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 GETTER function taking OBJECT in argument
#define CHECK_AND_FILL_VEC_BPLUS_FUNC(OBJECT, FEATURE, GETTER) \
case static_cast<uint8_t>(InputFeaturesBplusToD0Pi::FEATURE): { \
inputFeatures.emplace_back(GETTER(OBJECT)); \
break; \
}

// Specific case of CHECK_AND_FILL_VEC_BPLUS_FULL(OBJECT, FEATURE, GETTER)
// where OBJECT is named candidate and FEATURE = GETTER
#define CHECK_AND_FILL_VEC_BPLUS(GETTER) \
case static_cast<uint8_t>(InputFeaturesBplusToD0Pi::GETTER): { \
inputFeatures.emplace_back(candidate.GETTER()); \
break; \
}

namespace o2::pid_tpc_tof_utils
{
template <typename T1>
float getTpcTofNSigmaPi1(const T1& prong1)
{
float defaultNSigma = -999.f; // -999.f is the default value set in TPCPIDResponse.h and PIDTOF.h

bool hasTpc = prong1.hasTPC();
bool hasTof = prong1.hasTOF();

if (hasTpc && hasTof) {
float tpcNSigma = prong1.tpcNSigmaPi();
float tofNSigma = prong1.tofNSigmaPi();
return sqrt(.5f * tpcNSigma * tpcNSigma + .5f * tofNSigma * tofNSigma);
}
if (hasTpc) {
return abs(prong1.tpcNSigmaPi());
}
if (hasTof) {
return abs(prong1.tofNSigmaPi());
}
return defaultNSigma;
}
} // namespace o2::pid_tpc_tof_utils

namespace o2::analysis
{

enum class InputFeaturesBplusToD0Pi : uint8_t {
ptProng0 = 0,
ptProng1,
impactParameter0,
impactParameter1,
impactParameterProduct,
chi2PCA,
decayLength,
decayLengthXY,
decayLengthNormalised,
decayLengthXYNormalised,
cpa,
cpaXY,
maxNormalisedDeltaIP,
prong0MlScoreBkg,
prong0MlScorePrompt,
prong0MlScoreNonprompt,
tpcNSigmaPi1,
tofNSigmaPi1,
tpcTofNSigmaPi1
};

template <typename TypeOutputScore = float>
class HfMlResponseBplusToD0Pi : public HfMlResponse<TypeOutputScore>
{
public:
/// Default constructor
HfMlResponseBplusToD0Pi() = default;
/// Default destructor
virtual ~HfMlResponseBplusToD0Pi() = default;

/// Method to get the input features vector needed for ML inference
/// \param candidate is the B+ candidate
/// \param prong1 is the candidate's prong1
/// \return inputFeatures vector
template <bool withDmesMl, typename T1, typename T2>
std::vector<float> getInputFeatures(T1 const& candidate,
T2 const& prong1)
{
std::vector<float> inputFeatures;

for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
if constexpr (withDmesMl) {
switch (idx) {
CHECK_AND_FILL_VEC_BPLUS(ptProng0);
CHECK_AND_FILL_VEC_BPLUS(ptProng1);
CHECK_AND_FILL_VEC_BPLUS(impactParameter0);
CHECK_AND_FILL_VEC_BPLUS(impactParameter1);
CHECK_AND_FILL_VEC_BPLUS(impactParameterProduct);
CHECK_AND_FILL_VEC_BPLUS(chi2PCA);
CHECK_AND_FILL_VEC_BPLUS(decayLength);
CHECK_AND_FILL_VEC_BPLUS(decayLengthXY);
CHECK_AND_FILL_VEC_BPLUS(decayLengthNormalised);
CHECK_AND_FILL_VEC_BPLUS(decayLengthXYNormalised);
CHECK_AND_FILL_VEC_BPLUS(cpa);
CHECK_AND_FILL_VEC_BPLUS(cpaXY);
CHECK_AND_FILL_VEC_BPLUS(maxNormalisedDeltaIP);
CHECK_AND_FILL_VEC_BPLUS(prong0MlScoreBkg);
CHECK_AND_FILL_VEC_BPLUS(prong0MlScorePrompt);
CHECK_AND_FILL_VEC_BPLUS(prong0MlScoreNonprompt);
// TPC PID variable
CHECK_AND_FILL_VEC_BPLUS_FULL(prong1, tpcNSigmaPi1, tpcNSigmaPi);
// TOF PID variable
CHECK_AND_FILL_VEC_BPLUS_FULL(prong1, tofNSigmaPi1, tofNSigmaPi);
// Combined PID variables
CHECK_AND_FILL_VEC_BPLUS_FUNC(prong1, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
}
} else {
switch (idx) {
CHECK_AND_FILL_VEC_BPLUS(ptProng0);
CHECK_AND_FILL_VEC_BPLUS(ptProng1);
CHECK_AND_FILL_VEC_BPLUS(impactParameter0);
CHECK_AND_FILL_VEC_BPLUS(impactParameter1);
CHECK_AND_FILL_VEC_BPLUS(impactParameterProduct);
CHECK_AND_FILL_VEC_BPLUS(chi2PCA);
CHECK_AND_FILL_VEC_BPLUS(decayLength);
CHECK_AND_FILL_VEC_BPLUS(decayLengthXY);
CHECK_AND_FILL_VEC_BPLUS(decayLengthNormalised);
CHECK_AND_FILL_VEC_BPLUS(decayLengthXYNormalised);
CHECK_AND_FILL_VEC_BPLUS(cpa);
CHECK_AND_FILL_VEC_BPLUS(cpaXY);
CHECK_AND_FILL_VEC_BPLUS(maxNormalisedDeltaIP);
// TPC PID variable
CHECK_AND_FILL_VEC_BPLUS_FULL(prong1, tpcNSigmaPi1, tpcNSigmaPi);
// TOF PID variable
CHECK_AND_FILL_VEC_BPLUS_FULL(prong1, tofNSigmaPi1, tofNSigmaPi);
// Combined PID variables
CHECK_AND_FILL_VEC_BPLUS_FUNC(prong1, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
}
}
}

return inputFeatures;
}

protected:
/// Method to fill the map of available input features
void setAvailableInputFeatures()
{
MlResponse<TypeOutputScore>::mAvailableInputFeatures = {
FILL_MAP_BPLUS(ptProng0),
FILL_MAP_BPLUS(ptProng1),
FILL_MAP_BPLUS(impactParameter0),
FILL_MAP_BPLUS(impactParameter1),
FILL_MAP_BPLUS(impactParameterProduct),
FILL_MAP_BPLUS(chi2PCA),
FILL_MAP_BPLUS(decayLength),
FILL_MAP_BPLUS(decayLengthXY),
FILL_MAP_BPLUS(decayLengthNormalised),
FILL_MAP_BPLUS(decayLengthXYNormalised),
FILL_MAP_BPLUS(cpa),
FILL_MAP_BPLUS(cpaXY),
FILL_MAP_BPLUS(maxNormalisedDeltaIP),
FILL_MAP_BPLUS(prong0MlScoreBkg),
FILL_MAP_BPLUS(prong0MlScorePrompt),
FILL_MAP_BPLUS(prong0MlScoreNonprompt),
// TPC PID variable
FILL_MAP_BPLUS(tpcNSigmaPi1),
// TOF PID variable
FILL_MAP_BPLUS(tofNSigmaPi1),
// Combined PID variable
FILL_MAP_BPLUS(tpcTofNSigmaPi1)};
}
};

} // namespace o2::analysis

#undef FILL_MAP_BPLUS
#undef CHECK_AND_FILL_VEC_BPLUS_FULL
#undef CHECK_AND_FILL_VEC_BPLUS_FUNC
#undef CHECK_AND_FILL_VEC_BPLUS

#endif // PWGHF_CORE_HFMLRESPONSEBPLUSTOD0PI_H_
29 changes: 29 additions & 0 deletions PWGHF/D2H/DataModel/ReducedDataModel.h
Original file line number Diff line number Diff line change
Expand Up @@ -238,11 +238,20 @@ namespace hf_cand_bplus_reduced
{
DECLARE_SOA_INDEX_COLUMN_FULL(Prong0, prong0, int, HfRed2Prongs, "_0"); //! Prong0 index
DECLARE_SOA_INDEX_COLUMN_FULL(Prong1, prong1, int, HfRedTrackBases, "_1"); //! Prong1 index
DECLARE_SOA_COLUMN(Prong0MlScoreBkg, prong0MlScoreBkg, float); //! Bkg ML score of the D daughter
DECLARE_SOA_COLUMN(Prong0MlScorePrompt, prong0MlScorePrompt, float); //! Prompt ML score of the D daughter
DECLARE_SOA_COLUMN(Prong0MlScoreNonprompt, prong0MlScoreNonprompt, float); //! Nonprompt ML score of the D daughter
} // namespace hf_cand_bplus_reduced

DECLARE_SOA_TABLE(HfRedBplusProngs, "AOD", "HFREDBPPRONG",
hf_cand_bplus_reduced::Prong0Id, hf_cand_bplus_reduced::Prong1Id);

DECLARE_SOA_TABLE(HfRedBplusD0Mls, "AOD", "HFREDBPLUSD0ML", //! Table with ML scores for the D+ daughter
hf_cand_bplus_reduced::Prong0MlScoreBkg,
hf_cand_bplus_reduced::Prong0MlScorePrompt,
hf_cand_bplus_reduced::Prong0MlScoreNonprompt,
o2::soa::Marker<1>);

using HfRedCandBplus = soa::Join<HfCandBplusExt, HfRedBplusProngs>;

namespace hf_b0_mc
Expand Down Expand Up @@ -340,6 +349,11 @@ DECLARE_SOA_COLUMN(EtaProng0, etaProng0, float); //! Pseudorapidity of the track
DECLARE_SOA_COLUMN(PtProng1, ptProng1, float); //! Transverse momentum of the track's prong1 in GeV/c
DECLARE_SOA_COLUMN(YProng1, yProng1, float); //! Rapidity of the track's prong1
DECLARE_SOA_COLUMN(EtaProng1, etaProng1, float); //! Pseudorapidity of the track's prong1

DECLARE_SOA_COLUMN(PdgCodeBeautyMother, pdgCodeBeautyMother, int); //! Pdg code of beauty mother
DECLARE_SOA_COLUMN(PdgCodeProng0, pdgCodeProng0, int); //! Pdg code of prong0
DECLARE_SOA_COLUMN(PdgCodeProng1, pdgCodeProng1, int); //! Pdg code of prong1
DECLARE_SOA_COLUMN(PdgCodeProng2, pdgCodeProng2, int); //! Pdg code of prong2
} // namespace hf_bplus_mc

// table with results of reconstruction level MC matching
Expand All @@ -350,12 +364,27 @@ DECLARE_SOA_TABLE(HfMcRecRedD0Pis, "AOD", "HFMCRECREDD0PI", //! Table with recon
hf_cand_bplus::DebugMcRec,
hf_bplus_mc::PtMother);

// DECLARE_SOA_EXTENDED_TABLE_USER(ExTable, Tracks, "EXTABLE",
DECLARE_SOA_TABLE(HfMcCheckD0Pis, "AOD", "HFMCCHECKD0PI", //! Table with reconstructed MC information on D0Pi(<-B0) pairs for MC checks in reduced workflow
hf_bplus_mc::PdgCodeBeautyMother,
hf_bplus_mc::PdgCodeProng0,
hf_bplus_mc::PdgCodeProng1,
hf_bplus_mc::PdgCodeProng2,
o2::soa::Marker<1>);

// Table with same size as HFCANDBPLUS
DECLARE_SOA_TABLE(HfMcRecRedBps, "AOD", "HFMCRECREDBP", //! Reconstruction-level MC information on B+ candidates for reduced workflow
hf_cand_bplus::FlagMcMatchRec,
hf_cand_bplus::DebugMcRec,
hf_bplus_mc::PtMother);

DECLARE_SOA_TABLE(HfMcCheckBps, "AOD", "HFMCCHECKBP", //! Table with reconstructed MC information on B+ candidates for MC checks in reduced workflow
hf_bplus_mc::PdgCodeBeautyMother,
hf_bplus_mc::PdgCodeProng0,
hf_bplus_mc::PdgCodeProng1,
hf_bplus_mc::PdgCodeProng2,
o2::soa::Marker<2>);

DECLARE_SOA_TABLE(HfMcGenRedBps, "AOD", "HFMCGENREDBP", //! Generation-level MC information on B+ candidates for reduced workflow
hf_cand_bplus::FlagMcMatchGen,
hf_bplus_mc::PtTrack,
Expand Down
2 changes: 1 addition & 1 deletion PWGHF/D2H/TableProducer/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ o2physics_add_dpl_workflow(candidate-selector-b0-to-d-pi-reduced

o2physics_add_dpl_workflow(candidate-selector-bplus-to-d0-pi-reduced
SOURCES candidateSelectorBplusToD0PiReduced.cxx
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
COMPONENT_NAME Analysis)

# Data creators
Expand Down
3 changes: 1 addition & 2 deletions PWGHF/D2H/TableProducer/candidateCreatorB0Reduced.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -193,8 +193,7 @@ struct HfCandidateCreatorB0Reduced {
pVecD[0], pVecD[1], pVecD[2],
pVecPion[0], pVecPion[1], pVecPion[2],
dcaD.getY(), dcaPion.getY(),
std::sqrt(dcaD.getSigmaY2()), std::sqrt(dcaPion.getSigmaY2()),
BIT(hf_cand_b0::DecayType::B0ToDPi));
std::sqrt(dcaD.getSigmaY2()), std::sqrt(dcaPion.getSigmaY2()));

rowCandidateProngs(candD.globalIndex(), trackPion.globalIndex());

Expand Down
Loading

0 comments on commit f347b6b

Please sign in to comment.