ITS GPU seeding vertexer: add clamps to adaptive low mult threshold, and propate cells towards beamline

Author: cima22

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h

@@ -211,6 +211,8 @@ struct TrackingKernels {
                                            int* lineSlots,
                                            const float beamX,
                                            const float beamY,
+                                           const float bz,
+                                           const float curvatureScale,
                                            const float maxZ,
                                            const float minPt,
                                            float* linesZs,

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

@@ -266,6 +266,8 @@ void TrackerTraitsGPU<NLayers>::computeVertexCandidates(const int iteration)
                                                          mTimeFrameGPU->getDeviceLineSlots(),
                                                          mTimeFrameGPU->getBeamX(),
                                                          mTimeFrameGPU->getBeamY(),
+                                                         this->getBz(),
+                                                         this->mTrkParams[iteration].VtxLineCurvatureScale,
                                                          this->mTrkParams[iteration].VtxMaxZPositionAllowed,
                                                          this->mTrkParams[iteration].VtxLineMinPt,
                                                          mTimeFrameGPU->getDeviceLineZs(),
@@ -451,13 +453,13 @@ void TrackerTraitsGPU<NLayers>::computeVertices(const int iteration)
         }
       }
     }
-    const float sigThreshold = goodSig > 0.f ? goodSig * std::sqrt(static_cast<float>(std::max(rofLoad, 1.))) : 0.f;
+    const float debrisThreshold = goodSig > 0.f ? getDebrisThreshold(goodSig, rofLoad, suppressLowMultDebris) : 0.f;
 
     for (const int p : accepted) {
       const auto& c = cands[p];
       if (!rofVertices[rofId].empty()) {
         if (goodSig > 0.f) {
-          if (c.nGood <= sigThreshold) {
+          if (c.nGood < debrisThreshold) {
             continue;
           }
         } else if (c.size < suppressLowMultDebris) {
@@ -503,6 +505,9 @@ void TrackerTraitsGPU<NLayers>::computeVertices(const int iteration)
       }
     }
   }
+  if (!this->mTrkParams[iteration].PassFlags[IterationStep::MarkVerticesAsUPC]) { // UPC ROFs are near-empty by construction
+    pruneOverpopulatedRofs(rofVertices, rofLabels, this->mTrkParams[iteration].VtxOverpopulatedRofNSigma, suppressLowMultDebris);
+  }
 
   for (int rofId = 0; rofId < nRofs; ++rofId) {
     for (auto& vertex : rofVertices[rofId]) {

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu

@@ -874,22 +874,23 @@ GPUg() void dedupCellsKernel(
   const int ownedClustersCut,
   const float beamX,
   const float beamY,
+  const float bz,
+  const float curvatureScale,
   const float maxZ,
   const float minPt,
   int* cellAccepted)
 {
   for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nCells; i += blockDim.x * gridDim.x) {
     const CellSeed& cell = cells[i];
     std::array<float, 3> origin, direction;
-    if (!cell.getPxPyPzGlo(direction)) {
+    if (!getCellLineAtBeam(cell, beamX, beamY, bz, curvatureScale, origin, direction)) {
       cellAccepted[i] = 0;
       continue;
     }
     const bool owned0 = static_cast<uint32_t>(clusterOwners[0][cell.getFirstClusterIndex()]) == static_cast<uint32_t>(i);
     const bool owned1 = static_cast<uint32_t>(clusterOwners[1][cell.getSecondClusterIndex()]) == static_cast<uint32_t>(i);
     const bool owned2 = static_cast<uint32_t>(clusterOwners[2][cell.getThirdClusterIndex()]) == static_cast<uint32_t>(i);
     const bool keepCell = (static_cast<int>(owned0) + static_cast<int>(owned1) + static_cast<int>(owned2)) >= 3 - ownedClustersCut;
-    cell.getXYZGlo(origin);
     const float dx = origin[0] - beamX;
     const float dy = origin[1] - beamY;
     const float den = direction[0] * direction[0] + direction[1] * direction[1];
@@ -910,6 +911,8 @@ GPUg() void linearizeCellsKernel(
   int* lineRof,
   const float beamX,
   const float beamY,
+  const float bz,
+  const float curvatureScale,
   float* lineZs,
   o2::its::TimeEstBC* lineTimes,
   int* lineClusters, // 3 per line (L0,L1,L2 cluster ids), for the host-side MC label derivation
@@ -923,8 +926,7 @@ GPUg() void linearizeCellsKernel(
     }
     const CellSeed& cell = cells[i];
     std::array<float, 3> origin, direction;
-    cell.getXYZGlo(origin);
-    cell.getPxPyPzGlo(direction);
+    getCellLineAtBeam(cell, beamX, beamY, bz, curvatureScale, origin, direction); // accepted by dedupCellsKernel: succeeds
     lines[slot] = o2::its::Line{origin.data(), direction.data(), cell.getTimeStamp()};
     lineRof[slot] = deviceUpperBound(rofFramesClustersL1, 0, nRofsL1 + 1, cell.getSecondClusterIndex()) - 1;
     float zAtBeam;
@@ -1577,6 +1579,8 @@ void TrackingKernels<NLayers>::linearizeCellsToLinesHandler(const int nCells,
                                                             int* lineSlots, // nCells + 1 scratch: accept flags, scanned in place into slots
                                                             const float beamX,
                                                             const float beamY,
+                                                            const float bz,
+                                                            const float curvatureScale,
                                                             const float maxZ,
                                                             const float minPt,
                                                             float* linesZs,
@@ -1593,6 +1597,8 @@ void TrackingKernels<NLayers>::linearizeCellsToLinesHandler(const int nCells,
     ownedClustersCut,
     beamX,
     beamY,
+    bz,
+    curvatureScale,
     maxZ,
     minPt,
     lineSlots);
@@ -1608,6 +1614,8 @@ void TrackingKernels<NLayers>::linearizeCellsToLinesHandler(const int nCells,
     lineRof,
     beamX,
     beamY,
+    bz,
+    curvatureScale,
     linesZs,
     lineTimes,
     lineClusters,

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

@@ -19,6 +19,7 @@
 #include <cmath>
 #endif
 #include "ITStracking/Cluster.h"
+#include "ITStracking/Cell.h"
 #include "ITSMFTTracking/Constants.h"
 #include "ITStracking/Tracklet.h"
 #include "GPUCommonDef.h"
@@ -121,6 +122,25 @@ struct Line final {
   TimeEstBC mTime;
 };
 
+GPUdi() bool getCellLineAtBeam(const CellSeed& cell, const float beamX, const float beamY, const float bz, const float curvatureScale, std::array<float, 3>& origin, std::array<float, 3>& direction)
+{
+  cell.getXYZGlo(origin);
+  if (!cell.getPxPyPzGlo(direction)) {
+    return false;
+  }
+  o2::track::TrackParametrization<float> par{cell};
+  par.setQ2Pt(par.getQ2Pt() * curvatureScale);
+  std::array<float, 3> dcaOrigin, dcaDirection;
+  if (par.propagateParamToDCA({beamX, beamY, 0.f}, bz)) {
+    par.getXYZGlo(dcaOrigin);
+    if (par.getPxPyPzGlo(dcaDirection)) {
+      origin = dcaOrigin;
+      direction = dcaDirection;
+    }
+  }
+  return true;
+}
+
 /// Least-squares vertex fit over a set of lines (the normal equations AX = -B).
 class ClusterLines final
 {

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

@@ -151,6 +151,7 @@ struct TrackingParameters {
   int VertPerRofThreshold = 0; // max vertices in a ROF for the UPC pass to still run on it
   float VtxPhiCut = -1.f;
   float VtxLineMinPt = -1.f;
+  float VtxLineCurvatureScale = 0.5f;
   float VtxMaxZPositionAllowed = -1.f;
   float VtxClusterCut = -1.f;
   float VtxPairCut = -1.f;
@@ -165,6 +166,7 @@ struct TrackingParameters {
   float VtxGoodContributorsSignificance = -1.f;
   int VtxClusterContributorsCut = -1;
   int VtxSuppressLowMultDebris = -1;
+  float VtxOverpopulatedRofNSigma = -1.f;
 };
 
 struct VertexingParameters {

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

@@ -19,9 +19,14 @@
 #include "SimulationDataFormat/MCCompLabel.h"
 #include "ITStracking/Configuration.h"
 
+#include "Framework/Logger.h"
+
+#include <algorithm>
+#include <cmath>
 #include <limits>
 #include <unordered_map>
 #include <utility>
+#include <vector>
 
 namespace o2::its
 {
@@ -60,6 +65,62 @@ inline Vertex makeDiamondVertex(const TrackingParameters& trkParam)
   return diamond;
 }
 
+/// Good lines a further vertex needs to not count as debris in a ROF that already has one: goodSig * sqrt(ROF load), clamped to
+/// [constants::VtxMinGoodThreshold, suppressLowMultDebris] unless the debris cut is off (UPC pass).
+inline float getDebrisThreshold(const float goodSig, const double rofLoad, const int suppressLowMultDebris)
+{
+  const float threshold = goodSig * std::sqrt(static_cast<float>(std::max(rofLoad, 1.)));
+  if (suppressLowMultDebris < constants::VtxMinGoodThreshold) {
+    return threshold;
+  }
+  return std::clamp(threshold, constants::VtxMinGoodThreshold, static_cast<float>(suppressLowMultDebris));
+}
+
+/// Caps ROFs whose vertex count is an outlier of the TF
+template <typename VtxVec, typename LabVec>
+void pruneOverpopulatedRofs(std::vector<VtxVec>& rofVertices, std::vector<LabVec>& rofLabels, const float nSigma, const int minContributors)
+{
+  const int nRofs = static_cast<int>(rofVertices.size());
+  if (nSigma <= 0.f || nRofs == 0) {
+    return;
+  }
+  std::vector<int> counts(nRofs);
+  for (int r = 0; r < nRofs; ++r) {
+    counts[r] = static_cast<int>(rofVertices[r].size());
+  }
+  std::vector<int> sorted(counts);
+  std::sort(sorted.begin(), sorted.end());
+  const int nUsed = std::max(1, nRofs - std::max(1, static_cast<int>(constants::VtxOverpopulatedRofTrimFraction * nRofs)));
+  double sum = 0.;
+  for (int i = 0; i < nUsed; ++i) {
+    sum += sorted[i];
+  }
+  const double mean = sum / nUsed;
+  const double threshold = mean + nSigma * std::sqrt(mean + 1.);
+  for (int r = 0; r < nRofs; ++r) {
+    if (counts[r] <= threshold) {
+      continue;
+    }
+    auto& vtx = rofVertices[r];
+    const bool withLabels = static_cast<int>(rofLabels.size()) == nRofs && rofLabels[r].size() == vtx.size();
+    size_t out = 1; // the largest vertex always stays
+    for (size_t i = 1; i < vtx.size(); ++i) {
+      if (vtx[i].getNContributors() >= minContributors) {
+        vtx[out] = vtx[i];
+        if (withLabels) {
+          rofLabels[r][out] = rofLabels[r][i];
+        }
+        ++out;
+      }
+    }
+    LOGP(info, "Seeding vertexer: overpopulated ROF {} pruned {} -> {} vertices (threshold {:.1f}, mean {:.2f} per ROF)", r, vtx.size(), out, threshold, mean);
+    vtx.erase(vtx.begin() + out, vtx.end());
+    if (withLabels) {
+      rofLabels[r].erase(rofLabels[r].begin() + out, rofLabels[r].end());
+    }
+  }
+}
+
 } // namespace o2::its
 
 #endif /* O2_ITS_TRACKING_VERTEXUTILS_H_ */

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

@@ -308,6 +308,7 @@ std::vector<TrackingParameters> TrackingMode::getTrackingParameters(TrackingMode
     seedingPass.VertPerRofThreshold = vc.vertPerRofThreshold;
     seedingPass.VtxPhiCut = vc.phiCut;
     seedingPass.VtxLineMinPt = vc.lineMinPt;
+    seedingPass.VtxLineCurvatureScale = vc.lineCurvatureScale;
     seedingPass.VtxMaxZPositionAllowed = vc.maxZPositionAllowed;
     seedingPass.VtxClusterCut = vc.clusterCut;
     seedingPass.VtxPairCut = vc.pairCut;
@@ -322,6 +323,7 @@ std::vector<TrackingParameters> TrackingMode::getTrackingParameters(TrackingMode
     seedingPass.VtxGoodContributorsSignificance = vc.goodContributorsSignificance;
     seedingPass.VtxClusterContributorsCut = vc.clusterContributorsCut;
     seedingPass.VtxSuppressLowMultDebris = vc.suppressLowMultDebris;
+    seedingPass.VtxOverpopulatedRofNSigma = vc.overpopulatedRofNSigma;
     std::vector<TrackingParameters> seedingPasses;
     seedingPasses.push_back(seedingPass);