[PWGDQ] Apply phi-dependent correction to track rotation (#17245)

Co-authored-by: ayatsuji <13807650+asakuratou@user.noreply.gitee.com>

Author: ayatsujiringo

PWGDQ/Core/VarManager.cxx

@@ -33,6 +33,7 @@
 #include <Rtypes.h>
 #include <RtypesCore.h>
 
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstddef>
@@ -75,7 +76,9 @@ int VarManager::fgCalibrationType = 0;                // 0 - no calibration, 1 -
 bool VarManager::fgUseInterpolatedCalibration = true; // use interpolated calibration histograms (default: true)
 int VarManager::fgEfficiencyType = 0;                 // type of efficiency to be applied, default is no efficiency
 TObject* VarManager::fgEfficiencyHist = nullptr;      // histogram for efficiency
-
+TObject* VarManager::fgPosiPhiMap = nullptr;
+TObject* VarManager::fgNegaPhiMap = nullptr;
+bool VarManager::fgUsePhiCorrection = false;
 //__________________________________________________________________
 VarManager::VarManager() : TObject()
 {
@@ -453,6 +456,76 @@ void VarManager::FillEfficiency(float* values)
   }
 }
 
+void VarManager::SetPhiMap(TObject* hposi, TObject* hnega, bool option)
+{
+  fgPosiPhiMap = hposi;
+  fgNegaPhiMap = hnega;
+  fgUsePhiCorrection = option;
+}
+
+double VarManager::SampleRotationPhi(double pt, double eta, int charge)
+{
+  // each type only alarm once
+  static bool warnedEmptyPhi = false;
+
+  if (!fgUsePhiCorrection) {
+    return gRandom->Uniform(0., o2::constants::math::TwoPI);
+  } else {
+
+    TH3D* hMap = nullptr;
+    if (charge > 0) {
+      hMap = dynamic_cast<TH3D*>(fgPosiPhiMap);
+    } else {
+      hMap = dynamic_cast<TH3D*>(fgNegaPhiMap);
+    }
+
+    if (!hMap) {
+      LOGF(fatal, "Phi map is not a TH3D");
+    }
+    // TH3 axes: X=pT, Y=phi, Z=eta
+    int ptBin = hMap->GetXaxis()->FindBin(pt);
+    int etaBin = hMap->GetZaxis()->FindBin(eta);
+
+    ptBin = std::clamp(ptBin, 1, hMap->GetNbinsX());
+    etaBin = std::clamp(etaBin, 1, hMap->GetNbinsZ());
+
+    TH1D* hPhi = hMap->ProjectionY(
+      Form("hTRPhi_tmp_charge%d_ptbin%d_etabin%d",
+           charge, ptBin, etaBin),
+      ptBin,
+      ptBin,
+      etaBin,
+      etaBin);
+
+    if (!hPhi || hPhi->Integral(1, hPhi->GetNbinsX()) <= 0.) {
+      if (!warnedEmptyPhi) {
+        LOGF(warn,
+             "Empty phi distribution for "
+             "pt=%f, eta=%f, charge=%d, ptBin=%d, etaBin=%d. "
+             "Falling back to uniform phi sampling.",
+             pt,
+             eta,
+             charge,
+             ptBin,
+             etaBin);
+
+        warnedEmptyPhi = true;
+      }
+
+      delete hPhi;
+
+      return gRandom->Uniform(
+        0., o2::constants::math::TwoPI);
+    }
+
+    const double phi =
+      RecoDecay::constrainAngle(hPhi->GetRandom());
+
+    delete hPhi;
+    return phi;
+  }
+}
+
 //__________________________________________________________________
 std::tuple<float, float, float, float, float> VarManager::BimodalityCoefficientUnbinned(const std::vector<float>& data)
 {

PWGDQ/Core/VarManager.h

@@ -1534,6 +1534,8 @@ class VarManager : public TObject
 
   static void SetEfficiencyObject(int type, TObject* obj);
   static void FillEfficiency(float* values = nullptr);
+  static void SetPhiMap(TObject* h1, TObject* h2, bool option);
+  static double SampleRotationPhi(double pT, double eta, int charge);
   static TObject* GetCalibrationObject(CalibObjects calib)
   {
     auto obj = fgCalibs.find(calib);
@@ -1618,6 +1620,10 @@ class VarManager : public TObject
   static int fgEfficiencyType;      // type of efficiency correction to apply
   static TObject* fgEfficiencyHist; // histogram for efficiency correction
 
+  static TObject* fgPosiPhiMap; // phi map to correct track rotation
+  static TObject* fgNegaPhiMap;
+  static bool fgUsePhiCorrection;
+
   VarManager& operator=(const VarManager& c);
   VarManager(const VarManager& c);
 };
@@ -3960,8 +3966,7 @@ void VarManager::FillPairRotation(T1 const& t1, T2 const& t2, float* values)
     m2 = o2::constants::physics::MassMuon;
   }
 
