[PWGDQ] Usage of ML models for MFT-Muon matching

PWGDQ/Core/MuonMatchingMlResponse.h

@@ -0,0 +1,377 @@
+// 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 MuonMatchingMlResponse.h
+/// \brief Class to compute the ML response for MFT-Muon matching
+/// \author Maurice Coquet <maurice.louis.coquet@cern.ch>
+
+#ifndef PWGDQ_CORE_MUONMATCHINGMLRESPONSE_H_
+#define PWGDQ_CORE_MUONMATCHINGMLRESPONSE_H_
+
+#include "Tools/ML/MlResponse.h"
+
+#include <map>
+#include <string>
+#include <vector>
+
+// 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_MFTMUON_MATCH(FEATURE)                                \
+  {                                                                    \
+    #FEATURE, static_cast<uint8_t>(InputFeaturesMFTMuonMatch::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_MUON_TRACK(FEATURE, GETTER)                 \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::FEATURE): { \
+    inputFeature = muon.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_MFT_TRACK(FEATURE, GETTER)                  \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::FEATURE): { \
+    inputFeature = mft.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_MUON_COV(FEATURE, GETTER)                   \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::FEATURE): { \
+    inputFeature = muoncov.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_MFT_COV(FEATURE, GETTER)                    \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::FEATURE): { \
+    inputFeature = mftcov.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 and GETTER2 from track.
+#define CHECK_AND_FILL_MFTMUON_DIFF(FEATURE, GETTER1, GETTER2)     \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::FEATURE): { \
+    inputFeature = (mft.GETTER2() - muon.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_MFTMUON_COLLISION(GETTER)                  \
+  case static_cast<uint8_t>(InputFeaturesMFTMuonMatch::GETTER): { \
+    inputFeature = collision.GETTER();                            \
+    break;                                                        \
+  }
+
+namespace o2::analysis
+{
+// possible input features for ML
+enum class InputFeaturesMFTMuonMatch : uint8_t {
+  zMatching,
+  xMFT,
+  yMFT,
+  qOverptMFT,
+  tglMFT,
+  phiMFT,
+  dcaXY,
+  dcaZ,
+  chi2MFT,
+  nClustersMFT,
+  xMCH,
+  yMCH,
+  qOverptMCH,
+  tglMCH,
+  phiMCH,
+  nClustersMCH,
+  chi2MCH,
+  pdca,
+  cXXMFT,
+  cXYMFT,
+  cYYMFT,
+  cPhiYMFT,
+  cPhiXMFT,
+  cPhiPhiMFT,
+  cTglYMFT,
+  cTglXMFT,
+  cTglPhiMFT,
+  cTglTglMFT,
+  c1PtYMFT,
+  c1PtXMFT,
+  c1PtPhiMFT,
+  c1PtTglMFT,
+  c1Pt21Pt2MFT,
+  cXXMCH,
+  cXYMCH,
+  cYYMCH,
+  cPhiYMCH,
+  cPhiXMCH,
+  cPhiPhiMCH,
+  cTglYMCH,
+  cTglXMCH,
+  cTglPhiMCH,
+  cTglTglMCH,
+  c1PtYMCH,
+  c1PtXMCH,
+  c1PtPhiMCH,
+  c1PtTglMCH,
+  c1Pt21Pt2MCH,
+  deltaX,
+  deltaY,
+  deltaPhi,
+  deltaEta,
+  deltaPt,
+  posX,
+  posY,
+  posZ,
+  numContrib,
+  trackOccupancyInTimeRange,
+  ft0cOccupancyInTimeRange,
+  multFT0A,
+  multFT0C,
+  multNTracksPV,
+  multNTracksPVeta1,
+  multNTracksPVetaHalf,
+  isInelGt0,
+  isInelGt1,
+  multFT0M,
+  centFT0M,
+  centFT0A,
+  centFT0C,
+  chi2MCHMFT
+};
+
+template <typename TypeOutputScore = float>
+class MlResponseMFTMuonMatch : public MlResponse<TypeOutputScore>
+{
+ public:
+  /// Default constructor
+  MlResponseMFTMuonMatch() = default;
+  /// Default destructor
+  virtual ~MlResponseMFTMuonMatch() = default;
+
+  template <typename T1, typename T2, typename C1, typename C2, typename U>
+  float returnFeature(uint8_t idx, T1 const& muon, T2 const& mft, C1 const& muoncov, C2 const& mftcov, U const& collision)
+  {
+    float inputFeature = 0.;
+    switch (idx) {
+      CHECK_AND_FILL_MFT_TRACK(zMatching, z);
+      CHECK_AND_FILL_MFT_TRACK(xMFT, x);
+      CHECK_AND_FILL_MFT_TRACK(yMFT, y);
+      CHECK_AND_FILL_MFT_TRACK(qOverptMFT, signed1Pt);
+      CHECK_AND_FILL_MFT_TRACK(tglMFT, tgl);
+      CHECK_AND_FILL_MFT_TRACK(phiMFT, phi);
+      CHECK_AND_FILL_MFT_TRACK(chi2MFT, chi2);
+      CHECK_AND_FILL_MFT_TRACK(nClustersMFT, nClusters);
+      CHECK_AND_FILL_MUON_TRACK(dcaXY, fwddcaXY);
+      CHECK_AND_FILL_MUON_TRACK(dcaZ, fwddcaz);
+      CHECK_AND_FILL_MUON_TRACK(xMCH, x);
+      CHECK_AND_FILL_MUON_TRACK(yMCH, y);
+      CHECK_AND_FILL_MUON_TRACK(qOverptMCH, signed1Pt);
+      CHECK_AND_FILL_MUON_TRACK(tglMCH, tgl);
+      CHECK_AND_FILL_MUON_TRACK(phiMCH, phi);
+      CHECK_AND_FILL_MUON_TRACK(nClustersMCH, nClusters);
+      CHECK_AND_FILL_MUON_TRACK(chi2MCH, chi2);
+      CHECK_AND_FILL_MUON_TRACK(pdca, pDca);
+      CHECK_AND_FILL_MFT_COV(cXXMFT, cXX);
+      CHECK_AND_FILL_MFT_COV(cXYMFT, cXY);
+      CHECK_AND_FILL_MFT_COV(cYYMFT, cYY);
+      CHECK_AND_FILL_MFT_COV(cPhiYMFT, cPhiY);
+      CHECK_AND_FILL_MFT_COV(cPhiXMFT, cPhiX);
+      CHECK_AND_FILL_MFT_COV(cPhiPhiMFT, cPhiPhi);
+      CHECK_AND_FILL_MFT_COV(cTglYMFT, cTglY);
+      CHECK_AND_FILL_MFT_COV(cTglXMFT, cTglX);
+      CHECK_AND_FILL_MFT_COV(cTglPhiMFT, cTglPhi);
+      CHECK_AND_FILL_MFT_COV(cTglTglMFT, cTglTgl);
+      CHECK_AND_FILL_MFT_COV(c1PtYMFT, c1PtY);
+      CHECK_AND_FILL_MFT_COV(c1PtXMFT, c1PtX);
+      CHECK_AND_FILL_MFT_COV(c1PtPhiMFT, c1PtPhi);
+      CHECK_AND_FILL_MFT_COV(c1PtTglMFT, c1PtTgl);
+      CHECK_AND_FILL_MFT_COV(c1Pt21Pt2MFT, c1Pt21Pt2);
+      CHECK_AND_FILL_MUON_COV(cXXMCH, cXX);
+      CHECK_AND_FILL_MUON_COV(cXYMCH, cXY);
+      CHECK_AND_FILL_MUON_COV(cYYMCH, cYY);
+      CHECK_AND_FILL_MUON_COV(cPhiYMCH, cPhiY);
+      CHECK_AND_FILL_MUON_COV(cPhiXMCH, cPhiX);
+      CHECK_AND_FILL_MUON_COV(cPhiPhiMCH, cPhiPhi);
+      CHECK_AND_FILL_MUON_COV(cTglYMCH, cTglY);
+      CHECK_AND_FILL_MUON_COV(cTglXMCH, cTglX);
+      CHECK_AND_FILL_MUON_COV(cTglPhiMCH, cTglPhi);
+      CHECK_AND_FILL_MUON_COV(cTglTglMCH, cTglTgl);
+      CHECK_AND_FILL_MUON_COV(c1PtYMCH, c1PtY);
+      CHECK_AND_FILL_MUON_COV(c1PtXMCH, c1PtX);
+      CHECK_AND_FILL_MUON_COV(c1PtPhiMCH, c1PtPhi);
+      CHECK_AND_FILL_MUON_COV(c1PtTglMCH, c1PtTgl);
+      CHECK_AND_FILL_MUON_COV(c1Pt21Pt2MCH, c1Pt21Pt2);
+      CHECK_AND_FILL_MFTMUON_COLLISION(posX);
+      CHECK_AND_FILL_MFTMUON_COLLISION(posY);
+      CHECK_AND_FILL_MFTMUON_COLLISION(posZ);
+      CHECK_AND_FILL_MFTMUON_COLLISION(numContrib);
+      CHECK_AND_FILL_MFTMUON_COLLISION(trackOccupancyInTimeRange);
+      CHECK_AND_FILL_MFTMUON_COLLISION(ft0cOccupancyInTimeRange);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multFT0A);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multFT0C);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multNTracksPV);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multNTracksPVeta1);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multNTracksPVetaHalf);
+      CHECK_AND_FILL_MFTMUON_COLLISION(isInelGt0);
+      CHECK_AND_FILL_MFTMUON_COLLISION(isInelGt1);
+      CHECK_AND_FILL_MFTMUON_COLLISION(multFT0M);
+      CHECK_AND_FILL_MFTMUON_COLLISION(centFT0M);
+      CHECK_AND_FILL_MFTMUON_COLLISION(centFT0A);
+      CHECK_AND_FILL_MFTMUON_COLLISION(centFT0C);
+      CHECK_AND_FILL_MUON_TRACK(chi2MCHMFT, chi2MatchMCHMFT);
+    }
+    return inputFeature;
+  }
+
+  template <typename T1>
+  float returnFeatureTest(uint8_t idx, T1 const& muon)
+  {
+    float inputFeature = 0.;
+    switch (idx) {
+      CHECK_AND_FILL_MUON_TRACK(chi2MCHMFT, chi2MatchMCHMFT);
+    }
+    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 <typename T1, typename T2, typename C1, typename C2, typename U>
+  std::vector<float> getInputFeatures(T1 const& muon, T2 const& mft, C1 const& muoncov, C2 const& mftcov, U const& collision)
+  {
+    std::vector<float> inputFeatures;
+    for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
+      float inputFeature = returnFeature(idx, muon, mft, muoncov, mftcov, collision);
+      inputFeatures.emplace_back(inputFeature);
+    }
+    return inputFeatures;
+  }
+
+  template <typename T1>
+  std::vector<float> getInputFeaturesTest(T1 const& muon)
+  {
+    std::vector<float> inputFeatures;
+    for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
+      float inputFeature = returnFeatureTest(idx, muon);
+      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 <typename T1, typename T2, typename C1, typename C2, typename U>
+  float getBinningFeature(T1 const& muon, T2 const& mft, C1 const& muoncov, C2 const& mftcov, U const& collision)
+  {
+    return returnFeature(mCachedIndexBinning, muon, mft, muoncov, mftcov, collision);
+  }
+
+  void cacheBinningIndex(std::string const& cfgBinningFeature)
+  {
+    setAvailableInputFeatures();
+    if (MlResponse<TypeOutputScore>::mAvailableInputFeatures.count(cfgBinningFeature)) {
+      mCachedIndexBinning = MlResponse<TypeOutputScore>::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<TypeOutputScore>::mAvailableInputFeatures = {
+      FILL_MAP_MFTMUON_MATCH(zMatching),
+      FILL_MAP_MFTMUON_MATCH(xMFT),
+      FILL_MAP_MFTMUON_MATCH(yMFT),
+      FILL_MAP_MFTMUON_MATCH(qOverptMFT),
+      FILL_MAP_MFTMUON_MATCH(tglMFT),
+      FILL_MAP_MFTMUON_MATCH(phiMFT),
+      FILL_MAP_MFTMUON_MATCH(dcaXY),
+      FILL_MAP_MFTMUON_MATCH(dcaZ),
+      FILL_MAP_MFTMUON_MATCH(chi2MFT),
+      FILL_MAP_MFTMUON_MATCH(nClustersMFT),
+      FILL_MAP_MFTMUON_MATCH(xMCH),
+      FILL_MAP_MFTMUON_MATCH(yMCH),
+      FILL_MAP_MFTMUON_MATCH(qOverptMCH),
+      FILL_MAP_MFTMUON_MATCH(tglMCH),
+      FILL_MAP_MFTMUON_MATCH(phiMCH),
+      FILL_MAP_MFTMUON_MATCH(nClustersMCH),
+      FILL_MAP_MFTMUON_MATCH(chi2MCH),
+      FILL_MAP_MFTMUON_MATCH(pdca),
+      FILL_MAP_MFTMUON_MATCH(cXXMFT),
+      FILL_MAP_MFTMUON_MATCH(cXYMFT),
+      FILL_MAP_MFTMUON_MATCH(cYYMFT),
+      FILL_MAP_MFTMUON_MATCH(cPhiYMFT),
+      FILL_MAP_MFTMUON_MATCH(cPhiXMFT),
+      FILL_MAP_MFTMUON_MATCH(cPhiPhiMFT),
+      FILL_MAP_MFTMUON_MATCH(cTglYMFT),
+      FILL_MAP_MFTMUON_MATCH(cTglXMFT),
+      FILL_MAP_MFTMUON_MATCH(cTglPhiMFT),
+      FILL_MAP_MFTMUON_MATCH(cTglTglMFT),
+      FILL_MAP_MFTMUON_MATCH(c1PtYMFT),
+      FILL_MAP_MFTMUON_MATCH(c1PtXMFT),
+      FILL_MAP_MFTMUON_MATCH(c1PtPhiMFT),
+      FILL_MAP_MFTMUON_MATCH(c1PtTglMFT),
+      FILL_MAP_MFTMUON_MATCH(c1Pt21Pt2MFT),
+      FILL_MAP_MFTMUON_MATCH(cXXMCH),
+      FILL_MAP_MFTMUON_MATCH(cXYMCH),
+      FILL_MAP_MFTMUON_MATCH(cYYMCH),
+      FILL_MAP_MFTMUON_MATCH(cPhiYMCH),
+      FILL_MAP_MFTMUON_MATCH(cPhiXMCH),
+      FILL_MAP_MFTMUON_MATCH(cPhiPhiMCH),
+      FILL_MAP_MFTMUON_MATCH(cTglYMCH),
+      FILL_MAP_MFTMUON_MATCH(cTglXMCH),
+      FILL_MAP_MFTMUON_MATCH(cTglPhiMCH),
+      FILL_MAP_MFTMUON_MATCH(cTglTglMCH),
+      FILL_MAP_MFTMUON_MATCH(c1PtYMCH),
+      FILL_MAP_MFTMUON_MATCH(c1PtXMCH),
+      FILL_MAP_MFTMUON_MATCH(c1PtPhiMCH),
+      FILL_MAP_MFTMUON_MATCH(c1PtTglMCH),
+      FILL_MAP_MFTMUON_MATCH(c1Pt21Pt2MCH),
+      FILL_MAP_MFTMUON_MATCH(chi2MCHMFT)};
+  }
+
+  uint8_t mCachedIndexBinning; // index correspondance between configurable and available input features
+};
+
+} // namespace o2::analysis
+
+#undef FILL_MAP_MFTMUON_MAP
+#undef CHECK_AND_FILL_MUON_TRACK
+#undef CHECK_AND_FILL_MFT_TRACK
+#undef CHECK_AND_FILL_MUON_COV
+#undef CHECK_AND_FILL_MFT_COV
+#undef CHECK_AND_FILL_MFTMUON_DIFF
+#undef CHECK_AND_FILL_MFTMUON_COLLISION
+
+#endif // PWGDQ_CORE_MUONMATCHINGMLRESPONSE_H_