Add more info to MCStudy w-flow

Author: shahor02

Detectors/GlobalTrackingWorkflow/study/include/GlobalTrackingStudy/TrackMCStudyConfig.h

@@ -28,6 +28,8 @@ struct TrackMCStudyConfig : o2::conf::ConfigurableParamHelper<TrackMCStudyConfig
   bool requireITSorTPCTrackRefs = true;
   bool requireTopBottomRefs = false;
   int minTPCRefsToExtractClRes = 2;
+  int nOccBinsDrift = 10; // number of bins for TPC max drift time, where we integrate the occupancies
+  int nTBPerOccBin = 48;  // number of TB per occ bin
   float rejectClustersResStat = 0.1;
   float maxTPCRefExtrap = 2;                   // max dX to extrapolate the track ref when extrapolating track true posions
   int decayPDG[5] = {310, 3122, 411, 421, -1}; // decays to study, must end by -1

Detectors/GlobalTrackingWorkflow/study/include/GlobalTrackingStudy/TrackMCStudyTypes.h

@@ -32,7 +32,7 @@ struct MCTrackInfo {
   int getNITSClusForAB() const;
   int getLowestITSLayer() const;
   int getHighestITSLayer() const;
-
+  std::vector<float> occTPCV{};
   o2::track::TrackPar track{};
   o2::MCCompLabel label{};
   float occTPC = -1.f;
@@ -52,8 +52,24 @@ struct MCTrackInfo {
   uint8_t maxTPCRowSect = -1;
   int8_t nITSCl = 0;
   int8_t pattITSCl = 0;
-  bool addedAtRecStage = false;
-  ClassDefNV(MCTrackInfo, 5);
+  uint8_t flags = 0;
+
+  enum Flags : uint32_t { Primary = 0,
+                          AddedAtRecStage = 2,
+                          BitMask = 0xff };
+
+  bool isPrimary() const { return isBitSet(Primary); }
+  bool isAddedAtRecStage() const { return isBitSet(AddedAtRecStage); }
+  void setPrimary() { setBit(Primary); }
+  void setAddedAtRecStage() { setBit(AddedAtRecStage); }
+
+  uint8_t getBits() const { return flags; }
+  bool isBitSet(int bit) const { return flags & (0xff & (0x1 << bit)); }
+  void setBits(std::uint8_t b) { flags = b; }
+  void setBit(int bit) { flags |= BitMask & (0x1 << bit); }
+  void resetBit(int bit) { flags &= ~(BitMask & (0x1 << bit)); }
+
+  ClassDefNV(MCTrackInfo, 7);
 };
 
 struct RecTrack {
@@ -272,7 +288,8 @@ struct MCVertex {
   int nTrackSel = 0; // number of selected MC charged tracks
   int ID = -1;
   std::vector<RecPV> recVtx{};
-  ClassDefNV(MCVertex, 1);
+  std::vector<float> occTPCV{};
+  ClassDefNV(MCVertex, 2);
 };
 
 } // namespace o2::trackstudy

Detectors/GlobalTrackingWorkflow/study/src/TrackMCStudy.cxx

@@ -109,7 +109,7 @@ class TrackMCStudy : public Task
   void updateTimeDependentParams(ProcessingContext& pc);
   float getDCAYCut(float pt) const;
 
-  gsl::span<const MCTrack> mCurrMCTracks;
+  const std::vector<o2::MCTrack>* mCurrMCTracks = nullptr;
   TVector3 mCurrMCVertex;
   o2::tpc::VDriftHelper mTPCVDriftHelper{};
   o2::tpc::CorrectionMapsLoader mTPCCorrMapsLoader{};
@@ -124,7 +124,7 @@ class TrackMCStudy : public Task
   bool mCheckSV = false;         //< check SV binding (apart from prongs availability)
   bool mRecProcStage = false;    //< flag that the MC particle was added only at the stage of reco tracks processing
   int mNTPCOccBinLength = 0;     ///< TPC occ. histo bin length in TBs
-  float mNTPCOccBinLengthInv;
+  float mNTPCOccBinLengthInv = -1.f;
   int mVerbose = 0;
   float mITSTimeBiasMUS = 0.f;
   float mITSROFrameLengthMUS = 0.f; ///< ITS RO frame in mus
@@ -182,7 +182,7 @@ void TrackMCStudy::run(ProcessingContext& pc)
   }
   mDecProdLblPool.clear();
   mMCVtVec.clear();
-  mCurrMCTracks = {};
+  mCurrMCTracks = nullptr;
 
   recoData.collectData(pc, *mDataRequest.get()); // select tracks of needed type, with minimal cuts, the real selected will be done in the vertexer
   updateTimeDependentParams(pc);                 // Make sure this is called after recoData.collectData, which may load some conditions
@@ -346,6 +346,21 @@ void TrackMCStudy::process(const o2::globaltracking::RecoContainer& recoData)
         }
         break;
       }
+      if (mNTPCOccBinLengthInv > 0.f) {
+        mcVtx.occTPCV.resize(params.nOccBinsDrift);
+        int grp = TMath::Max(1, TMath::Nint(params.nTBPerOccBin * mNTPCOccBinLengthInv));
+        for (int ib = 0; ib < params.nOccBinsDrift; ib++) {
+          float smb = 0;
+          int tbs = occBin + TMath::Nint(ib * params.nTBPerOccBin * mNTPCOccBinLengthInv);
+          for (int ig = 0; ig < grp; ig++) {
+            if (tbs >= 0 && tbs < int(mTBinClOccHist.size())) {
+              smb += mTBinClOccHist[tbs];
+            }
+            tbs++;
+          }
+          mcVtx.occTPCV[ib] = smb;
+        }
+      }
       if (rofCount >= ITSClusROFRec.size()) {
         mITSOcc.push_back(0); // IR after the last ROF
       }
