AOD: Add thinner option to producer and seedable

This MR adds the following features:
 - make the random number optionally seedable, good for debugging
 - allow to thin tpc-only tracks and produce same output as the
 aod-thinner
 - adds some const for correctness

Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>

Author: f3sch

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -227,17 +227,19 @@ class AODProducerWorkflowDPL : public Task
     return std::uint64_t(mStartIR.toLong()) + relativeTime_to_LocalBC(relativeTimeStampInNS);
   }
 
+  bool mThinTracks{false};
   bool mPropTracks{false};
   bool mPropMuons{false};
   float mTrackQCFraction{0.00};
   int64_t mTrackQCNTrCut{4};
   float mSqrtS{13860.};
-  std::mt19937 mGenerator; ///< random generator for trackQA sampling
+  std::mt19937 mGenerator{}; ///< random generator for trackQA sampling
   o2::base::Propagator::MatCorrType mMatCorr{o2::base::Propagator::MatCorrType::USEMatCorrLUT};
   o2::dataformats::MeanVertexObject mVtx;
   float mMinPropR{o2::constants::geom::XTPCInnerRef + 0.1f};
 
   std::unordered_set<GIndex> mGIDUsedBySVtx;
+  std::unordered_set<GIndex> mGIDUsedByStr;
 
   int mNThreads = 1;
   bool mUseMC = true;

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -455,7 +455,7 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
       tracksExtraCursor.reserve(nToReserve + tracksExtraCursor.lastIndex());
     }
     for (int ti = start; ti < end; ti++) {
-      auto& trackIndex = GIndices[ti];
+      const auto& trackIndex = GIndices[ti];
       if (GIndex::includesSource(src, mInputSources)) {
         if (src == GIndex::Source::MFT) {                                                                // MFT tracks are treated separately since they are stored in a different table
           if (trackIndex.isAmbiguous() && mGIDToTableMFTID.find(trackIndex) != mGIDToTableMFTID.end()) { // was it already stored ?
@@ -477,11 +477,12 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
           if (trackIndex.isAmbiguous() && mGIDToTableID.find(trackIndex) != mGIDToTableID.end()) { // was it already stored ?
             continue;
           }
-          auto extraInfoHolder = processBarrelTrack(collisionID, collisionBC, trackIndex, data, bcsMap);
 
           float weight = 0;
-          std::uniform_real_distribution<> distr(0., 1.);
+          static std::uniform_real_distribution<> distr(0., 1.);
           bool writeQAData = o2::math_utils::Tsallis::downsampleTsallisCharged(data.getTrackParam(trackIndex).getPt(), mTrackQCFraction, mSqrtS, weight, distr(mGenerator));
+          auto extraInfoHolder = processBarrelTrack(collisionID, collisionBC, trackIndex, data, bcsMap);
+
           if (writeQAData) {
             auto trackQAInfoHolder = processBarrelTrackQA(collisionID, collisionBC, trackIndex, data, bcsMap);
             if (std::bitset<8>(trackQAInfoHolder.tpcClusterByteMask).count() >= mTrackQCNTrCut) {
@@ -490,17 +491,26 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
               trackQAInfoHolder.tpcTime0 = (trackQAInfoHolder.tpcTime0 * 8 - extraInfoHolder.diffBCRef) * o2::constants::lhc::LHCBunchSpacingNS - extraInfoHolder.trackTime;
               // difference between TPC track time0 and stored track nominal time in ns instead of TF start
               addToTracksQATable(tracksQACursor, trackQAInfoHolder);
+            } else {
+              writeQAData = false;
             }
           }
 
+          if (mThinTracks && src == GIndex::Source::TPC && mGIDUsedBySVtx.find(trackIndex) == mGIDUsedBySVtx.end() && mGIDUsedByStr.find(trackIndex) == mGIDUsedByStr.end() && !writeQAData) {
+            mGIDToTableID.emplace(trackIndex, -1); // pretend skipped tracks are stored; this is safe since they are are not written to disk and -1 indicates to all users to not use this track
+            continue;
+          }
+
           if (!extraInfoHolder.isTPConly && extraInfoHolder.trackTimeRes < 0.f) { // failed or rejected?
             LOG(warning) << "Barrel track " << trackIndex << " has no time set, rejection is not expected : time=" << extraInfoHolder.trackTime
                          << " timeErr=" << extraInfoHolder.trackTimeRes << " BCSlice: " << extraInfoHolder.bcSlice[0] << ":" << extraInfoHolder.bcSlice[1];
             continue;
           }
           const auto& trOrig = data.getTrackParam(trackIndex);
           bool isProp = false;
-          if (mPropTracks && trOrig.getX() < mMinPropR && mGIDUsedBySVtx.find(trackIndex) == mGIDUsedBySVtx.end()) { // Do not propagate track assoc. to V0s
+          if (mPropTracks && trOrig.getX() < mMinPropR &&
+              mGIDUsedBySVtx.find(trackIndex) == mGIDUsedBySVtx.end() &&
+              mGIDUsedByStr.find(trackIndex) == mGIDUsedByStr.end()) { // Do not propagate track assoc. to V0s and str. tracking
             auto trackPar(trOrig);
             isProp = propagateTrackToPV(trackPar, data, collisionID);
             if (isProp) {
@@ -511,7 +521,7 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
             addToTracksTable(tracksCursor, tracksCovCursor, trOrig, collisionID, aod::track::TrackIU);
           }
           addToTracksExtraTable(tracksExtraCursor, extraInfoHolder);
-          // addToTracksQATable(tracksQACursor, trackQAInfoHolder);
+
           //  collecting table indices of barrel tracks for V0s table
           if (extraInfoHolder.bcSlice[0] >= 0 && collisionID < 0) {
             ambigTracksCursor(mTableTrID, extraInfoHolder.bcSlice);
@@ -1334,7 +1344,7 @@ void AODProducerWorkflowDPL::fillStrangenessTrackingTables(const o2::globaltrack
   int nCasc = 0;
   int nD3Body = 0;
 
-  for (auto& sTrk : recoData.getStrangeTracks()) {
+  for (const auto& sTrk : recoData.getStrangeTracks()) {
     if (sTrk.mPartType == dataformats::kStrkV0) {
       nV0++;
     } else if (sTrk.mPartType == dataformats::kStrkCascade) {
@@ -1348,13 +1358,13 @@ void AODProducerWorkflowDPL::fillStrangenessTrackingTables(const o2::globaltrack
   cascCurs.reserve(nCasc);
   d3BodyCurs.reserve(nD3Body);
 
-  for (auto& sTrk : recoData.getStrangeTracks()) {
+  for (const auto& sTrk : recoData.getStrangeTracks()) {
     auto ITSIndex = GIndex{sTrk.mITSRef, GIndex::ITS};
     auto item = mGIDToTableID.find(ITSIndex);
     if (item != mGIDToTableID.end()) {
       itsTableIdx = item->second;
     } else {
-      LOG(warn) << "Could not find a ITS strange track index";
+      LOG(warn) << "Could not find a ITS strange track index " << ITSIndex;
       continue;
     }
     if (sTrk.mPartType == dataformats::kStrkV0) {
@@ -1649,11 +1659,22 @@ void AODProducerWorkflowDPL::init(InitContext& ic)
   mCTPReadout = ic.options().get<int>("ctpreadout-create");
   mNThreads = std::max(1, ic.options().get<int>("nthreads"));
   mEMCselectLeading = ic.options().get<bool>("emc-select-leading");
+  mThinTracks = ic.options().get<bool>("thin-tracks");
   mPropTracks = ic.options().get<bool>("propagate-tracks");
   mPropMuons = ic.options().get<bool>("propagate-muons");
   mTrackQCFraction = ic.options().get<float>("trackqc-fraction");
   mTrackQCNTrCut = ic.options().get<int64_t>("trackqc-NTrCut");
-  mGenerator = std::mt19937(std::random_device{}());
+  if (auto seed = ic.options().get<int>("seed"); seed == 0) {
+    LOGP(info, "Using random device for seeding");
+    std::random_device rd;
+    std::array<int, std::mt19937::state_size> seed_data{};
+    std::generate(std::begin(seed_data), std::end(seed_data), std::ref(rd));
+    std::seed_seq seq(std::begin(seed_data), std::end(seed_data));
+    mGenerator = std::mt19937(seq);
+  } else {
+    LOGP(info, "Using seed {} for sampling", seed);
+    mGenerator.seed(seed);
+  }
 #ifdef WITH_OPENMP
   LOGP(info, "Multi-threaded parts will run with {} OpenMP threads", mNThreads);
 #else
@@ -2094,7 +2115,7 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   mGIDToTableFwdID.clear(); // reset the tables to be used by 'fillTrackTablesPerCollision'
   mGIDToTableMFTID.clear();
 
-  if (mPropTracks) {
+  if (mPropTracks || mThinTracks) {
     auto v0s = recoData.getV0sIdx();
     auto cascades = recoData.getCascadesIdx();
     auto decays3Body = recoData.getDecays3BodyIdx();
@@ -2111,6 +2132,11 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
       mGIDUsedBySVtx.insert(id3Body.getProngID(1));
       mGIDUsedBySVtx.insert(id3Body.getProngID(2));
     }
+
+    mGIDUsedByStr.reserve(recoData.getStrangeTracks().size());
+    for (const auto& sTrk : recoData.getStrangeTracks()) {
+      mGIDUsedByStr.emplace(sTrk.mITSRef, GIndex::ITS);
+    }
   }
 
   // filling unassigned tracks first
@@ -2284,6 +2310,7 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   mBCLookup.clear();
 
   mGIDUsedBySVtx.clear();
+  mGIDUsedByStr.clear();
 
   originCursor(tfNumber);
 
@@ -3003,8 +3030,10 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, boo
       ConfigParamSpec{"propagate-tracks", VariantType::Bool, false, {"Propagate tracks (not used for secondary vertices) to IP"}},
       ConfigParamSpec{"hepmc-update", VariantType::String, "always", {"When to update HepMC Aux tables: always - force update, never - never update, all - if all keys are present, any - when any key is present (not valid yet)"}},
       ConfigParamSpec{"propagate-muons", VariantType::Bool, false, {"Propagate muons to IP"}},
+      ConfigParamSpec{"thin-tracks", VariantType::Bool, false, {"Produce thinned track tables"}},
       ConfigParamSpec{"trackqc-fraction", VariantType::Float, float(0.1), {"Fraction of tracks to QC"}},
       ConfigParamSpec{"trackqc-NTrCut", VariantType::Int64, 4L, {"Minimal length of the track - in amount of tracklets"}},
+      ConfigParamSpec{"seed", VariantType::Int, 0, {"Set seed for random generator used for sampling (0 (default) means using a random_device)"}},
     }};
 }