[PWGEM/PhotonMeson] Add option for fast analytical propagation of tracks

jokonig · jokonig · commit 7c6d3aa6eafc · 2026-08-03T17:13:26.000+02:00
- Photonconversion builder is slow due to many propagation of the tracks to the primary vertex etc.
- Add a fast routine to calcuate DCA and Phiv, Psi-Pair from the track helix.
- Add histograms to compare the results of both methods
diff --git a/PWGEM/PhotonMeson/TableProducer/photonconversionbuilder.cxx b/PWGEM/PhotonMeson/TableProducer/photonconversionbuilder.cxx
@@ -94,6 +94,12 @@ enum MatCorrType {
   LUT = 2
 };
 
+enum TrackPropMode {
+  kProper = 0,
+  kFast = 1,
+  kBoth = 2
+};
+
 struct PhotonConversionBuilder {
   Produces<aod::V0PhotonsKF> v0photonskf;
   Produces<aod::V0Legs> v0legs;
@@ -114,6 +120,7 @@ struct PhotonConversionBuilder {
   // Operation and minimisation criteria
   Configurable<double> d_bz_input{"d_bz", -999, "bz field, -999 is automatic"};
   Configurable<int> useMatCorrType{"useMatCorrType", 0, "0: none, 1: TGeo, 2: LUT"};
+  Configurable<int> modeTrackPropagation{"modeTrackPropagation", 0, "0: use real track propagation, including material, 1: use fast approximation using only geometry, 2: Use real track propagation and make comparison to fast propagation (only for debugging and testing)"};
 
   // single track cuts
   Configurable<int> min_ncluster_tpc{"min_ncluster_tpc", 0, "min ncluster tpc"};
@@ -325,6 +332,14 @@ struct PhotonConversionBuilder {
         registry.add("V0/hBDTScoreAfterCutVsPt", "BDT score after cut vs pT; pT (GeV/c); BDT score", {HistType::kTH2F, {{1000, 0.0f, 20.0f}, {1000, 0.0f, 1.0f}}});
       }
     }
+
+    // Compare proper propagation and fast geometrical propagation
+    if (modeTrackPropagation == TrackPropMode::kBoth) {
+      registry.add("V0Leg/hDCAxyPropagationCompare", "Comparison of DCA_{xy} propagation;DCA_{xy} (cm) proper propagation; DCA_{xy} (cm) geom. propagation", {HistType::kTH2F, {{200, -10., 10.}, {200, -10., 10.}}});
+      registry.add("V0Leg/hDCAzPropagationCompare", "Comparison of DCA_{z} propagation;DCA_{z} (cm) proper propagation; DCA_{z} (cm) geom. propagation", {HistType::kTH2F, {{200, -10., 10.}, {200, -10., 10.}}});
+      registry.add("V0/hPhivPropagationCompare", "Comparison of #phi_{v};#phi_{v} proper propagation; #phi_{v} geom. propagation", {HistType::kTH2F, {{100, 0., 1.6}, {100, 0., 1.6}}});
+      registry.add("V0/hPsiPairPropagationCompare", "Comparison of #Psi_{pair};#Psi_{pair} proper propagation; #Psi_{pair} geom. propagation", {HistType::kTH2F, {{100, 0., 1.6}, {100, 0., 1.6}}});
+    }
   }
 
   void initCCDB(aod::BCsWithTimestamps::iterator const& bc)
@@ -555,8 +570,23 @@ struct PhotonConversionBuilder {
       return;
     }
     auto pTrackC = pTrack;
+    o2::math_utils::Point3D<float> vtxPrim{
+      collision.posX(),
+      collision.posY(),
+      collision.posZ()};
+    if (modeTrackPropagation != TrackPropMode::kProper) {
+      dcaInfo = CalculateDCAFast(pTrackC, vtxPrim, d_bz);
+    }
     pTrackC.setPID(o2::track::PID::Electron);
