add TPC occupancy est. to SVStudy output

shahor02 · shahor02 · commit 9849cbcfaafa · 2025-01-22T19:41:14.000+01:00
diff --git a/Detectors/GlobalTrackingWorkflow/study/include/GlobalTrackingStudy/SVStudy.h b/Detectors/GlobalTrackingWorkflow/study/include/GlobalTrackingStudy/SVStudy.h
@@ -22,7 +22,7 @@
 namespace o2::svstudy
 {
 /// create a processor spec
-o2::framework::DataProcessorSpec getSVStudySpec(o2::dataformats::GlobalTrackID::mask_t srcTracks, bool useMC);
+o2::framework::DataProcessorSpec getSVStudySpec(o2::dataformats::GlobalTrackID::mask_t srcTracks, o2::dataformats::GlobalTrackID::mask_t srcCls, bool useMC);
 
 } // namespace o2::svstudy
 
diff --git a/Detectors/GlobalTrackingWorkflow/study/src/SVStudy.cxx b/Detectors/GlobalTrackingWorkflow/study/src/SVStudy.cxx
@@ -44,6 +44,12 @@
 #include "DCAFitter/DCAFitterN.h"
 #include "MathUtils/fit.h"
 #include "GlobalTrackingStudy/V0Ext.h"
+#include "GPUO2InterfaceConfiguration.h"
+// #include "GPUSettingsO2.h"
+#include "GPUParam.h"
+#include "GPUParam.inc"
+#include "GPUO2InterfaceRefit.h"
+#include "GPUO2InterfaceUtils.h"
 
 namespace o2::svstudy
 {
@@ -64,8 +70,8 @@ using timeEst = o2::dataformats::TimeStampWithError<float, float>;
 class SVStudySpec : public Task
 {
  public:
-  SVStudySpec(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, GTrackID::mask_t src, bool useMC)
-    : mDataRequest(dr), mGGCCDBRequest(gr), mTracksSrc(src), mUseMC(useMC) {}
+  SVStudySpec(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, GTrackID::mask_t src, bool useTPCCl, bool useMC)
+    : mDataRequest(dr), mGGCCDBRequest(gr), mTracksSrc(src), mUseTPCCl(useTPCCl), mUseMC(useMC) {}
   ~SVStudySpec() final = default;
   void init(InitContext& ic) final;
   void run(ProcessingContext& pc) final;
@@ -83,11 +89,18 @@ class SVStudySpec : public Task
   std::unique_ptr<o2::utils::TreeStreamRedirector> mDBGOut;
   float mSelK0 = -1;
   bool mRefit = false;
+  bool mUseTPCCl = false;
   float mMaxEta = 0.8;
   float mBz = 0;
+  int mNHBPerTF = 0;
+  int mNTPCOccBinLength = 0; ///< TPC occ. histo bin length in TBs
+  float mNTPCOccBinLengthInv;
+  float mTPCTBinMUSInv = 0.f;
   GTrackID::mask_t mTracksSrc{};
   o2::vertexing::DCAFitterN<2> mFitterV0;
+  std::vector<float> mTBinClOccAft, mTBinClOccBef;
   std::unique_ptr<o2::steer::MCKinematicsReader> mcReader; // reader of MC information
+  std::shared_ptr<o2::gpu::GPUParam> mParam = nullptr;
 };
 
 void SVStudySpec::init(InitContext& ic)
@@ -107,6 +120,48 @@ void SVStudySpec::run(ProcessingContext& pc)
   o2::globaltracking::RecoContainer recoData;
   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
+
+  size_t occupancyMapSizeBytes = o2::gpu::GPUO2InterfaceRefit::fillOccupancyMapGetSize(mNHBPerTF, mParam.get());
+  gsl::span<const unsigned int> TPCRefitterOccMap = recoData.occupancyMapTPC;
+  o2::gpu::GPUO2InterfaceUtils::paramUseExternalOccupancyMap(mParam.get(), mNHBPerTF, TPCRefitterOccMap.data(), occupancyMapSizeBytes);
+
+  mTBinClOccBef.resize(1);
+  mTBinClOccAft.resize(1);
+  if (recoData.inputsTPCclusters && mUseTPCCl) {
+    mNTPCOccBinLength = mParam->rec.tpc.occupancyMapTimeBins;
+    mTBinClOccBef.clear();
+    mTBinClOccAft.clear();
+    // prepare TPC occupancy data
+    if (mNTPCOccBinLength > 1 && recoData.occupancyMapTPC.size()) {
+      mNTPCOccBinLengthInv = 1. / mNTPCOccBinLength;
+      int nTPCBins = mNHBPerTF * o2::constants::lhc::LHCMaxBunches / 8, ninteg = 0;
+      int nTPCOccBins = nTPCBins * mNTPCOccBinLengthInv, sumBins = std::max(1, int(o2::constants::lhc::LHCMaxBunches / 8 * mNTPCOccBinLengthInv));
+      mTBinClOccAft.resize(nTPCOccBins);
+      mTBinClOccBef.resize(nTPCOccBins);
+      float sm = 0., tb = 0.5 * mNTPCOccBinLength;
+      std::vector<float> mltHistTB(nTPCOccBins);
+      for (int i = 0; i < nTPCOccBins; i++) {
+        mltHistTB[i] = mParam->GetUnscaledMult(tb);
+        tb += mNTPCOccBinLength;
+      }
+      for (int i = nTPCOccBins; i--;) {
+        sm += mltHistTB[i];
+        if (i + sumBins < nTPCOccBins) {
+          sm -= mltHistTB[i + sumBins];
+        }
+        mTBinClOccAft[i] = sm;
+      }
+      sm = 0;
+      for (int i = 0; i < nTPCOccBins; i++) {
+        sm += mltHistTB[i];
+        if (i - sumBins > 0) {
+          sm -= mltHistTB[i - sumBins];
+        }
+        mTBinClOccBef[i] = sm;
+      }
+    }
+  }
+
   process(recoData);
 }
 
