Make track clustering algo for vertexing configurable (cms-sw#347)

fwyzard · May 15, 2019 · c7d66fd · c7d66fd
1 parent d3dde7a
commit c7d66fd
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 9 deletions.
diff --git a/RecoPixelVertexing/PixelVertexFinding/src/PixelVertexHeterogeneousProducer.cc b/RecoPixelVertexing/PixelVertexFinding/src/PixelVertexHeterogeneousProducer.cc
@@ -1,6 +1,6 @@
+#include <cmath>
 #include <memory>
 #include <string>
-#include <cmath>
 
 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
 #include "DataFormats/GeometryCommonDetAlgo/interface/GlobalError.h"
@@ -67,6 +67,13 @@ class PixelVertexHeterogeneousProducer
 
 void PixelVertexHeterogeneousProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
   edm::ParameterSetDescription desc;
+
+  // Only one of these three algos can be used at once.
+  // Maybe this should become a Plugin Factory
+  desc.add<bool>("useDensity", true);
+  desc.add<bool>("useDBSCAN", false);
+  desc.add<bool>("useIterative", false);
+
   desc.add<int>("minT", 2);          // min number of neighbours to be "core"
   desc.add<double>("eps", 0.07);     // max absolute distance to cluster
   desc.add<double>("errmax", 0.01);  // max error to be "seed"
@@ -91,7 +98,10 @@ PixelVertexHeterogeneousProducer::PixelVertexHeterogeneousProducer(const edm::Pa
       enableConversion_(conf.getParameter<bool>("gpuEnableConversion")),
       enableTransfer_(enableConversion_ || conf.getParameter<bool>("gpuEnableTransfer")),
       gpuToken_(consumes<HeterogeneousProduct>(conf.getParameter<edm::InputTag>("src"))),
-      m_gpuAlgo(conf.getParameter<int>("minT"),
+      m_gpuAlgo(conf.getParameter<bool>("useDensity"),
+                conf.getParameter<bool>("useDBSCAN"),
+                conf.getParameter<bool>("useIterative"),
+                conf.getParameter<int>("minT"),
                 conf.getParameter<double>("eps"),
                 conf.getParameter<double>("errmax"),
                 conf.getParameter<double>("chi2max"),

diff --git a/RecoPixelVertexing/PixelVertexFinding/src/gpuClusterTracks.h b/RecoPixelVertexing/PixelVertexFinding/src/gpuClusterTracks.h
diff --git a/RecoPixelVertexing/PixelVertexFinding/src/gpuVertexFinder.cu b/RecoPixelVertexing/PixelVertexFinding/src/gpuVertexFinder.cu
@@ -1,4 +1,7 @@
-#include "gpuClusterTracks.h"
+#include "RecoPixelVertexing/PixelVertexFinding/src/gpuClusterTracksByDensity.h"
+#include "RecoPixelVertexing/PixelVertexFinding/src/gpuClusterTracksDBSCAN.h"
+#include "RecoPixelVertexing/PixelVertexFinding/src/gpuClusterTracksIterative.h"
+
 #include "gpuFitVertices.h"
 #include "gpuSortByPt2.h"
 #include "gpuSplitVertices.h"
@@ -58,8 +61,13 @@ namespace gpuVertexFinder {
     auto numberOfBlocks = (CAConstants::maxTuples() + blockSize - 1) / blockSize;
     loadTracks<<<numberOfBlocks, blockSize, 0, stream>>>(tracks.gpu_d, ws_d, ptMin);
     cudaCheck(cudaGetLastError());
-
-    CLUSTERIZE<<<1, 1024 - 256, 0, stream>>>(gpu_d, ws_d, minT, eps, errmax, chi2max);
+    if (useDensity_) {
+      clusterTracksByDensity<<<1, 1024 - 256, 0, stream>>>(gpu_d, ws_d, minT, eps, errmax, chi2max);
+    } else if (useDBSCAN_) {
+      clusterTracksDBSCAN<<<1, 1024 - 256, 0, stream>>>(gpu_d, ws_d, minT, eps, errmax, chi2max);
+    } else if (useIterative_) {
+      clusterTracksIterative<<<1, 1024 - 256, 0, stream>>>(gpu_d, ws_d, minT, eps, errmax, chi2max);
+    }
     cudaCheck(cudaGetLastError());
     fitVertices<<<1, 1024 - 256, 0, stream>>>(gpu_d, ws_d, 50.);
     cudaCheck(cudaGetLastError());

diff --git a/RecoPixelVertexing/PixelVertexFinding/src/gpuVertexFinder.h b/RecoPixelVertexing/PixelVertexFinding/src/gpuVertexFinder.h
@@ -79,12 +79,22 @@ namespace gpuVertexFinder {
     using ZVertices = gpuVertexFinder::ZVertices;
     using WorkSpace = gpuVertexFinder::WorkSpace;
 
-    Producer(int iminT,       // min number of neighbours to be "core"
+    Producer(bool useDensity,
+             bool useDBSCAN,
+             bool useIterative,
+             int iminT,       // min number of neighbours to be "core"
              float ieps,      // max absolute distance to cluster
              float ierrmax,   // max error to be "seed"
              float ichi2max,  // max normalized distance to cluster
              bool ienableTransfer)
-        : minT(iminT), eps(ieps), errmax(ierrmax), chi2max(ichi2max), enableTransfer(ienableTransfer) {}
+        : useDensity_(useDensity),
+          useDBSCAN_(useDBSCAN),
+          useIterative_(useIterative),
+          minT(iminT),
+          eps(ieps),
+          errmax(ierrmax),
+          chi2max(ichi2max),
+          enableTransfer(ienableTransfer) {}
 
     ~Producer() { deallocate(); }
 
@@ -100,6 +110,10 @@ namespace gpuVertexFinder {
     ZVertices* gpu_d = nullptr;
     WorkSpace* ws_d = nullptr;
 
+    const bool useDensity_;
+    const bool useDBSCAN_;
+    const bool useIterative_;
+
     int minT;       // min number of neighbours to be "core"
     float eps;      // max absolute distance to cluster
     float errmax;   // max error to be "seed"