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, constants::VtxOverpopulatedRofTrimFraction);
+  }
 
   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,65 @@ 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>
+int pruneOverpopulatedRofs(std::vector<VtxVec>& rofVertices, std::vector<LabVec>& rofLabels, const float nSigma, const int minContributors, const float trimFraction)
+{
+  const int nRofs = static_cast<int>(rofVertices.size());
+  if (nSigma <= 0.f || nRofs == 0) {
+    return 0;
+  }
+  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>(trimFraction * 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.);
+  int removed = 0;
+  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);
+    removed += static_cast<int>(vtx.size() - out);
+    vtx.erase(vtx.begin() + out, vtx.end());
+    if (withLabels) {
+      rofLabels[r].erase(rofLabels[r].begin() + out, rofLabels[r].end());
+    }
+  }
+  return removed;
+}
+
 } // 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);
 

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

@@ -305,6 +305,7 @@ void TrackerTraits<NLayers>::computeVertexCandidates(const int iteration)
   const float lineMinPt = mTrkParams[iteration].VtxLineMinPt;
   const float beamX = mTimeFrame->getBeamX();
   const float beamY = mTimeFrame->getBeamY();
+  const float lineCurvatureScale = mTrkParams[iteration].VtxLineCurvatureScale;
   auto makeKey = [](float attribute, int cellIdx) -> size_t {
     const uint32_t attributeInt = std::bit_cast<uint32_t>(attribute);
     return (static_cast<size_t>(attributeInt) << 32) | static_cast<uint32_t>(cellIdx);
@@ -327,8 +328,7 @@ void TrackerTraits<NLayers>::computeVertexCandidates(const int iteration)
       kCl1[k] = c1;
       kCl2[k] = cell.getThirdClusterIndex();
       std::array<float, 3> origin, direction;
-      cell.getXYZGlo(origin);
-      if (!cell.getPxPyPzGlo(direction)) {
+      if (!getCellLineAtBeam(cell, beamX, beamY, getBz(), lineCurvatureScale, origin, direction)) {
         return;
       }
       kGeomOk[k] = 1;
@@ -711,11 +711,11 @@ void TrackerTraits<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) {
       if (!rofVertices[rofId].empty()) {
         if (goodSig > 0.f) {
-          if (nGoodCand[p] <= sigThreshold) {
+          if (nGoodCand[p] < debrisThreshold) {
             continue;
           }
         } else if (static_cast<int>(cand[p].getSize()) < suppressLowMultDebris) {
@@ -758,6 +758,9 @@ void TrackerTraits<NLayers>::computeVertices(const int iteration)
       });
     });
   }
+  if (!tp.PassFlags[IterationStep::MarkVerticesAsUPC]) { // UPC ROFs are near-empty by construction
+    pruneOverpopulatedRofs(rofVertices, rofLabels, tp.VtxOverpopulatedRofNSigma, suppressLowMultDebris, constants::VtxOverpopulatedRofTrimFraction);
+  }
   for (int rofId{0}; rofId < nRofs; ++rofId) {
     for (auto& vertex : rofVertices[rofId]) {
       mTimeFrame->addPrimaryVertex(vertex);

Detectors/ITSMFT/common/tracking/include/ITSMFTTracking/Constants.h

@@ -41,8 +41,10 @@ constexpr int MinNumberOfConcurrentSeeds = (1 << 8);  // minimum chunk size for
 constexpr int MaxNumberOfConcurrentSeeds = (1 << 12); // maximum chunk size for a worker for the final track fit/extraploation step
 constexpr float MaxTrackSeedQ2Pt = 1.e3f;             // maximum q/pt for track seeds
 
-constexpr int MaxBootstrapPasses = 5;               // beam bootstrap: cap on the re-trackleting passes
-constexpr float BeamConvergence2 = 5.e-3f * 5.e-3f; // beam bootstrap: stop below a (50 um)^2 beam shift
+constexpr int MaxBootstrapPasses = 5;                    // beam bootstrap: cap on the re-trackleting passes
+constexpr float BeamConvergence2 = 5.e-3f * 5.e-3f;      // beam bootstrap: stop below a (50 um)^2 beam shift
+constexpr float VtxMinGoodThreshold = 2.f;               // seeding emit: floor of the k*sqrt(ROF load) debris threshold (bounds included)
+constexpr float VtxOverpopulatedRofTrimFraction = 0.02f; // overpopulated-ROF pruning: busiest fraction of ROFs left out of the per-ROF mean
 
 namespace helpers
 {

Detectors/ITSMFT/common/tracking/include/ITSMFTTracking/ITSTrackingConfigParam.h

@@ -42,15 +42,17 @@ struct VertexerParamConfig : public o2::conf::ConfigurableParamHelper<VertexerPa
   float maxZPositionAllowed = 25.f; // 4x sZ of the beam
 
   // Artefacts selections
-  int clusterContributorsCut = 2; // minimum number of contributors for an accepted final vertex
-  int suppressLowMultDebris = 16; // suppress all vertices below this threshold if a vertex was already found in a rof
-  float lineMinPt = 0.10f;        // drop soft lines before the density scan
-  float fineZWindow = 0.010f;     // second, narrow density pass (dip search); <=0 disables
+  int clusterContributorsCut = 2;  // minimum number of contributors for an accepted final vertex
+  int suppressLowMultDebris = 16;  // suppress all vertices below this threshold if a vertex was already found in a rof
+  float lineMinPt = 0.10f;         // drop soft lines before the density scan
+  float lineCurvatureScale = 0.5f; // q/pT scale when propagating a seeding line to its xy-DCA to the beam
+  float fineZWindow = 0.010f;      // second, narrow density pass (dip search); <=0 disables
   int fineMinDensity = 8;
   float fineMaxDrift = 0.005f; // |z_fit - z_seed| cap on fine-only candidates; <=0 disables
   float goodLineChi2Cut = 5.f;
   float goodLinePtCut = 0.5f;
-  float goodContributorsSignificance = 0.070f; // emit threshold k, scaled by sqrt(ROF load); <=0 disables
+  float goodContributorsSignificance = 0.070f; // emit threshold k, scaled by sqrt(ROF load) and clamped to [constants::VtxMinGoodThreshold (2), suppressLowMultDebris], bounds included; <=0 disables
+  float overpopulatedRofNSigma = -1.f;         // a ROF with more seeding vertices than mean + overpopulatedRofNSigma*sqrt(mean+1) of its TF keeps its largest vertex and those with >= suppressLowMultDebris contributors (not in the UPC pass); <=0 disables
   float duplicateZScale = 0.7f;                // per-candidate dedup radius scale/sqrt(size); <=0 uses duplicateZCut
   int seedMemberRadiusTime = 0;
   int seedMemberRadiusZ = 2;