[PWGDQ] Allowing independent shifts of global tracks for MFT top or bottom halves (#17689)

Co-authored-by: Maurice Coquet <mcoquet@alicecerno2.cern.ch>

Author: mcoquet642

PWGDQ/Core/VarManager.cxx

@@ -59,6 +59,10 @@ float VarManager::fgzMatching = -77.5;
 float VarManager::fgxShiftFwd = 0.0;
 float VarManager::fgyShiftFwd = 0.0;
 float VarManager::fgzShiftFwd = 0.0;
+bool VarManager::fgUseTopBottomShift = false;
+float VarManager::fgxShiftFwdBottom = 0.0;
+float VarManager::fgyShiftFwdBottom = 0.0;
+float VarManager::fgzShiftFwdBottom = 0.0;
 float VarManager::fgValues[VarManager::kNVars] = {0.0f};
 float VarManager::fgTPCInterSectorBoundary = 1.0; // cm
 int VarManager::fgITSROFbias = 0;
@@ -147,7 +151,7 @@ void VarManager::ResetValues(int startValue, int endValue, float* values)
 }
 
 //__________________________________________________________________
-void VarManager::SetCollisionSystem(TString system, float energy)
+void VarManager::SetCollisionSystem(const TString& system, float energy)
 {
   //
   // Set the collision system and the center of mass energy
@@ -200,8 +204,8 @@ void VarManager::SetCollisionSystem(TString system, float energy)
   // TO Do: add more systems
 
   // set the beam 4-momentum vectors
-  float beamAEnergy = energy / 2.0 * sqrt(NumberOfProtonsA * NumberOfProtonsC / NumberOfProtonsC / NumberOfProtonsA); // GeV
-  float beamCEnergy = energy / 2.0 * sqrt(NumberOfProtonsC * NumberOfProtonsA / NumberOfProtonsA / NumberOfProtonsC); // GeV
+  float beamAEnergy = energy / 2.0f * std::sqrt(static_cast<float>(NumberOfProtonsA) * NumberOfProtonsC / NumberOfProtonsC / NumberOfProtonsA); // GeV
+  float beamCEnergy = energy / 2.0f * std::sqrt(static_cast<float>(NumberOfProtonsC) * NumberOfProtonsA / NumberOfProtonsA / NumberOfProtonsC); // GeV
   float beamAMomentum = std::sqrt(beamAEnergy * beamAEnergy - NumberOfNucleonsA * NumberOfNucleonsA * MassProton * MassProton);
   float beamCMomentum = std::sqrt(beamCEnergy * beamCEnergy - NumberOfNucleonsC * NumberOfNucleonsC * MassProton * MassProton);
   fgBeamA.SetPxPyPzE(0, 0, beamAMomentum, beamAEnergy);
@@ -243,7 +247,7 @@ void VarManager::FillTrackDerived(float* values)
 }
 
 //__________________________________________________________________
-float VarManager::calculateCosPA(KFParticle kfp, KFParticle PV)
+float VarManager::calculateCosPA(const KFParticle& kfp, const KFParticle& PV)
 {
   return cpaFromKF(kfp, PV);
 }
