diff --git a/PWGEM/Dilepton/Core/Dilepton.h b/PWGEM/Dilepton/Core/Dilepton.h index d5d63c90b8f..bd8d07a952d 100644 --- a/PWGEM/Dilepton/Core/Dilepton.h +++ b/PWGEM/Dilepton/Core/Dilepton.h @@ -707,7 +707,7 @@ struct Dilepton { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; const std::vector labelsClasses = {"Signal", "Background"}; const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; const std::vector labelsBins(nBinsMl, "bin"); diff --git a/PWGEM/Dilepton/Core/DileptonMC.h b/PWGEM/Dilepton/Core/DileptonMC.h index b7e24e7d4d5..e23b18d56db 100644 --- a/PWGEM/Dilepton/Core/DileptonMC.h +++ b/PWGEM/Dilepton/Core/DileptonMC.h @@ -672,7 +672,7 @@ struct DileptonMC { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; const std::vector labelsClasses = {"Signal", "Background"}; const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; const std::vector labelsBins(nBinsMl, "bin"); diff --git a/PWGEM/Dilepton/Core/PhotonHBT.h b/PWGEM/Dilepton/Core/PhotonHBT.h index cc8251fb0cd..99be337aa67 100644 --- a/PWGEM/Dilepton/Core/PhotonHBT.h +++ b/PWGEM/Dilepton/Core/PhotonHBT.h @@ -589,7 +589,7 @@ struct PhotonHBT { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; const std::vector labelsClasses = {"Signal", "Background"}; const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; const std::vector labelsBins(nBinsMl, "bin"); diff --git a/PWGEM/Dilepton/Core/SingleTrackQC.h b/PWGEM/Dilepton/Core/SingleTrackQC.h index 799e276d790..2d3389354d9 100644 --- a/PWGEM/Dilepton/Core/SingleTrackQC.h +++ b/PWGEM/Dilepton/Core/SingleTrackQC.h @@ -415,7 +415,7 @@ struct SingleTrackQC { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; const std::vector labelsClasses = {"Signal", "Background"}; const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; const std::vector labelsBins(nBinsMl, "bin"); diff --git a/PWGEM/Dilepton/Core/SingleTrackQCMC.h b/PWGEM/Dilepton/Core/SingleTrackQCMC.h index e817e1b30fd..3b76e8f8b93 100644 --- a/PWGEM/Dilepton/Core/SingleTrackQCMC.h +++ b/PWGEM/Dilepton/Core/SingleTrackQCMC.h @@ -474,7 +474,7 @@ struct SingleTrackQCMC { if (dielectroncuts.cfg_pid_scheme == static_cast(DielectronCut::PIDSchemes::kPIDML)) { // please call this at the end of DefineDileptonCut static constexpr int nClassesMl = 2; - const std::vector cutDirMl = {o2::cuts_ml::CutSmaller, o2::cuts_ml::CutNot}; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; const std::vector labelsClasses = {"Signal", "Background"}; const uint32_t nBinsMl = dielectroncuts.binsMl.value.size() - 1; const std::vector labelsBins(nBinsMl, "bin"); diff --git a/PWGEM/Dilepton/TableProducer/CMakeLists.txt b/PWGEM/Dilepton/TableProducer/CMakeLists.txt index c95928b3eb1..e756d0a9fc7 100644 --- a/PWGEM/Dilepton/TableProducer/CMakeLists.txt +++ b/PWGEM/Dilepton/TableProducer/CMakeLists.txt @@ -22,7 +22,7 @@ o2physics_add_dpl_workflow(tree-creator-electron-ml-dda o2physics_add_dpl_workflow(skimmer-primary-electron SOURCES skimmerPrimaryElectron.cxx - PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore + PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::AnalysisCore O2Physics::MLCore COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(skimmer-primary-muon diff --git a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx index a4eac761725..ad5976b9850 100644 --- a/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx +++ b/PWGEM/Dilepton/TableProducer/skimmerPrimaryElectron.cxx @@ -12,30 +12,34 @@ /// \brief write relevant information about primary electrons. /// \author daiki.sekihata@cern.ch -#include -#include -#include -#include +#include "PWGEM/Dilepton/DataModel/dileptonTables.h" +#include "PWGEM/Dilepton/Utils/MlResponseO2Track.h" +#include "PWGEM/Dilepton/Utils/PairUtilities.h" -#include "Math/Vector4D.h" -#include "Framework/runDataProcessing.h" -#include "Framework/AnalysisTask.h" -#include "Framework/AnalysisDataModel.h" -#include "DetectorsBase/Propagator.h" -#include "DetectorsBase/GeometryManager.h" -#include "DataFormatsParameters/GRPObject.h" -#include "DataFormatsParameters/GRPMagField.h" -#include "DataFormatsCalibration/MeanVertexObject.h" -#include "CCDB/BasicCCDBManager.h" +#include "Common/Core/TableHelper.h" #include "Common/Core/trackUtilities.h" -#include "CommonConstants/PhysicsConstants.h" #include "Common/DataModel/CollisionAssociationTables.h" -#include "Common/Core/TableHelper.h" #include "Common/DataModel/PIDResponse.h" #include "Common/DataModel/PIDResponseITS.h" +#include "Tools/ML/MlResponse.h" -#include "PWGEM/Dilepton/DataModel/dileptonTables.h" -#include "PWGEM/Dilepton/Utils/PairUtilities.h" +#include "CCDB/BasicCCDBManager.h" +#include "CommonConstants/PhysicsConstants.h" +#include "DataFormatsCalibration/MeanVertexObject.h" +#include "DataFormatsParameters/GRPMagField.h" +#include "DataFormatsParameters/GRPObject.h" +#include "DetectorsBase/GeometryManager.h" +#include "DetectorsBase/Propagator.h" +#include "Framework/AnalysisDataModel.h" +#include "Framework/AnalysisTask.h" +#include "Framework/runDataProcessing.h" + +#include "Math/Vector4D.h" + +#include +#include +#include +#include using namespace o2; using namespace o2::soa; @@ -70,7 +74,7 @@ struct skimmerPrimaryElectron { // Operation and minimisation criteria Configurable fillQAHistogram{"fillQAHistogram", false, "flag to fill QA histograms"}; Configurable d_bz_input{"d_bz_input", -999, "bz field in kG, -999 is automatic"}; - Configurable min_ncluster_tpc{"min_ncluster_tpc", 10, "min ncluster tpc"}; + Configurable min_ncluster_tpc{"min_ncluster_tpc", 0, "min ncluster tpc"}; Configurable mincrossedrows{"mincrossedrows", 70, "min. crossed rows"}; Configurable min_tpc_cr_findable_ratio{"min_tpc_cr_findable_ratio", 0.8, "min. TPC Ncr/Nf ratio"}; Configurable min_ncluster_its{"min_ncluster_its", 4, "min ncluster its"}; @@ -79,8 +83,8 @@ struct skimmerPrimaryElectron { Configurable maxchi2its{"maxchi2its", 6.0, "max. chi2/NclsITS"}; Configurable minpt{"minpt", 0.15, "min pt for track"}; Configurable maxeta{"maxeta", 0.9, "eta acceptance"}; - Configurable dca_xy_max{"dca_xy_max", 0.3f, "max DCAxy in cm"}; - Configurable dca_z_max{"dca_z_max", 0.3f, "max DCAz in cm"}; + Configurable dca_xy_max{"dca_xy_max", 1.0, "max DCAxy in cm"}; + Configurable dca_z_max{"dca_z_max", 1.0, "max DCAz in cm"}; Configurable dca_3d_sigma_max{"dca_3d_sigma_max", 1e+10, "max DCA 3D in sigma"}; Configurable minTPCNsigmaEl{"minTPCNsigmaEl", -2.5, "min. TPC n sigma for electron inclusion"}; Configurable maxTPCNsigmaEl{"maxTPCNsigmaEl", 3.5, "max. TPC n sigma for electron inclusion"}; @@ -96,7 +100,20 @@ struct skimmerPrimaryElectron { Configurable max_pin_for_pion_rejection{"max_pin_for_pion_rejection", 0.5, "pion rejection is applied below this pin"}; Configurable max_frac_shared_clusters_tpc{"max_frac_shared_clusters_tpc", 999.f, "max fraction of shared clusters in TPC"}; + // configuration for PID ML + Configurable usePIDML{"usePIDML", false, "Flag to use PID ML"}; + Configurable> onnxFileNames{"onnxFileNames", std::vector{"filename"}, "ONNX file names for each bin (if not from CCDB full path)"}; + Configurable> onnxPathsCCDB{"onnxPathsCCDB", std::vector{"path"}, "Paths of models on CCDB"}; + Configurable> binsMl{"binsMl", std::vector{-999999., 999999.}, "Bin limits for ML application"}; + Configurable> cutsMl{"cutsMl", std::vector{0.95}, "ML cuts per bin"}; + Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature"}, "Names of ML model input features"}; + Configurable nameBinningFeature{"nameBinningFeature", "pt", "Names of ML model binning feature"}; + Configurable timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB. Exceptions: > 0 for the specific timestamp, 0 gets the run dependent timestamp"}; + Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + Configurable enableOptimizations{"enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; + HistogramRegistry fRegistry{"output", {}, OutputObjHandlingPolicy::AnalysisObject, false, false}; + o2::analysis::MlResponseO2Track mlResponseSingleTrack; int mRunNumber; float d_bz; @@ -106,6 +123,7 @@ struct skimmerPrimaryElectron { o2::dataformats::VertexBase mVtx; const o2::dataformats::MeanVertexObject* mMeanVtx = nullptr; o2::base::MatLayerCylSet* lut = nullptr; + o2::ccdb::CcdbApi ccdbApi; void init(InitContext&) { @@ -116,6 +134,7 @@ struct skimmerPrimaryElectron { ccdb->setCaching(true); ccdb->setLocalObjectValidityChecking(); ccdb->setFatalWhenNull(false); + ccdbApi.init(ccdburl); if (fillQAHistogram) { fRegistry.add("Track/hPt", "pT;p_{T} (GeV/c)", kTH1F, {{1000, 0.0f, 10}}, false); @@ -154,6 +173,31 @@ struct skimmerPrimaryElectron { fRegistry.add("Track/hITSNsigmaKa", "ITS n sigma ka;p_{pv} (GeV/c);n #sigma_{K}^{ITS}", kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); fRegistry.add("Track/hITSNsigmaPr", "ITS n sigma pr;p_{pv} (GeV/c);n #sigma_{p}^{ITS}", kTH2F, {{1000, 0, 10}, {100, -5, +5}}, false); } + + if (usePIDML) { + static constexpr int nClassesMl = 2; + const std::vector cutDirMl = {o2::cuts_ml::CutGreater, o2::cuts_ml::CutNot}; + const std::vector labelsClasses = {"Signal", "Background"}; + const uint32_t nBinsMl = binsMl.value.size() - 1; + const std::vector labelsBins(nBinsMl, "bin"); + double cutsMlArr[nBinsMl][nClassesMl]; + for (uint32_t i = 0; i < nBinsMl; i++) { + cutsMlArr[i][0] = cutsMl.value[i]; + cutsMlArr[i][1] = 0.; + } + o2::framework::LabeledArray cutsMl = {cutsMlArr[0], nBinsMl, nClassesMl, labelsBins, labelsClasses}; + + mlResponseSingleTrack.configure(binsMl.value, cutsMl, cutDirMl, nClassesMl); + if (loadModelsFromCCDB) { + ccdbApi.init(ccdburl); + mlResponseSingleTrack.setModelPathsCCDB(onnxFileNames.value, ccdbApi, onnxPathsCCDB.value, timestampCCDB.value); + } else { + mlResponseSingleTrack.setModelPathsLocal(onnxFileNames.value); + } + mlResponseSingleTrack.cacheInputFeaturesIndices(namesInputFeatures); + mlResponseSingleTrack.cacheBinningIndex(nameBinningFeature); + mlResponseSingleTrack.init(enableOptimizations.value); + } // end of PID ML } void initCCDB(aod::BCsWithTimestamps::iterator const& bc) @@ -299,10 +343,32 @@ struct skimmerPrimaryElectron { return true; } - template - bool isElectron(TTrack const& track) + template + bool isElectron(TCollision const& collision, TTrack const& track) { - return isElectron_TPChadrej(track) || isElectron_TOFreq(track); + if (usePIDML) { + if (track.tpcNSigmaEl() < minTPCNsigmaEl || maxTPCNsigmaEl < track.tpcNSigmaEl()) { + return false; + } + if (track.hasTOF() && (maxTOFNsigmaEl < std::fabs(track.tofNSigmaEl()))) { + return false; + } + + // return false; + o2::dataformats::DCA mDcaInfoCov; + mDcaInfoCov.set(999, 999, 999, 999, 999); + auto trackParCov = getTrackParCov(track); + trackParCov.setPID(o2::track::PID::Electron); + mVtx.setPos({collision.posX(), collision.posY(), collision.posZ()}); + mVtx.setCov(collision.covXX(), collision.covXY(), collision.covYY(), collision.covXZ(), collision.covYZ(), collision.covZZ()); + o2::base::Propagator::Instance()->propagateToDCABxByBz(mVtx, trackParCov, 2.f, matCorr, &mDcaInfoCov); + + std::vector inputFeatures = mlResponseSingleTrack.getInputFeatures(track, trackParCov, collision); + float binningFeature = mlResponseSingleTrack.getBinningFeature(track, trackParCov, collision); + return mlResponseSingleTrack.isSelectedMl(inputFeatures, binningFeature); + } else { + return isElectron_TPChadrej(track) || isElectron_TOFreq(track); + } } template @@ -460,7 +526,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracksWithITSPid.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); @@ -491,7 +557,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { // auto track = trackId.template track_as(); auto track = tracksWithITSPid.rawIteratorAt(trackId.trackId()); - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); @@ -522,7 +588,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracksWithITSPid.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); @@ -556,7 +622,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { // auto track = trackId.template track_as(); auto track = tracksWithITSPid.rawIteratorAt(trackId.trackId()); - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); @@ -591,7 +657,7 @@ struct skimmerPrimaryElectron { auto tracks_per_coll = tracksWithITSPid.sliceBy(perCol, collision.globalIndex()); for (const auto& track : tracks_per_coll) { - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); @@ -624,7 +690,7 @@ struct skimmerPrimaryElectron { for (const auto& trackId : trackIdsThisCollision) { // auto track = trackId.template track_as(); auto track = tracksWithITSPid.rawIteratorAt(trackId.trackId()); - if (!checkTrack(collision, track) || !isElectron(track)) { + if (!checkTrack(collision, track) || !isElectron(collision, track)) { continue; } fillTrackTable(collision, track); diff --git a/PWGEM/Dilepton/Utils/MlResponseDielectronSingleTrack.h b/PWGEM/Dilepton/Utils/MlResponseDielectronSingleTrack.h index ef1c7b33474..3db865a69e4 100644 --- a/PWGEM/Dilepton/Utils/MlResponseDielectronSingleTrack.h +++ b/PWGEM/Dilepton/Utils/MlResponseDielectronSingleTrack.h @@ -17,12 +17,12 @@ #ifndef PWGEM_DILEPTON_UTILS_MLRESPONSEDIELECTRONSINGLETRACK_H_ #define PWGEM_DILEPTON_UTILS_MLRESPONSEDIELECTRONSINGLETRACK_H_ +#include "Tools/ML/MlResponse.h" + #include #include #include -#include "Tools/ML/MlResponse.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 @@ -76,6 +76,16 @@ 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 corresponding GETTER1 and GETTER2 from track. +#define CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_RELDIFF(FEATURE, GETTER1, GETTER2) \ + case static_cast(InputFeaturesDielectronSingleTrack::FEATURE): { \ + inputFeature = (track.GETTER2() - track.GETTER1()) / track.GETTER1(); \ + 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 @@ -104,6 +114,7 @@ enum class InputFeaturesDielectronSingleTrack : uint8_t { tpcNClsShared, tpcChi2NCl, tpcInnerParam, + reldiffp, tpcSignal, tpcNSigmaEl, tpcNSigmaMu, @@ -225,6 +236,7 @@ class MlResponseDielectronSingleTrack : public MlResponse CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcNClsShared); CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcChi2NCl); CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcInnerParam); + CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_RELDIFF(reldiffp, p, tpcInnerParam); CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcSignal); CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcNSigmaEl); CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK(tpcNSigmaMu); @@ -372,6 +384,7 @@ class MlResponseDielectronSingleTrack : public MlResponse FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcNClsShared), FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcChi2NCl), FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcInnerParam), + FILL_MAP_DIELECTRON_SINGLE_TRACK(reldiffp), FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcSignal), FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcNSigmaEl), FILL_MAP_DIELECTRON_SINGLE_TRACK(tpcNSigmaMu), @@ -475,6 +488,7 @@ class MlResponseDielectronSingleTrack : public MlResponse #undef CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_SQRT #undef CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_COS #undef CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_TPCTOF +#undef CHECK_AND_FILL_DIELECTRON_SINGLE_TRACK_RELDIFF #undef CHECK_AND_FILL_DIELECTRON_COLLISION #endif // PWGEM_DILEPTON_UTILS_MLRESPONSEDIELECTRONSINGLETRACK_H_ diff --git a/PWGEM/Dilepton/Utils/MlResponseO2Track.h b/PWGEM/Dilepton/Utils/MlResponseO2Track.h new file mode 100644 index 00000000000..cd9be049af6 --- /dev/null +++ b/PWGEM/Dilepton/Utils/MlResponseO2Track.h @@ -0,0 +1,333 @@ +// 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 MlResponseO2Track.h +/// \brief Class to compute the ML response for dielectron analyses at the single track level +/// \author Daniel Samitz , SMI Vienna +/// Elisa Meninno, , SMI Vienna + +#ifndef PWGEM_DILEPTON_UTILS_MLRESPONSEO2TRACK_H_ +#define PWGEM_DILEPTON_UTILS_MLRESPONSEO2TRACK_H_ + +#include "Tools/ML/MlResponse.h" + +#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_O2_TRACK(FEATURE) \ + { \ + #FEATURE, static_cast(InputFeaturesO2Track::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=FEATURE from track +#define CHECK_AND_FILL_O2_TRACK(GETTER) \ + case static_cast(InputFeaturesO2Track::GETTER): { \ + inputFeature = track.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 corresponding GETTER=FEATURE from track +#define CHECK_AND_FILL_O2_TRACKPARCOV(FEATURE, GETTER) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + inputFeature = trackParCov.GETTER(); \ + break; \ + } + +// mean ITS cluster size +#define CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE(FEATURE, v1, v2) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + int nsize = 0; \ + int ncls = 0; \ + for (int il = v1; il < v2; il++) { \ + nsize += track.itsClsSizeInLayer(il); \ + if (nsize > 0) { \ + ncls++; \ + } \ + } \ + inputFeature = static_cast(nsize) / static_cast(ncls); \ + break; \ + } + +// mean ITS cluster size x cos(lambda) +#define CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE_COS(FEATURE, v1, v2) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + int nsize = 0; \ + int ncls = 0; \ + for (int il = v1; il < v2; il++) { \ + nsize += track.itsClsSizeInLayer(il); \ + if (nsize > 0) { \ + ncls++; \ + } \ + } \ + inputFeature = static_cast(nsize) / static_cast(ncls) * std::cos(std::atan(trackParCov.getTgl())); \ + break; \ + } + +// TPC+TOF combined nSigma +#define CHECK_AND_FILL_O2_TRACK_TPCTOF(FEATURE, GETTER1, GETTER2, GETTER3) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + if (!track.GETTER3()) { \ + inputFeature = track.GETTER1(); \ + } else { \ + if (track.GETTER1() > 0) { \ + inputFeature = sqrt((pow(track.GETTER1(), 2) + pow(track.GETTER2(), 2)) / 2.); \ + } else { \ + inputFeature = (-1) * sqrt((pow(track.GETTER1(), 2) + pow(track.GETTER2(), 2)) / 2.); \ + } \ + } \ + 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 corresponding GETTER from track and applying a sqrt +#define CHECK_AND_FILL_O2_TRACK_SQRT(FEATURE, GETTER) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + inputFeature = sqrt(track.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 corresponding GETTER1 from track and multiplying with cos(atan(GETTER2)) +#define CHECK_AND_FILL_O2_TRACK_COS(FEATURE, GETTER1, GETTER2) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + inputFeature = track.GETTER1() * std::cos(std::atan(track.GETTER2())); \ + 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 corresponding GETTER1 and GETTER2 from track. +#define CHECK_AND_FILL_O2_TRACK_RELDIFF(FEATURE, GETTER1, GETTER2) \ + case static_cast(InputFeaturesO2Track::FEATURE): { \ + inputFeature = (track.GETTER2() - trackParCov.GETTER1()) / trackParCov.GETTER1(); \ + 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 corresponding GETTER=FEATURE from collision +#define CHECK_AND_FILL_DIELECTRON_COLLISION(GETTER) \ + case static_cast(InputFeaturesO2Track::GETTER): { \ + inputFeature = collision.GETTER(); \ + break; \ + } + +namespace o2::analysis +{ +// possible input features for ML +enum class InputFeaturesO2Track : uint8_t { + tpcInnerParam, + reldiffp, + tpcSignal, + tpcNSigmaEl, + tpcNSigmaMu, + tpcNSigmaPi, + tpcNSigmaKa, + tpcNSigmaPr, + beta, + tofNSigmaEl, + tofNSigmaMu, + tofNSigmaPi, + tofNSigmaKa, + tofNSigmaPr, + tpctofNSigmaEl, + tpctofNSigmaMu, + tpctofNSigmaPi, + tpctofNSigmaKa, + tpctofNSigmaPr, + tpcNClsFound, + tpcNClsCrossedRows, + hasITS, + hasTPC, + hasTRD, + hasTOF, + tgl, + p, + itsClusterSizes, + meanClusterSizeITS, + meanClusterSizeITSib, + meanClusterSizeITSob, + meanClusterSizeITSCos, + meanClusterSizeITSibCos, + meanClusterSizeITSobCos, + posZ, + numContrib, + trackOccupancyInTimeRange, + ft0cOccupancyInTimeRange, +}; + +template +class MlResponseO2Track : public MlResponse +{ + public: + /// Default constructor + MlResponseO2Track() = default; + /// Default destructor + virtual ~MlResponseO2Track() = default; + + template + float return_feature(uint8_t idx, T const& track, U const& trackParCov, V const& collision) + { + float inputFeature = 0.; + switch (idx) { + CHECK_AND_FILL_O2_TRACK(tpcInnerParam); + CHECK_AND_FILL_O2_TRACK_RELDIFF(reldiffp, getP, tpcInnerParam); + CHECK_AND_FILL_O2_TRACK(tpcSignal); + CHECK_AND_FILL_O2_TRACK(tpcNSigmaEl); + CHECK_AND_FILL_O2_TRACK(tpcNSigmaMu); + CHECK_AND_FILL_O2_TRACK(tpcNSigmaPi); + CHECK_AND_FILL_O2_TRACK(tpcNSigmaKa); + CHECK_AND_FILL_O2_TRACK(tpcNSigmaPr); + CHECK_AND_FILL_O2_TRACK(beta); + CHECK_AND_FILL_O2_TRACK(tofNSigmaEl); + CHECK_AND_FILL_O2_TRACK(tofNSigmaMu); + CHECK_AND_FILL_O2_TRACK(tofNSigmaPi); + CHECK_AND_FILL_O2_TRACK(tofNSigmaKa); + CHECK_AND_FILL_O2_TRACK(tofNSigmaPr); + CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaEl, tpcNSigmaEl, tofNSigmaEl, hasTOF); + CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaMu, tpcNSigmaMu, tofNSigmaMu, hasTOF); + CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaPi, tpcNSigmaPi, tofNSigmaPi, hasTOF); + CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaKa, tpcNSigmaKa, tofNSigmaKa, hasTOF); + CHECK_AND_FILL_O2_TRACK_TPCTOF(tpctofNSigmaPr, tpcNSigmaPr, tofNSigmaPr, hasTOF); + CHECK_AND_FILL_O2_TRACK(tpcNClsFound); + CHECK_AND_FILL_O2_TRACK(tpcNClsCrossedRows); + CHECK_AND_FILL_O2_TRACK(hasITS); + CHECK_AND_FILL_O2_TRACK(hasTPC); + CHECK_AND_FILL_O2_TRACK(hasTRD); + CHECK_AND_FILL_O2_TRACK(hasTOF); + CHECK_AND_FILL_O2_TRACKPARCOV(tgl, getTgl); + CHECK_AND_FILL_O2_TRACKPARCOV(p, getP); + CHECK_AND_FILL_O2_TRACK(itsClusterSizes); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE(meanClusterSizeITS, 0, 7); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE(meanClusterSizeITSib, 0, 3); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE(meanClusterSizeITSob, 3, 7); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE_COS(meanClusterSizeITSCos, 0, 7); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE_COS(meanClusterSizeITSibCos, 0, 3); + CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE_COS(meanClusterSizeITSobCos, 3, 7); + CHECK_AND_FILL_DIELECTRON_COLLISION(posZ); + CHECK_AND_FILL_DIELECTRON_COLLISION(numContrib); + CHECK_AND_FILL_DIELECTRON_COLLISION(trackOccupancyInTimeRange); + CHECK_AND_FILL_DIELECTRON_COLLISION(ft0cOccupancyInTimeRange); + } + return inputFeature; + } + + /// Method to get the input features vector needed for ML inference + /// \param track is the single track, \param collision is the collision + /// \return inputFeatures vector + template + std::vector getInputFeatures(T const& track, U const& trackParCov, V const& collision) + { + std::vector inputFeatures; + for (const auto& idx : MlResponse::mCachedIndices) { + float inputFeature = return_feature(idx, track, trackParCov, collision); + inputFeatures.emplace_back(inputFeature); + } + return inputFeatures; + } + + /// Method to get the value of variable chosen for binning + /// \param track is the single track, \param collision is the collision + /// \return binning variable + template + float getBinningFeature(T const& track, U const& trackParCov, V const& collision) + { + return return_feature(mCachedIndexBinning, track, trackParCov, collision); + } + + void cacheBinningIndex(std::string const& cfgBinningFeature) + { + setAvailableInputFeatures(); + if (MlResponse::mAvailableInputFeatures.count(cfgBinningFeature)) { + mCachedIndexBinning = MlResponse::mAvailableInputFeatures[cfgBinningFeature]; + } else { + LOG(fatal) << "Binning feature " << cfgBinningFeature << " not available! Please check your configurables."; + } + } + + protected: + /// Method to fill the map of available input features + void setAvailableInputFeatures() + { + MlResponse::mAvailableInputFeatures = { + FILL_MAP_O2_TRACK(tpcInnerParam), + FILL_MAP_O2_TRACK(reldiffp), + FILL_MAP_O2_TRACK(tpcSignal), + FILL_MAP_O2_TRACK(tpcNSigmaEl), + FILL_MAP_O2_TRACK(tpcNSigmaMu), + FILL_MAP_O2_TRACK(tpcNSigmaPi), + FILL_MAP_O2_TRACK(tpcNSigmaKa), + FILL_MAP_O2_TRACK(tpcNSigmaPr), + FILL_MAP_O2_TRACK(beta), + FILL_MAP_O2_TRACK(tofNSigmaEl), + FILL_MAP_O2_TRACK(tofNSigmaMu), + FILL_MAP_O2_TRACK(tofNSigmaPi), + FILL_MAP_O2_TRACK(tofNSigmaKa), + FILL_MAP_O2_TRACK(tofNSigmaPr), + FILL_MAP_O2_TRACK(tpctofNSigmaEl), + FILL_MAP_O2_TRACK(tpctofNSigmaMu), + FILL_MAP_O2_TRACK(tpctofNSigmaPi), + FILL_MAP_O2_TRACK(tpctofNSigmaKa), + FILL_MAP_O2_TRACK(tpctofNSigmaPr), + FILL_MAP_O2_TRACK(tpcNClsFound), + FILL_MAP_O2_TRACK(tpcNClsCrossedRows), + FILL_MAP_O2_TRACK(hasITS), + FILL_MAP_O2_TRACK(hasTPC), + FILL_MAP_O2_TRACK(hasTRD), + FILL_MAP_O2_TRACK(hasTOF), + FILL_MAP_O2_TRACK(tgl), + FILL_MAP_O2_TRACK(p), + FILL_MAP_O2_TRACK(itsClusterSizes), + FILL_MAP_O2_TRACK(meanClusterSizeITS), + FILL_MAP_O2_TRACK(meanClusterSizeITSib), + FILL_MAP_O2_TRACK(meanClusterSizeITSob), + FILL_MAP_O2_TRACK(meanClusterSizeITSCos), + FILL_MAP_O2_TRACK(meanClusterSizeITSibCos), + FILL_MAP_O2_TRACK(meanClusterSizeITSobCos), + FILL_MAP_O2_TRACK(posZ), + FILL_MAP_O2_TRACK(numContrib), + FILL_MAP_O2_TRACK(trackOccupancyInTimeRange), + FILL_MAP_O2_TRACK(ft0cOccupancyInTimeRange)}; + } + + uint8_t mCachedIndexBinning; // index correspondance between configurable and available input features +}; + +} // namespace o2::analysis + +#undef FILL_MAP_O2_TRACK +#undef CHECK_AND_FILL_O2_TRACK +#undef CHECK_AND_FILL_O2_TRACKPARCOV +#undef CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE +#undef CHECK_AND_FILL_O2_TRACK_MEAN_ITSCLUSTER_SIZE_COS +#undef CHECK_AND_FILL_O2_TRACK_SQRT +#undef CHECK_AND_FILL_O2_TRACK_COS +#undef CHECK_AND_FILL_O2_TRACK_TPCTOF +#undef CHECK_AND_FILL_O2_TRACK_RELDIFF +#undef CHECK_AND_FILL_DIELECTRON_COLLISION + +#endif // PWGEM_DILEPTON_UTILS_MLRESPONSEO2TRACK_H_