GPU TPC: Interpolate missing pad rows when rebuilding track

davidrohr · davidrohr · commit 6b2534cc32b2 · 2025-11-02T15:28:00.000+01:00
diff --git a/GPU/GPUTracking/Merger/GPUTPCGMMergerTypes.h b/GPU/GPUTracking/Merger/GPUTPCGMMergerTypes.h
@@ -34,11 +34,13 @@ enum attachTypes { attachProtect = 0x80000000,
 struct InterpolationErrorHit {
   float posY, posZ;
   GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A errorY, errorZ;
+  bool isValid() const { return errorY >= (GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A)0; }
+  void markInvalid() { errorY = (GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A)-1; }
 };
 
 struct InterpolationErrors {
-  InterpolationErrorHit hit[GPUCA_MERGER_MAX_TRACK_CLUSTERS];
   static constexpr size_t size = GPUCA_MERGER_MAX_TRACK_CLUSTERS;
+  InterpolationErrorHit hit[size];
 };
 
 struct GPUResolveSharedMemory : public GPUKernelTemplate::GPUSharedMemoryScan64<int16_t, GPUCA_GET_THREAD_COUNT(GPUCA_LB_GPUTPCGMMergerResolve_step3)> {
diff --git a/GPU/GPUTracking/Merger/GPUTPCGMTrackParam.cxx b/GPU/GPUTracking/Merger/GPUTPCGMTrackParam.cxx
@@ -67,12 +67,11 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger& GPUrestrict() merger, int32_
   prop.SetPolynomialField(&param.polynomialField);
   prop.SetMaxSinPhi(maxSinPhi);
   if (param.rec.tpc.mergerInterpolateErrors && !rebuilt) {
-    for (uint32_t i = 0; i < interpolation.size; i++) { // TODO: Tune the zeroing size
-      interpolation.hit[i].errorY = -1;
+    for (uint32_t i = 0; i < interpolation.size; i++) { // TODO: Tune the zeroing range
+      interpolation.hit[i].markInvalid();
     }
     if (param.rec.tpc.rebuildTrackInFit) {
-      merger.TrackRebuildHelper()[iTrk].highInclRowLow = 255;
-      merger.TrackRebuildHelper()[iTrk].highInclRowHigh = 255;
+      merger.TrackRebuildHelper()[iTrk].highInclRowLow = merger.TrackRebuildHelper()[iTrk].highInclRowHigh = 255;
     }
   }
 
@@ -140,22 +139,24 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger& GPUrestrict() merger, int32_
       CADEBUG(printf("\tHit %3d/%3d Row %3d: Cluster Alpha %8.3f %3d, X %8.3f - Y %8.3f, Z %8.3f (Missed %d)\n", ihit, maxN, (int32_t)clusters[ihit].row, clAlpha, (int32_t)clusters[ihit].sector, xx, yy, zz, nMissed));
 
       uint8_t dEdxSubThresholdRow = 255;
-      bool doInterpolate = param.rec.tpc.rebuildTrackInFit && (iWay == nWays - 3 || iWay == nWays - 2);
       if (lastPropagateRow != 255 && CAMath::Abs(cluster.row - lastPropagateRow) > 1) {
-        bool dodEdx = param.dodEdxEnabled && param.rec.tpc.adddEdxSubThresholdClusters && finalFit && CAMath::Abs(cluster.row - lastUpdateRow) == 2 && cluster.sector == lastSector && currentClusterStatus == 0;
-        bool doAttach = allowChangeClusters && !param.rec.tpc.rebuildTrackInFit && !(merger.Param().rec.tpc.disableRefitAttachment & 2);
-        if (dodEdx || doAttach || doInterpolate) {
+        const bool dodEdx = param.dodEdxEnabled && param.rec.tpc.adddEdxSubThresholdClusters && finalFit && CAMath::Abs(cluster.row - lastUpdateRow) == 2 && cluster.sector == lastSector && currentClusterStatus == 0;
+        const bool doAttach = allowChangeClusters && !param.rec.tpc.rebuildTrackInFit && !(merger.Param().rec.tpc.disableRefitAttachment & 2);
+        const uint8_t doInterpolate = (param.rec.tpc.rebuildTrackInFit && iWay == nWays - 3) ? 1 : ((param.rec.tpc.rebuildTrackInFit && iWay == nWays - 2) ? 2 : 0);
+        if (lastUpdateRow != 255 && ( dodEdx || doAttach || doInterpolate)) {
           int32_t step = cluster.row > lastPropagateRow ? 1 : -1;
-          for (int32_t iRow = lastPropagateRow + step; iRow != cluster.row; iRow += step) {
+          uint8_t sector = lastSector;
+
+          for (int32_t iRow = lastPropagateRow + step, index = interpolationIndex; iRow != cluster.row; iRow += step, index += wayDirection) {
             float tmpX, tmpY, tmpZ;
-            if (prop.GetPropagatedYZ(mX - GPUTPCGeometry::Row2X(iRow - step) + GPUTPCGeometry::Row2X(iRow), tmpY, tmpZ)) {
+            if (prop.GetPropagatedYZ(GPUTPCGeometry::Row2X(iRow), tmpY, tmpZ)) {
               break;
             }
-            merger.GetConstantMem()->calibObjects.fastTransformHelper->InverseTransformYZtoX(cluster.sector, iRow, tmpY, tmpZ, tmpX);
+            merger.GetConstantMem()->calibObjects.fastTransformHelper->InverseTransformYZtoX(sector, iRow, tmpY, tmpZ, tmpX);
             if (prop.PropagateToXAlpha(tmpX, prop.GetAlpha(), inFlyDirection)) {
               break;
             }
-            FitAddRow(iRow, cluster.sector, iTrk, track, prop, inFlyDirection, merger, &dEdxSubThresholdRow, dodEdx, doAttach, doInterpolate);
+            FitAddRow(iRow, sector, iTrk, track, prop, inFlyDirection, merger, &dEdxSubThresholdRow, dodEdx, doAttach, doInterpolate, interpolation.hit[index], deltaZ, sumInvSqrtCharge, nAvgCharge);
           }
         }
         interpolationIndex += (CAMath::Abs(cluster.row - lastPropagateRow) - 1) * wayDirection;
@@ -239,6 +240,7 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger& GPUrestrict() merger, int32_
               merger.TrackRebuildHelper()[iTrk].highInclRowHigh = cluster.row;
             }
           }
+          // TODO: We can perhaps break here, if we pick up remaining rows
         }
         continue;
       }
@@ -304,7 +306,7 @@ GPUd() bool GPUTPCGMTrackParam::Fit(GPUTPCGMMerger& GPUrestrict() merger, int32_
   return true;
 }
 
-GPUdii() void GPUTPCGMTrackParam::FitAddRow(const int32_t iRow, const uint8_t sector, const int32_t iTrk, const GPUTPCGMMergedTrack& GPUrestrict() track, GPUTPCGMPropagator& GPUrestrict() prop, const bool inFlyDirection, GPUTPCGMMerger& GPUrestrict() merger, uint8_t* GPUrestrict() dEdxSubThresholdRow, const bool dodEdx, const bool doAttach, const bool doInterpolate)
+GPUdii() void GPUTPCGMTrackParam::FitAddRow(const int32_t iRow, const uint8_t sector, const int32_t iTrk, const GPUTPCGMMergedTrack& GPUrestrict() track, GPUTPCGMPropagator& GPUrestrict() prop, const bool inFlyDirection, GPUTPCGMMerger& GPUrestrict() merger, uint8_t* GPUrestrict() dEdxSubThresholdRow, const bool dodEdx, const bool doAttach, const uint8_t doInterpolate, gputpcgmmergertypes::InterpolationErrorHit& GPUrestrict() inter, const float deltaZ, const float sumInvSqrtCharge, const int32_t nAvgCharge)
 {
   if (CAMath::Abs(mP[2]) > GPUCA_MAX_SIN_PHI || CAMath::Abs(mP[0]) > CAMath::Abs(mX) * CAMath::Tan(GPUTPCGeometry::kSectAngle() / 2.f)) {
     return;
@@ -323,6 +325,16 @@ GPUdii() void GPUTPCGMTrackParam::FitAddRow(const int32_t iRow, const uint8_t se
   if (doAttach) {
     AttachClusters(merger, sector, iRow, iTrk, track.Leg() == 0, prop);
   }
+  if (doInterpolate) {
+    if (doInterpolate == 2) {
+      if (inter.isValid() && CAMath::Abs(inter.posY - mP[0]) < 4.f) {
+        inter.markInvalid();
+      }
+      FindBestInterpolatedHit(merger, inter, sector, iRow, deltaZ, sumInvSqrtCharge, nAvgCharge, prop, iTrk, false);
+    } else {
+      prop.InterpolateFill(&inter);
+    }
+  }
 }
 
 GPUdii() void GPUTPCGMTrackParam::HandleCrossCE(const GPUParam& GPUrestrict() param, const uint8_t sector, const uint8_t& lastSector)
@@ -369,10 +381,10 @@ GPUdii() int32_t GPUTPCGMTrackParam::FitHit(GPUTPCGMMerger& GPUrestrict() merger
     if (param.rec.tpc.mergerInterpolateErrors) {
       if (iWay == nWays - 2) {
         if (!param.rec.tpc.rebuildTrackInFit) {
-          if (inter.errorY < (GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A)0) {
-            rejectChi2 = true;
-          } else {
+          if (inter.isValid()) {
             retValInt = prop.InterpolateReject(param, yy, zz, clusterState, &inter, err2Y, err2Z, deltaZ);
+          } else {
+            rejectChi2 = true;
           }
         }
       } else if (iWay == nWays - 1) {
@@ -435,7 +447,7 @@ GPUdii() float GPUTPCGMTrackParam::FindBestInterpolatedHit(GPUTPCGMMerger& GPUre
   GPUglobalref() const cahit2* hits = tracker.HitData(rowData);
   GPUglobalref() const calink* firsthit = tracker.FirstHitInBin(rowData);
   float uncorrectedY = -1e6f, uncorrectedZ;
-  if (rowData.NHits() && (inter.errorY >= (GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A)0 || (!interOnly && param.rec.tpc.rebuildTrackMaxNonIntCov > 0 && mC[0] < param.rec.tpc.rebuildTrackMaxNonIntCov && mC[2] < param.rec.tpc.rebuildTrackMaxNonIntCov))) {
+  if (rowData.NHits() && (inter.isValid() || (!interOnly && param.rec.tpc.rebuildTrackMaxNonIntCov > 0 && mC[0] < param.rec.tpc.rebuildTrackMaxNonIntCov && mC[2] < param.rec.tpc.rebuildTrackMaxNonIntCov))) {
     const float zOffset = param.par.continuousTracking ? merger.GetConstantMem()->calibObjects.fastTransformHelper->getCorrMap()->convVertexTimeToZOffset(sector, mTOffset, param.continuousMaxTimeBin) : 0;
     const float y0 = rowData.Grid().YMin();
     const float stepY = rowData.HstepY();
@@ -449,7 +461,7 @@ GPUdii() float GPUTPCGMTrackParam::FindBestInterpolatedHit(GPUTPCGMMerger& GPUre
       ImP1 = (float)inter.posZ + deltaZ;
       ImC0 = (float)inter.errorY;
       ImC2 = (float)inter.errorZ;
-    } else if (inter.errorY >= (GPUCA_PAR_MERGER_INTERPOLATION_ERROR_TYPE_A)0) {
+    } else if (inter.isValid()) {
       const float Iz0 = (float)inter.posY - mP[0];
       const float Iz1 = (float)inter.posZ + deltaZ - mP[1];
       const float Iw0 = 1.f / (mC[0] + (float)inter.errorY);
@@ -526,7 +538,7 @@ GPUdii() float GPUTPCGMTrackParam::FindBestInterpolatedHit(GPUTPCGMMerger& GPUre
               if (insert < param.rec.tpc.rebuildTrackInFitClusterCandidates) {
                 for (int32_t c = CAMath::Min(nCandidates, param.rec.tpc.rebuildTrackInFitClusterCandidates - 1); c > insert; c--) {
                   merger.ClusterCandidates()[(iTrk * GPUCA_ROW_COUNT + row) * param.rec.tpc.rebuildTrackInFitClusterCandidates + c] = merger.ClusterCandidates()[(iTrk * GPUCA_ROW_COUNT + row) * param.rec.tpc.rebuildTrackInFitClusterCandidates + c - 1];
-                }
+                } // TODO: Downweight non-interpolated candidates
                 merger.ClusterCandidates()[(iTrk * GPUCA_ROW_COUNT + row) * param.rec.tpc.rebuildTrackInFitClusterCandidates + insert] = {.id = (uint32_t)(idOffset + ids[ih] + 2), .row = row, .sector = sector, .error = err, .weight = 0, .best = 0};
                 nCandidates += (nCandidates < param.rec.tpc.rebuildTrackInFitClusterCandidates);
               }
diff --git a/GPU/GPUTracking/Merger/GPUTPCGMTrackParam.h b/GPU/GPUTracking/Merger/GPUTPCGMTrackParam.h
@@ -157,7 +157,7 @@ class GPUTPCGMTrackParam
   GPUd() bool Fit(GPUTPCGMMerger& merger, int32_t iTrk, int32_t& N, int32_t& NTolerated, float& Alpha, GPUTPCGMMergedTrack& track, bool rebuilt);
   GPUd() void DodEdx(GPUdEdx& dEdx, GPUdEdx& dEdxAlt, GPUTPCGMMerger& merger, bool finalFit, int ihit, int ihitMergeFirst, int wayDirection, const GPUTPCGMMergedTrackHit* clusters, uint8_t clusterState, float zz, uint8_t dEdxSubThresholdRow);
   GPUd() int32_t FitHit(GPUTPCGMMerger& merger, const int32_t iTrk, const GPUTPCGMMergedTrack& track, const float xx, const float yy, const float zz, const uint8_t clusterState, const float clAlpha, const int32_t iWay, const bool inFlyDirection, float& deltaZ, float& lastUpdateX, GPUTPCGMMergedTrackHit* clusters, GPUTPCGMPropagator& prop, gputpcgmmergertypes::InterpolationErrorHit& inter, GPUdEdx& dEdx, GPUdEdx& dEdxAlt, float& sumInvSqrtCharge, int32_t& nAvgCharge, const int32_t ihit, const int32_t ihitMergeFirst, const bool allowChangeClusters, const bool refit, const bool finalFit, int32_t& nMissed, int32_t& nMissed2, int32_t& resetT0, float uncorrectedY);
-  GPUd() void FitAddRow(const int32_t iRow, const uint8_t sector, const int32_t iTrk, const GPUTPCGMMergedTrack& track, GPUTPCGMPropagator& prop, const bool inFlyDirection, GPUTPCGMMerger& merger, uint8_t* dEdxSubThresholdRow, const bool dodEdx, const bool doAttach, const bool doInterpolate);
+  GPUd() void FitAddRow(const int32_t iRow, const uint8_t sector, const int32_t iTrk, const GPUTPCGMMergedTrack& track, GPUTPCGMPropagator& prop, const bool inFlyDirection, GPUTPCGMMerger& merger, uint8_t* dEdxSubThresholdRow, const bool dodEdx, const bool doAttach, const uint8_t doInterpolate, gputpcgmmergertypes::InterpolationErrorHit& inter, const float deltaZ, const float sumInvSqrtCharge, const int32_t nAvgCharge);
   GPUd() void HandleCrossCE(const GPUParam& param, const uint8_t sector, const uint8_t& lastSector);
   GPUd() static void RefitTrack(GPUTPCGMMergedTrack& track, int32_t iTrk, GPUTPCGMMerger& merger, bool rebuilt);
   GPUd() void MoveToReference(GPUTPCGMPropagator& prop, const GPUParam& param, float& alpha);
