ITS: add truth seeding

Signed-off-by: Felix Schlepper <[email protected]>

Author: f3sch

Detectors/ITSMFT/ITS/tracking/CMakeLists.txt

@@ -36,7 +36,9 @@ o2_add_library(ITStracking
                        O2::ITSReconstruction
                        O2::ITSMFTReconstruction
                        O2::DataFormatsITS
-               PRIVATE_LINK_LIBRARIES TBB::tbb)
+               PRIVATE_LINK_LIBRARIES
+                       O2::Steer
+                       TBB::tbb)
 
 o2_add_library(ITSTrackingInterface
                TARGETVARNAME targetName

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

@@ -112,6 +112,8 @@ struct VertexingParameters {
   int zSpan = -1;
   bool SaveTimeBenchmarks = false;
 
+  bool useTruthSeeding = false; // overwrite found vertices with MC events
+
   int nThreads = 1;
   bool PrintMemory = false; // print allocator usage in epilog report
   size_t MaxMemory = std::numeric_limits<size_t>::max();

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

@@ -48,6 +48,8 @@ struct VertexerParamConfig : public o2::conf::ConfigurableParamHelper<VertexerPa
   int ZBins = 1;     // z-phi index table configutation: number of z bins
   int PhiBins = 128; // z-phi index table configutation: number of phi bins
 
+  bool useTruthSeeding{false}; // overwrite seeding vertices with MC truth
+
   int nThreads = 1;
   bool printMemory = false;
   size_t maxMemory = std::numeric_limits<size_t>::max();

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

@@ -75,7 +75,8 @@ class Vertexer
   void validateTracklets(T&&... args);
   template <typename... T>
   void findVertices(T&&... args);
-  void findHistVertices();
+
+  void addTruthSeeds() { mTraits->addTruthSeedingVertices(); }
 
   template <typename... T>
   void initialiseVertexer(T&&... args);
@@ -108,10 +109,11 @@ class Vertexer
     Trackleting,
     Validating,
     Finding,
+    TruthSeeding,
     NStates,
   };
   State mCurState{Init};
-  static constexpr std::array<const char*, NStates> StateNames{"Initialisation", "Tracklet finding", "Tracklet validation", "Vertex finding"};
+  static constexpr std::array<const char*, NStates> StateNames{"Initialisation", "Tracklet finding", "Tracklet validation", "Vertex finding", "Truth seeding"};
 };
 
 template <typename... T>

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

@@ -73,6 +73,8 @@ class VertexerTraits
   virtual void computeVertices(const int iteration = 0);
   virtual void adoptTimeFrame(TimeFrame7* tf) noexcept { mTimeFrame = tf; }
   virtual void updateVertexingParameters(const std::vector<VertexingParameters>& vrtPar, const TimeFrameGPUParameters& gpuTfPar);
