From 29d5eaed52946c28504401775b035e0acf820c6e Mon Sep 17 00:00:00 2001
From: mmusich <marco.musich@cern.ch>
Date: Wed, 5 Jun 2024 22:26:14 +0200
Subject: [PATCH] move PhotonXGBoostEstimator to use pat::XGBooster

---
 RecoEgamma/PhotonIdentification/BuildFile.xml | 16 ++---
 .../interface/PhotonXGBoostEstimator.h        |  4 +-
 .../src/PhotonXGBoostEstimator.cc             | 66 +++++++------------
 3 files changed, 32 insertions(+), 54 deletions(-)

diff --git a/RecoEgamma/PhotonIdentification/BuildFile.xml b/RecoEgamma/PhotonIdentification/BuildFile.xml
index 5e23ba0376cda..27e2a21485113 100644
--- a/RecoEgamma/PhotonIdentification/BuildFile.xml
+++ b/RecoEgamma/PhotonIdentification/BuildFile.xml
@@ -1,9 +1,7 @@
-<use name="FWCore/Framework"/>
-<use name="FWCore/MessageLogger"/>
-<use name="FWCore/ParameterSet"/>
+<use name="CommonTools/MVAUtils" />
+<use name="CondFormats/DataRecord"/>
 <use name="CondFormats/EcalObjects"/>
 <use name="CondFormats/HcalObjects"/>
-<use name="CondFormats/DataRecord"/>
 <use name="DataFormats/BeamSpot"/>
 <use name="DataFormats/DetId"/>
 <use name="DataFormats/EcalDetId"/>
@@ -12,17 +10,19 @@
 <use name="DataFormats/EgammaReco"/>
 <use name="DataFormats/HcalRecHit"/>
 <use name="DataFormats/TrackReco"/>
+<use name="FWCore/Framework"/>
+<use name="FWCore/MessageLogger"/>
+<use name="FWCore/ParameterSet"/>
 <use name="Geometry/CaloGeometry"/>
-<use name="Geometry/Records"/>
 <use name="Geometry/CaloTopology"/>
+<use name="Geometry/Records"/>
+<use name="PhysicsTools/TensorFlow" />
+<use name="PhysicsTools/XGBoost" />
 <use name="RecoEcal/EgammaCoreTools"/>
 <use name="RecoEgamma/EgammaIsolationAlgos"/>
 <use name="RecoEgamma/EgammaTools"/>
 <use name="RecoLocalCalo/EcalRecAlgos"/>
 <use name="RecoLocalCalo/HcalRecAlgos"/>
-<use name="PhysicsTools/TensorFlow" />
-<use name="CommonTools/MVAUtils" />
-<use name="xgboost"/>
 <export>
   <lib name="1"/>
 </export>
diff --git a/RecoEgamma/PhotonIdentification/interface/PhotonXGBoostEstimator.h b/RecoEgamma/PhotonIdentification/interface/PhotonXGBoostEstimator.h
index 7252e5116d948..e21f03e515476 100644
--- a/RecoEgamma/PhotonIdentification/interface/PhotonXGBoostEstimator.h
+++ b/RecoEgamma/PhotonIdentification/interface/PhotonXGBoostEstimator.h
@@ -2,7 +2,7 @@
 #define ReciEgamma_PhotonIdentification_PhotonXGBoostEstimator_h
 
 #include "FWCore/ParameterSet/interface/FileInPath.h"
-#include "xgboost/c_api.h"
+#include "PhysicsTools/XGBoost/interface/XGBooster.h"
 
 class PhotonXGBoostEstimator {
 public:
@@ -20,7 +20,7 @@ class PhotonXGBoostEstimator {
                    float ecalPFIsoIn) const;
 
 private:
-  BoosterHandle booster_;
+  std::unique_ptr<pat::XGBooster> booster_;
   int best_ntree_limit_ = -1;
   std::string config_;
 };