@@ -133,6 +188,12 @@ void SVStudySpec::updateTimeDependentParams(ProcessingContext& pc)
     mFitterV0.setMaxStep(svparam.maxStep);
     mFitterV0.setMaxSnp(svparam.maxSnp);
     mFitterV0.setMinXSeed(svparam.minXSeed);
+
+    mNHBPerTF = o2::base::GRPGeomHelper::instance().getGRPECS()->getNHBFPerTF();
+    if (!mParam) {
+      // for occupancy estimator
+      mParam = o2::gpu::GPUO2InterfaceUtils::getFullParamShared(0.f, mNHBPerTF);
+    }
   }
   mBz = o2::base::Propagator::Instance()->getNominalBz();
   mFitterV0.setBz(mBz);
@@ -268,8 +329,13 @@ void SVStudySpec::process(o2::globaltracking::RecoContainer& recoData)
     }
     if (v0extVec.size()) {
       const auto& pv = recoData.getPrimaryVertex(pvID);
+      float tpcOccBef = 0., tpcOccAft = 0.;
+      int tb = pv.getTimeStamp().getTimeStamp() * mTPCTBinMUSInv * mNTPCOccBinLengthInv;
+      tpcOccBef = tb < 0 ? mTBinClOccBef[0] : (tb >= mTBinClOccBef.size() ? mTBinClOccBef.back() : mTBinClOccBef[tb]);
+      tpcOccAft = tb < 0 ? mTBinClOccAft[0] : (tb >= mTBinClOccAft.size() ? mTBinClOccAft.back() : mTBinClOccAft[tb]);
+
       (*mDBGOut) << "v0"
-                 << "orbit=" << recoData.startIR.orbit << "tfID=" << tfID
+                 << "orbit=" << recoData.startIR.orbit << "tfID=" << tfID << "tpcOccBef=" << tpcOccBef << "tpcOccAft=" << tpcOccAft
                  << "v0Ext=" << v0extVec
                  << "pv=" << pv
                  << "\n";
@@ -334,12 +400,13 @@ void SVStudySpec::finaliseCCDB(ConcreteDataMatcher& matcher, void* obj)
   }
 }
 
-DataProcessorSpec getSVStudySpec(GTrackID::mask_t srcTracks, bool useMC)
+DataProcessorSpec getSVStudySpec(GTrackID::mask_t srcTracks, GTrackID::mask_t srcCls, bool useMC)
 {
   std::vector<OutputSpec> outputs;
   auto dataRequest = std::make_shared<DataRequest>();
 
   dataRequest->requestTracks(srcTracks, useMC);
+  dataRequest->requestClusters(srcCls, false);
   dataRequest->requestPrimaryVertices(useMC);
   dataRequest->requestSecondaryVertices(useMC);
   dataRequest->inputs.emplace_back("meanvtx", "GLO", "MEANVERTEX", 0, Lifetime::Condition, ccdbParamSpec("GLO/Calib/MeanVertex", {}, 1));
@@ -351,12 +418,12 @@ DataProcessorSpec getSVStudySpec(GTrackID::mask_t srcTracks, bool useMC)
                                                               o2::base::GRPGeomRequest::None, // geometry
                                                               dataRequest->inputs,
                                                               true);
