Add barrel track propagation to PV at AOD creation stage

nburmaso · nburmaso · commit 38887fc13829 · 2023-06-08T14:46:26.000+03:00
diff --git a/Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h b/Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h
@@ -28,6 +28,7 @@
 #include "DataFormatsEMCAL/EventHandler.h"
 #include "DataFormatsPHOS/EventHandler.h"
 #include "DetectorsBase/GRPGeomHelper.h"
+#include "DetectorsBase/Propagator.h"
 #include "Framework/AnalysisDataModel.h"
 #include "Framework/AnalysisHelpers.h"
 #include "Framework/DataProcessorSpec.h"
@@ -248,6 +249,19 @@ class AODProducerWorkflowDPL : public Task
     return c;
   }
 
+  bool mPropTracks{false};
+  bool mPropTracksCov{false};
+  std::string mCCDBUrl{"http://alice-ccdb.cern.ch"};
+  std::string mVtxPath{"GLO/Calib/MeanVertex"};
+
+  o2::ccdb::CcdbApi mCCDBAPI;
+
+  o2::base::Propagator::MatCorrType mMatCorr{o2::base::Propagator::MatCorrType::USEMatCorrLUT};
+  o2::dataformats::MeanVertexObject* mVtx{nullptr};
+  float mMinPropR{o2::constants::geom::XTPCInnerRef + 0.1f};
+
+  std::unordered_set<GIndex> mGIDUsedBySVtx;
+
   int mNThreads = 1;
   bool mUseMC = true;
   bool mEnableSV = true; // enable secondary vertices
@@ -323,7 +337,6 @@ class AODProducerWorkflowDPL : public Task
   // production metadata
   std::vector<TString> mMetaDataKeys;
   std::vector<TString> mMetaDataVals;
-  bool mIsMDSent{false};
 
   std::shared_ptr<DataRequest> mDataRequest;
   std::shared_ptr<o2::base::GRPGeomRequest> mGGCCDBRequest;
@@ -456,6 +469,19 @@ class AODProducerWorkflowDPL : public Task
   };
   std::vector<TPCCounters> mTPCCounters;
 
+  void fetchMeanVtxCCDB(int runNumber)
+  {
+    static int prevRunNumber = -1;
+    if (runNumber == prevRunNumber)
+      return;
+    auto& ccdbMgr = o2::ccdb::BasicCCDBManager::instance();
+    std::map<std::string, std::string> headers, metadataRCT;
+    headers = mCCDBAPI.retrieveHeaders(Form("RCT/Info/RunInformation/%i", runNumber), metadataRCT, -1);
+    int64_t ts = std::atol(headers["SOR"].c_str()); // timestamp in ms
+    mVtx = ccdbMgr.getForTimeStamp<o2::dataformats::MeanVertexObject>(mVtxPath, ts);
+    prevRunNumber = runNumber;
+  }
+
   void updateTimeDependentParams(ProcessingContext& pc);
 
   void addRefGlobalBCsForTOF(const o2::dataformats::VtxTrackRef& trackRef, const gsl::span<const GIndex>& GIndices,
@@ -465,7 +491,6 @@ class AODProducerWorkflowDPL : public Task
                   const std::vector<o2::InteractionTimeRecord>& mcRecords,
                   std::map<uint64_t, int>& bcsMap);
 
-  uint64_t getTFNumber(const o2::InteractionRecord& tfStartIR, int runNumber);
   template <typename TracksCursorType, typename TracksCovCursorType>
   void addToTracksTable(TracksCursorType& tracksCursor, TracksCovCursorType& tracksCovCursor,
                         const o2::track::TrackParCov& track, int collisionID, aod::track::TrackTypeEnum type = aod::track::TrackIU);
@@ -483,6 +508,7 @@ class AODProducerWorkflowDPL : public Task
                            GIndex trackID, const o2::globaltracking::RecoContainer& data, int collisionID, std::uint64_t collisionBC, const std::map<uint64_t, int>& bcsMap);
 
   TrackExtraInfo processBarrelTrack(int collisionID, std::uint64_t collisionBC, GIndex trackIndex, const o2::globaltracking::RecoContainer& data, const std::map<uint64_t, int>& bcsMap);
+  bool propagateTrackToPV(o2::track::TrackParametrizationWithError<float>& trackPar, const o2::globaltracking::RecoContainer& data, int colID);
   void extrapolateToCalorimeters(TrackExtraInfo& extraInfoHolder, const o2::track::TrackPar& track);
   void cacheTriggers(const o2::globaltracking::RecoContainer& recoData);
 