-    o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, pTrackC, 2.f, matCorr, &dcaInfo);
+
+    std::array<float, 2> dcaInfoFast = dcaInfo;
+    if (modeTrackPropagation != TrackPropMode::kFast) {
+      o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, pTrackC, 2.f, matCorr, &dcaInfo);
+      if (modeTrackPropagation == TrackPropMode::kBoth) {
+        registry.fill(HIST("V0Leg/hDCAxyPropagationCompare"), dcaInfo[0], dcaInfoFast[0]);
+        registry.fill(HIST("V0Leg/hDCAzPropagationCompare"), dcaInfo[1], dcaInfoFast[1]);
+      }
+    }
     auto posdcaXY = dcaInfo[0];
     auto posdcaZ = dcaInfo[1];
 
@@ -566,8 +596,19 @@ struct PhotonConversionBuilder {
       return;
     }
     auto nTrackC = nTrack;
+    if (modeTrackPropagation != TrackPropMode::kProper) {
+      dcaInfo = CalculateDCAFast(nTrackC, vtxPrim, d_bz);
+    }
     nTrackC.setPID(o2::track::PID::Electron);
-    o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, nTrackC, 2.f, matCorr, &dcaInfo);
+
+    dcaInfoFast = dcaInfo;
+    if (modeTrackPropagation != TrackPropMode::kFast) {
+      o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, nTrackC, 2.f, matCorr, &dcaInfo);
+      if (modeTrackPropagation == TrackPropMode::kBoth) {
+        registry.fill(HIST("V0Leg/hDCAxyPropagationCompare"), dcaInfo[0], dcaInfoFast[0]);
+        registry.fill(HIST("V0Leg/hDCAzPropagationCompare"), dcaInfo[1], dcaInfoFast[1]);
+      }
+    }
     auto eledcaXY = dcaInfo[0];
     auto eledcaZ = dcaInfo[1];
 