-
+  bool useTPCcl = srcCls[GTrackID::TPC];
   return DataProcessorSpec{
     "sv-study",
     dataRequest->inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<SVStudySpec>(dataRequest, ggRequest, srcTracks, useMC)},
+    AlgorithmSpec{adaptFromTask<SVStudySpec>(dataRequest, ggRequest, srcTracks, useTPCcl, useMC)},
     Options{
       {"refit", VariantType::Bool, false, {"refit SVertices"}},
       {"sel-k0", VariantType::Float, -1.f, {"If positive, select K0s with this mass margin"}},
diff --git a/Detectors/GlobalTrackingWorkflow/study/src/TrackingStudy.cxx b/Detectors/GlobalTrackingWorkflow/study/src/TrackingStudy.cxx
@@ -49,6 +49,7 @@
 #include "GPUParam.inc"
 #include "Steer/MCKinematicsReader.h"
 #include "MathUtils/fit.h"
+#include <TF1.h>
 
 namespace o2::trackstudy
 {
@@ -93,7 +94,8 @@ class TrackingStudySpec : public Task
   std::unique_ptr<o2::utils::TreeStreamRedirector> mDBGOut;
   std::unique_ptr<o2::utils::TreeStreamRedirector> mDBGOutVtx;
   std::unique_ptr<o2::gpu::GPUO2InterfaceRefit> mTPCRefitter; ///< TPC refitter used for TPC tracks refit during the reconstruction
-  std::vector<float> mTBinClOccAft, mTBinClOccBef;            ///< TPC occupancy histo: i-th entry is the integrated occupancy for ~1 orbit starting/preceding from the TB = i*mNTPCOccBinLength
+  std::vector<float> mTBinClOccAft, mTBinClOccBef, mTBinClOccWgh; ///< TPC occupancy histo: i-th entry is the integrated occupancy for ~1 orbit starting/preceding from the TB = i*mNTPCOccBinLength
+  std::unique_ptr<TF1> mOccWghFun;
   float mITSROFrameLengthMUS = 0.f;
   float mTPCTBinMUS = 0.f; // TPC bin in microseconds
   float mTPCTBinMUSInv = 0.f;
@@ -139,6 +141,10 @@ void TrackingStudySpec::init(InitContext& ic)
   mDCAYFormula = ic.options().get<std::string>("dcay-vs-pt");
   mDCAZFormula = ic.options().get<std::string>("dcaz-vs-pt");
   mDoPairsCorr = ic.options().get<bool>("pair-correlations");
+  auto str = ic.options().get<std::string>("occ-weight-fun");
+  if (!str.empty()) {
+    mOccWghFun = std::make_unique<TF1>("occFun", str.c_str(), -100., 100.);
+  }
 }
 
 void TrackingStudySpec::run(ProcessingContext& pc)
@@ -154,16 +160,37 @@ void TrackingStudySpec::run(ProcessingContext& pc)
     mNTPCOccBinLength = mTPCRefitter->getParam()->rec.tpc.occupancyMapTimeBins;
     mTBinClOccBef.clear();
     mTBinClOccAft.clear();
+    mTBinClOccWgh.clear();
   }
