Update TrackUtilities.h

Author: njacazio

ALICE3/Core/TrackUtilities.h

@@ -19,6 +19,8 @@
 #ifndef ALICE3_CORE_TRACKUTILITIES_H_
 #define ALICE3_CORE_TRACKUTILITIES_H_
 
+#include <vector>
+
 #include "ReconstructionDataFormats/Track.h"
 #include "Framework/O2DatabasePDGPlugin.h"
 #include "Framework/AnalysisHelpers.h"
@@ -27,50 +29,37 @@
 namespace o2::upgrade
 {
 
-/// Function to convert a McParticle into a perfect Track
-/// \param particle the particle to convert (mcParticle)
+/// Function to convert a TLorentzVector into a perfect Track
+/// \param charge particle charge (integer)
+/// \param particle the particle to convert (TLorentzVector)
+/// \param productionVertex where the particle was produced
 /// \param o2track the address of the resulting TrackParCov
-template <typename McParticleType>
-void convertMCParticleToO2Track(McParticleType& particle,
-                                o2::track::TrackParCov& o2track,
-                                const o2::framework::Service<o2::framework::O2DatabasePDG>& pdg)
+void convertTLorentzVectorToO2Track(const int charge,
+                                    const TLorentzVector particle,
+                                    const std::vector<double> productionVertex,
+                                    o2::track::TrackParCov& o2track)
 {
-  const auto pdgInfo = pdg->GetParticle(particle.pdgCode());
-  int charge = 0;
-  if (pdgInfo != nullptr) {
-    charge = pdgInfo->Charge() / 3;
-  }
   std::array<float, 5> params;
   std::array<float, 15> covm = {0.};
   float s, c, x;
-  o2::math_utils::sincos(particle.phi(), s, c);
-  o2::math_utils::rotateZInv(particle.vx(), particle.vy(), x, params[0], s, c);
-  params[1] = particle.vz();
+  o2::math_utils::sincos(static_cast<float>(particle.Phi()), s, c);
+  o2::math_utils::rotateZInv(static_cast<float>(productionVertex[0]), static_cast<float>(productionVertex[1]), x, params[0], s, c);
+  params[1] = static_cast<float>(productionVertex[2]);
   params[2] = 0.; // since alpha = phi
-  const auto theta = 2. * std::atan(std::exp(-particle.eta()));
+  const auto theta = 2. * std::atan(std::exp(-particle.PseudoRapidity()));
   params[3] = 1. / std::tan(theta);
-  params[4] = charge / particle.pt();
+  params[4] = charge / particle.Pt();
 
   // Initialize TrackParCov in-place
-  new (&o2track)(o2::track::TrackParCov)(x, particle.phi(), params, covm);
-}
-
-/// Function to convert a McParticle into a perfect Track
-/// \param particle the particle to convert (mcParticle)
-/// \param o2track the address of the resulting TrackParCov
-template <typename McParticleType>
-o2::track::TrackParCov convertMCParticleToO2Track(McParticleType& particle,
-                                                  const o2::framework::Service<o2::framework::O2DatabasePDG>& pdg)
-{
-  o2::track::TrackParCov o2track;
-  return convertMCParticleToO2Track(particle, o2track, pdg);
+  new (&o2track)(o2::track::TrackParCov)(x, particle.Phi(), params, covm);
 }
 
 /// Function to convert a TLorentzVector into a perfect Track
 /// \param pdgCode particle pdg
 /// \param particle the particle to convert (TLorentzVector)
 /// \param productionVertex where the particle was produced
 /// \param o2track the address of the resulting TrackParCov
+/// \param pdg the pdg service
 void convertTLorentzVectorToO2Track(int pdgCode,
                                     TLorentzVector particle,
                                     std::vector<double> productionVertex,
@@ -82,19 +71,35 @@ void convertTLorentzVectorToO2Track(int pdgCode,
   if (pdgInfo != nullptr) {
     charge = pdgInfo->Charge() / 3;
   }
-  std::array<float, 5> params;
-  std::array<float, 15> covm = {0.};
-  float s, c, x;
-  o2::math_utils::sincos(static_cast<float>(particle.Phi()), s, c);
-  o2::math_utils::rotateZInv(static_cast<float>(productionVertex[0]), static_cast<float>(productionVertex[1]), x, params[0], s, c);
-  params[1] = static_cast<float>(productionVertex[2]);
-  params[2] = 0;
-  const auto theta = 2. * std::atan(std::exp(-particle.PseudoRapidity()));
-  params[3] = 1. / std::tan(theta);
-  params[4] = charge / particle.Pt();
+  convertTLorentzVectorToO2Track(charge, particle, productionVertex, o2track);
+}
 
-  // Initialize TrackParCov in-place
-  new (&o2track)(o2::track::TrackParCov)(x, particle.Phi(), params, covm);
+/// Function to convert a McParticle into a perfect Track
+/// \param particle the particle to convert (mcParticle)
+/// \param o2track the address of the resulting TrackParCov
+/// \param pdg the pdg service
+template <typename McParticleType>
+void convertMCParticleToO2Track(McParticleType& particle,
+                                o2::track::TrackParCov& o2track,
+                                const o2::framework::Service<o2::framework::O2DatabasePDG>& pdg)
+{
+  static TLorentzVector tlv;
+  tlv.SetPxPyPzE(particle.px(), particle.py(), particle.pz(), particle.e());
+  tlv.SetXYZT(particle.vx(), particle.vy(), particle.vz(), particle.t());
+  convertTLorentzVectorToO2Track(particle.pdgCode(), tlv, {particle.vx(), particle.vy(), particle.vz()}, o2track, pdg);
+}
+
+/// Function to convert a McParticle into a perfect Track
+/// \param particle the particle to convert (mcParticle)
+/// \param o2track the address of the resulting TrackParCov
+/// \param pdg the pdg service
+template <typename McParticleType>
+o2::track::TrackParCov convertMCParticleToO2Track(McParticleType& particle,
+                                                  const o2::framework::Service<o2::framework::O2DatabasePDG>& pdg)
+{
+  o2::track::TrackParCov o2track;
+  convertMCParticleToO2Track(particle, o2track, pdg);
+  return o2track;
 }
 
 } // namespace o2::upgrade