-  double dphi = gRandom->Uniform(0., o2::constants::math::TwoPI);
-  double rotationphi2 = RecoDecay::constrainAngle(t2.phi() + dphi);
+  double rotationphi2 = SampleRotationPhi(t2.pt(), t2.eta(), t2.sign());
 
   values[kCharge] = t1.sign() + t2.sign();
   values[kCharge1] = t1.sign();

PWGDQ/Tasks/tableReader_withAssoc.cxx

@@ -1362,6 +1362,7 @@ struct AnalysisSameEventPairing {
     Configurable<std::string> GrpLhcIfPath{"grplhcif", "GLO/Config/GRPLHCIF", "Path on the CCDB for the GRPLHCIF object"};
     Configurable<std::string> efficiencyPath{"effHistPath", "Users/z/zhxiong/efficiency", "Path on the CCDB for the efficiency histograms"};
     Configurable<std::string> flowPath{"flowPath", "Users/y/yiping/FlowResolution", "Path to the flow resolution object"};
+    Configurable<std::string> phiPath{"phiPath", "Users/h/hxiaoyu/TrackPhi/LHC23", "Path to load phi distribution for track rotation"};
   } fConfigCCDB;
 
   struct : ConfigurableGroup {
@@ -1381,6 +1382,7 @@ struct AnalysisSameEventPairing {
     Configurable<bool> useEfficiencyWeighting{"cfgUseEfficiencyWeighting", false, "Apply efficiency weighting to the pairs from CCDB"};
     Configurable<int> efficiencyType{"cfgEfficiencyType", 0, "Type of efficiency to apply from CCDB: 0 no efficiency, 1 pt-cent-costhetastar"};
     Configurable<bool> useFlowReso{"cfgUseFlowReso", false, "Use remote flow information from CCDB"};
+    Configurable<bool> usePhiDistribution{"cfgUsePhiDistribution", false, "Use phi distribution to correct track rotation"};
   } fConfigOptions;
   struct : ConfigurableGroup {
     Configurable<bool> applyBDT{"applyBDT", false, "Flag to apply ML selections"};
@@ -1855,6 +1857,18 @@ struct AnalysisSameEventPairing {
         LOGF(fatal, "Flow resolution histograms not available in CCDB at timestamp=%llu", timestamp);
       }
     }
+
+    if (fConfigOptions.usePhiDistribution) {
+      TString PathPhi = fConfigCCDB.phiPath.value;
+      TString ccdbPathPhiPosi = Form("%s/hPtPhiPositive", PathPhi.Data());
+      TString ccdbPathPhiNega = Form("%s/hPtPhiNegative", PathPhi.Data());
+      auto phiPosi = fCCDB->getForTimeStamp<TH3D>(ccdbPathPhiPosi.Data(), timestamp);
+      auto phiNega = fCCDB->getForTimeStamp<TH3D>(ccdbPathPhiNega.Data(), timestamp);
+      if (phiPosi == nullptr || phiNega == nullptr) {
+        LOGF(fatal, "Phi distribution histograms not available in CCDB at timestamp=%llu", timestamp);
+      }
+      VarManager::SetPhiMap(phiPosi, phiNega, fConfigOptions.usePhiDistribution.value);
+    }
   }
 
   // Template function to run same event pairing (barrel-barrel, muon-muon, barrel-muon)