+
   // prepare TPC occupancy data
   if (mNTPCOccBinLength > 1 && recoData.occupancyMapTPC.size()) {
     mNTPCOccBinLengthInv = 1. / mNTPCOccBinLength;
     int nTPCBins = mNHBPerTF * o2::constants::lhc::LHCMaxBunches / 8, ninteg = 0;
     int nTPCOccBins = nTPCBins * mNTPCOccBinLengthInv, sumBins = std::max(1, int(o2::constants::lhc::LHCMaxBunches / 8 * mNTPCOccBinLengthInv));
     mTBinClOccAft.resize(nTPCOccBins);
     mTBinClOccBef.resize(nTPCOccBins);
-    std::vector<float> mltHistTB(nTPCOccBins);
     float sm = 0., tb = 0.5 * mNTPCOccBinLength;
+    /* // at the moment not used
+    if (mOccWghFun) {
+      mTBinClOccWgh.resize(nTPCBins);
+      float occBin2MUS = 8 * o2::constants::lhc::LHCBunchSpacingMUS;
+      int covWghTB = TMath::NInt(100./occBin2MUS); // coverage of weighted occ. in TBins
+      for (int i = 0; i < nTPCBins; i++) {
+  sm = 0.;
+  for (int j=-covWghTB;j<covWghTB;j++) {
+    if (j+i<0 || j+i>=nTPCBins) {
+      continue;
+    }
+    sm += mOccWghFun->Eval(j*occBin2MUS)*mTPCRefitter->getParam()->GetUnscaledMult(j+i);
+  }
+  mTBinClOccWgh[i] = sm;
+      }
+    } else {
+      mTBinClOccWgh.resize(1);
+    }
+    */
+    std::vector<float> mltHistTB(nTPCOccBins);
     for (int i = 0; i < nTPCOccBins; i++) {
       mltHistTB[i] = mTPCRefitter->getParam()->GetUnscaledMult(tb);
       tb += mNTPCOccBinLength;
@@ -719,6 +746,7 @@ DataProcessorSpec getTrackingStudySpec(GTrackID::mask_t srcTracks, GTrackID::mas
     {"min-pt", VariantType::Float, 0.1f, {"Cut on track pT"}},
     {"with-its-only", VariantType::Bool, false, {"Store tracks with ITS only"}},
     {"pair-correlations", VariantType::Bool, false, {"Do pairs correlation"}},
+    {"occ-weight-fun", VariantType::String, "(x>=-40&&x<-5) ? (1./1225*pow(x+40,2)) : ((x>-5&&x<15) ? 1. : ((x>=15&&x<40) ? (-0.4/25*x+1.24 ) : ( (x>40&&x<100) ? -0.4/60*x+0.6+0.8/3 : 0)))", {"Occupancy weighting f-n vs time in musec"}},
     {"min-x-prop", VariantType::Float, 100.f, {"track should be propagated to this X at least"}},
   };
   o2::tpc::VDriftHelper::requestCCDBInputs(dataRequest->inputs);
diff --git a/Detectors/GlobalTrackingWorkflow/study/src/sv-study-workflow.cxx b/Detectors/GlobalTrackingWorkflow/study/src/sv-study-workflow.cxx
@@ -39,6 +39,7 @@ void customize(std::vector<ConfigParamSpec>& workflowOptions)
     {"disable-mc", o2::framework::VariantType::Bool, false, {"disable MC propagation"}},
     {"track-sources", VariantType::String, std::string{GID::ALL}, {"comma-separated list of track sources to use"}},
     {"disable-root-input", VariantType::Bool, false, {"disable root-files input reader"}},
+    {"ignore-tpc-occ", VariantType::Bool, false, {"do not fill TPC occupancy (needs TPC clusters)"}},
     {"configKeyValues", VariantType::String, "", {"Semicolon separated key=value strings ..."}}};
   o2::raw::HBFUtilsInitializer::addConfigOption(options);
   std::swap(workflowOptions, options);
@@ -61,10 +62,14 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
 
   GID::mask_t srcTrc = allowedSourcesTrc & GID::getSourcesMask(configcontext.options().get<std::string>("track-sources"));
   GID::mask_t srcCls{};
+  bool fillTPCOcc = !configcontext.options().get<bool>("ignore-tpc-occ");
+  if (fillTPCOcc) {
+    srcCls = srcCls | GID::getSourcesMask("TPC");
+  }
   o2::globaltracking::InputHelper::addInputSpecs(configcontext, specs, srcCls, srcTrc, srcTrc, useMC);
   o2::globaltracking::InputHelper::addInputSpecsPVertex(configcontext, specs, useMC); // P-vertex is always needed
   o2::globaltracking::InputHelper::addInputSpecsSVertex(configcontext, specs);        // S-vertex is always needed
-  specs.emplace_back(o2::svstudy::getSVStudySpec(srcTrc, useMC));
+  specs.emplace_back(o2::svstudy::getSVStudySpec(srcTrc, srcCls, useMC));
 
   // configure dpl timer to inject correct firstTForbit: start from the 1st orbit of TF containing 1st sampled orbit
   o2::raw::HBFUtilsInitializer hbfIni(configcontext, specs);

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@`
`22`	`22`	`namespace o2::svstudy`
`23`	`23`	`{`
`24`	`24`	`/// create a processor spec`
`25`		`-o2::framework::DataProcessorSpec getSVStudySpec(o2::dataformats::GlobalTrackID::mask_t srcTracks, bool useMC);`
	`25`	`+o2::framework::DataProcessorSpec getSVStudySpec(o2::dataformats::GlobalTrackID::mask_t srcTracks, o2::dataformats::GlobalTrackID::mask_t srcCls, bool useMC);`
`26`	`26`
`27`	`27`	`} // namespace o2::svstudy`
`28`	`28`