From a872cfab965110c2e1c07a4706a6d2046f7e32dd Mon Sep 17 00:00:00 2001 From: Felix Schlepper Date: Tue, 24 Jun 2025 22:58:15 +0200 Subject: [PATCH] ITS: opt. vertex cont. output Signed-off-by: Felix Schlepper --- .../include/ITStracking/Configuration.h | 1 + .../tracking/include/ITStracking/TimeFrame.h | 20 +++++++++++++++++ .../include/ITStracking/TrackingConfigParam.h | 3 ++- .../ITSMFT/ITS/tracking/src/Configuration.cxx | 1 + .../ITSMFT/ITS/tracking/src/TimeFrame.cxx | 22 +++++++++++++++++++ .../ITS/tracking/src/TrackingInterface.cxx | 16 +++++++++++++- .../ITS/tracking/src/VertexerTraits.cxx | 10 +++++++++ .../ITS/workflow/src/TrackWriterSpec.cxx | 7 ++++++ .../ITSMFT/ITS/workflow/src/TrackerSpec.cxx | 1 + 9 files changed, 79 insertions(+), 2 deletions(-) diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/Configuration.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/Configuration.h index 8c46b2e72078a..19c4617426304 100644 --- a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/Configuration.h +++ b/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 diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h index 230a46681385d..9434fc2292750 100644 --- a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h +++ b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TimeFrame.h @@ -73,15 +73,18 @@ struct TimeFrame { gsl::span getPrimaryVertices(int rofId) const; gsl::span getPrimaryVertices(int romin, int romax) const; gsl::span> getPrimaryVerticesMCRecInfo(const int rofId) const; + gsl::span getPrimaryVerticesContributors(const int rofId) const; gsl::span> getPrimaryVerticesXAlpha(int rofId) const; void fillPrimaryVerticesXandAlpha(); int getPrimaryVerticesNum(int rofId = -1) const; void addPrimaryVertices(const bounded_vector& vertices); void addPrimaryVerticesLabels(bounded_vector>& labels); + void addPrimaryVerticesContributorLabels(bounded_vector& labels); void addPrimaryVertices(const bounded_vector& vertices, const int rofId, const int iteration); void addPrimaryVertices(const gsl::span& vertices, const int rofId, const int iteration); void addPrimaryVerticesInROF(const bounded_vector& vertices, const int rofId, const int iteration); void addPrimaryVerticesLabelsInROF(const bounded_vector>& labels, const int rofId); + void addPrimaryVerticesContributorLabelsInROF(const bounded_vector& labels, const int rofId); void removePrimaryVerticesInROf(const int rofId); int loadROFrameData(const o2::itsmft::ROFRecord& rof, gsl::span clusters, const dataformats::MCTruthContainer* mcLabels = nullptr); @@ -342,6 +345,7 @@ struct TimeFrame { std::array, 2> mTrackletsIndexROF; std::vector> mLinesLabels; std::vector> mVerticesMCRecInfo; + bounded_vector mVerticesContributorLabels; std::array mTotalTracklets = {0, 0}; unsigned int mNoVertexROF = 0; bounded_vector mTotVertPerIteration; @@ -371,6 +375,22 @@ inline gsl::span> TimeFrame::getPri return {&(mVerticesMCRecInfo[start]), static_cast>::size_type>(delta)}; } +template +inline gsl::span TimeFrame::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::size_type>(delta)}; +} + template inline gsl::span TimeFrame::getPrimaryVertices(int romin, int romax) const { diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackingConfigParam.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackingConfigParam.h index 8de80d5e4cd07..dacc2019356f4 100644 --- a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackingConfigParam.h +++ b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackingConfigParam.h @@ -48,7 +48,8 @@ struct VertexerParamConfig : public o2::conf::ConfigurableParamHelper 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; diff --git a/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx b/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx index 9c683112791ab..ea57e5fa8e3b9 100644 --- a/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx +++ b/Detectors/ITSMFT/ITS/tracking/src/TimeFrame.cxx @@ -85,6 +85,12 @@ void TimeFrame::addPrimaryVerticesLabels(bounded_vector +void TimeFrame::addPrimaryVerticesContributorLabels(bounded_vector& labels) +{ + mVerticesContributorLabels.insert(mVerticesContributorLabels.end(), labels.begin(), labels.end()); +} + template void TimeFrame::addPrimaryVerticesInROF(const bounded_vector& vertices, const int rofId, const int iteration) { @@ -101,6 +107,18 @@ void TimeFrame::addPrimaryVerticesLabelsInROF(const bounded_vector +void TimeFrame::addPrimaryVerticesContributorLabelsInROF(const bounded_vector& 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 void TimeFrame::addPrimaryVertices(const gsl::span& vertices, const int rofId, const int iteration) { @@ -295,6 +313,7 @@ void TimeFrame::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::setMemoryPool(std::shared_ptr& p initVector(mClusterSize); initVector(mPValphaX); initVector(mBogusClusters); + initVector(mVerticesContributorLabels); initArrays(mTrackletsIndexROF); initVectors(mTracks); initVectors(mTracklets); @@ -689,6 +709,8 @@ void TimeFrame::wipe() deepVectorClear(mBogusClusters); deepVectorClear(mTrackletsIndexROF); deepVectorClear(mPrimaryVertices); + deepVectorClear(mTrackletClusters); + deepVectorClear(mVerticesContributorLabels); } template class TimeFrame<7>; diff --git a/Detectors/ITSMFT/ITS/tracking/src/TrackingInterface.cxx b/Detectors/ITSMFT/ITS/tracking/src/TrackingInterface.cxx index 31dc68d03a7e8..f5fa195983240 100644 --- a/Detectors/ITSMFT/ITS/tracking/src/TrackingInterface.cxx +++ b/Detectors/ITSMFT/ITS/tracking/src/TrackingInterface.cxx @@ -117,6 +117,8 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc) static pmr::vector dummyMCPurVerts; auto& allTrackLabels = mIsMC ? pc.outputs().make>(Output{"ITS", "TRACKSMCTR", 0}) : dummyMCLabTracks; auto& allVerticesLabels = mIsMC ? pc.outputs().make>(Output{"ITS", "VERTICESMCTR", 0}) : dummyMCLabVerts; + bool writeContLabels = mIsMC && o2::its::VertexerParamConfig::Instance().outputContLabels; + auto& allVerticesContLabels = writeContLabels ? pc.outputs().make>(Output{"ITS", "VERTICESMCTRCONT", 0}) : dummyMCLabVerts; auto& allVerticesPurities = mIsMC ? pc.outputs().make>(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> vMCRecInfo; + gsl::span vMCContLabels; for (auto iRof{0}; iRof < trackROFspan.size(); ++iRof) { std::vector 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,9 +192,10 @@ 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; @@ -196,7 +203,11 @@ void ITSTrackingInterface::run(framework::ProcessingContext& pc) 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()); } } diff --git a/Detectors/ITSMFT/ITS/tracking/src/VertexerTraits.cxx b/Detectors/ITSMFT/ITS/tracking/src/VertexerTraits.cxx index a0f044c5f62ca..34c3ebfcb26ca 100644 --- a/Detectors/ITSMFT/ITS/tracking/src/VertexerTraits.cxx +++ b/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 vertices(mMemoryPool.get()); bounded_vector> polls(mMemoryPool.get()); + bounded_vector contLabels(mMemoryPool.get()); #ifdef VTX_DEBUG std::vector> dbg_clusLines(mTimeFrame->getNrof()); #endif @@ -531,6 +532,9 @@ 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()); + } } } } @@ -538,11 +542,17 @@ void VertexerTraits::computeVertices(const int 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)) { diff --git a/Detectors/ITSMFT/ITS/workflow/src/TrackWriterSpec.cxx b/Detectors/ITSMFT/ITS/workflow/src/TrackWriterSpec.cxx index 5cb6aa199ab64..9fbb138b376b2 100644 --- a/Detectors/ITSMFT/ITS/workflow/src/TrackWriterSpec.cxx +++ b/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(0); auto tracksSizeGetter = [tracksSize](std::vector const& tracks) { *tracksSize = tracks.size(); @@ -69,6 +71,11 @@ DataProcessorSpec getTrackWriterSpec(bool useMC) "ITSVertexMCTruth", (useMC ? 1 : 0), // one branch if mc labels enabled ""}, + BranchDefinition{InputSpec{"labelsVerticesContributors", "ITS", "VERTICESMCTRCONT", 0}, + "ITSVertexMCTruthCont", + (writeContLabels ? 1 : 0), // one branch if + // requested + ""}, BranchDefinition{InputSpec{"MC2ROframes", "ITS", "ITSTrackMC2ROF", 0}, "ITSTracksMC2ROF", (useMC ? 1 : 0), // one branch if mc labels enabled diff --git a/Detectors/ITSMFT/ITS/workflow/src/TrackerSpec.cxx b/Detectors/ITSMFT/ITS/workflow/src/TrackerSpec.cxx index fedc42c017f7e..c8a785b6a3627 100644 --- a/Detectors/ITSMFT/ITS/workflow/src/TrackerSpec.cxx +++ b/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);