Store ITS and TPC labels for matched tracks

Author: shahor02

Detectors/GlobalTracking/include/GlobalTracking/MatchTPCITS.h

@@ -138,6 +138,7 @@ struct matchRecord {
 
 class MatchTPCITS
 {
+  using MCLabCont = o2::dataformats::MCTruthContainer<o2::MCCompLabel>;
 
  public:
   ///< perform matching for provided input
@@ -176,6 +177,8 @@ class MatchTPCITS
   void setITSMCTruthBranchName(const std::string& nm) { mITSMCTruthBranchName = nm; }
   void setTPCMCTruthBranchName(const std::string& nm) { mTPCMCTruthBranchName = nm; }
   void setOutTPCITSTracksBranchName(const std::string& nm) { mOutTPCITSTracksBranchName = nm; }
+  void setOutTPCMCTruthBranchName(const std::string& nm) { mOutTPCMCTruthBranchName = nm; }
+  void setOutITSMCTruthBranchName(const std::string& nm) { mOutITSMCTruthBranchName = nm; }
 
   ///< get input branch names for the input from the tree
   const std::string& getITSTrackBranchName() const { return mITSTrackBranchName; }
@@ -184,6 +187,8 @@ class MatchTPCITS
   const std::string& getITSMCTruthBranchName() const { return mITSMCTruthBranchName; }
   const std::string& getTPCMCTruthBranchName() const { return mTPCMCTruthBranchName; }
   const std::string& getOutTPCITSTracksBranchName() const { return mOutTPCITSTracksBranchName; }
+  const std::string& getOutTPCMCTruthBranchName() const { return mOutTPCMCTruthBranchName; }
+  const std::string& getOutITSMCTruthBranchName() const { return mOutITSMCTruthBranchName; }
 
   ///< print settings
   void print() const;
@@ -268,7 +273,7 @@ class MatchTPCITS
   void doMatching(int sec);
 
   void refitWinners();
-  bool refitTrackITSTPC(const TrackLocITS& tITS);
+  bool refitTrackITSTPC(int iITS);
   void selectBestMatches();
   void buildMatch2TrackTables();
   bool validateTPCMatch(int mtID);
@@ -439,6 +444,9 @@ class MatchTPCITS
 
   ///<outputs tracks container
   std::vector<o2::dataformats::TrackTPCITS> mMatchedTracks;
+  std::vector<o2::MCCompLabel> mOutITSLabels; ///< ITS label of matched track
+  std::vector<o2::MCCompLabel> mOutTPCLabels; ///< TPC label of matched track
+
   int mMaxOutputTracksPerEntry = 500; ///< max number of output tracks to store per entry
 
   std::string mITSTrackBranchName = "ITSTrack";          ///< name of branch containing input ITS tracks
@@ -447,6 +455,8 @@ class MatchTPCITS
   std::string mITSMCTruthBranchName = "ITSTrackMCTruth"; ///< name of branch containing ITS MC labels
   std::string mTPCMCTruthBranchName = "TracksMCTruth";   ///< name of branch containing input TPC tracks
   std::string mOutTPCITSTracksBranchName = "TPCITS";     ///< name of branch containing output matched tracks
+  std::string mOutTPCMCTruthBranchName = "MatchTPCMCTruth"; ///< name of branch for output matched tracks TPC MC
+  std::string mOutITSMCTruthBranchName = "MatchITSMCTruth"; ///< name of branch for output matched tracks ITS MC
 
 #ifdef _ALLOW_DEBUG_TREES_
   std::unique_ptr<o2::utils::TreeStreamRedirector> mDBGOut;

Detectors/GlobalTracking/src/MatchTPCITS.cxx

@@ -15,8 +15,6 @@
 #include "Field/MagFieldFast.h"
 #include "ITSBase/GeometryTGeo.h"
 
-#include "SimulationDataFormat/MCTruthContainer.h"
-
 #include "DetectorsBase/Propagator.h"
 #include "CommonUtils/TreeStream.h"
 
@@ -35,7 +33,7 @@
 #include <TFile.h>
 #include <TGeoGlobalMagField.h>
 #include "DataFormatsParameters/GRPObject.h"
-
+#include "SimulationDataFormat/MCTruthContainer.h"
 #include "GlobalTracking/MatchTPCITS.h"
 
 using namespace o2::globaltracking;
@@ -131,9 +129,15 @@ void MatchTPCITS::init()
 
   // create output branch
   if (mOutputTree) {
+    LOG(INFO) << "ITS-TPC Matching results will be stored in the tree " << mOutputTree->GetName();
     mOutputTree->Branch(mOutTPCITSTracksBranchName.data(), &mMatchedTracks);
-    LOG(INFO) << "Matched tracks will be stored in " << mOutTPCITSTracksBranchName << " branch of tree "
-              << mOutputTree->GetName() << FairLogger::endl;
+    LOG(INFO) << "Matched tracks branch: " << mOutTPCITSTracksBranchName;
+    if (mMCTruthON) {
+      mOutputTree->Branch(mOutTPCMCTruthBranchName.data(), &mOutITSLabels);
+      LOG(INFO) << "ITS Tracks Labels branch: " << mOutITSMCTruthBranchName;
+      mOutputTree->Branch(mOutITSMCTruthBranchName.data(), &mOutTPCLabels);
+      LOG(INFO) << "TPC Tracks Labels branch: " << mOutTPCMCTruthBranchName;
+    }
   } else {
     LOG(ERROR) << "Output tree is not attached, matched tracks will not be stored" << FairLogger::endl;
   }
@@ -1068,39 +1072,48 @@ void MatchTPCITS::refitWinners()
   mWinnerChi2Refit.resize(mITSWork.size(), -1.f);
   mCurrITSTracksTreeEntry = -1;
   mCurrITSClustersTreeEntry = -1;
-  for (int i = 0; i < mITSWork.size(); i++) {
-    const auto& its = mITSWork[i];
-    if (!refitTrackITSTPC(its)) {
+  for (int iITS = 0; iITS < mITSWork.size(); iITS++) {
+    if (!refitTrackITSTPC(iITS)) {
       continue;
     }
-    mWinnerChi2Refit[i] = mMatchedTracks.back().getChi2Refit();
+    mWinnerChi2Refit[iITS] = mMatchedTracks.back().getChi2Refit();
     if (mMatchedTracks.size() == mMaxOutputTracksPerEntry) {
       if (mOutputTree) {
         mTimerRefit.Stop();
         mOutputTree->Fill();
         mTimerRefit.Start(false);
       }
       mMatchedTracks.clear();
+      if (mMCTruthON) {
+        mOutITSLabels.clear();
+        mOutTPCLabels.clear();
+      }
     }
   }
   // flush last tracks
   if (mMatchedTracks.size() && mOutputTree) {
     mOutputTree->Fill();
   }
   mMatchedTracks.clear();
+  if (mMCTruthON) {
+    mOutITSLabels.clear();
+    mOutTPCLabels.clear();
+  }
   mTimerRefit.Stop();
 }
 
 //______________________________________________
-bool MatchTPCITS::refitTrackITSTPC(const TrackLocITS& tITS)
+bool MatchTPCITS::refitTrackITSTPC(int iITS)
 {
   ///< refit in inward direction the pair of TPC and ITS tracks
+  const auto& tITS = mITSWork[iITS];
   if (tITS.matchID < 0 || isDisabledITS(mMatchesITS[tITS.matchID])) {
     return false; // no match
   }
   const auto& itsMatch = mMatchesITS[tITS.matchID];
   const auto& itsMatchRec = mMatchRecordsITS[itsMatch.first];
-  const auto& tTPC = mTPCWork[mTPCMatch2Track[itsMatchRec.matchID]];
+  int iTPC = mTPCMatch2Track[itsMatchRec.matchID];
+  const auto& tTPC = mTPCWork[iTPC];
 
   loadITSClustersChunk(tITS.source.getEvent());
   loadITSTracksChunk(tITS.source.getEvent());
@@ -1162,6 +1175,12 @@ bool MatchTPCITS::refitTrackITSTPC(const TrackLocITS& tITS)
   trfit.setTimeMUS(time, timeErr);
   trfit.setRefTPC(tTPC.source);
   trfit.setRefITS(tITS.source);
+
+  if (mMCTruthON) { // store MC info
+    mOutITSLabels.emplace_back(mITSLblWork[iITS]);
+    mOutTPCLabels.emplace_back(mTPCLblWork[iTPC]);
+  }
+
   return true;
 }