@@ -362,13 +377,15 @@ void TrackMCStudy::process(const o2::globaltracking::RecoContainer& recoData)
     if (nev != (int)mMCVtVec.size()) {
       LOGP(debug, "source {} has {} events while {} MC vertices were booked", curSrcMC, nev, mMCVtVec.size());
       okAccVtx = false;
+      if (nev > (int)mMCVtVec.size()) { // QED
+        continue;
+      }
     }
     for (curEvMC = 0; curEvMC < nev; curEvMC++) {
       if (mVerbose > 1) {
         LOGP(info, "Event {}", curEvMC);
       }
-      const auto& mt = mcReader.getTracks(curSrcMC, curEvMC);
-      mCurrMCTracks = gsl::span<const MCTrack>(mt.data(), mt.size());
+      mCurrMCTracks = &mcReader.getTracks(curSrcMC, curEvMC);
       const_cast<o2::dataformats::MCEventHeader&>(mcReader.getMCEventHeader(curSrcMC, curEvMC)).GetVertex(mCurrMCVertex);
       if (okAccVtx) {
         auto& pos = mMCVtVec[curEvMC].pos;
@@ -378,7 +395,7 @@ void TrackMCStudy::process(const o2::globaltracking::RecoContainer& recoData)
           pos[2] = mCurrMCVertex.Z();
         }
       }
-      for (int itr = 0; itr < mCurrMCTracks.size(); itr++) {
+      for (int itr = 0; itr < mCurrMCTracks->size(); itr++) {
         processMCParticle(curSrcMC, curEvMC, itr);
       }
     }
@@ -425,11 +442,10 @@ void TrackMCStudy::process(const o2::globaltracking::RecoContainer& recoData)
             if (lbl.getSourceID() != curSrcMC || lbl.getEventID() != curEvMC) {
               curSrcMC = lbl.getSourceID();
               curEvMC = lbl.getEventID();
-              const auto& mt = mcReader.getTracks(curSrcMC, curEvMC);
-              mCurrMCTracks = gsl::span<const MCTrack>(mt.data(), mt.size());
+              mCurrMCTracks = &mcReader.getTracks(curSrcMC, curEvMC);
               const_cast<o2::dataformats::MCEventHeader&>(mcReader.getMCEventHeader(curSrcMC, curEvMC)).GetVertex(mCurrMCVertex);
             }
-            if (!acceptMCCharged(mCurrMCTracks[lbl.getTrackID()], lbl)) {
+            if (!acceptMCCharged((*mCurrMCTracks)[lbl.getTrackID()], lbl)) {
               continue;
             }
             entry = mSelMCTracks.find(lbl);
@@ -977,7 +993,7 @@ float TrackMCStudy::getDCAYCut(float pt) const
 
 bool TrackMCStudy::processMCParticle(int src, int ev, int trid)
 {
-  const auto& mcPart = mCurrMCTracks[trid];
+  const auto& mcPart = (*mCurrMCTracks)[trid];
   int pdg = mcPart.GetPdgCode();
   bool res = false;
   while (true) {
@@ -999,7 +1015,7 @@ bool TrackMCStudy::processMCParticle(int src, int ev, int trid)
           break;
         }
         for (int idd = idd0; idd <= idd1; idd++) {
-          const auto& product = mCurrMCTracks[idd];
+          const auto& product = (*mCurrMCTracks)[idd];
           auto lbld = o2::MCCompLabel(idd, ev, src);
           if (!acceptMCCharged(product, lbld, decay)) {
             decay = -1; // discard decay
@@ -1097,11 +1113,17 @@ bool TrackMCStudy::addMCParticle(const MCTrack& mcPart, const o2::MCCompLabel& l
   mcEntry.mcTrackInfo.bcInTF = mIntBC[lb.getEventID()];
   mcEntry.mcTrackInfo.occTPC = mTPCOcc[lb.getEventID()];
   mcEntry.mcTrackInfo.occITS = mITSOcc[lb.getEventID()];
-  mcEntry.mcTrackInfo.addedAtRecStage = mRecProcStage;
+  mcEntry.mcTrackInfo.occTPCV = mMCVtVec[lb.getEventID()].occTPCV;
+  if (mRecProcStage) {
+    mcEntry.mcTrackInfo.setAddedAtRecStage();
+  }
+  if (o2::mcutils::MCTrackNavigator::isPhysicalPrimary(mcPart, *mCurrMCTracks)) {
+    mcEntry.mcTrackInfo.setPrimary();
+  }
   int moth = -1;
   o2::MCCompLabel mclbPar;
   if ((moth = mcPart.getMotherTrackId()) >= 0) {
-    const auto& mcPartPar = mCurrMCTracks[moth];
+    const auto& mcPartPar = (*mCurrMCTracks)[moth];
     mcEntry.mcTrackInfo.pdgParent = mcPartPar.GetPdgCode();
   }
   if (mcPart.isPrimary() && mcReader.getNEvents(lb.getSourceID()) == mMCVtVec.size()) {