@@ -587,30 +628,69 @@ struct PhotonConversionBuilder {
 
     float phiv = 999.f;
     float psipair = 999.f;
+    float phivFast = 999.f;
+    float psipairFast = 999.f;
     float baseR = std::hypot(xyz[0], xyz[1]);
-    std::array<float, 3> offsetsR = {propV0LegsRadius, 30.f, 10.f};
-    bool pPropagatedSuccess = false;
-    bool nPropagatedSuccess = false;
-    auto pTrackProp = pTrack;
-    auto nTrackProp = nTrack;
-    for (const float& offsetR : offsetsR) {
-      pTrackProp = pTrack;
-      pTrackProp.setPID(o2::track::PID::Electron);
-      nTrackProp = nTrack;
-      nTrackProp.setPID(o2::track::PID::Electron);
-      o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, pTrackProp, 2.f, matCorr, &dcaInfo);
-      o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, nTrackProp, 2.f, matCorr, &dcaInfo);
-      pPropagatedSuccess = o2::base::Propagator::Instance()->propagateToR(pTrackProp, baseR + offsetR);
-      nPropagatedSuccess = o2::base::Propagator::Instance()->propagateToR(nTrackProp, baseR + offsetR);
-      if (pPropagatedSuccess && nPropagatedSuccess) {
-        KFPTrack kfp_track_posProp = createKFPTrackFromTrackParCov(pTrackProp, pos.sign(), pos.tpcNClsFound(), pos.tpcChi2NCl());
-        KFPTrack kfp_track_eleProp = createKFPTrackFromTrackParCov(nTrackProp, ele.sign(), ele.tpcNClsFound(), ele.tpcChi2NCl());
-        phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(kfp_track_posProp.GetPx(), kfp_track_posProp.GetPy(), kfp_track_posProp.GetPz(), kfp_track_eleProp.GetPx(), kfp_track_eleProp.GetPy(), kfp_track_eleProp.GetPz(), pos.sign(), ele.sign(), d_bz);
-        psipair = o2::aod::pwgem::dilepton::utils::pairutil::getPsiPair(kfp_track_posProp.GetPx(), kfp_track_posProp.GetPy(), kfp_track_posProp.GetPz(), kfp_track_eleProp.GetPx(), kfp_track_eleProp.GetPy(), kfp_track_eleProp.GetPz());
-        break;
+    // This method uses the track Helix instead of the full propagation.
+    // Hence, it is only an approximation but much faster
+    if (modeTrackPropagation != TrackPropMode::kProper) {
+
+      o2::track::TrackAuxPar helixPosEle(nTrack, d_bz);
+      o2::track::TrackAuxPar helixPosPos(pTrack, d_bz);
+
+      float diffX = helixPosEle.xC - helixPosPos.xC;
+      float diffY = helixPosEle.yC - helixPosPos.yC;
+      auto phiHelix = RecoDecay::constrainAngle<float, float>(std::atan2(diffY, diffX) - o2::constants::math::PI / 2.);
+
+      // Electron
+      float arcLenghtEle = helixPosEle.rC * 0.9 > propV0LegsRadius ? std::asin(propV0LegsRadius / helixPosEle.rC) * helixPosEle.rC : o2::constants::math::PI / 2.2 * helixPosEle.rC; // This assumes that the photon momentum vector is a tangent of the circle
+      auto propTrackEle = getPropMomentumFromTrackHelix(arcLenghtEle, ele, helixPosEle, d_bz / 10., phiHelix - ele.phi());
+      // Positron
+      float arcLenghtPos = helixPosPos.rC * 0.9 > propV0LegsRadius ? std::asin(propV0LegsRadius / helixPosPos.rC) * helixPosPos.rC : o2::constants::math::PI / 2.2 * helixPosPos.rC; // This assumes that the photon momentum vector is a tangent of the circle
+      auto propTrackPos = getPropMomentumFromTrackHelix(arcLenghtPos, pos, helixPosPos, d_bz / 10., phiHelix - pos.phi());
+
+      phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(propTrackPos[0], propTrackPos[1], propTrackPos[2], propTrackEle[0], propTrackEle[1], propTrackEle[2], pos.sign(), ele.sign(), d_bz);
+      psipair = o2::aod::pwgem::dilepton::utils::pairutil::getPsiPair(propTrackPos[0], propTrackPos[1], propTrackPos[2], propTrackEle[0], propTrackEle[1], propTrackEle[2]);
+
+      // Store values for later comparison
+      phivFast = phiv;
+      psipairFast = psipair;
+    }
+    // This uses the full propagation including material effects
+    if (modeTrackPropagation != TrackPropMode::kFast) {
+      std::array<float, 3> offsetsR = {propV0LegsRadius, 30.f, 10.f};
+      bool pPropagatedSuccess = false;
+      bool nPropagatedSuccess = false;
+      auto pTrackProp = pTrack;
+      auto nTrackProp = nTrack;
+      for (const auto& offsetR : offsetsR) {
+        pTrackProp = pTrack;
+        pTrackProp.setPID(o2::track::PID::Electron);
+        nTrackProp = nTrack;
+        nTrackProp.setPID(o2::track::PID::Electron);
+
+        o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, pTrackProp, 2.f, matCorr, &dcaInfo);
+        o2::base::Propagator::Instance()->propagateToDCABxByBz({collision.posX(), collision.posY(), collision.posZ()}, nTrackProp, 2.f, matCorr, &dcaInfo);
+
+        pPropagatedSuccess = o2::base::Propagator::Instance()->propagateToR(pTrackProp, baseR + offsetR);
+        nPropagatedSuccess = o2::base::Propagator::Instance()->propagateToR(nTrackProp, baseR + offsetR);
+
+        if (pPropagatedSuccess && nPropagatedSuccess) {
+          KFPTrack kfp_track_posProp = createKFPTrackFromTrackParCov(pTrackProp, pos.sign(), pos.tpcNClsFound(), pos.tpcChi2NCl());
+          KFPTrack kfp_track_eleProp = createKFPTrackFromTrackParCov(nTrackProp, ele.sign(), ele.tpcNClsFound(), ele.tpcChi2NCl());
+          phiv = o2::aod::pwgem::dilepton::utils::pairutil::getPhivPair(kfp_track_posProp.GetPx(), kfp_track_posProp.GetPy(), kfp_track_posProp.GetPz(), kfp_track_eleProp.GetPx(), kfp_track_eleProp.GetPy(), kfp_track_eleProp.GetPz(), pos.sign(), ele.sign(), d_bz);
+          psipair = o2::aod::pwgem::dilepton::utils::pairutil::getPsiPair(kfp_track_posProp.GetPx(), kfp_track_posProp.GetPy(), kfp_track_posProp.GetPz(), kfp_track_eleProp.GetPx(), kfp_track_eleProp.GetPy(), kfp_track_eleProp.GetPz());
+          break;
+        } else {
+          LOG(debug) << "Propagation to offset" << offsetR << " cm failed for " << (pPropagatedSuccess ? "negative" : "positive") << " track. Trying smaller offset.";
+        }
+      }
+      if (modeTrackPropagation == TrackPropMode::kBoth) {
+        registry.fill(HIST("V0/hPhivPropagationCompare"), phiv, phivFast);
+        registry.fill(HIST("V0/hPsiPairPropagationCompare"), psipair, psipairFast);
       }
-      LOG(debug) << "Propagation to offset" << offsetR << " cm failed for " << (pPropagatedSuccess ? "negative" : "positive") << " track. Trying smaller offset.";
     }
