Add phi-map correction for track rotation

Author: ayatsuji

PWGDQ/Core/VarManager.cxx

@@ -75,7 +75,8 @@ 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
-
+TH3D* VarManager::fgPosiPhiMap = nullptr;
+TH3D* VarManager::fgNegaPhiMap = nullptr;
 //__________________________________________________________________
 VarManager::VarManager() : TObject()
 {
@@ -453,6 +454,129 @@ void VarManager::FillEfficiency(float* values)
   }
 }
 
+void VarManager::SetPhiMap(TH3D* hposi, TH3D* hnega)
+{
+  fgPosiPhiMap = hposi;
+  fgNegaPhiMap = hnega;
+}
+
+double VarManager::SampleRotationPhi(double pt, double eta, int charge)
+{
+  // each type only alarm once
+  static bool warnedInvalidCharge = false;
+  static bool warnedMissingPositiveMap = false;
+  static bool warnedMissingNegativeMap = false;
+  static bool warnedOutOfRange = false;
+  static bool warnedEmptyPhi = false;
+
+  TH3D* hMap = nullptr;
+
+  if (charge > 0) {
+    hMap = fgPosiPhiMap;
+
+    if (!hMap) {
+      if (!warnedMissingPositiveMap) {
+        LOGF(warn,
+             "Positive phi correction map is not available. "
+             "Falling back to uniform phi sampling.");
+        warnedMissingPositiveMap = true;
+      }
+
+      return gRandom->Uniform(
+          0., o2::constants::math::TwoPI);
+    }
+  } else if (charge < 0) {
+    hMap = fgNegaPhiMap;
+
+    if (!hMap) {
+      if (!warnedMissingNegativeMap) {
+        LOGF(warn,
+             "Negative phi correction map is not available. "
+             "Falling back to uniform phi sampling.");
+        warnedMissingNegativeMap = true;
+      }
+
+      return gRandom->Uniform(
+          0., o2::constants::math::TwoPI);
+    }
+  } else {
+    if (!warnedInvalidCharge) {
+      LOGF(warn,
+           "Track with charge=0 passed to SampleRotationPhi. "
+           "Falling back to uniform phi sampling.");
+      warnedInvalidCharge = true;
+    }
+
+    return gRandom->Uniform(
+        0., o2::constants::math::TwoPI);
+  }
+
+  // TH3 axes: X=pT, Y=phi, Z=eta
+  const int ptBin = hMap->GetXaxis()->FindBin(pt);
+  const int etaBin = hMap->GetZaxis()->FindBin(eta);
+
+  if (ptBin < 1 || ptBin > hMap->GetNbinsX() ||
+      etaBin < 1 || etaBin > hMap->GetNbinsZ()) {
+    if (!warnedOutOfRange) {
+      LOGF(warn,
+           "Track outside phi correction map: "
+           "pt=%f, eta=%f, charge=%d, ptBin=%d, etaBin=%d. "
+           "Map ranges: pt=[%f,%f], eta=[%f,%f]. "
+           "Falling back to uniform phi sampling.",
+           pt,
+           eta,
+           charge,
+           ptBin,
+           etaBin,
+           hMap->GetXaxis()->GetXmin(),
+           hMap->GetXaxis()->GetXmax(),
+           hMap->GetZaxis()->GetXmin(),
+           hMap->GetZaxis()->GetXmax());
+
+      warnedOutOfRange = true;
+    }
+
+    return gRandom->Uniform(
+        0., o2::constants::math::TwoPI);
+  }
+
+  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

@@ -1530,6 +1530,8 @@ class VarManager : public TObject
 
   static void SetEfficiencyObject(int type, TObject* obj);
   static void FillEfficiency(float* values = nullptr);
+  static void SetPhiMap(TH3D* h1, TH3D* h2);
+  static double SampleRotationPhi(double pT, int charge);
   static TObject* GetCalibrationObject(CalibObjects calib)
   {
     auto obj = fgCalibs.find(calib);
@@ -3956,7 +3958,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 dphi = SampleRotationPhi(t2.pt(),t2.eta(),t2.sign());
   double rotationphi2 = RecoDecay::constrainAngle(t2.phi() + dphi);
 
   values[kCharge] = t1.sign() + t2.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);
+    }
   }
 
   // Template function to run same event pairing (barrel-barrel, muon-muon, barrel-muon)