MFT-MCH matching: Implementing save of N best matching candidates (#13478)

* Implementing saving N matching candidates

* Clang format

* Adding comments

* Fixing mistake in merging

Author: mcoquet642

Detectors/GlobalTracking/include/GlobalTracking/MatchGlobalFwd.h

@@ -333,11 +333,14 @@ class MatchGlobalFwd
   std::vector<o2::mft::TrackMFT> mMFTMatchPlaneParams;         ///< MFT track parameters at matching plane
   std::vector<o2::track::TrackParCovFwd> mMCHMatchPlaneParams; ///< MCH track parameters at matching plane
 
+  std::map<int, std::vector<std::pair<int, int>>> mCandidates; ///< map each MCH track id to vector of best match candidates
+
   const o2::itsmft::TopologyDictionary* mMFTDict{nullptr}; // cluster patterns dictionary
   o2::itsmft::ChipMappingMFT mMFTMapping;
   bool mMCTruthON = false;      ///< Flag availability of MC truth
   bool mUseMIDMCHMatch = false; ///< Flag for using MCHMID matches (TrackMCHMID)
-  int mSaveMode = 0;            ///< Output mode [0 = SaveBestMatch; 1 = SaveAllMatches; 2 = SaveTrainingData]
+  int mSaveMode = 0;            ///< Output mode [0 = SaveBestMatch; 1 = SaveAllMatches; 2 = SaveTrainingData; 3 = SaveNCandidates]
+  int mNCandidates = 5;         ///< Numbers of matching candidates to save in savemode=3
   MatchingType mMatchingType = MATCHINGUNDEFINED;
   TGeoManager* mGeoManager;
 };

Detectors/GlobalTracking/include/GlobalTracking/MatchGlobalFwdParam.h

@@ -26,7 +26,8 @@ namespace globaltracking
 
 enum SaveMode { kBestMatch = 0,
                 kSaveAll,
-                kSaveTrainingData };
+                kSaveTrainingData,
+                kSaveNCandidates };
 
 struct GlobalFwdMatchingParam : public o2::conf::ConfigurableParamHelper<GlobalFwdMatchingParam> {
 
@@ -41,6 +42,7 @@ struct GlobalFwdMatchingParam : public o2::conf::ConfigurableParamHelper<GlobalF
   Int_t saveMode = kBestMatch;                            ///< Global Forward Tracks save mode
   float MFTRadLength = 0.042;                             ///< MFT thickness in radiation length
   float alignResidual = 1.;                               ///< Alignment residual for cluster position uncertainty
+  int nCandidates = 5;                                    ///< Number of best matching candidates to save in savemode=3
 
   bool
     isMatchUpstream() const

Detectors/GlobalTracking/src/MatchGlobalFwd.cxx

@@ -10,6 +10,7 @@
 // or submit itself to any jurisdiction.
 
 #include "GlobalTracking/MatchGlobalFwd.h"
+#include <queue>
 
 using namespace o2::globaltracking;
 
@@ -66,6 +67,8 @@ void MatchGlobalFwd::init()
 
   mSaveMode = matchingParam.saveMode;
   LOG(info) << "Save mode MFTMCH candidates = " << mSaveMode;
+
+  mNCandidates = matchingParam.nCandidates;
 }
 
 //_________________________________________________________
@@ -98,6 +101,9 @@ void MatchGlobalFwd::run(const o2::globaltracking::RecoContainer& inp)
           case kSaveTrainingData:
             doMatching<kSaveTrainingData>();
             break;
+          case kSaveNCandidates:
+            doMatching<kSaveNCandidates>();
+            break;
           default:
             LOG(fatal) << "Invalid MFTMCH save mode";
         }
@@ -132,6 +138,7 @@ void MatchGlobalFwd::clear()
   mMatchLabels.clear();
   mMFTTrackROFContMapping.clear();
   mMatchingInfo.clear();
+  mCandidates.clear();
 }
 
 //_________________________________________________________
@@ -400,6 +407,25 @@ void MatchGlobalFwd::doMatching()
     if (mMCTruthON) {
       LOG(info) << "  MFT-MCH Matching: nFakes = " << nFakes << " nTrue = " << nTrue;
     }
+  } else if constexpr (saveAllMode == SaveMode::kSaveNCandidates) {
+    int nFakes = 0, nTrue = 0;
+    auto& matchAllChi2 = mMatchingFunctionMap["matchALL"];
+    for (auto MCHId = 0; MCHId < mMCHWork.size(); MCHId++) {
+      auto& thisMCHTrack = mMCHWork[MCHId];
+      for (auto& pairCandidate : mCandidates[MCHId]) {
+        thisMCHTrack.setMFTTrackID(pairCandidate.second);
+        auto& thisMFTTrack = mMFTWork[pairCandidate.second];
+        auto chi2 = matchAllChi2(thisMCHTrack, thisMFTTrack); // Matching chi2 is stored independently
+        thisMCHTrack.setMFTMCHMatchingScore(pairCandidate.first);
+        thisMCHTrack.setMFTMCHMatchingChi2(chi2);
+        mMatchedTracks.emplace_back(thisMCHTrack);
+        mMatchingInfo.emplace_back(thisMCHTrack);
+        if (mMCTruthON) {
+          mMatchLabels.push_back(computeLabel(MCHId, pairCandidate.second));
+          mMatchLabels.back().isFake() ? nFakes++ : nTrue++;
+        }
+      }
+    }
   }
 }
 
@@ -413,6 +439,10 @@ void MatchGlobalFwd::ROFMatch(int MFTROFId, int firstMCHROFId, int lastMCHROFId)
   const auto& lastMCHROF = mMCHTrackROFRec[lastMCHROFId];
   int nFakes = 0, nTrue = 0;
 
+  auto compare = [](const std::pair<int, int>& a, const std::pair<int, int>& b) {
+    return a.first < b.first;
+  };
+
   auto firstMFTTrackID = thisMFTROF.getFirstEntry();
   auto lastMFTTrackID = firstMFTTrackID + thisMFTROF.getNEntries() - 1;
 
@@ -464,12 +494,21 @@ void MatchGlobalFwd::ROFMatch(int MFTROFId, int firstMCHROFId, int lastMCHROFId)
           }
         }
 
+        if constexpr (saveAllMode == SaveMode::kSaveNCandidates) { // In saveAllmode save all pairs to output container
+          auto score = mMatchFunc(thisMCHTrack, thisMFTTrack);
+          std::pair<int, int> scoreID = {score, MFTId};
+          mCandidates[MCHId].push_back(scoreID);
+          std::sort(mCandidates[MCHId].begin(), mCandidates[MCHId].end(), compare);
+          if (mCandidates[MCHId].size() > mNCandidates) {
+            mCandidates[MCHId].pop_back();
+          }
+        }
+
         if constexpr (saveAllMode == SaveMode::kSaveTrainingData) { // In save training data mode store track parameters at matching plane
           thisMCHTrack.setMFTTrackID(MFTId);
           mMatchingInfo.emplace_back(thisMCHTrack);
           mMCHMatchPlaneParams.emplace_back(thisMCHTrack);
           mMFTMatchPlaneParams.emplace_back(static_cast<o2::mft::TrackMFT>(thisMFTTrack));
-
           if (mMCTruthON) {
             mMatchLabels.push_back(matchLabel);
             mMatchLabels.back().isFake() ? nFakes++ : nTrue++;