@@ -256,7 +260,8 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
 
   if (fgCalibrationType == 1) {
     // get the calibration histograms
-    CalibObjects calibMean, calibSigma;
+    CalibObjects calibMean = kTPCElectronMean;
+    CalibObjects calibSigma = kTPCElectronSigma;
     switch (species) {
       case 0:
         calibMean = kTPCElectronMean;
@@ -279,8 +284,8 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
         return -999.0; // Return zero if species is invalid
     }
 
-    TH3F* calibMeanHist = reinterpret_cast<TH3F*>(fgCalibs[calibMean]);
-    TH3F* calibSigmaHist = reinterpret_cast<TH3F*>(fgCalibs[calibSigma]);
+    TH3F* calibMeanHist = dynamic_cast<TH3F*>(fgCalibs[calibMean]);
+    TH3F* calibSigmaHist = dynamic_cast<TH3F*>(fgCalibs[calibSigma]);
     if (!calibMeanHist || !calibSigmaHist) {
       LOG(fatal) << "Calibration histograms not found for species: " << species;
       return -999.0; // Return zero if histograms are not found
@@ -302,7 +307,9 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
     return (nSigmaValue - mean) / sigma; // Return the calibrated nSigma value
   } else if (fgCalibrationType == 2) {
     // get the calibration histograms
-    CalibObjects calibMean, calibSigma, calibStatus;
+    CalibObjects calibMean = kTPCElectronMean;
+    CalibObjects calibSigma = kTPCElectronSigma;
+    CalibObjects calibStatus = kTPCElectronStatus;
     switch (species) {
       case 0:
         calibMean = kTPCElectronMean;
@@ -329,9 +336,9 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
         return -999.0; // Return zero if species is invalid
     }
 
-    THnF* calibMeanHist = reinterpret_cast<THnF*>(fgCalibs[calibMean]);
-    THnF* calibSigmaHist = reinterpret_cast<THnF*>(fgCalibs[calibSigma]);
-    THnF* calibStatusHist = reinterpret_cast<THnF*>(fgCalibs[calibStatus]);
+    THnF* calibMeanHist = dynamic_cast<THnF*>(fgCalibs[calibMean]);
+    THnF* calibSigmaHist = dynamic_cast<THnF*>(fgCalibs[calibSigma]);
+    THnF* calibStatusHist = dynamic_cast<THnF*>(fgCalibs[calibStatus]);
     if (!calibMeanHist || !calibSigmaHist || !calibStatusHist) {
       LOG(fatal) << "Calibration histograms not found for species: " << species;
       return -999.0; // Return zero if histograms are not found
@@ -351,17 +358,18 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
     binTlong = (binTlong == 0 ? 1 : binTlong);
     binTlong = (binTlong > calibMeanHist->GetAxis(3)->GetNbins() ? calibMeanHist->GetAxis(3)->GetNbins() : binTlong);
 
-    int bin[4] = {binEta, binNpv, binNlong, binTlong};
-    int status = static_cast<int>(calibStatusHist->GetBinContent(bin));
-    double mean = calibMeanHist->GetBinContent(bin);
-    double sigma = calibSigmaHist->GetBinContent(bin);
+    std::array<int, 4> bin{binEta, binNpv, binNlong, binTlong};
+    int status = static_cast<int>(calibStatusHist->GetBinContent(bin.data()));
+    double mean = calibMeanHist->GetBinContent(bin.data());
+    double sigma = calibSigmaHist->GetBinContent(bin.data());
     switch (status) {
       case 0:
         // good calibration, return the calibrated nSigma value
         return (nSigmaValue - mean) / sigma;
         break;
       case 1:
-        // calibration not valid, return the original nSigma value
+      case 4:
+        // calibration not valid or interpolation failed, return the original nSigma value
         return nSigmaValue;
         break;
       case 2: // calibration constant has poor stat uncertainty, consider the user option for what to do
@@ -374,10 +382,6 @@ double VarManager::ComputePIDcalibration(int species, double nSigmaValue)
           return nSigmaValue;
         }
         break;
-      case 4:
-        // calibration constants interpolation failed, return the original nSigma value
-        return nSigmaValue;
-        break;
       default:
         return nSigmaValue; // unknown status, return the original nSigma value
         break;
@@ -419,7 +423,7 @@ void VarManager::FillEfficiency(float* values)
       LOG(fatal) << "efficiency histogram not set";
       return;
     }
-    TH3F* efficiencyHist = reinterpret_cast<TH3F*>(fgEfficiencyHist);
+    TH3F* efficiencyHist = dynamic_cast<TH3F*>(fgEfficiencyHist);
     // Get the bin indices for the efficiency histogram
     int binPt = efficiencyHist->GetXaxis()->FindBin(values[kPt]);
     binPt = (binPt == 0 ? 1 : binPt);
@@ -439,7 +443,7 @@ void VarManager::FillEfficiency(float* values)
       LOG(fatal) << "efficiency histogram not set";
       return;
     }
-    TH3F* efficiencyHist = reinterpret_cast<TH3F*>(fgEfficiencyHist);
+    TH3F* efficiencyHist = dynamic_cast<TH3F*>(fgEfficiencyHist);
     // Get the bin indices for the efficiency histogram
     int binPt = efficiencyHist->GetXaxis()->FindBin(values[kPt]);
     binPt = (binPt == 0 ? 1 : binPt);
@@ -541,10 +545,9 @@ std::tuple<float, float, float, float, float> VarManager::BimodalityCoefficientU
   float mean = std::accumulate(data.begin(), data.end(), 0.0) / n;
 
   float m2 = 0.0, m3 = 0.0, m4 = 0.0;
-  float diff, diff2;
-  for (float x : data) {
-    diff = x - mean;
-    diff2 = diff * diff;
+  for (const float& x : data) {
+    const float diff = x - mean;
+    const float diff2 = diff * diff;
     m2 += diff2;
     m3 += diff2 * diff;
     m4 += diff2 * diff2;
@@ -581,7 +584,7 @@ std::tuple<float, float, float, float, float, int> VarManager::BimodalityCoeffic
   int nBins = static_cast<int>((max - min) / binWidth);
   std::vector<int> counts(nBins, 0.0);
 
-  for (float x : data) {
+  for (const float& x : data) {
     if (x < min || x >= max) {
       continue; // skip out-of-range values
     }
@@ -688,14 +691,13 @@ std::tuple<float, float, float, float, float, int> VarManager::BimodalityCoeffic
 
   // then compute the second, third, and fourth central moments
   float m2 = 0.0, m3 = 0.0, m4 = 0.0;
-  float diff, diff2, binCenter;
   for (int i = 0; i < nBins; ++i) {
     if (counts[i] == 0) {
       continue; // skip empty bins
     }
-    binCenter = min + (i + 0.5) * binWidth;
-    diff = binCenter - mean;
-    diff2 = diff * diff;
+    const float binCenter = min + (i + 0.5f) * binWidth;
+    const float diff = binCenter - mean;
+    const float diff2 = diff * diff;
     m2 += counts[i] * diff2;
     m3 += counts[i] * diff2 * diff;
     m4 += counts[i] * diff2 * diff2;

PWGDQ/Core/VarManager.h

@@ -1274,7 +1274,7 @@ class VarManager : public TObject
   }
 
   // Setup the collision system
-  static void SetCollisionSystem(TString system, float energy);
+  static void SetCollisionSystem(const TString& system, float energy);
   static void SetCollisionSystem(o2::parameters::GRPLHCIFData* grplhcif);
 
   static void SetMagneticField(float magField)
@@ -1297,6 +1297,7 @@ class VarManager : public TObject
   static void SetZShift(float z)
   {
     fgzShiftFwd = z;
+    fgUseTopBottomShift = false;
   }
 
   // Set x, y and z shifts for forward tracks
@@ -1305,6 +1306,33 @@ class VarManager : public TObject
     fgxShiftFwd = x;
     fgyShiftFwd = y;
     fgzShiftFwd = z;
+    fgUseTopBottomShift = false;
+  }
+
+  // Set separate x, y, z shifts for top (y >= 0) and bottom (y < 0) forward tracks
+  // Top shifts are stored in fgx/y/zShiftFwd; bottom shifts in fgx/y/zShiftFwdBottom
+  static void SetTopBottom3DShift(float xTop, float yTop, float zTop, float xBottom, float yBottom, float zBottom)
+  {
+    fgxShiftFwd = xTop;
+    fgyShiftFwd = yTop;
+    fgzShiftFwd = zTop;
+    fgxShiftFwdBottom = xBottom;
+    fgyShiftFwdBottom = yBottom;
+    fgzShiftFwdBottom = zBottom;
+    fgUseTopBottomShift = true;
+  }
+
+  static void GetFwdShiftForY(float y, float& xShift, float& yShift, float& zShift)
+  {
+    if (fgUseTopBottomShift && y < 0.f) {
+      xShift = fgxShiftFwdBottom;
+      yShift = fgyShiftFwdBottom;
+      zShift = fgzShiftFwdBottom;
+    } else {
+      xShift = fgxShiftFwd;
+      yShift = fgyShiftFwd;
+      zShift = fgzShiftFwd;
+    }
   }
 
   // Setup the 2 prong KFParticle
@@ -1619,6 +1647,10 @@ class VarManager : public TObject
   static float fgxShiftFwd;
   static float fgyShiftFwd;
   static float fgzShiftFwd;
+  static bool fgUseTopBottomShift;
+  static float fgxShiftFwdBottom;
+  static float fgyShiftFwdBottom;
+  static float fgzShiftFwdBottom;
   static float fgCenterOfMassEnergy;        // collision energy
   static float fgMassofCollidingParticle;   // mass of the colliding particle
   static float fgTPCInterSectorBoundary;    // TPC inter-sector border size at the TPC outer radius, in cm
@@ -1641,7 +1673,7 @@ class VarManager : public TObject
   static KFPTrack createKFPFwdTrackFromFwdTrack(const T& muon);
   template <typename T>
   static KFPVertex createKFPVertexFromCollision(const T& collision);
-  static float calculateCosPA(KFParticle kfp, KFParticle PV);
+  static float calculateCosPA(const KFParticle& kfp, const KFParticle& PV);
   template <int pairType, typename T1, typename T2>
   static float calculatePhiV(const T1& t1, const T2& t2);
   template <typename T1, typename T2>
@@ -1797,7 +1829,11 @@ o2::dataformats::VertexBase VarManager::RecalculatePrimaryVertex(T const& track0
 template <typename T, typename C>
 o2::dataformats::GlobalFwdTrack VarManager::PropagateMuon(const T& muon, const C& collision, const int endPoint)
 {
-  o2::track::TrackParCovFwd fwdtrack = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(muon, fgxShiftFwd, fgyShiftFwd, fgzShiftFwd, muon);
+  float xShift = 0.f;
+  float yShift = 0.f;
+  float zShift = 0.f;
+  GetFwdShiftForY(muon.y(), xShift, yShift, zShift);
+  o2::track::TrackParCovFwd fwdtrack = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(muon, xShift, yShift, zShift, muon);
   o2::dataformats::GlobalFwdTrack propmuon;
   if (static_cast<int>(muon.trackType()) > 2) {
     o2::dataformats::GlobalFwdTrack track;
@@ -1924,12 +1960,16 @@ void VarManager::FillGlobalMuonRefit(T1 const& muontrack, T2 const& mfttrack, co
     values = fgValues;
   }
   if constexpr ((fillMap & MuonCov) > 0 || (fillMap & ReducedMuonCov) > 0) {
-    o2::dataformats::GlobalFwdTrack propmuon = PropagateMuon(muontrack, collision);
+    float xShift = 0.f;
+    float yShift = 0.f;
+    float zShift = 0.f;
+    GetFwdShiftForY(mfttrack.y(), xShift, yShift, zShift);
+    o2::dataformats::GlobalFwdTrack propmuon = PropagateMuon(muontrack, collision, kToVertex);
     double px = propmuon.getP() * std::sin(o2::constants::math::PIHalf - std::atan(mfttrack.tgl())) * std::cos(mfttrack.phi());
     double py = propmuon.getP() * std::sin(o2::constants::math::PIHalf - std::atan(mfttrack.tgl())) * std::sin(mfttrack.phi());
     double pz = propmuon.getP() * std::cos(o2::constants::math::PIHalf - std::atan(mfttrack.tgl()));
     double pt = std::sqrt(std::pow(px, 2) + std::pow(py, 2));
-    auto mftprop = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(mfttrack, fgxShiftFwd, fgyShiftFwd, fgzShiftFwd);
+    auto mftprop = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(mfttrack, xShift, yShift, zShift);
     values[kX] = mftprop.getX();
     values[kY] = mftprop.getY();
     values[kZ] = mftprop.getZ();
@@ -1949,8 +1989,12 @@ void VarManager::FillGlobalMuonRefitCov(T1 const& muontrack, T2 const& mfttrack,
   }
   if constexpr ((MuonfillMap & MuonCov) > 0) {
     if constexpr ((MFTfillMap & MFTCov) > 0) {
-      o2::dataformats::GlobalFwdTrack propmuon = PropagateMuon(muontrack, collision);
-      auto mft = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(mfttrack, fgxShiftFwd, fgyShiftFwd, fgzShiftFwd, mftcov);
+      float xShift = 0.f;
+      float yShift = 0.f;
+      float zShift = 0.f;
+      GetFwdShiftForY(mfttrack.y(), xShift, yShift, zShift);
+      o2::dataformats::GlobalFwdTrack propmuon = PropagateMuon(muontrack, collision, kToVertex);
+      auto mft = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(mfttrack, xShift, yShift, zShift, mftcov);
 
       o2::dataformats::GlobalFwdTrack globalRefit = o2::aod::fwdtrackutils::refitGlobalMuonCov(propmuon, mft);
       values[kX] = globalRefit.getX();
@@ -3369,7 +3413,11 @@ void VarManager::FillTrack(T const& track, float* values)
     values[kMuonC1Pt21Pt2] = track.c1Pt21Pt2();
   }
   if constexpr ((fillMap & MuonCov) > 0 || (fillMap & MuonCovRealign) > 0) {
-    auto muonTrack = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(track, fgxShiftFwd, fgyShiftFwd, fgzShiftFwd, track);
+    float xShift = 0.f;
+    float yShift = 0.f;
+    float zShift = 0.f;
+    GetFwdShiftForY(track.y(), xShift, yShift, zShift);
+    auto muonTrack = o2::aod::fwdtrackutils::getTrackParCovFwd3DShift(track, xShift, yShift, zShift, track);
     auto muonCov = muonTrack.getCovariances();
     values[kX] = muonTrack.getX();
     values[kY] = muonTrack.getY();
@@ -4052,11 +4100,7 @@ void VarManager::FillPairRotation(T1 const& t1, T2 const& t2, int rotation, floa
     rotationphi2 = 2 * values[kPsi2A] - t2.phi() + o2::constants::math::PI;
   }
 
-  if (rotationphi2 >= o2::constants::math::TwoPI) {
-    rotationphi2 -= o2::constants::math::TwoPI;
-  } else if (rotationphi2 < 0) {
-    rotationphi2 += o2::constants::math::TwoPI;
-  }
+  rotationphi2 = RecoDecay::constrainAngle(rotationphi2);
 
   values[kCharge] = t1.sign() + t2.sign();
   values[kCharge1] = t1.sign();

PWGDQ/TableProducer/tableMaker_withAssoc.cxx

@@ -288,7 +288,7 @@ struct TableMaker {
     Configurable<int64_t> fConfigNoLaterThan{"ccdb-no-later-than", std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count(), "latest acceptable timestamp of creation for the object"};
     Configurable<std::string> fConfigGeoPath{"geoPath", "GLO/Config/GeometryAligned", "Path of the geometry file"};
     Configurable<std::string> fConfigGrpMagPath{"grpmagPath", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object"};
-    Configurable<std::string> fFwdShiftPath{"fwdShiftPath", "Users/m/mcoquet/ZShift", "CCDB path for the shift to apply to forward tracks, either 1 value (z) or 3 values (x, y, z)"};
+    Configurable<std::string> fFwdShiftPath{"fwdShiftPath", "Users/m/mcoquet/ZShift", "CCDB path for the shift to apply to forward tracks: 1 (z), 3 (x,y,z), or 10 (x,y,z,slopeX,slopeY for top then bottom; slopes unused)"};
     Configurable<bool> fUseRemoteFwdShift{"cfgUseRemoteFwdShift", false, "Enable getting the forward track shift from ccdb"};
     Configurable<float> fManualZShift{"cfgManualZShift", 0.f, "Manual value for the Zshift for muons."};
     Configurable<std::string> fConfigGrpMagPathRun2{"grpmagPathRun2", "GLO/GRP/GRP", "CCDB path of the GRPObject (Usage for Run 2)"};
@@ -1870,8 +1870,15 @@ struct TableMaker {
             VarManager::SetZShift((*fFwdShift)[0]);
           } else if (fFwdShift->size() == 3) {
             VarManager::Set3DShift((*fFwdShift)[0], (*fFwdShift)[1], (*fFwdShift)[2]);
+          } else if (fFwdShift->size() == 10) {
+            // x_top, y_top, z_top, slopeX_top, slopeY_top, x_bottom, y_bottom, z_bottom, slopeX_bottom, slopeY_bottom
+            // Slopes are unused for now; shift is selected from track y (top: y >= 0, bottom: y < 0)
+            VarManager::SetTopBottom3DShift((*fFwdShift)[0], (*fFwdShift)[1], (*fFwdShift)[2],
+                                            (*fFwdShift)[5], (*fFwdShift)[6], (*fFwdShift)[7]);
+            LOG(info) << "Loaded top/bottom forward track shifts from CCDB: top=(" << (*fFwdShift)[0] << ", " << (*fFwdShift)[1] << ", " << (*fFwdShift)[2]
+                      << "), bottom=(" << (*fFwdShift)[5] << ", " << (*fFwdShift)[6] << ", " << (*fFwdShift)[7] << ")";
           } else {
-            LOG(fatal) << "Unexpected number of shift values from CCDB: " << fFwdShift->size() << ", expected 1 (z) or 3 (x, y, z)";
+            LOG(fatal) << "Unexpected number of shift values from CCDB: " << fFwdShift->size() << ", expected 1 (z), 3 (x, y, z) or 10 (top/bottom x,y,z + slopes)";
           }
         } else {
           VarManager::SetZShift(fConfigCCDB.fManualZShift.value);