+
     if (phiv == 999.f || psipair == 999.f) {
       LOG(debug) << "Propagation failed for all radii (" << propV0LegsRadius << ", 30, 10 cm). Using default values for phiv and psipair (999.f).";
     }
@@ -985,14 +1065,14 @@ struct PhotonConversionBuilder {
       fillV0Table<isMC, TBCs, TCollisions, TTracks>(v0, true);
     } // end of fullv0Id loop
 
-    for (const auto& collision : collisions) {
-      if constexpr (isMC) {
-        if (!collision.has_mcCollision()) {
-          continue;
-        }
-      }
-      // events_ngpcm(nv0_map[collision.globalIndex()]);
-    } // end of collision loop
+    // for (const auto& collision : collisions) {
+    //   if constexpr (isMC) {
+    //     if (!collision.has_mcCollision()) {
+    //       continue;
+    //     }
+    //   }
+    //   // events_ngpcm(nv0_map[collision.globalIndex()]);
+    // } // end of collision loop
 
     pca_map.clear();
     cospa_map.clear();
diff --git a/PWGEM/PhotonMeson/Utils/PCMUtilities.h b/PWGEM/PhotonMeson/Utils/PCMUtilities.h
@@ -22,6 +22,7 @@
 #include <CommonConstants/MathConstants.h>
 #include <DetectorsBase/Propagator.h>
 #include <Framework/Concepts.h>
+#include <Framework/ASoA.h>
 #include <ReconstructionDataFormats/HelixHelper.h>
 #include <ReconstructionDataFormats/TrackParametrizationWithError.h>
 
@@ -101,10 +102,92 @@ inline void Vtx_recalculationParCov(o2::base::Propagator* prop, const o2::track:
 
   xyz[2] = (trackPosInformationCopy.getZ() * helixNeg.rC + trackNegInformationCopy.getZ() * helixPos.rC) / (helixPos.rC + helixNeg.rC);
 }