diff --git a/Detectors/AOD/src/AODProducerWorkflowSpec.cxx b/Detectors/AOD/src/AODProducerWorkflowSpec.cxx
@@ -281,37 +281,6 @@ void AODProducerWorkflowDPL::collectBCs(const o2::globaltracking::RecoContainer&
   }
 }
 
-uint64_t AODProducerWorkflowDPL::getTFNumber(const o2::InteractionRecord& tfStartIR, int runNumber)
-{
-  auto& mgr = o2::ccdb::BasicCCDBManager::instance();
-  o2::ccdb::CcdbApi ccdb_api;
-  const std::string rct_path = "RCT/Info/RunInformation/";
-  const std::string start_orbit_path = "Trigger/StartOrbit";
-
-  mgr.setURL(o2::base::NameConf::getCCDBServer());
-  ccdb_api.init(o2::base::NameConf::getCCDBServer());
-
-  std::map<int, int>* mapStartOrbit = mgr.get<std::map<int, int>>(start_orbit_path);
-  int64_t ts = 0;
-  std::map<std::string, std::string> metadata;
-  std::map<std::string, std::string> headers;
-  const std::string run_path = Form("%s/%i", rct_path.data(), runNumber);
-  headers = ccdb_api.retrieveHeaders(run_path, metadata, -1);
-  ts = atol(headers["SOR"].c_str());
-
-  // ccdb returns timestamp in mus
-  // mus to ms
-  ts = ts / 1000;
-
-  uint32_t initialOrbit = mapStartOrbit->at(runNumber);
-  uint16_t firstRecBC = tfStartIR.bc;
-  uint32_t firstRecOrbit = tfStartIR.orbit;
-  const o2::InteractionRecord firstRec(firstRecBC, firstRecOrbit);
-  ts += firstRec.bc2ns() / 1000000;
-
-  return ts;
-};
-
 template <typename TracksCursorType, typename TracksCovCursorType>
 void AODProducerWorkflowDPL::addToTracksTable(TracksCursorType& tracksCursor, TracksCovCursorType& tracksCovCursor,
                                               const o2::track::TrackParCov& track, int collisionID, aod::track::TrackTypeEnum type)
@@ -472,7 +441,14 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
                          << " timeErr=" << extraInfoHolder.trackTimeRes << " BCSlice: " << extraInfoHolder.bcSlice[0] << ":" << extraInfoHolder.bcSlice[1];
             continue;
           }
-          addToTracksTable(tracksCursor, tracksCovCursor, data.getTrackParam(trackIndex), collisionID);
+          o2::track::TrackParametrizationWithError<float> trackPar = data.getTrackParam(trackIndex);
+          bool isProp = false;
+          if (mPropTracks) {
+            if (mGIDUsedBySVtx.find(trackIndex) == mGIDUsedBySVtx.end())
+              isProp = propagateTrackToPV(trackPar, data, collisionID);
+          }
+          auto trkType = isProp ? aod::track::Track : aod::track::TrackIU;
+          addToTracksTable(tracksCursor, tracksCovCursor, trackPar, collisionID, trkType);
           addToTracksExtraTable(tracksExtraCursor, extraInfoHolder);
           // collecting table indices of barrel tracks for V0s table
           if (extraInfoHolder.bcSlice[0] >= 0 && collisionID < 0) {
@@ -1586,6 +1562,8 @@ 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");
+  mPropTracks = ic.options().get<bool>("propagate-tracks");
+  mPropTracksCov = ic.options().get<bool>("propagate-tracks-cov");
 #ifdef WITH_OPENMP
   LOGP(info, "Multi-threaded parts will run with {} OpenMP threads", mNThreads);
 #else
@@ -1648,6 +1626,14 @@ void AODProducerWorkflowDPL::init(InitContext& ic)
     mZDCTDCMap[ic] = -std::numeric_limits<float>::infinity();
   }
 
+  if (mPropTracks) {
+    mCCDBAPI.init(mCCDBUrl);
+    auto& ccdbMgr = o2::ccdb::BasicCCDBManager::instance();
+    ccdbMgr.setURL(mCCDBUrl);
+    ccdbMgr.setCaching(true);
+    ccdbMgr.setLocalObjectValidityChecking();
+  }
+
   mTimer.Reset();
 }
 
@@ -1752,6 +1738,10 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
     tfNumber = mTFNumber;
   }
 
+  if (mPropTracks) {
+    fetchMeanVtxCCDB(runNumber);
+  }
+
   std::vector<float> aAmplitudes;
   std::vector<uint8_t> aChannels;
   fv0aCursor.reserve(fv0RecPoints.size());