+  // truth tracking
+  void addTruthSeedingVertices();
 
   void computeVerticesInRof(int,
                             gsl::span<const o2::its::Line>&,

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

@@ -282,6 +282,8 @@ std::vector<VertexingParameters> TrackingMode::getVertexingParameters(TrackingMo
     p.nThreads = vc.nThreads;
     p.ZBins = vc.ZBins;
     p.PhiBins = vc.PhiBins;
+
+    p.useTruthSeeding = vc.useTruthSeeding;
   }
   // set for now outside to not disturb status quo
   vertParams[0].vertNsigmaCut = vc.vertNsigmaCut;

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

@@ -187,13 +187,13 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
       vtxROF.setNEntries(vtxSpan.size());
       bool selROF = vtxSpan.empty();
       for (auto iV{0}; iV < vtxSpan.size(); ++iV) {
-        auto& v = vtxSpan[iV];
+        const auto& v = vtxSpan[iV];
         if (multEstConf.isVtxMultCutRequested() && !multEstConf.isPassingVtxMultCut(v.getNContributors())) {
           continue; // skip vertex of unwanted multiplicity
         }
         selROF = true;
         vertices.push_back(v);
-        if (mIsMC) {
+        if (mIsMC && !VertexerParamConfig::Instance().useTruthSeeding) {
           allVerticesLabels.push_back(vMCRecInfo[iV].first);
           allVerticesPurities.push_back(vMCRecInfo[iV].second);
         }

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

@@ -36,6 +36,11 @@ Vertexer::Vertexer(VertexerTraits* traits) : mTraits(traits)
 float Vertexer::clustersToVertices(LogFunc logger)
 {
   LogFunc evalLog = [](const std::string&) {};
+
+  if (mTimeFrame->hasMCinformation() && mVertParams[0].useTruthSeeding) {
+    return evaluateTask(&Vertexer::addTruthSeeds, StateNames[mCurState = TruthSeeding], 0, evalLog);
+  }
+
   TrackingParameters trkPars;
   TimeFrameGPUParameters tfGPUpar;
   mTraits->updateVertexingParameters(mVertParams, tfGPUpar);
@@ -58,14 +63,11 @@ float Vertexer::clustersToVertices(LogFunc logger)
       logger(fmt::format("=== ITS {} Seeding vertexer iteration {} summary:", mTraits->getName(), iteration));
       trkPars.PhiBins = mTraits->getVertexingParameters()[0].PhiBins;
       trkPars.ZBins = mTraits->getVertexingParameters()[0].ZBins;
-      auto timeInitIteration = evaluateTask(
-        &Vertexer::initialiseVertexer, StateNames[mCurState = Init], iteration, evalLog, trkPars, iteration);
-      auto timeTrackletIteration = evaluateTask(
-        &Vertexer::findTracklets, StateNames[mCurState = Trackleting], iteration, evalLog, iteration);
+      auto timeInitIteration = evaluateTask(&Vertexer::initialiseVertexer, StateNames[mCurState = Init], iteration, evalLog, trkPars, iteration);
+      auto timeTrackletIteration = evaluateTask(&Vertexer::findTracklets, StateNames[mCurState = Trackleting], iteration, evalLog, iteration);
       nTracklets01 = mTimeFrame->getTotalTrackletsTF(0);
       nTracklets12 = mTimeFrame->getTotalTrackletsTF(1);
-      auto timeSelectionIteration = evaluateTask(
-        &Vertexer::validateTracklets, StateNames[mCurState = Validating], iteration, evalLog, iteration);
+      auto timeSelectionIteration = evaluateTask(&Vertexer::validateTracklets, StateNames[mCurState = Validating], iteration, evalLog, iteration);
       auto timeVertexingIteration = evaluateTask(&Vertexer::findVertices, StateNames[mCurState = Finding], iteration, evalLog, iteration);
       printEpilog(logger, nTracklets01, nTracklets12, mTimeFrame->getNLinesTotal(), mTimeFrame->getTotVertIteration()[iteration], timeInitIteration, timeTrackletIteration, timeSelectionIteration, timeVertexingIteration);
       timeInit += timeInitIteration;

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

@@ -10,10 +10,10 @@
 // or submit itself to any jurisdiction.
 ///
 
-#include <iostream>
 #include <memory>
-#include <string>
-#include <chrono>
+#include <ranges>
+#include <map>
+#include <algorithm>
 
 #include <oneapi/tbb/blocked_range.h>
 #include <oneapi/tbb/parallel_for.h>
@@ -22,6 +22,9 @@
 #include "ITStracking/BoundedAllocator.h"
 #include "ITStracking/ClusterLines.h"
 #include "ITStracking/Tracklet.h"
+#include "SimulationDataFormat/DigitizationContext.h"
+#include "Steer/MCKinematicsReader.h"
+#include "ITSMFTBase/DPLAlpideParam.h"
 
 #ifdef VTX_DEBUG
 #include "TTree.h"
@@ -693,6 +696,59 @@ void VertexerTraits::computeVerticesInRof(int rofId,
   verticesInRof.push_back(foundVertices);
 }
 
+void VertexerTraits::addTruthSeedingVertices()
+{
+  LOGP(info, "Using truth seeds as vertices; will skip computations");
+  mTimeFrame->resetRofPV();
+  const auto dc = o2::steer::DigitizationContext::loadFromFile("collisioncontext.root");
+  const auto irs = dc->getEventRecords();
+  int64_t roFrameBiasInBC = o2::itsmft::DPLAlpideParam<o2::detectors::DetID::ITS>::Instance().roFrameBiasInBC;
+  int64_t roFrameLengthInBC = o2::itsmft::DPLAlpideParam<o2::detectors::DetID::ITS>::Instance().roFrameLengthInBC;
+  o2::steer::MCKinematicsReader mcReader(dc);
+  std::map<int, bounded_vector<Vertex>> vertices;
+  for (int iSrc{0}; iSrc < mcReader.getNSources(); ++iSrc) {
+    auto eveId2colId = dc->getCollisionIndicesForSource(iSrc);
+    for (int iEve{0}; iEve < mcReader.getNEvents(iSrc); ++iEve) {
+      const auto& ir = irs[eveId2colId[iEve]];
+      if (!ir.isDummy()) { // do we need this, is this for diffractive events?
+        const auto& eve = mcReader.getMCEventHeader(iSrc, iEve);
+        int rofId = (ir.toLong() - roFrameBiasInBC) / roFrameLengthInBC;
+        if (!vertices.contains(rofId)) {
+          vertices[rofId] = bounded_vector<Vertex>(mMemoryPool.get());
+        }
+        Vertex vert;
+        vert.setTimeStamp(rofId);
+        vert.setNContributors(std::ranges::count_if(mcReader.getTracks(iSrc, iEve), [](const auto& trk) {
+          return trk.isPrimary() && trk.GetPt() > 0.2 && std::abs(trk.GetEta()) < 1.3;
+        }));
+        vert.setXYZ((float)eve.GetX(), (float)eve.GetY(), (float)eve.GetZ());
+        vert.setChi2(1);
+        constexpr float cov = 50e-9;
+        vert.setCov(cov, cov, cov, cov, cov, cov);
+        vertices[rofId].push_back(vert);
+      }
+    }
+  }
+  size_t nVerts{0};
+  for (int iROF{0}; iROF < mTimeFrame->getNrof(); ++iROF) {
+    bounded_vector<Vertex> verts(mMemoryPool.get());
+    bounded_vector<std::pair<o2::MCCompLabel, float>> polls(mMemoryPool.get());
+    if (vertices.contains(iROF)) {
+      verts = vertices[iROF];
+      nVerts += verts.size();
+      for (size_t i{0}; i < verts.size(); ++i) {
+        o2::MCCompLabel lbl; // unset label for now
+        polls.emplace_back(lbl, 1.f);
+      }
+    } else {
+      mTimeFrame->getNoVertexROF()++;
+    }
+    mTimeFrame->addPrimaryVertices(verts, iROF, 0);
+    mTimeFrame->addPrimaryVerticesLabels(polls);
+  }
+  LOGP(info, "Found {}/{} ROFs with {} vertices -> <NV>={:.2f}", vertices.size(), mTimeFrame->getNrof(), nVerts, (float)nVerts / (float)vertices.size());
+}
+
 void VertexerTraits::setNThreads(int n, std::shared_ptr<tbb::task_arena>& arena)
 {
 #if defined(VTX_DEBUG)