+
+//_______________________________________________________________________
+/// \brief Calculate DCA for tracks using the track helix and the inclination angl tan(lambda)
+/// \param trk track parameterization to obtain helix parameters
+/// \param vtx primary vertex position
+/// \param magField magnetic field strenght of L3
+/// \return DCAxy, DCAz
+template <typename TrackPrecision = float>
+std::array<float, 2> CalculateDCAFast(const o2::track::TrackParametrizationWithError<TrackPrecision>& trk, const o2::math_utils::Point3D<float>& vtx, const float magField)
+{
+
+  std::array<float, 2> dca;
+
+  // obtain circle from track in x-y plane
+  const o2::track::TrackAuxPar helixPos(trk, magField);
+
+  // obtain position in global coordinates and tan(lambda)
+  const auto posTrack = trk.getXYZGlo();
+  const float trX = posTrack.X();
+  const float trY = posTrack.Y();
+  const float trZ = posTrack.Z();
+  const float tangentLambda = trk.getTgl();
+
+  // Calculate DCAxy
+  // Use distance in x and y between circle center and vtx, afterwards subtract radius of circle
+  const float distX = helixPos.xC - vtx.X();
+  const float distY = helixPos.yC - vtx.Y();
+  const float trackCircCenter = std::sqrt(distX * distX + distY * distY);
+  dca[0] = helixPos.rC - trackCircCenter;
+
+  // Calculate DCAz
+  // First step: Calculate arc lenght of circle between current position and primary vertex in x-y
+  const float theta0 = std::atan2(trY - helixPos.yC, trX - helixPos.xC);
+  const float thetav = std::atan2(vtx.Y() - helixPos.yC, vtx.X() - helixPos.xC);
+
+  // Make sure angle is between -pi and pi
+  const auto dtheta = RecoDecay::constrainAngle<float>(thetav - theta0, -o2::constants::math::PI);
+
+  // arc-lenght along helix
+  const float arcLenght = std::fabs(helixPos.rC * dtheta);
+
+  // get global z-position at DCA
+  const float ZPosGlo = trZ - arcLenght * tangentLambda;
+
+  // DCA calculated from
+  dca[1] = ZPosGlo - vtx.Z();
+
+  return dca;
+}
+
+//_______________________________________________________________________
+/// \brief Calculate the track momentum at a different place of the track Helix.
+/// \param s arc-lenght where track should be propagated to
+/// \param track particle track
+/// \param trHelix track helix param. for circle approximation
+/// \param bz magnetic field strenght of L3
+/// \param addPhi optional additional rotation of the track in phi-direction
+/// \return track momentum vector at new position
+template <o2::soa::is_iterator TTrack>
+inline std::array<float, 3> getPropMomentumFromTrackHelix(const float s, const TTrack& track, const o2::track::TrackAuxPar& trHelix, const float bz, float addPhi = 0.f)
+{
+
+  // Calculate the change in phi considering the track radius and the track arc lenght s
+  const float dphi = -track.sign() * 0.3f * bz * s / 100. / track.pt(); // s in cm
+  const float dotProd = std::cos(track.phi() + addPhi) * track.pt() * trHelix.xC + std::sin(track.phi() + addPhi) * track.pt() * trHelix.yC;
+  if (dotProd < 0) {
+    addPhi -= o2::constants::math::PI;
+  }
+
+  // Calculate the phi at the secondary vertex
+  const auto phi = RecoDecay::constrainAngle<float>(track.phi() + dphi + addPhi);
+
+  // Calculate px,y,z at the new propagated vertex
+  std::array<float, 3> trackP;
+  trackP[0] = std::cos(phi) * track.pt();
+  trackP[1] = std::sin(phi) * track.pt();
+  trackP[2] = track.tgl() * track.pt();
+
+  return trackP;
+}
+
 //_______________________________________________________________________
 template <typename TrackPrecision = float, o2::soa::is_iterator T1, o2::soa::is_iterator T2>
 inline void Vtx_recalculation(o2::base::Propagator* prop, T1 lTrackPos, T2 lTrackNeg, std::array<float, 3>& xyz, o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrNONE)
 {
+
   // o2::track::TrackParametrizationWithError<TrackPrecision> = TrackParCov, I use the full version to have control over the data type
   o2::track::TrackParametrizationWithError<TrackPrecision> trackPosInformation = getTrackParCov(lTrackPos); // first get an object that stores Track information (positive)
   o2::track::TrackParametrizationWithError<TrackPrecision> trackNegInformation = getTrackParCov(lTrackNeg); // first get an object that stores Track information (negative)
