[PWGEM] Add and update ML-based photon cuts

PWGEM/PhotonMeson/Core/Pi0EtaToGammaGamma.h

@@ -147,7 +147,6 @@ struct Pi0EtaToGammaGamma {
     o2::framework::Configurable<bool> cfg_load_ml_models_from_ccdb{"cfg_load_ml_models_from_ccdb", true, "flag to load ML models from CCDB"};
     o2::framework::Configurable<int> cfg_timestamp_ccdb{"cfg_timestamp_ccdb", -1, "timestamp for CCDB"};
     o2::framework::Configurable<int> cfg_nclasses_ml{"cfg_nclasses_ml", static_cast<int>(o2::analysis::em_cuts_ml::NCutScores), "number of classes for ML"};
-    o2::framework::Configurable<std::string> cfg_cent_type_ml{"cfg_cent_type_ml", "CentFT0C", "centrality type for 2D ML application: CentFT0C, CentFT0M, or CentFT0A"};
     o2::framework::Configurable<std::vector<int>> cfg_cut_dir_ml{"cfg_cut_dir_ml", std::vector<int>{o2::analysis::em_cuts_ml::vecCutDir}, "cut direction for ML"};
     o2::framework::Configurable<std::vector<std::string>> cfg_input_feature_names{"cfg_input_feature_names", std::vector<std::string>{"feature1", "feature2"}, "input feature names for ML models"};
     o2::framework::Configurable<std::vector<std::string>> cfg_model_paths_ccdb{"cfg_model_paths_ccdb", std::vector<std::string>{"path_ccdb/BDT_PCM/"}, "CCDB paths for ML models"};
@@ -508,7 +507,9 @@ struct Pi0EtaToGammaGamma {
     fV0PhotonCut.SetNClassesMl(pcmcuts.cfg_nclasses_ml);
     fV0PhotonCut.SetMlTimestampCCDB(pcmcuts.cfg_timestamp_ccdb);
     fV0PhotonCut.SetCcdbUrl(ccdburl);
-    fV0PhotonCut.SetCentralityTypeMl(pcmcuts.cfg_cent_type_ml);
+    CentType mCentralityTypeMlEnum;
+    mCentralityTypeMlEnum = static_cast<CentType>(cfgCentEstimator.value);
+    fV0PhotonCut.SetCentralityTypeMl(mCentralityTypeMlEnum);
     fV0PhotonCut.SetCutDirMl(pcmcuts.cfg_cut_dir_ml);
     fV0PhotonCut.SetMlModelPathsCCDB(pcmcuts.cfg_model_paths_ccdb);
     fV0PhotonCut.SetMlOnnxFileNames(pcmcuts.cfg_onnx_file_names);
@@ -702,7 +703,6 @@ struct Pi0EtaToGammaGamma {
   {
     for (const auto& collision : collisions) {
       initCCDB(collision);
-      fV0PhotonCut.SetCentrality(collision.centFT0A(), collision.centFT0C(), collision.centFT0M());
       int ndiphoton = 0;
       if ((pairtype == o2::aod::pwgem::photonmeson::photonpair::PairType::kPHOSPHOS || pairtype == o2::aod::pwgem::photonmeson::photonpair::PairType::kPCMPHOS) && !collision.alias_bit(triggerAliases::kTVXinPHOS)) {
         continue;
@@ -718,6 +718,7 @@ struct Pi0EtaToGammaGamma {
       }
 
       const float centralities[3] = {collision.centFT0M(), collision.centFT0A(), collision.centFT0C()};
+      fV0PhotonCut.SetCentrality(centralities[cfgCentEstimator]);
       if (centralities[cfgCentEstimator] < cfgCentMin || cfgCentMax < centralities[cfgCentEstimator]) {
         continue;
       }

PWGEM/PhotonMeson/Core/Pi0EtaToGammaGammaMC.h

@@ -136,7 +136,6 @@ struct Pi0EtaToGammaGammaMC {
     o2::framework::Configurable<bool> cfg_load_ml_models_from_ccdb{"cfg_load_ml_models_from_ccdb", true, "flag to load ML models from CCDB"};
     o2::framework::Configurable<int> cfg_timestamp_ccdb{"cfg_timestamp_ccdb", -1, "timestamp for CCDB"};
     o2::framework::Configurable<int> cfg_nclasses_ml{"cfg_nclasses_ml", static_cast<int>(o2::analysis::em_cuts_ml::NCutScores), "number of classes for ML"};
-    o2::framework::Configurable<std::string> cfg_cent_type_ml{"cfg_cent_type_ml", "CentFT0C", "centrality type for 2D ML application: CentFT0C, CentFT0M, or CentFT0A"};
     o2::framework::Configurable<std::vector<int>> cfg_cut_dir_ml{"cfg_cut_dir_ml", std::vector<int>{o2::analysis::em_cuts_ml::vecCutDir}, "cut direction for ML"};
     o2::framework::Configurable<std::vector<std::string>> cfg_input_feature_names{"cfg_input_feature_names", std::vector<std::string>{"feature1", "feature2"}, "input feature names for ML models"};
     o2::framework::Configurable<std::vector<std::string>> cfg_model_paths_ccdb{"cfg_model_paths_ccdb", std::vector<std::string>{"path_ccdb/BDT_PCM/"}, "CCDB paths for ML models"};
@@ -348,7 +347,9 @@ struct Pi0EtaToGammaGammaMC {
     fV0PhotonCut.SetNClassesMl(pcmcuts.cfg_nclasses_ml);
     fV0PhotonCut.SetMlTimestampCCDB(pcmcuts.cfg_timestamp_ccdb);
     fV0PhotonCut.SetCcdbUrl(ccdburl);
-    fV0PhotonCut.SetCentralityTypeMl(pcmcuts.cfg_cent_type_ml);
+    CentType mCentralityTypeMlEnum;
+    mCentralityTypeMlEnum = static_cast<CentType>(cfgCentEstimator.value);
+    fV0PhotonCut.SetCentralityTypeMl(mCentralityTypeMlEnum);
     fV0PhotonCut.SetCutDirMl(pcmcuts.cfg_cut_dir_ml);
     fV0PhotonCut.SetMlModelPathsCCDB(pcmcuts.cfg_model_paths_ccdb);
     fV0PhotonCut.SetMlOnnxFileNames(pcmcuts.cfg_onnx_file_names);
@@ -560,7 +561,6 @@ struct Pi0EtaToGammaGammaMC {
   {
     for (auto& collision : collisions) {
       initCCDB(collision);
-      fV0PhotonCut.SetCentrality(collision.centFT0A(), collision.centFT0C(), collision.centFT0M());
       if ((pairtype == o2::aod::pwgem::photonmeson::photonpair::PairType::kPHOSPHOS || pairtype == o2::aod::pwgem::photonmeson::photonpair::PairType::kPCMPHOS) && !collision.alias_bit(triggerAliases::kTVXinPHOS)) {
         continue;
       }
@@ -575,6 +575,7 @@ struct Pi0EtaToGammaGammaMC {
       }
 
       const float centralities[3] = {collision.centFT0M(), collision.centFT0A(), collision.centFT0C()};
+      fV0PhotonCut.SetCentrality(centralities[cfgCentEstimator]);
       if (centralities[cfgCentEstimator] < cfgCentMin || cfgCentMax < centralities[cfgCentEstimator]) {
         continue;
       }

PWGEM/PhotonMeson/Core/V0PhotonCandidate.h

@@ -23,11 +23,19 @@
 
 #include <KFParticle.h>
 
+enum CentType : uint8_t {
+  CentFT0M = 0,
+  CentFT0A = 1,
+  CentFT0C = 2
+};
+
 struct V0PhotonCandidate {
 
  public:
-  // Constructor for photonconversionbuilder
-  V0PhotonCandidate(const KFParticle& v0, const KFParticle& pos, const KFParticle& ele, const auto& collision, float cospa, float d_bz) : cospa(cospa)
+  // Empty Constructor
+  V0PhotonCandidate() = default;
+  // Set method for photonconversionbuilder
+  void setPhotonCandidate(const KFParticle& v0, const KFParticle& pos, const KFParticle& ele, const auto& collision, float cospa, float psipair, float phiv, CentType centType)
   {
     px = v0.GetPx();
     py = v0.GetPy();
@@ -57,18 +65,27 @@ struct V0PhotonCandidate {
 
     alpha = v0_alpha(posPx, posPy, posPz, elePx, elePy, elePz);
     qt = v0_qt(posPx, posPy, posPz, elePx, elePy, elePz);
-    int posSign = (pos.GetQ() > 0) - (pos.GetQ() < 0);
-    int eleSign = (ele.GetQ() > 0) - (ele.GetQ() < 0);
-    phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(posPx, posPy, posPz, elePx, elePy, elePz, posSign, eleSign, d_bz);
-    psipair = o2::aod::pwgem::dilepton::utils::pairutil::getPsiPair(posPx, posPy, posPz, elePx, elePy, elePz);
-
-    centFT0M = collision.centFT0M();
-    centFT0C = collision.centFT0C();
-    centFT0A = collision.centFT0A();
+
+    this->cospa = cospa;
+    this->psipair = psipair;
+    this->phiv = phiv;
+    this->centType = centType;
+
+    switch (centType) {
+      case CentType::CentFT0A:
+        cent = collision.centFT0A();
+        break;
+      case CentType::CentFT0C:
+        cent = collision.centFT0C();
+        break;
+      case CentType::CentFT0M:
+        cent = collision.centFT0M();
+        break;
+    }
   }
 
-  // Constructor for V0PhotonCut
-  V0PhotonCandidate(const auto& v0, const auto& pos, const auto& ele, float centFT0A, float centFT0C, float centFT0M, float d_bz) : centFT0A(centFT0A), centFT0C(centFT0C), centFT0M(centFT0M)
+  // Set-Method for V0PhotonCut
+  void setPhoton(const auto& v0, const auto& pos, const auto& ele, float cent, CentType centType)
   {
     px = v0.px();
     py = v0.py();
@@ -93,9 +110,14 @@ struct V0PhotonCandidate {
     cospa = v0.cospa();
     alpha = v0.alpha();
     qt = v0.qtarm();
-
-    phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(posPx, posPy, posPz, elePx, elePy, elePz, pos.sign(), ele.sign(), d_bz);
-    psipair = o2::aod::pwgem::dilepton::utils::pairutil::getPsiPair(posPx, posPy, posPz, elePx, elePy, elePz);
+    psipair = 999.f; // default if V0PhotonPhiVPsi table is not included
+    phiv = 999.f;    // default if V0PhotonPhiVPsi table is not included
+    if constexpr (requires { v0.psipair(); v0.phiv(); }) {
+      psipair = v0.psipair();
+      phiv = v0.phiv();
+    }
+    this->cent = cent;
+    this->centType = centType;
   }
 
   // Getter functions
@@ -119,10 +141,9 @@ struct V0PhotonCandidate {
   float getElePx() const { return elePx; }
   float getElePy() const { return elePy; }
   float getElePz() const { return elePz; }
-  float getCentFT0M() const { return centFT0M; }
-  float getCentFT0C() const { return centFT0C; }
-  float getCentFT0A() const { return centFT0A; }
+  float getCent() const { return cent; }
   float getPCA() const { return pca; }
+  CentType getCentType() const { return centType; }
 
  private:
   float px;
@@ -145,10 +166,9 @@ struct V0PhotonCandidate {
   float psipair;
   float cospa;
   float chi2ndf;
-  float centFT0A;
-  float centFT0C;
-  float centFT0M;
+  float cent;
   float pca;
+  CentType centType;
 };
 
 #endif // PWGEM_PHOTONMESON_CORE_V0PHOTONCANDIDATE_H_

PWGEM/PhotonMeson/Core/V0PhotonCut.cxx

@@ -306,20 +306,20 @@ void V0PhotonCut::SetCutsMl(const std::vector<double>& cuts)
 void V0PhotonCut::SetNClassesMl(int nClasses)
 {
   mNClassesMl = nClasses;
+  mOutputML.reserve(mNClassesMl);
   LOG(info) << "V0 Photon Cut, set number of classes ML: " << mNClassesMl;
 }
 
 void V0PhotonCut::SetNamesInputFeatures(const std::vector<std::string>& featureNames)
 {
   mNamesInputFeatures = featureNames;
+  mMlInputFeatures.reserve(mNamesInputFeatures.size());
   LOG(info) << "V0 Photon Cut, set ML input feature names with size:" << mNamesInputFeatures.size();
 }
 
-void V0PhotonCut::SetCentrality(float centFT0A, float centFT0C, float centFT0M)
+void V0PhotonCut::SetCentrality(float cent)
 {
-  mCentFT0A = centFT0A;
-  mCentFT0C = centFT0C;
-  mCentFT0M = centFT0M;
+  mCent = cent;
 }
 void V0PhotonCut::SetD_Bz(float d_bz)
 {
@@ -332,10 +332,10 @@ void V0PhotonCut::SetCutDirMl(const std::vector<int>& cutDirMl)
   LOG(info) << "V0 Photon Cut, set ML cut directions with size:" << mCutDirMl.size();
 }
 
-void V0PhotonCut::SetCentralityTypeMl(const std::string& centType)
+void V0PhotonCut::SetCentralityTypeMl(CentType centType)
 {
   mCentralityTypeMl = centType;
-  LOG(info) << "V0 Photon Cut, set centrality type ML: " << mCentralityTypeMl;
+  LOG(info) << "V0 Photon Cut, set centrality type ML: " << mCentralityTypeMl << " (0: CentFT0M, 1: CentFT0A, 2: CentFT0C)";
 }
 
 void V0PhotonCut::SetLabelsBinsMl(const std::vector<std::string>& labelsBins)

PWGEM/PhotonMeson/Core/V0PhotonCut.h

@@ -39,6 +39,7 @@
 #include <cstdint>
 #include <functional>
 #include <set>
+#include <span>
 #include <string>
 #include <utility>
 #include <vector>
@@ -144,7 +145,7 @@ static const std::vector<std::string> labelsCent = {
   "Cent bin 10"};
 
 // column labels
-static const std::vector<std::string> labelsCutScore = {"score primary photons", "score background"};
+static const std::vector<std::string> labelsCutScore = {"score background", "score primary photons"};
 } // namespace em_cuts_ml
 
 } // namespace o2::analysis
@@ -573,30 +574,22 @@ class V0PhotonCut : public TNamed
       }
     }
     if (mApplyMlCuts) {
-      if (!mEmMlResponse) {
+      if (mEmMlResponse == nullptr) {
         LOG(error) << "EM ML Response is not initialized!";
         return false;
       }
-      bool mIsSelectedMl = false;
-      std::vector<float> mOutputML;
-      V0PhotonCandidate v0photoncandidate(v0, pos, ele, mCentFT0A, mCentFT0C, mCentFT0M, mD_Bz);
-      std::vector<float> mlInputFeatures = mEmMlResponse->getInputFeatures(v0photoncandidate, pos, ele);
+      mIsSelectedMl = false;
+      mV0PhotonForMl.setPhoton(v0, pos, ele, mCent, mCentralityTypeMl);
+      mMlInputFeatures = mEmMlResponse->getInputFeatures(mV0PhotonForMl, pos, ele);
       if (mUse2DBinning) {
-        if (mCentralityTypeMl == "CentFT0C") {
-          mIsSelectedMl = mEmMlResponse->isSelectedMl(mlInputFeatures, v0photoncandidate.getPt(), v0photoncandidate.getCentFT0C(), mOutputML);
-        } else if (mCentralityTypeMl == "CentFT0A") {
-          mIsSelectedMl = mEmMlResponse->isSelectedMl(mlInputFeatures, v0photoncandidate.getPt(), v0photoncandidate.getCentFT0A(), mOutputML);
-        } else if (mCentralityTypeMl == "CentFT0M") {
-          mIsSelectedMl = mEmMlResponse->isSelectedMl(mlInputFeatures, v0photoncandidate.getPt(), v0photoncandidate.getCentFT0M(), mOutputML);
-        } else {
-          LOG(fatal) << "Unsupported centTypePCMMl: " << mCentralityTypeMl << " , please choose from CentFT0C, CentFT0A, CentFT0M.";
-        }
+        mIsSelectedMl = mEmMlResponse->isSelectedMl(mMlInputFeatures, mV0PhotonForMl.getPt(), mV0PhotonForMl.getCent(), mOutputML);
       } else {
-        mIsSelectedMl = mEmMlResponse->isSelectedMl(mlInputFeatures, v0photoncandidate.getPt(), mOutputML);
+        mIsSelectedMl = mEmMlResponse->isSelectedMl(mMlInputFeatures, mV0PhotonForMl.getPt(), mOutputML);
       }
       if (!mIsSelectedMl) {
         return false;
       }
+      mMlBDTScores = std::span<float>(mOutputML.data(), mOutputML.size());
     }
     if (doQA) {
       fillAfterPhotonHistogram(v0, pos, ele, fRegistry);
@@ -845,7 +838,7 @@ class V0PhotonCut : public TNamed
 
   void initV0MlModels(o2::ccdb::CcdbApi& ccdbApi)
   {
-    if (!mEmMlResponse) {
+    if (mEmMlResponse == nullptr) {
       mEmMlResponse = new o2::analysis::EmMlResponsePCM<float>();
     }
     if (mUse2DBinning) {
@@ -899,6 +892,11 @@ class V0PhotonCut : public TNamed
     mEmMlResponse->init();
   }
 
+  const std::span<float> getBDTValue() const
+  {
+    return mMlBDTScores;
+  }
+
   template <o2::soa::is_iterator TMCPhoton>
   bool IsConversionPointInAcceptance(TMCPhoton const& mcphoton, float convRadius) const
   {
@@ -968,10 +966,10 @@ class V0PhotonCut : public TNamed
   void SetLoadMlModelsFromCCDB(bool flag = true);
   void SetNClassesMl(int nClasses);
   void SetMlTimestampCCDB(int timestamp);
-  void SetCentrality(float centFT0A, float centFT0C, float centFT0M);
+  void SetCentralityTypeMl(CentType centType);
+  void SetCentrality(float cent);
   void SetD_Bz(float d_bz);
   void SetCcdbUrl(const std::string& url = "http://alice-ccdb.cern.ch");
-  void SetCentralityTypeMl(const std::string& centType);
   void SetCutDirMl(const std::vector<int>& cutDirMl);
   void SetMlModelPathsCCDB(const std::vector<std::string>& modelPaths);
   void SetMlOnnxFileNames(const std::vector<std::string>& onnxFileNamesVec);
@@ -1011,22 +1009,25 @@ class V0PhotonCut : public TNamed
   bool mLoadMlModelsFromCCDB{true};
   int mTimestampCCDB{-1};
   int mNClassesMl{static_cast<int>(o2::analysis::em_cuts_ml::NCutScores)};
-  float mCentFT0A{0.f};
-  float mCentFT0C{0.f};
-  float mCentFT0M{0.f};
+  float mCent{0.f};
   float mD_Bz{0.f};
   std::string mCcdbUrl{"http://alice-ccdb.cern.ch"};
-  std::string mCentralityTypeMl{"CentFT0C"};
   std::vector<int> mCutDirMl{std::vector<int>{o2::analysis::em_cuts_ml::vecCutDir}};
   std::vector<std::string> mModelPathsCCDB{std::vector<std::string>{"path_ccdb/BDT_PCM/"}};
   std::vector<std::string> mOnnxFileNames{std::vector<std::string>{"ModelHandler_onnx_PCM.onnx"}};
   std::vector<std::string> mNamesInputFeatures{std::vector<std::string>{"feature1", "feature2"}};
   std::vector<std::string> mLabelsBinsMl{std::vector<std::string>{"bin 0", "bin 1"}};
-  std::vector<std::string> mLabelsCutScoresMl{std::vector<std::string>{"score primary photons", "score background"}};
+  std::vector<std::string> mLabelsCutScoresMl{std::vector<std::string>{o2::analysis::em_cuts_ml::labelsCutScore}};
   std::vector<double> mBinsPtMl{std::vector<double>{o2::analysis::em_cuts_ml::vecBinsPt}};
   std::vector<double> mBinsCentMl{std::vector<double>{o2::analysis::em_cuts_ml::vecBinsCent}};
   std::vector<double> mCutsMlFlat{std::vector<double>{0.5}};
   o2::analysis::EmMlResponsePCM<float>* mEmMlResponse{nullptr};
+  mutable bool mIsSelectedMl{false};
+  mutable std::vector<float> mOutputML{};
+  mutable std::vector<float> mMlInputFeatures{};
+  mutable std::span<float> mMlBDTScores{};
+  CentType mCentralityTypeMl{CentType::CentFT0C};
+  mutable V0PhotonCandidate mV0PhotonForMl;
 
   // pid cuts
   float mMinTPCNsigmaEl{-5}, mMaxTPCNsigmaEl{+5};

PWGEM/PhotonMeson/DataModel/gammaTables.h

@@ -416,15 +416,16 @@ using EMPrimaryElectronsFromDalitz = EMPrimaryElectronsFromDalitz_001;
 // iterators
 using EMPrimaryElectronFromDalitz = EMPrimaryElectronsFromDalitz::iterator;
 
-namespace v0photonsphiv
+namespace v0photonsphivpsi
 {
 DECLARE_SOA_INDEX_COLUMN(EMEvent, emevent); //!
 DECLARE_SOA_COLUMN(PhiV, phiv, float);      //!
-} // namespace v0photonsphiv
-DECLARE_SOA_TABLE(V0PhotonsPhiV, "AOD", "V0PHOTONPHIV", //!
-                  o2::soa::Index<>, v0photonsphiv::PhiV);
+DECLARE_SOA_COLUMN(PsiPair, psipair, float);
+} // namespace v0photonsphivpsi
+DECLARE_SOA_TABLE(V0PhotonsPhiVPsi, "AOD", "V0PHOTONPHIVPSI", //!
+                  o2::soa::Index<>, v0photonsphivpsi::PhiV, v0photonsphivpsi::PsiPair);
 // iterators
-using V0PhotonsPhiV = V0PhotonsPhiV;
+using V0PhotonsPhiVPsi = V0PhotonsPhiVPsi;
 
 namespace dalitzee
 {