diff --git a/RecoEgamma/PhotonIdentification/src/PhotonXGBoostEstimator.cc b/RecoEgamma/PhotonIdentification/src/PhotonXGBoostEstimator.cc
index 6bca23e1e2ff8..27c3e4ef1dbc0 100644
--- a/RecoEgamma/PhotonIdentification/src/PhotonXGBoostEstimator.cc
+++ b/RecoEgamma/PhotonIdentification/src/PhotonXGBoostEstimator.cc
@@ -2,34 +2,21 @@
 #include <sstream>
 
 PhotonXGBoostEstimator::PhotonXGBoostEstimator(const edm::FileInPath& weightsFile, int best_ntree_limit) {
-  XGBoosterCreate(NULL, 0, &booster_);
-  // Set number of threads to 1, to avoid spawning hundreds of OpenMP threads
-  // See https://github.com/cms-sw/cmssw/issues/44923 for details
-  XGBoosterSetParam(booster_, "nthread", "1");
-  XGBoosterLoadModel(booster_, weightsFile.fullPath().c_str());
-  best_ntree_limit_ = best_ntree_limit;
+  booster_ = std::make_unique<pat::XGBooster>(weightsFile.fullPath());
+  booster_->addFeature("rawEnergy");
+  booster_->addFeature("r9");
+  booster_->addFeature("sigmaIEtaIEta");
+  booster_->addFeature("etaWidth");
+  booster_->addFeature("phiWidth");
+  booster_->addFeature("s4");
+  booster_->addFeature("eta");
+  booster_->addFeature("hOvrE");
+  booster_->addFeature("ecalPFIso");
 
-  std::stringstream config;
-  config << "{\"training\": false, \"type\": 0, \"iteration_begin\": 0, \"iteration_end\": " << best_ntree_limit_
-         << ", \"strict_shape\": false}";
-  config_ = config.str();
+  best_ntree_limit_ = best_ntree_limit;
 }
 
-PhotonXGBoostEstimator::~PhotonXGBoostEstimator() { XGBoosterFree(booster_); }
-
-namespace {
-  enum inputIndexes {
-    rawEnergy = 0,      // 0
-    r9 = 1,             // 1
-    sigmaIEtaIEta = 2,  // 2
-    etaWidth = 3,       // 3
-    phiWidth = 4,       // 4
-    s4 = 5,             // 5
-    eta = 6,            // 6
-    hOvrE = 7,          // 7
-    ecalPFIso = 8,      // 8
-  };
-}  // namespace
+PhotonXGBoostEstimator::~PhotonXGBoostEstimator() {}
 
 float PhotonXGBoostEstimator::computeMva(float rawEnergyIn,
                                          float r9In,
@@ -40,24 +27,15 @@ float PhotonXGBoostEstimator::computeMva(float rawEnergyIn,
                                          float etaIn,
                                          float hOvrEIn,
                                          float ecalPFIsoIn) const {
-  float var[9];
-  var[rawEnergy] = rawEnergyIn;
-  var[r9] = r9In;
-  var[sigmaIEtaIEta] = sigmaIEtaIEtaIn;
-  var[etaWidth] = etaWidthIn;
-  var[phiWidth] = phiWidthIn;
-  var[s4] = s4In;
-  var[eta] = etaIn;
-  var[hOvrE] = hOvrEIn;
-  var[ecalPFIso] = ecalPFIsoIn;
+  booster_->set("rawEnergy", rawEnergyIn);
+  booster_->set("r9", r9In);
+  booster_->set("sigmaIEtaIEta", sigmaIEtaIEtaIn);
+  booster_->set("etaWidth", etaWidthIn);
+  booster_->set("phiWidth", phiWidthIn);
+  booster_->set("s4", s4In);
+  booster_->set("eta", etaIn);
+  booster_->set("hOvrE", hOvrEIn);
+  booster_->set("ecalPFIso", ecalPFIsoIn);
 
-  DMatrixHandle dmat;
-  XGDMatrixCreateFromMat(var, 1, 9, -999.9f, &dmat);
-  uint64_t const* out_shape;
-  uint64_t out_dim;
-  const float* out_result = NULL;
-  XGBoosterPredictFromDMatrix(booster_, dmat, config_.c_str(), &out_shape, &out_dim, &out_result);
-  float ret = out_result[0];
-  XGDMatrixFree(dmat);
-  return ret;
+  return booster_->predict(best_ntree_limit_);
 }