ITS: add selections on tracks sharing clusters (#15406)

* ITS: add selections on tracks sharing clusters

* Refactor selection function, return simple int instead of reference to int

* Use isPhiDifferenceBelow, avoid variable narrowing

* Improve memory management for selection of tracks with shared clusters

Author: fchinu

DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TrackITS.h

@@ -192,7 +192,13 @@ class TrackITSExt : public TrackITS
     getClusterRefs().setEntries(ncl);
   }
 
-  GPUhdi() const int& getClusterIndex(int lr) const { return mIndex[lr]; }
+  GPUhdi() const int getClusterIndex(int lr) const { return mIndex[lr]; }
+
+  GPUh() const int getFirstLayerClusterIndex() const
+  {
+    int firstLayer = getFirstClusterLayer();
+    return getClusterIndex(firstLayer);
+  }
 
   GPUhdi() void setExternalClusterIndex(int layer, int idx, bool newCluster = false)
   {

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx

@@ -387,10 +387,10 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
     mTimeFrameGPU->downloadTrackITSExtDevice();
 
     auto& tracks = mTimeFrameGPU->getTrackITSExt();
-    this->acceptTracks(iteration, tracks, firstClusters, sharedFirstClusters);
+    this->acceptTracks(iteration, tracks, firstClusters);
     mTimeFrameGPU->loadUsedClustersDevice();
   }
-  this->markTracks(iteration, sharedFirstClusters);
+  this->markTracks(iteration);
   // wipe the artefact memory
   mTimeFrameGPU->popMemoryStack(iteration);
 };

Detectors/ITSMFT/ITS/tracking/include/ITStracking/Configuration.h

