ITS: opt. vertex cont. output

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

@@ -113,6 +113,7 @@ struct VertexingParameters {
   bool SaveTimeBenchmarks = false;
 
   bool useTruthSeeding = false; // overwrite found vertices with MC events
+  bool outputContLabels = false;
 
   int nThreads = 1;
   bool PrintMemory = false; // print allocator usage in epilog report

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

@@ -73,15 +73,18 @@ struct TimeFrame {
   gsl::span<const Vertex> getPrimaryVertices(int rofId) const;
   gsl::span<const Vertex> getPrimaryVertices(int romin, int romax) const;
   gsl::span<const std::pair<MCCompLabel, float>> getPrimaryVerticesMCRecInfo(const int rofId) const;
+  gsl::span<const MCCompLabel> getPrimaryVerticesContributors(const int rofId) const;
   gsl::span<const std::array<float, 2>> getPrimaryVerticesXAlpha(int rofId) const;
   void fillPrimaryVerticesXandAlpha();
   int getPrimaryVerticesNum(int rofId = -1) const;
   void addPrimaryVertices(const bounded_vector<Vertex>& vertices);
   void addPrimaryVerticesLabels(bounded_vector<std::pair<MCCompLabel, float>>& labels);
+  void addPrimaryVerticesContributorLabels(bounded_vector<MCCompLabel>& labels);
   void addPrimaryVertices(const bounded_vector<Vertex>& vertices, const int rofId, const int iteration);
   void addPrimaryVertices(const gsl::span<const Vertex>& vertices, const int rofId, const int iteration);
   void addPrimaryVerticesInROF(const bounded_vector<Vertex>& vertices, const int rofId, const int iteration);
   void addPrimaryVerticesLabelsInROF(const bounded_vector<std::pair<MCCompLabel, float>>& labels, const int rofId);
+  void addPrimaryVerticesContributorLabelsInROF(const bounded_vector<MCCompLabel>& labels, const int rofId);
   void removePrimaryVerticesInROf(const int rofId);
   int loadROFrameData(const o2::itsmft::ROFRecord& rof, gsl::span<const itsmft::Cluster> clusters,
                       const dataformats::MCTruthContainer<MCCompLabel>* mcLabels = nullptr);
@@ -342,6 +345,7 @@ struct TimeFrame {
   std::array<bounded_vector<int>, 2> mTrackletsIndexROF;
   std::vector<bounded_vector<MCCompLabel>> mLinesLabels;
   std::vector<std::pair<MCCompLabel, float>> mVerticesMCRecInfo;
+  bounded_vector<MCCompLabel> mVerticesContributorLabels;
   std::array<uint32_t, 2> mTotalTracklets = {0, 0};
   unsigned int mNoVertexROF = 0;
   bounded_vector<int> mTotVertPerIteration;
@@ -371,6 +375,22 @@ inline gsl::span<const std::pair<MCCompLabel, float>> TimeFrame<nLayers>::getPri
   return {&(mVerticesMCRecInfo[start]), static_cast<gsl::span<const std::pair<MCCompLabel, float>>::size_type>(delta)};
 }
 
+template <int nLayers>
+inline gsl::span<const MCCompLabel> TimeFrame<nLayers>::getPrimaryVerticesContributors(const int rofId) const
+{
+  // count the number of cont. in rofs before target rof
+  unsigned int start{0}, delta{0};
+  const auto& pvsBefore = getPrimaryVertices(0, rofId - 1);
+  for (const auto& pv : pvsBefore) {
+    start += pv.getNContributors();
+  }
+  const auto& pvsIn = getPrimaryVertices(rofId);
+  for (const auto& pv : pvsIn) {
+    delta += pv.getNContributors();
+  }
+  return {&(mVerticesContributorLabels[start]), static_cast<gsl::span<const MCCompLabel>::size_type>(delta)};
+}
+
 template <int nLayers>
 inline gsl::span<const Vertex> TimeFrame<nLayers>::getPrimaryVertices(int romin, int romax) const
 {

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

@@ -48,7 +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
+  bool useTruthSeeding{false};  // overwrite seeding vertices with MC truth
+  bool outputContLabels{false}; // output additioanlly for each vertex its contributing line labels
 
   int nThreads = 1;
   bool printMemory = false;

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

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

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

@@ -85,6 +85,12 @@ void TimeFrame<nLayers>::addPrimaryVerticesLabels(bounded_vector<std::pair<MCCom
   mVerticesMCRecInfo.insert(mVerticesMCRecInfo.end(), labels.begin(), labels.end());
 }
 
+template <int nLayers>
+void TimeFrame<nLayers>::addPrimaryVerticesContributorLabels(bounded_vector<MCCompLabel>& labels)
+{
+  mVerticesContributorLabels.insert(mVerticesContributorLabels.end(), labels.begin(), labels.end());
+}
+
 template <int nLayers>
 void TimeFrame<nLayers>::addPrimaryVerticesInROF(const bounded_vector<Vertex>& vertices, const int rofId, const int iteration)
 {
@@ -101,6 +107,18 @@ void TimeFrame<nLayers>::addPrimaryVerticesLabelsInROF(const bounded_vector<std:
   mVerticesMCRecInfo.insert(mVerticesMCRecInfo.begin() + mROFramesPV[rofId], labels.begin(), labels.end());
 }
 
+template <int nLayers>
+void TimeFrame<nLayers>::addPrimaryVerticesContributorLabelsInROF(const bounded_vector<MCCompLabel>& labels, const int rofId)
+{
+  // count the number of cont. in rofs before and including the target rof
+  unsigned int n{0};
+  const auto& pvs = getPrimaryVertices(0, rofId);
+  for (const auto& pv : pvs) {
+    n += pv.getNContributors();
+  }
+  mVerticesContributorLabels.insert(mVerticesContributorLabels.begin() + n, labels.begin(), labels.end());
+}
+
 template <int nLayers>
 void TimeFrame<nLayers>::addPrimaryVertices(const gsl::span<const Vertex>& vertices, const int rofId, const int iteration)
 {
@@ -295,6 +313,7 @@ void TimeFrame<nLayers>::initialise(const int iteration, const TrackingParameter
     deepVectorClear(mLinesLabels);
     if (resetVertices) {
       deepVectorClear(mVerticesMCRecInfo);
+      deepVectorClear(mVerticesContributorLabels);
     }
     clearResizeBoundedVector(mTracks, mNrof, mMemoryPool.get());
     clearResizeBoundedVector(mTracksLabel, mNrof, mMemoryPool.get());
@@ -646,6 +665,7 @@ void TimeFrame<nLayers>::setMemoryPool(std::shared_ptr<BoundedMemoryResource>& p
   initVector(mClusterSize);
   initVector(mPValphaX);
   initVector(mBogusClusters);
+  initVector(mVerticesContributorLabels);
   initArrays(mTrackletsIndexROF);
   initVectors(mTracks);
   initVectors(mTracklets);
@@ -689,6 +709,8 @@ void TimeFrame<nLayers>::wipe()
   deepVectorClear(mBogusClusters);
   deepVectorClear(mTrackletsIndexROF);
   deepVectorClear(mPrimaryVertices);
+  deepVectorClear(mTrackletClusters);
+  deepVectorClear(mVerticesContributorLabels);
 }
 
 template class TimeFrame<7>;

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

@@ -117,6 +117,8 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
   static pmr::vector<float> dummyMCPurVerts;
   auto& allTrackLabels = mIsMC ? pc.outputs().make<std::vector<o2::MCCompLabel>>(Output{"ITS", "TRACKSMCTR", 0}) : dummyMCLabTracks;
   auto& allVerticesLabels = mIsMC ? pc.outputs().make<std::vector<o2::MCCompLabel>>(Output{"ITS", "VERTICESMCTR", 0}) : dummyMCLabVerts;
+  bool writeContLabels = mIsMC && o2::its::VertexerParamConfig::Instance().outputContLabels;
+  auto& allVerticesContLabels = writeContLabels ? pc.outputs().make<std::vector<o2::MCCompLabel>>(Output{"ITS", "VERTICESMCTRCONT", 0}) : dummyMCLabVerts;
   auto& allVerticesPurities = mIsMC ? pc.outputs().make<std::vector<float>>(Output{"ITS", "VERTICESMCPUR", 0}) : dummyMCPurVerts;
 
   std::uint32_t roFrame = 0;
@@ -159,6 +161,7 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
   }
   const auto& multEstConf = FastMultEstConfig::Instance(); // parameters for mult estimation and cuts
   gsl::span<const std::pair<MCCompLabel, float>> vMCRecInfo;
+  gsl::span<const MCCompLabel> vMCContLabels;
   for (auto iRof{0}; iRof < trackROFspan.size(); ++iRof) {
     std::vector<Vertex> vtxVecLoc;
     auto& vtxROF = vertROFvec.emplace_back(trackROFspan[iRof]);
@@ -167,6 +170,9 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
       auto vtxSpan = mTimeFrame->getPrimaryVertices(iRof);
       if (mIsMC) {
         vMCRecInfo = mTimeFrame->getPrimaryVerticesMCRecInfo(iRof);
+        if (o2::its::VertexerParamConfig::Instance().outputContLabels) {
+          vMCContLabels = mTimeFrame->getPrimaryVerticesContributors(iRof);
+        }
       }
       if (o2::its::TrackerParamConfig::Instance().doUPCIteration) {
         if (!vtxSpan.empty()) {
@@ -186,17 +192,22 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
       }
       vtxROF.setNEntries(vtxSpan.size());
       bool selROF = vtxSpan.empty();
-      for (auto iV{0}; iV < vtxSpan.size(); ++iV) {
+      for (int iV{0}, iVC{0}; iV < vtxSpan.size(); ++iV) {
         const auto& v = vtxSpan[iV];
         if (multEstConf.isVtxMultCutRequested() && !multEstConf.isPassingVtxMultCut(v.getNContributors())) {
+          iVC += v.getNContributors();
           continue; // skip vertex of unwanted multiplicity
         }
         selROF = true;
         vertices.push_back(v);
         if (mIsMC && !VertexerParamConfig::Instance().useTruthSeeding) {
           allVerticesLabels.push_back(vMCRecInfo[iV].first);
           allVerticesPurities.push_back(vMCRecInfo[iV].second);
+          if (o2::its::VertexerParamConfig::Instance().outputContLabels) {
+            allVerticesContLabels.insert(allVerticesContLabels.end(), vMCContLabels.begin() + iVC, vMCContLabels.begin() + iVC + v.getNContributors());
+          }
         }
+        iVC += v.getNContributors();
       }
       if (processingMask[iRof] && !selROF) { // passed selection in clusters and not in vertex multiplicity
         LOGP(info, "ROF {} rejected by the vertex multiplicity selection [{},{}]", iRof, multEstConf.cutMultVtxLow, multEstConf.cutMultVtxHigh);
@@ -291,6 +302,9 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc)
     if (mIsMC) {
       LOGP(info, "ITSTracker pushed {} track labels", allTrackLabels.size());
       LOGP(info, "ITSTracker pushed {} vertex labels", allVerticesLabels.size());
+      if (!allVerticesContLabels.empty()) {
+        LOGP(info, "ITSTracker pushed {} vertex contributor labels", allVerticesContLabels.size());
+      }
       LOGP(info, "ITSTracker pushed {} vertex purities", allVerticesPurities.size());
     }
   }

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

@@ -408,6 +408,7 @@ void VertexerTraits::computeVertices(const int iteration)
   auto nsigmaCut{std::min(mVrtParams[iteration].vertNsigmaCut * mVrtParams[iteration].vertNsigmaCut * (mVrtParams[iteration].vertRadiusSigma * mVrtParams[iteration].vertRadiusSigma + mVrtParams[iteration].trackletSigma * mVrtParams[iteration].trackletSigma), 1.98f)};
   bounded_vector<Vertex> vertices(mMemoryPool.get());
   bounded_vector<std::pair<o2::MCCompLabel, float>> polls(mMemoryPool.get());
+  bounded_vector<o2::MCCompLabel> contLabels(mMemoryPool.get());
 #ifdef VTX_DEBUG
   std::vector<std::vector<ClusterLines>> dbg_clusLines(mTimeFrame->getNrof());
 #endif
@@ -531,18 +532,27 @@ void VertexerTraits::computeVertices(const int iteration)
             labels.push_back(mTimeFrame->getLinesLabel(rofId)[index]); // then we can use nContributors from vertices to get the labels
           }
           polls.push_back(computeMain(labels));
+          if (mVrtParams[iteration].outputContLabels) {
+            contLabels.insert(contLabels.end(), labels.begin(), labels.end());
+          }
         }
       }
     }
     if (!iteration) {
       mTimeFrame->addPrimaryVertices(vertices, rofId, iteration);
       if (mTimeFrame->hasMCinformation()) {
         mTimeFrame->addPrimaryVerticesLabels(polls);
+        if (mVrtParams[iteration].outputContLabels) {
+          mTimeFrame->addPrimaryVerticesContributorLabels(contLabels);
+        }
       }
     } else {
       mTimeFrame->addPrimaryVerticesInROF(vertices, rofId, iteration);
       if (mTimeFrame->hasMCinformation()) {
         mTimeFrame->addPrimaryVerticesLabelsInROF(polls, rofId);
+        if (mVrtParams[iteration].outputContLabels) {
+          mTimeFrame->addPrimaryVerticesContributorLabelsInROF(contLabels, rofId);
+        }
       }
     }
     if (vertices.empty() && !(iteration && (int)mTimeFrame->getPrimaryVertices(rofId).size() > mVrtParams[iteration].vertPerRofThreshold)) {

Detectors/ITSMFT/ITS/workflow/src/TrackWriterSpec.cxx

@@ -20,6 +20,7 @@
 #include "SimulationDataFormat/MCCompLabel.h"
 #include "SimulationDataFormat/MCTruthContainer.h"
 #include "ReconstructionDataFormats/Vertex.h"
+#include "ITStracking/TrackingConfigParam.h"
 
 using namespace o2::framework;
 
@@ -39,6 +40,7 @@ DataProcessorSpec getTrackWriterSpec(bool useMC)
 {
   // Spectators for logging
   // this is only to restore the original behavior
+  const auto writeContLabels = VertexerParamConfig::Instance().outputContLabels && useMC;
   auto tracksSize = std::make_shared<int>(0);
   auto tracksSizeGetter = [tracksSize](std::vector<o2::its::TrackITS> const& tracks) {
     *tracksSize = tracks.size();
@@ -69,6 +71,11 @@ DataProcessorSpec getTrackWriterSpec(bool useMC)
                                                              "ITSVertexMCTruth",
                                                              (useMC ? 1 : 0), // one branch if mc labels enabled
                                                              ""},
+                                BranchDefinition<LabelsType>{InputSpec{"labelsVerticesContributors", "ITS", "VERTICESMCTRCONT", 0},
+                                                             "ITSVertexMCTruthCont",
+                                                             (writeContLabels ? 1 : 0), // one branch if
+                                                                                        // requested
+                                                             ""},
                                 BranchDefinition<ROFRecLblT>{InputSpec{"MC2ROframes", "ITS", "ITSTrackMC2ROF", 0},
                                                              "ITSTracksMC2ROF",
                                                              (useMC ? 1 : 0), // one branch if mc labels enabled

Detectors/ITSMFT/ITS/workflow/src/TrackerSpec.cxx

@@ -120,6 +120,7 @@ DataProcessorSpec getTrackerSpec(bool useMC, bool useGeom, int trgType, Tracking
     inputs.emplace_back("itsmclabels", "ITS", "CLUSTERSMCTR", 0, Lifetime::Timeframe);
     inputs.emplace_back("ITSMC2ROframes", "ITS", "CLUSTERSMC2ROF", 0, Lifetime::Timeframe);
     outputs.emplace_back("ITS", "VERTICESMCTR", 0, Lifetime::Timeframe);
+    outputs.emplace_back("ITS", "VERTICESMCTRCONT", 0, Lifetime::Timeframe);
     outputs.emplace_back("ITS", "VERTICESMCPUR", 0, Lifetime::Timeframe);
     outputs.emplace_back("ITS", "TRACKSMCTR", 0, Lifetime::Timeframe);
     outputs.emplace_back("ITS", "ITSTrackMC2ROF", 0, Lifetime::Timeframe);