[PWGEM/Dilepton,Common] update muon analyses (#13348)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

Author: dsekihat

Common/Core/fwdtrackUtilities.h

@@ -18,6 +18,7 @@
 #ifndef COMMON_CORE_FWDTRACKUTILITIES_H_
 #define COMMON_CORE_FWDTRACKUTILITIES_H_
 
+#include "Framework/AnalysisDataModel.h"
 #include <DetectorsBase/GeometryManager.h>
 #include <Field/MagneticField.h>
 #include <GlobalTracking/MatchGlobalFwd.h>
@@ -29,6 +30,7 @@
 #include <Math/SMatrix.h>
 #include <TGeoGlobalMagField.h>
 
+#include <type_traits>
 #include <vector>
 
 namespace o2::aod
@@ -40,30 +42,71 @@ enum class propagationPoint : int {
   kToVertex = 0,
   kToDCA = 1,
   kToRabs = 2,
+  kToMatchingPlane = 3,
 };
 using SMatrix55 = ROOT::Math::SMatrix<double, 5, 5, ROOT::Math::MatRepSym<double, 5>>;
 using SMatrix55Std = ROOT::Math::SMatrix<double, 5>;
 using SMatrix5 = ROOT::Math::SVector<double, 5>;
 
-/// propagate fwdtrack to a certain point.
-template <typename TFwdTrack, typename TCollision>
-o2::dataformats::GlobalFwdTrack propagateMuon(TFwdTrack const& muon, TCollision const& collision, const propagationPoint endPoint)
+template <typename TFwdTrack, typename TFwdTrackCov>
+o2::track::TrackParCovFwd getTrackParCovFwd(TFwdTrack const& track, TFwdTrackCov const& cov)
 {
-  double chi2 = muon.chi2();
-  SMatrix5 tpars(muon.x(), muon.y(), muon.phi(), muon.tgl(), muon.signed1Pt());
-  std::vector<double> v1{muon.cXX(), muon.cXY(), muon.cYY(), muon.cPhiX(), muon.cPhiY(),
-                         muon.cPhiPhi(), muon.cTglX(), muon.cTglY(), muon.cTglPhi(), muon.cTglTgl(),
-                         muon.c1PtX(), muon.c1PtY(), muon.c1PtPhi(), muon.c1PtTgl(), muon.c1Pt21Pt2()};
+  // This function works for both (saMuon, saMuon) and (MFTTrack, MFTTrackCov).
+  // Don't use covariant matrix of global muons stored in AO2D.root.
+
+  double chi2 = track.chi2();
+  if constexpr (std::is_same_v<std::decay_t<TFwdTrackCov>, aod::MFTTracksCov::iterator>) {
+    chi2 = track.chi2();
+  } else {
+    if (track.trackType() == o2::aod::fwdtrack::ForwardTrackTypeEnum::GlobalMuonTrack) {
+      chi2 = track.chi2();
+    } else if (track.trackType() == o2::aod::fwdtrack::ForwardTrackTypeEnum::MuonStandaloneTrack) {
+      chi2 = track.chi2() * (2.f * track.nClusters() - 5.f);
+    }
+  }
+
+  SMatrix5 tpars(track.x(), track.y(), track.phi(), track.tgl(), track.signed1Pt());
+  std::vector<double> v1{cov.cXX(), cov.cXY(), cov.cYY(), cov.cPhiX(), cov.cPhiY(),
+                         cov.cPhiPhi(), cov.cTglX(), cov.cTglY(), cov.cTglPhi(), cov.cTglTgl(),
+                         cov.c1PtX(), cov.c1PtY(), cov.c1PtPhi(), cov.c1PtTgl(), cov.c1Pt21Pt2()};
   SMatrix55 tcovs(v1.begin(), v1.end());
-  o2::track::TrackParCovFwd fwdtrack{muon.z(), tpars, tcovs, chi2};
+  o2::track::TrackParCovFwd trackparCov{track.z(), tpars, tcovs, chi2}; // this is chi2! Not chi2/ndf.
+  v1.clear();
+  v1.shrink_to_fit();
+  return trackparCov;
+}
+
+/// propagate fwdtrack to a certain point.
+template <typename TFwdTrack, typename TFwdTrackCov, typename TCollision>
+o2::dataformats::GlobalFwdTrack propagateMuon(TFwdTrack const& muon, TFwdTrackCov const& cov, TCollision const& collision, const propagationPoint endPoint, const float matchingZ, const float bzkG)
+{
+  o2::track::TrackParCovFwd trackParCovFwd;
+  if (muon.trackType() == o2::aod::fwdtrack::ForwardTrackTypeEnum::GlobalMuonTrack) {
+    trackParCovFwd = getTrackParCovFwd(muon, cov);
+  } else if (muon.trackType() == o2::aod::fwdtrack::ForwardTrackTypeEnum::MuonStandaloneTrack) {
+    trackParCovFwd = getTrackParCovFwd(muon, muon);
+  } else {
+    trackParCovFwd = getTrackParCovFwd(muon, muon);
+  }
+
+  o2::dataformats::GlobalFwdTrack propmuon = propagateTrackParCovFwd(trackParCovFwd, muon.trackType(), collision, endPoint, matchingZ, bzkG);
+  return propmuon;
+}
+
+template <typename TFwdTrackParCov, typename TCollision>
+o2::dataformats::GlobalFwdTrack propagateTrackParCovFwd(TFwdTrackParCov const& fwdtrackORG, uint8_t trackType, TCollision const& collision, const propagationPoint endPoint, const float matchingZ, const float bzkG)
+{
+  // TFwdTrackParCov is o2::track::TrackParCovFwd
+
+  o2::track::TrackParCovFwd fwdtrack(fwdtrackORG);
   o2::dataformats::GlobalFwdTrack propmuon;
   o2::globaltracking::MatchGlobalFwd mMatching;
 
-  if (static_cast<int>(muon.trackType()) > 2) { // MCH-MID or MCH standalone
+  if (trackType > 2) { // MCH-MID or MCH standalone
     o2::dataformats::GlobalFwdTrack track;
-    track.setParameters(tpars);
+    track.setParameters(fwdtrack.getParameters());
     track.setZ(fwdtrack.getZ());
-    track.setCovariances(tcovs);
+    track.setCovariances(fwdtrack.getCovariances());
     auto mchTrack = mMatching.FwdtoMCH(track);
 
     if (endPoint == propagationPoint::kToVertex) {
@@ -72,37 +115,42 @@ o2::dataformats::GlobalFwdTrack propagateMuon(TFwdTrack const& muon, TCollision
       o2::mch::TrackExtrap::extrapToVertexWithoutBranson(mchTrack, collision.posZ());
     } else if (endPoint == propagationPoint::kToRabs) {
       o2::mch::TrackExtrap::extrapToZ(mchTrack, -505.);
+    } else if (endPoint == propagationPoint::kToMatchingPlane) {
+      o2::mch::TrackExtrap::extrapToVertexWithoutBranson(mchTrack, matchingZ);
     }
 
     auto proptrack = mMatching.MCHtoFwd(mchTrack);
     propmuon.setParameters(proptrack.getParameters());
     propmuon.setZ(proptrack.getZ());
     propmuon.setCovariances(proptrack.getCovariances());
-  } else if (static_cast<int>(muon.trackType()) < 2) { // MFT-MCH-MID
-    const double centerMFT[3] = {0, 0, -61.4};
-    o2::field::MagneticField* field = static_cast<o2::field::MagneticField*>(TGeoGlobalMagField::Instance()->GetField());
-    auto Bz = field->getBz(centerMFT); // Get field at centre of MFT
-    auto geoMan = o2::base::GeometryManager::meanMaterialBudget(muon.x(), muon.y(), muon.z(), collision.posX(), collision.posY(), collision.posZ());
+  } else if (trackType < 2) { // MFT-MCH-MID
+    // const double centerMFT[3] = {0, 0, -61.4};
+    // o2::field::MagneticField* field = static_cast<o2::field::MagneticField*>(TGeoGlobalMagField::Instance()->GetField());
+    // auto Bz = field->getBz(centerMFT); // Get field at centre of MFT in kG.
+
+    auto geoMan = o2::base::GeometryManager::meanMaterialBudget(fwdtrack.getX(), fwdtrack.getY(), fwdtrack.getZ(), collision.posX(), collision.posY(), collision.posZ());
     auto x2x0 = static_cast<float>(geoMan.meanX2X0);
+
     if (endPoint == propagationPoint::kToVertex) {
-      fwdtrack.propagateToVtxhelixWithMCS(collision.posZ(), {collision.posX(), collision.posY()}, {collision.covXX(), collision.covYY()}, Bz, x2x0);
+      fwdtrack.propagateToVtxhelixWithMCS(collision.posZ(), {collision.posX(), collision.posY()}, {collision.covXX(), collision.covYY()}, bzkG, x2x0);
     } else if (endPoint == propagationPoint::kToDCA) {
-      fwdtrack.propagateToZhelix(collision.posZ(), Bz);
+      fwdtrack.propagateToZhelix(collision.posZ(), bzkG);
+    } else if (endPoint == propagationPoint::kToMatchingPlane) {
+      fwdtrack.propagateToZhelix(matchingZ, bzkG);
     }
     propmuon.setParameters(fwdtrack.getParameters());
     propmuon.setZ(fwdtrack.getZ());
     propmuon.setCovariances(fwdtrack.getCovariances());
   }
 
-  v1.clear();
-  v1.shrink_to_fit();
-
   return propmuon;
 }
 
 template <typename TFwdTrack, typename TMFTTrack>
 o2::dataformats::GlobalFwdTrack refitGlobalMuonCov(TFwdTrack const& muon, TMFTTrack const& mft)
 {
+  // TFwdTrack and TMFTTrack are o2::track::TrackParCovFwd.
+
   auto muonCov = muon.getCovariances();
   auto mftCov = mft.getCovariances();
 
@@ -135,6 +183,7 @@ o2::dataformats::GlobalFwdTrack refitGlobalMuonCov(TFwdTrack const& muon, TMFTTr
 
   o2::dataformats::GlobalFwdTrack globalTrack;
   globalTrack.setParameters(mft.getParameters());
+  globalTrack.setZ(mft.getZ());
   globalTrack.setInvQPt(invQPtGlob);
   globalTrack.setCovariances(globalCov);
 

PWGEM/Dilepton/DataModel/dileptonTables.h

@@ -706,12 +706,17 @@ DECLARE_SOA_COLUMN(MFTTrackId, mfttrackId, int);
 DECLARE_SOA_COLUMN(MCHTrackId, mchtrackId, int);                                     //!
 DECLARE_SOA_SELF_ARRAY_INDEX_COLUMN(GlobalMuonsWithSameMFT, globalMuonsWithSameMFT); //! self indices to global muons that have the same MFTTrackId
 DECLARE_SOA_SELF_ARRAY_INDEX_COLUMN(AmbiguousMuons, ambiguousMuons);
-DECLARE_SOA_COLUMN(CXXatDCA, cXXatDCA, float);                  //! DCAx resolution squared at DCA
-DECLARE_SOA_COLUMN(CYYatDCA, cYYatDCA, float);                  //! DCAy resolution squared at DCA
-DECLARE_SOA_COLUMN(CXYatDCA, cXYatDCA, float);                  //! correlation term of DCAx,y resolution at DCA
-DECLARE_SOA_COLUMN(PtMatchedMCHMID, ptMatchedMCHMID, float);    //! pt of MCH-MID track in MFT-MCH-MID track at PV
-DECLARE_SOA_COLUMN(EtaMatchedMCHMID, etaMatchedMCHMID, float);  //! eta of MCH-MID track in MFT-MCH-MID track at PV
-DECLARE_SOA_COLUMN(PhiMatchedMCHMID, phiMatchedMCHMID, float);  //! phi of MCH-MID track in MFT-MCH-MID track at PV
+DECLARE_SOA_COLUMN(CXXatDCA, cXXatDCA, float);                         //! DCAx resolution squared at DCA
+DECLARE_SOA_COLUMN(CYYatDCA, cYYatDCA, float);                         //! DCAy resolution squared at DCA
+DECLARE_SOA_COLUMN(CXYatDCA, cXYatDCA, float);                         //! correlation term of DCAx,y resolution at DCA
+DECLARE_SOA_COLUMN(PtMatchedMCHMID, ptMatchedMCHMID, float);           //! pt of MCH-MID track in MFT-MCH-MID track at PV
+DECLARE_SOA_COLUMN(EtaMatchedMCHMID, etaMatchedMCHMID, float);         //! eta of MCH-MID track in MFT-MCH-MID track at PV
+DECLARE_SOA_COLUMN(PhiMatchedMCHMID, phiMatchedMCHMID, float);         //! phi of MCH-MID track in MFT-MCH-MID track at PV
+DECLARE_SOA_COLUMN(EtaMatchedMCHMIDatMP, etaMatchedMCHMIDatMP, float); //! eta of MCH-MID track in MFT-MCH-MID track at matching plane
+DECLARE_SOA_COLUMN(PhiMatchedMCHMIDatMP, phiMatchedMCHMIDatMP, float); //! phi of MCH-MID track in MFT-MCH-MID track at matching plane
+DECLARE_SOA_COLUMN(EtaMatchedMFTatMP, etaMatchedMFTatMP, float);       //! eta of MFT track in MFT-MCH-MID track at matching plane
+DECLARE_SOA_COLUMN(PhiMatchedMFTatMP, phiMatchedMFTatMP, float);       //! phi of MFT track in MFT-MCH-MID track at matching plane
+
 DECLARE_SOA_COLUMN(IsAssociatedToMPC, isAssociatedToMPC, bool); //! is associated to most probable collision
 DECLARE_SOA_COLUMN(IsAmbiguous, isAmbiguous, bool);             //! is ambiguous
 DECLARE_SOA_COLUMN(Sign, sign, int8_t);                         //!
@@ -742,7 +747,7 @@ DECLARE_SOA_DYNAMIC_COLUMN(MFTClusterMap, mftClusterMap, //! MFT cluster map, on
                              return clmap;
                            });
 } // namespace emprimarymuon
-DECLARE_SOA_TABLE(EMPrimaryMuons, "AOD", "EMPRIMARYMU", //!
+DECLARE_SOA_TABLE(EMPrimaryMuons_000, "AOD", "EMPRIMARYMU", //!
                   o2::soa::Index<>, emprimarymuon::CollisionId,
                   emprimarymuon::FwdTrackId, emprimarymuon::MFTTrackId, emprimarymuon::MCHTrackId, fwdtrack::TrackType,
                   fwdtrack::Pt, fwdtrack::Eta, fwdtrack::Phi, emprimarymuon::Sign,
@@ -763,6 +768,31 @@ DECLARE_SOA_TABLE(EMPrimaryMuons, "AOD", "EMPRIMARYMU", //!
                   emprimarymuon::Px<fwdtrack::Pt, fwdtrack::Phi>,
                   emprimarymuon::Py<fwdtrack::Pt, fwdtrack::Phi>,
                   emprimarymuon::Pz<fwdtrack::Pt, fwdtrack::Eta>);
+
+DECLARE_SOA_TABLE_VERSIONED(EMPrimaryMuons_001, "AOD", "EMPRIMARYMU", 1, //!
+                            o2::soa::Index<>, emprimarymuon::CollisionId,
+                            emprimarymuon::FwdTrackId, emprimarymuon::MFTTrackId, emprimarymuon::MCHTrackId, fwdtrack::TrackType,
+                            fwdtrack::Pt, fwdtrack::Eta, fwdtrack::Phi, emprimarymuon::Sign,
+                            fwdtrack::FwdDcaX, fwdtrack::FwdDcaY, emprimarymuon::CXXatDCA, emprimarymuon::CYYatDCA, emprimarymuon::CXYatDCA,
+                            emprimarymuon::PtMatchedMCHMID, emprimarymuon::EtaMatchedMCHMID, emprimarymuon::PhiMatchedMCHMID,
+                            emprimarymuon::EtaMatchedMCHMIDatMP, emprimarymuon::PhiMatchedMCHMIDatMP,
+                            emprimarymuon::EtaMatchedMFTatMP, emprimarymuon::PhiMatchedMFTatMP,
+
+                            fwdtrack::NClusters, fwdtrack::PDca, fwdtrack::RAtAbsorberEnd,
+                            fwdtrack::Chi2, fwdtrack::Chi2MatchMCHMID, fwdtrack::Chi2MatchMCHMFT,
+                            fwdtrack::MCHBitMap, fwdtrack::MIDBitMap, fwdtrack::MIDBoards,
+                            fwdtrack::MFTClusterSizesAndTrackFlags, emprimarymuon::Chi2MFT, emprimarymuon::IsAssociatedToMPC, emprimarymuon::IsAmbiguous,
+
+                            // dynamic column
+                            emprimarymuon::Signed1Pt<fwdtrack::Pt, emprimarymuon::Sign>,
+                            emprimarymuon::NClustersMFT<fwdtrack::MFTClusterSizesAndTrackFlags>,
+                            emprimarymuon::MFTClusterMap<fwdtrack::MFTClusterSizesAndTrackFlags>,
+                            emprimarymuon::P<fwdtrack::Pt, fwdtrack::Eta>,
+                            emprimarymuon::Px<fwdtrack::Pt, fwdtrack::Phi>,
+                            emprimarymuon::Py<fwdtrack::Pt, fwdtrack::Phi>,
+                            emprimarymuon::Pz<fwdtrack::Pt, fwdtrack::Eta>);
+
+using EMPrimaryMuons = EMPrimaryMuons_001;
 // iterators
 using EMPrimaryMuon = EMPrimaryMuons::iterator;