@@ -70,7 +70,6 @@ struct TrackingParameters {
   float DiamondCov[6] = {25.e-6f, 0.f, 0.f, 25.e-6f, 0.f, 36.f};
 
   /// General parameters
-  bool AllowSharingFirstCluster = false;
   int ClusterSharing = 0;
   int MinTrackLength = 7;
   int MaxHoles = 0;
@@ -98,6 +97,12 @@ struct TrackingParameters {
   bool PrintMemory = false; // print allocator usage in epilog report
   size_t MaxMemory = std::numeric_limits<size_t>::max();
   bool DropTFUponFailure = false;
+
+  // Selections on tracks sharing clusters
+  bool AllowSharingFirstCluster = false;
+  float SharedClusterMaxDeltaPhi = 0.05f; // For tracks sharing clusters, maximum allowed delta phi at the cluster position
+  float SharedClusterMaxDeltaEta = 0.03f; // For tracks sharing clusters, maximum allowed delta eta at the cluster position
+  bool SharedClusterOppositeSign = false; // For tracks sharing clusters, require opposite sign of the tracklets
 };
 
 struct VertexingParameters {

Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackerTraits.h

@@ -55,8 +55,8 @@ class TrackerTraits
   template <typename InputSeed>
   void processNeighbours(int iteration, int defaultCellTopologyId, int iLevel, const bounded_vector<InputSeed>& currentCellSeed, const bounded_vector<int>& currentCellId, const bounded_vector<int>& currentCellTopologyId, bounded_vector<TrackSeedN>& updatedCellSeed, bounded_vector<int>& updatedCellId, bounded_vector<int>& updatedCellTopologyId);
 
-  void acceptTracks(int iteration, bounded_vector<TrackITSExt>& tracks, bounded_vector<bounded_vector<int>>& firstClusters, bounded_vector<bounded_vector<int>>& sharedFirstClusters);
-  void markTracks(int iteration, bounded_vector<bounded_vector<int>>& sharedFirstClusters);
+  void acceptTracks(int iteration, bounded_vector<TrackITSExt>& tracks, bounded_vector<bounded_vector<int>>& firstClusters);
+  void markTracks(int iteration);
 
   void updateTrackingParameters(const std::vector<TrackingParameters>& trkPars)
   {

Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackingConfigParam.h

@@ -102,7 +102,12 @@ struct TrackerParamConfig : public o2::conf::ConfigurableParamHelper<TrackerPara
   size_t maxMemory = std::numeric_limits<size_t>::max();
   bool dropTFUponFailure = false;
   bool fataliseUponFailure = true;       // granular management of the fatalisation in async mode
+
+  // Selections on tracks sharing clusters
   bool allowSharingFirstCluster = false; // allow first cluster sharing among tracks
+  float sharedClusterMaxDeltaPhi = 0.05f; // Maximum allowed delta phi at the cluster position
+  float sharedClusterMaxDeltaEta = 0.03f; // Maximum allowed delta eta at the cluster position
+  bool sharedClusterOppositeSign = false; // Require opposite sign of the tracklets
 
   O2ParamDef(TrackerParamConfig, "ITSCATrackerParam");
 };

Detectors/ITSMFT/ITS/tracking/src/Configuration.cxx

@@ -204,6 +204,9 @@ std::vector<TrackingParameters> TrackingMode::getTrackingParameters(TrackingMode
     p.SaveTimeBenchmarks = tc.saveTimeBenchmarks;
     p.FataliseUponFailure = tc.fataliseUponFailure;
     p.AllowSharingFirstCluster = tc.allowSharingFirstCluster;
+    p.SharedClusterMaxDeltaPhi = tc.sharedClusterMaxDeltaPhi;
+    p.SharedClusterMaxDeltaEta = tc.sharedClusterMaxDeltaEta;
+    p.SharedClusterOppositeSign = tc.sharedClusterOppositeSign;
     const auto iter = &p - trackParams.data();
     if (iter < constants::MaxIter) {
       p.MaxHoles = tc.maxHolesIter[iter];

Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx

@@ -661,9 +661,7 @@ template <int NLayers>
 void TrackerTraits<NLayers>::findRoads(const int iteration)
 {
   bounded_vector<bounded_vector<int>> firstClusters(mTrkParams[iteration].NLayers, bounded_vector<int>(mMemoryPool.get()), mMemoryPool.get());
-  bounded_vector<bounded_vector<int>> sharedFirstClusters(mTrkParams[iteration].NLayers, bounded_vector<int>(mMemoryPool.get()), mMemoryPool.get());
   firstClusters.resize(mTrkParams[iteration].NLayers);
-  sharedFirstClusters.resize(mTrkParams[iteration].NLayers);
   const auto propagator = o2::base::Propagator::Instance();
   const TrackingFrameInfo* tfInfos[NLayers]{};
   const Cluster* unsortedClusters[NLayers]{};
@@ -787,13 +785,13 @@ void TrackerTraits<NLayers>::findRoads(const int iteration)
       return track::isBetter(a, b);
     });
 
-    acceptTracks(iteration, tracks, firstClusters, sharedFirstClusters);
+    acceptTracks(iteration, tracks, firstClusters);
   }
-  markTracks(iteration, sharedFirstClusters);
+  markTracks(iteration);
 }
 
 template <int NLayers>
-void TrackerTraits<NLayers>::acceptTracks(int iteration, bounded_vector<TrackITSExt>& tracks, bounded_vector<bounded_vector<int>>& firstClusters, bounded_vector<bounded_vector<int>>& sharedFirstClusters)
+void TrackerTraits<NLayers>::acceptTracks(int iteration, bounded_vector<TrackITSExt>& tracks, bounded_vector<bounded_vector<int>>& firstClusters)
 {
   auto& trks = mTimeFrame->getTracks();
   trks.reserve(trks.size() + tracks.size());
@@ -860,34 +858,51 @@ void TrackerTraits<NLayers>::acceptTracks(int iteration, bounded_vector<TrackITS
 
     if (mTrkParams[iteration].AllowSharingFirstCluster) {
       firstClusters[firstLayer].push_back(firstCluster);
-      if (isFirstShared) {
-        sharedFirstClusters[firstLayer].push_back(firstCluster);
-      }
     }
   }
 }
 
 template <int NLayers>
-void TrackerTraits<NLayers>::markTracks(int iteration, bounded_vector<bounded_vector<int>>& sharedFirstClusters)
+void TrackerTraits<NLayers>::markTracks(int iteration)
 {
   if (mTrkParams[iteration].AllowSharingFirstCluster) {
     /// Now we have to set the shared cluster flag
-    for (int iLayer{0}; iLayer < mTrkParams[iteration].NLayers; ++iLayer) {
-      std::sort(sharedFirstClusters[iLayer].begin(), sharedFirstClusters[iLayer].end());
-    }
+    auto& tracks = mTimeFrame->getTracks();
 
-    for (auto& track : mTimeFrame->getTracks()) {
-      int firstLayer{mTrkParams[iteration].NLayers}, firstCluster{constants::UnusedIndex};
-      for (int iLayer{0}; iLayer < mTrkParams[iteration].NLayers; ++iLayer) {
-        if (track.getClusterIndex(iLayer) == constants::UnusedIndex) {
-          continue;
-        }
-        firstLayer = iLayer;
-        firstCluster = track.getClusterIndex(iLayer);
-        break;
+    bounded_vector<int> fclusSort(tracks.size(), mMemoryPool.get());
+    std::iota(fclusSort.begin(), fclusSort.end(), 0);
+    std::sort(fclusSort.begin(), fclusSort.end(), [&tracks](int a, int b) {
+      return tracks[a].getFirstLayerClusterIndex() < tracks[b].getFirstLayerClusterIndex();
+    });
+
+    auto areTracksSelected = [this, iteration](const TrackITSExt& t1, const TrackITSExt& t2) {
+      const auto t1FirstLayer{t1.getFirstClusterLayer()}, t2FirstLayer{t2.getFirstClusterLayer()};
+      if (t1FirstLayer != t2FirstLayer) {
+        return false;
+      }
+      if (mTimeFrame->getClusterROF(t1FirstLayer, t1.getClusterIndex(t1FirstLayer)) != mTimeFrame->getClusterROF(t2FirstLayer, t2.getClusterIndex(t2FirstLayer))) {
+        return false;
+      }
+      if (!math_utils::isPhiDifferenceBelow(t1.getPhi(), t2.getPhi(), mTrkParams[iteration].SharedClusterMaxDeltaPhi)) {
+        return false;
+      }
+      if (std::abs(t1.getEta() - t2.getEta()) > mTrkParams[iteration].SharedClusterMaxDeltaEta) {
+        return false;
       }
-      if (std::binary_search(sharedFirstClusters[firstLayer].begin(), sharedFirstClusters[firstLayer].end(), firstCluster)) {
-        track.setSharedClusters();
+      if (mTrkParams[iteration].SharedClusterOppositeSign && t1.getSign() == t2.getSign()) {
+        return false;
+      }
+      return true;
+    };
+
+    for (int i{0}; i < static_cast<int>(fclusSort.size()); ++i) {
+      auto& track = tracks[fclusSort[i]];
+      for (int j{i + 1}; j < static_cast<int>(fclusSort.size()) && tracks[fclusSort[j]].getFirstLayerClusterIndex() == track.getFirstLayerClusterIndex(); ++j) {
+        auto& track2 = tracks[fclusSort[j]];
+        if (areTracksSelected(track, track2)) {
+          track.setSharedClusters();
+          track2.setSharedClusters();
+        }
       }
     }
   }