GPU: Protect also clusters rejected during the fit, use large tube for protection, and small tube for removal

GPU/GPUTracking/DataCompression/GPUTPCClusterRejection.h

@@ -25,18 +25,19 @@ struct GPUTPCClusterRejection {
   {
     (void)counts; // FIXME: Avoid incorrect -Wunused-but-set-parameter warning
     (void)mev200;
+    bool retVal = false;
     if (attach == 0) {
-      return false;
+      retVal = false;
     } else if ((attach & gputpcgmmergertypes::attachGoodLeg) == 0) {
       if constexpr (C) {
         counts->nLoopers++;
       }
-      return true;
+      retVal = true;
     } else if (attach & gputpcgmmergertypes::attachHighIncl) {
       if constexpr (C) {
         counts->nHighIncl++;
       }
-      return true;
+      retVal = true;
     } else if (attach & gputpcgmmergertypes::attachTube) {
       protect = true;
       if constexpr (C) {
@@ -46,17 +47,23 @@ struct GPUTPCClusterRejection {
           counts->nTube++;
         }
       }
-      return false;
+      retVal = false;
     } else if ((attach & gputpcgmmergertypes::attachGood) == 0) {
       protect = true;
       if constexpr (C) {
         counts->nRejected++;
       }
-      return false;
+      retVal = false;
     } else {
       physics = true;
-      return false;
+      retVal = false;
     }
+
+    if (attach & gputpcgmmergertypes::attachProtect) {
+      protect = true;
+      retVal = false;
+    }
+    return retVal;
   }
 
   static constexpr inline bool GetIsRejected(int32_t attach)

GPU/GPUTracking/DataCompression/GPUTPCCompressionKernels.cxx

@@ -56,10 +56,12 @@ GPUdii() void GPUTPCCompressionKernels::Thread<GPUTPCCompressionKernels::step0at
       if ((attach & gputpcgmmergertypes::attachTrackMask) != i) {
         continue; // Main attachment to different track
       }
-      bool rejectCluster = processors.param.rec.tpc.rejectionStrategy >= GPUSettings::RejectionStrategyA && (rejectTrk || GPUTPCClusterRejection::GetIsRejected(attach));
+      bool rejectCluster = processors.param.rec.tpc.rejectionStrategy >= GPUSettings::RejectionStrategyA && !(attach & gputpcgmmergertypes::attachProtect) && (rejectTrk || GPUTPCClusterRejection::GetIsRejected(attach));
       if (rejectCluster) {
         compressor.mClusterStatus[hitId] = 1; // Cluster rejected, do not store
         continue;
+      } else if (processors.param.rec.tpc.rejectionStrategy >= GPUSettings::RejectionStrategyA && rejectTrk) {
+        continue;
       }
 
       if (!(param.rec.tpc.compressionTypeMask & GPUSettings::CompressionTrackModel)) {
@@ -199,6 +201,9 @@ GPUd() bool GPUTPCCompression::rejectCluster(int32_t idx, GPUParam& GPUrestrict(
     if (GPUTPCClusterRejection::GetIsRejected(attach)) {
       return true;
     }
+    if (attach & gputpcgmmergertypes::attachProtect) {
+      return false;
+    }
     int32_t id = attach & gputpcgmmergertypes::attachTrackMask;
     auto& trk = ioPtrs.mergedTracks[id];
     if (CAMath::Abs(trk.GetParam().GetQPt() * param.qptB5Scaler) > param.rec.tpc.rejectQPtB5 || trk.MergedLooper()) {

GPU/GPUTracking/Definitions/GPUSettingsList.h

@@ -68,8 +68,11 @@ AddOptionRTC(minNClustersFinalTrack, int32_t, -1, "", 0, "required min number of
 AddOptionRTC(searchWindowDZDR, float, 2.5f, "", 0, "Use DZDR window for seeding instead of neighboursSearchArea")
 AddOptionRTC(trackReferenceX, float, 1000.f, "", 0, "Transport all tracks to this X after tracking (disabled if > 500, auto = 1000)")
 AddOptionRTC(zsThreshold, float, 2.0f, "", 0, "Zero-Suppression threshold")
-AddOptionRTC(tubeChi2, float, 5.f * 5.f, "", 0, "Max chi2 to mark cluster adjacent to track")
-AddOptionRTC(tubeMaxSize2, float, 2.5f * 2.5f, "", 0, "Square of max tube size (normally derrived from tpcTubeChi2)")
+AddOptionRTC(tubeProtectSigma2, float, 5.f * 5.f, "", 0, "Max sigma2 to mark adjacent cluster for protection")
+AddOptionRTC(tubeProtectMaxSize2, float, 3.5f * 3.5f, "", 0, "Square of max tube size (if smaller than tubeProtectChi2)")
+AddOptionRTC(tubeProtectMinSize2, float, 1.0f * 1.0f, "", 0, "Square of min tube size (if larger than tubeProtectChi2)")
+AddOptionRTC(tubeRemoveSigma2, float, 1.f * 1.f, "", 0, "Max sigma2 to mark adjacent cluster for removal")
+AddOptionRTC(tubeRemoveMaxSize2, float, 1.5f * 1.5f, "", 0, "Square of max tube size (if smaller than tubeRejectChi2)")
 AddOptionRTC(clustersShiftTimebins, float, 0, "", 0, "Shift of TPC clusters (applied during CTF cluster decoding)")
 AddOptionRTC(clustersShiftTimebinsClusterizer, float, 0, "", 0, "Shift of TPC clusters (applied during CTF clusterization)")
 AddOptionRTC(clustersEdgeFixDistance, float, 0.f, "", 0, "If >0, revert cluster.flag edge bit distance to edge exceeds this parameter (fixed during CTF decoding)")

GPU/GPUTracking/Merger/GPUTPCGMMerger.cxx

@@ -1848,7 +1848,11 @@ GPUd() void GPUTPCGMMerger::PrepareForFit1(int32_t nBlocks, int32_t nThreads, in
     GPUTPCGMMergedTrack& trk = mMergedTracks[i];
     if (trk.OK()) {
       for (uint32_t j = 0; j < trk.NClusters(); j++) {
-        mClusterAttachment[mClusters[trk.FirstClusterRef() + j].num] = attachAttached | attachGood;
+        uint32_t weight = attachAttached | attachGood;
+        if (CAMath::Abs(trk.GetParam().GetQPt() * Param().qptB5Scaler) <= Param().rec.tpc.rejectQPtB5 && !trk.MergedLooper() && trk.Leg() == 0) {
+          weight |= attachProtect;
+        }
+        mClusterAttachment[mClusters[trk.FirstClusterRef() + j].num] = weight;
         CAMath::AtomicAdd(&mSharedCount[mClusters[trk.FirstClusterRef() + j].num], 1u);
       }
       if (!trk.CCE() && !trk.MergedLooper()) {
@@ -1896,7 +1900,9 @@ GPUd() void GPUTPCGMMerger::Finalize0(int32_t nBlocks, int32_t nThreads, int32_t
     mTrackSort[mTrackOrderAttach[i]] = i;
   }
   for (uint32_t i = iBlock * nThreads + iThread; i < mMemory->nMergedTrackClusters; i += nThreads * nBlocks) {
-    mClusterAttachment[mClusters[i].num] = 0; // Reset adjacent attachment for attached clusters, set correctly below
+    if (!(mClusterAttachment[mClusters[i].num] & attachProtect)) {
+      mClusterAttachment[mClusters[i].num] = 0; // Reset adjacent attachment for attached clusters, set correctly below
+    }
   }
 }
 
@@ -1919,6 +1925,9 @@ GPUd() void GPUTPCGMMerger::Finalize1(int32_t nBlocks, int32_t nThreads, int32_t
       if (trk.Leg() == 0) {
         weight |= attachGoodLeg;
       }
+      if (CAMath::Abs(trk.GetParam().GetQPt() * Param().qptB5Scaler) <= Param().rec.tpc.rejectQPtB5 && !trk.MergedLooper() && trk.Leg() == 0) {
+        weight |= attachProtect;
+      }
       CAMath::AtomicMax(&mClusterAttachment[id], weight);
     }
   }

GPU/GPUTracking/Merger/GPUTPCGMMergerTypes.h

@@ -21,7 +21,8 @@
 namespace o2::gpu::gputpcgmmergertypes
 {
 
-enum attachTypes { attachAttached = 0x40000000,
+enum attachTypes { attachProtect = 0x80000000,
+                   attachAttached = 0x40000000,
                    attachGood = 0x20000000,
                    attachGoodLeg = 0x10000000,
                    attachTube = 0x08000000,

GPU/GPUTracking/Merger/GPUTPCGMTrackParam.cxx

@@ -514,10 +514,15 @@ GPUd() float GPUTPCGMTrackParam::AttachClusters(const GPUTPCGMMerger* GPUrestric
   const float stepZ = row.HstepZ();
   int32_t bin, ny, nz;
 
+  bool protect = CAMath::Abs(GetQPt() * Merger->Param().qptB5Scaler) <= Merger->Param().rec.tpc.rejectQPtB5 && goodLeg;
+
   float err2Y, err2Z;
   Merger->Param().GetClusterErrors2(sector, iRow, Z, mP[2], mP[3], -1.f, 0.f, 0.f, err2Y, err2Z);                                       // TODO: Use correct time/avgCharge
-  const float sy2 = CAMath::Min(Merger->Param().rec.tpc.tubeMaxSize2, Merger->Param().rec.tpc.tubeChi2 * (err2Y + CAMath::Abs(mC[0]))); // Cov can be bogus when following circle
-  const float sz2 = CAMath::Min(Merger->Param().rec.tpc.tubeMaxSize2, Merger->Param().rec.tpc.tubeChi2 * (err2Z + CAMath::Abs(mC[2]))); // In that case we should provide the track error externally
+  const float tubeMaxSize2 = protect ? Merger->Param().rec.tpc.tubeProtectMaxSize2 : Merger->Param().rec.tpc.tubeRemoveMaxSize2;
+  const float tubeMinSize2 = protect ? Merger->Param().rec.tpc.tubeProtectMinSize2 : 0.f;
+  const float tubeSigma2 = protect ? Merger->Param().rec.tpc.tubeProtectSigma2 : Merger->Param().rec.tpc.tubeRemoveSigma2;
+  const float sy2 = CAMath::Max(tubeMinSize2, CAMath::Min(tubeMaxSize2, tubeSigma2 * (err2Y + CAMath::Abs(mC[0])))); // Cov can be bogus when following circle
+  const float sz2 = CAMath::Max(tubeMinSize2, CAMath::Min(tubeMaxSize2, tubeSigma2 * (err2Z + CAMath::Abs(mC[2])))); // In that case we should provide the track error externally
   const float tubeY = CAMath::Sqrt(sy2);
   const float tubeZ = CAMath::Sqrt(sz2);
   const float sy21 = 1.f / sy2;
@@ -538,6 +543,10 @@ GPUd() float GPUTPCGMTrackParam::AttachClusters(const GPUTPCGMMerger* GPUrestric
   if (goodLeg) {
     myWeight |= gputpcgmmergertypes::attachGoodLeg;
   }
+  if (protect) {
+    myWeight |= gputpcgmmergertypes::attachProtect;
+  }
+
   for (int32_t k = 0; k <= nz; k++) {
     const int32_t mybin = bin + k * nBinsY;
     const uint32_t hitFst = firsthit[mybin];