@@ -1992,6 +1982,25 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   mGIDToTableFwdID.clear(); // reset the tables to be used by 'fillTrackTablesPerCollision'
   mGIDToTableMFTID.clear();
 
+  if (mPropTracks) {
+    auto v0s = recoData.getV0s();
+    auto cascades = recoData.getCascades();
+    auto decays3Body = recoData.getDecays3Body();
+    mGIDUsedBySVtx.reserve(v0s.size() * 2 + cascades.size() + decays3Body.size() * 3);
+    for (const auto& v0 : v0s) {
+      mGIDUsedBySVtx.insert(v0.getProngID(0));
+      mGIDUsedBySVtx.insert(v0.getProngID(1));
+    }
+    for (const auto& cascade : cascades) {
+      mGIDUsedBySVtx.insert(cascade.getBachelorID());
+    }
+    for (const auto& id3Body : decays3Body) {
+      mGIDUsedBySVtx.insert(id3Body.getProngID(0));
+      mGIDUsedBySVtx.insert(id3Body.getProngID(1));
+      mGIDUsedBySVtx.insert(id3Body.getProngID(2));
+    }
+  }
+
   // filling unassigned tracks first
   // so that all unassigned tracks are stored in the beginning of the table together
   auto& trackRef = primVer2TRefs.back(); // references to unassigned tracks are at the end
@@ -2154,6 +2163,8 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
 
   mBCLookup.clear();
 
+  mGIDUsedBySVtx.clear();
+
   originCursor(tfNumber);
 
   // sending metadata to writer
@@ -2344,6 +2355,35 @@ AODProducerWorkflowDPL::TrackExtraInfo AODProducerWorkflowDPL::processBarrelTrac
   return extraInfoHolder;
 }
 
+bool AODProducerWorkflowDPL::propagateTrackToPV(o2::track::TrackParametrizationWithError<float>& trackPar,
+                                                const o2::globaltracking::RecoContainer& data,
+                                                int colID)
+{
+  if (trackPar.getX() > mMinPropR)
+    return false;
+  o2::dataformats::DCA dcaInfo;
+  dcaInfo.set(999.f, 999.f, 999.f, 999.f, 999.f);
+  o2::dataformats::VertexBase v;
+  if (colID >= 0) {
+    const auto& vtx = data.getPrimaryVertex(colID);
+    v.setPos({vtx.getXYZ()});
+    if (mPropTracksCov) {
+      v.setCov({vtx.getCov()});
+    }
+  } else {
+    v.setPos({mVtx->getXYZ()});
+    if (mPropTracksCov) {
+      v.setCov(mVtx->getSigmaX() * mVtx->getSigmaX(),
+               0.f,
+               mVtx->getSigmaY() * mVtx->getSigmaY(),
+               0.f,
+               0.f,
+               mVtx->getSigmaZ() * mVtx->getSigmaZ());
+    }
+  }
+  return o2::base::Propagator::Instance()->propagateToDCABxByBz(v, trackPar, 2.f, mMatCorr, &dcaInfo);
+};
+
 void AODProducerWorkflowDPL::extrapolateToCalorimeters(TrackExtraInfo& extraInfoHolder, const o2::track::TrackPar& track)
 {
   constexpr float XEMCAL = 440.f, XPHOS = 460.f, XEMCAL2 = XEMCAL * XEMCAL;
@@ -2734,7 +2774,9 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, boo
       ConfigParamSpec{"nthreads", VariantType::Int, std::max(1, int(std::thread::hardware_concurrency() / 2)), {"Number of threads"}},
       ConfigParamSpec{"reco-mctracks-only", VariantType::Int, 0, {"Store only reconstructed MC tracks and their mothers/daughters. 0 -- off, != 0 -- on"}},
       ConfigParamSpec{"ctpreadout-create", VariantType::Int, 0, {"Create CTP digits from detector readout and CTP inputs. !=1 -- off, 1 -- on"}},
-      ConfigParamSpec{"emc-select-leading", VariantType::Bool, false, {"Flag to select if only the leading contributing particle for an EMCal cell should be stored"}}}};
+      ConfigParamSpec{"emc-select-leading", VariantType::Bool, false, {"Flag to select if only the leading contributing particle for an EMCal cell should be stored"}},
+      ConfigParamSpec{"propagate-tracks", VariantType::Bool, false, {"Propagate tracks (not used for secondary vertices) to IP"}},
+      ConfigParamSpec{"propagate-tracks-cov", VariantType::Bool, false, {"Propagate track (not used for secondary vertices) to IP using covariance matrix"}}}};
 }
 
 } // namespace o2::aodproducer
