Create TOF-constrained TPC tracks, optionally refit

+ clean-up of TOF-related workflow from unnecessaty initializations

Author: shahor02

DataFormats/Detectors/TPC/include/DataFormatsTPC/WorkflowHelper.h

@@ -62,6 +62,21 @@ struct getWorkflowTPCInput_ret {
 };
 } // namespace internal
 
+static auto getWorkflowTPCInputEmpryPtr()
+{
+  /*
+  In case the need of TPCInput is conditional, this will allow to define an empty unique_ptr to this complex structure outside 
+  of the "if" scope, i.e. 
+  auto inp = getWorkflowTPCInputEmpryPtr();
+  if (want_TPCInput) {
+    inp = getWorkflowTPCInput(...);
+    consumer->setInp(inp);
+  }
+*/
+  std::unique_ptr<internal::getWorkflowTPCInput_ret> ptr;
+  return ptr;
+}
+
 static auto getWorkflowTPCInput(o2::framework::ProcessingContext& pc, int verbosity = 0, bool do_mcLabels = false, bool do_clusters = true, unsigned long tpcSectorMask = 0xFFFFFFFFF, bool do_digits = false)
 {
   auto retVal = std::make_unique<internal::getWorkflowTPCInput_ret>();

Detectors/GlobalTracking/include/GlobalTracking/MatchTOF.h

@@ -23,6 +23,7 @@
 #include <TStopwatch.h>
 #include "ReconstructionDataFormats/Track.h"
 #include "ReconstructionDataFormats/TrackTPCITS.h"
+#include "ReconstructionDataFormats/TrackTPCTOF.h"
 #include "ReconstructionDataFormats/MatchInfoTOF.h"
 #include "ReconstructionDataFormats/GlobalTrackID.h"
 #include "DataFormatsTOF/CalibInfoTOF.h"
@@ -34,7 +35,7 @@
 #include "GlobalTracking/MatchTPCITS.h"
 #include "DataFormatsTPC/TrackTPC.h"
 #include "ReconstructionDataFormats/PID.h"
-#include <gsl/span>
+#include "TPCFastTransform.h"
 
 // from FIT
 #include "DataFormatsFT0/RecPoints.h"
@@ -199,8 +200,41 @@ class MatchTOF
     mFITRecPoints = recpoints;
   }
 
+  ///< set input TPC tracks cluster indices
+  void setTPCTrackClusIdxInp(const gsl::span<const o2::tpc::TPCClRefElem> inp)
+  {
+    mTPCTrackClusIdx = inp;
+  }
+
+  ///< set input TPC cluster sharing map
+  void setTPCClustersSharingMap(const gsl::span<const unsigned char> inp)
+  {
+    mTPCRefitterShMap = inp;
+  }
+
+  ///< set input TPC clusters
+  void setTPCClustersInp(const o2::tpc::ClusterNativeAccess* inp)
+  {
+    mTPCClusterIdxStruct = inp;
+  }
+
   int findFITIndex(int bc);
 
+  ///< populate externally provided container by TOF-time-constrained TPC tracks
+  template <typename V>
+  void makeConstrainedTPCTracks(V& container)
+  {
+    checkRefitter();
+    int nmatched = mMatchedTracks.size(), nconstrained = 0;
+    container.resize(nmatched);
+    for (unsigned i = 0; i < nmatched; i++) {
+      if (makeConstrainedTPCTrack(i, container[nconstrained])) {
+        nconstrained++;
+      }
+    }
+    container.resize(nconstrained);
+  }
+
  private:
   void attachInputTrees();
   void attachInputTreesTPConly();
@@ -217,6 +251,10 @@ class MatchTOF
   bool propagateToRefX(o2::track::TrackParCov& trc, float xRef /*in cm*/, float stepInCm /*in cm*/, o2::track::TrackLTIntegral& intLT);
   bool propagateToRefXWithoutCov(o2::track::TrackParCov& trc, float xRef /*in cm*/, float stepInCm /*in cm*/, float bz);
 
+  void updateTimeDependentParams();
+  void checkRefitter();
+  bool makeConstrainedTPCTrack(int matchedID, o2::dataformats::TrackTPCTOF& trConstr);
+
   //================================================================
 
   // Data members
@@ -232,8 +270,12 @@ class MatchTOF
   bool mMCTruthON = false; ///< flag availability of MC truth
 
   ///========== Parameters to be set externally, e.g. from CCDB ====================
+  float mBz = 0; ///< nominal Bz
 
   // to be done later
+  float mTPCTBinMUS = 0.;    ///< TPC time bin duration in microseconds
+  float mTPCTBinMUSInv = 0.; ///< inverse TPC time bin duration in microseconds
+  float mTPCBin2Z = 0.;      ///< conversion coeff from TPC time-bin to Z
 
   float mTimeTolerance = 1e3; ///<tolerance in ns for track-TOF time bracket matching
   float mSpaceTolerance = 10; ///<tolerance in cm for track-TOF time bracket matching
@@ -258,6 +300,13 @@ class MatchTOF
   gsl::span<const Cluster> mTOFClustersArrayInp;                  ///< input TOF clusters
   std::vector<Cluster>* mTOFClustersArrayInpVect;                 ///< input TOF clusters (vector to read from tree)
 
+  /// data needed for refit of time-constrained TPC tracks
+  gsl::span<const o2::tpc::TPCClRefElem> mTPCTrackClusIdx;            ///< input TPC track cluster indices span
+  gsl::span<const unsigned char> mTPCRefitterShMap;                   ///< externally set TPC clusters sharing map
+  const o2::tpc::ClusterNativeAccess* mTPCClusterIdxStruct = nullptr; ///< struct holding the TPC cluster indices
+  std::unique_ptr<o2::gpu::TPCFastTransform> mTPCTransform;           ///< TPC cluster transformation
+  std::unique_ptr<o2::gpu::GPUO2InterfaceRefit> mTPCRefitter;         ///< TPC refitter used for TPC tracks refit during the reconstruction
+
   o2::dataformats::MCTruthContainer<o2::MCCompLabel> mTOFClusLabels;     ///< input TOF clusters MC labels
   o2::dataformats::MCTruthContainer<o2::MCCompLabel>* mTOFClusLabelsPtr; ///< input TOF clusters MC labels (pointer to read from tree)
   std::vector<o2::MCCompLabel> mTracksLblWork;                           ///<TPCITS track labels

Detectors/GlobalTracking/src/MatchTOF.cxx

@@ -38,6 +38,7 @@
 
 #include "TPCBase/ParameterGas.h"
 #include "TPCBase/ParameterElectronics.h"
+#include "TPCReconstruction/TPCFastTransformHelperO2.h"
 
 using namespace o2::globaltracking;
 using evGIdx = o2::dataformats::EvIndex<int, o2::dataformats::GlobalTrackID>;
@@ -49,6 +50,7 @@ ClassImp(MatchTOF);
 void MatchTOF::run()
 {
   ///< running the matching
+  updateTimeDependentParams();
 
   if (!mWFInputAttached && !mSAInitDone) {
     LOG(ERROR) << "run called with mSAInitDone=" << mSAInitDone << " and mWFInputAttached=" << mWFInputAttached;
@@ -449,10 +451,7 @@ bool MatchTOF::prepareTracks()
     mTracksSectIndexCache[sec].reserve(100 + 1.2 * mNumOfTracks / o2::constants::math::NSectors);
   }
 
-  // getting Bz (mag field)
-  auto o2field = static_cast<o2::field::MagneticField*>(TGeoGlobalMagField::Instance()->GetField());
-  float bzField = o2field->solenoidField(); // magnetic field in kGauss
-  float maxInvPt = abs(bzField) > 0.1 ? 1. / (abs(bzField) * 0.05) : 999.;
+  float maxInvPt = abs(mBz) > 0.1 ? 1. / (abs(mBz) * 0.05) : 999.;
 
   LOG(DEBUG) << "\n\nWe have %d tracks to try to match to TOF: " << mNumOfTracks;
   int nNotPropagatedToTOF = 0;
@@ -479,7 +478,7 @@ bool MatchTOF::prepareTracks()
     LOG(DEBUG) << "Global coordinates Before propagating to 371 cm: globalPos[0] = " << globalPos[0] << ", globalPos[1] = " << globalPos[1] << ", globalPos[2] = " << globalPos[2];
     LOG(DEBUG) << "Radius xy Before propagating to 371 cm = " << TMath::Sqrt(globalPos[0] * globalPos[0] + globalPos[1] * globalPos[1]);
     LOG(DEBUG) << "Radius xyz Before propagating to 371 cm = " << TMath::Sqrt(globalPos[0] * globalPos[0] + globalPos[1] * globalPos[1] + globalPos[2] * globalPos[2]);
-    if (!propagateToRefXWithoutCov(trc, mXRef, 2, bzField)) { // we first propagate to 371 cm without considering the covariance matrix
+    if (!propagateToRefXWithoutCov(trc, mXRef, 2, mBz)) { // we first propagate to 371 cm without considering the covariance matrix
       nNotPropagatedToTOF++;
       continue;
     }
@@ -567,10 +566,7 @@ bool MatchTOF::prepareTPCTracks()
     mTPCTracksSectIndexCache[sec].reserve(100 + 1.2 * mNumOfTracks / o2::constants::math::NSectors);
   }
 
-  // getting Bz (mag field)
-  auto o2field = static_cast<o2::field::MagneticField*>(TGeoGlobalMagField::Instance()->GetField());
-  float bzField = o2field->solenoidField(); // magnetic field in kGauss
-  float maxInvPt = abs(bzField) > 0.1 ? 1. / (abs(bzField) * 0.05) : 999.;
+  float maxInvPt = abs(mBz) > 0.1 ? 1. / (abs(mBz) * 0.05) : 999.;
   int nclustersMin = 0;
   LOG(INFO) << "Max track Inv pT allowed = " << maxInvPt;
   LOG(INFO) << "Min track Nclusters allowed = " << nclustersMin;
@@ -584,10 +580,10 @@ bool MatchTOF::prepareTPCTracks()
     // create working copy of track param
     timeEst timeInfo;
     // set
-    timeInfo.setTimeStamp(trcOrig.getTime0() * o2::tpc::ParameterElectronics::Instance().ZbinWidth);
-    timeInfo.setTimeStampError((trcOrig.getDeltaTBwd() + 5) * o2::tpc::ParameterElectronics::Instance().ZbinWidth);
+    timeInfo.setTimeStamp(trcOrig.getTime0() * mTPCTBinMUS);
+    timeInfo.setTimeStampError((trcOrig.getDeltaTBwd() + 5) * mTPCTBinMUS);
     mSideTPC.push_back(trcOrig.hasASideClustersOnly() ? 1 : (trcOrig.hasCSideClustersOnly() ? -1 : 0));
-    mExtraTPCFwdTime.push_back((trcOrig.getDeltaTFwd() + 5) * o2::tpc::ParameterElectronics::Instance().ZbinWidth);
+    mExtraTPCFwdTime.push_back((trcOrig.getDeltaTFwd() + 5) * mTPCTBinMUS);
 
     o2::track::TrackLTIntegral intLT0; //mTPCTracksWork.back().getLTIntegralOut(); // we get the integrated length from TPC-ITC outward propagation
     // make a copy of the TPC track that we have to propagate
@@ -617,7 +613,7 @@ bool MatchTOF::prepareTPCTracks()
     // the "very rough" propagation worked; now we can propagate considering also the cov matrix
 #endif
 
-    if (!propagateToRefXWithoutCov(trc, mXRef, 10, bzField)) { // we first propagate to 371 cm without considering the covariance matrix
+    if (!propagateToRefXWithoutCov(trc, mXRef, 10, mBz)) { // we first propagate to 371 cm without considering the covariance matrix
       nNotPropagatedToTOF++;
       continue;
     }
@@ -1137,7 +1133,6 @@ void MatchTOF::doMatching(int sec)
 //______________________________________________
 void MatchTOF::doMatchingForTPC(int sec)
 {
-  printf("here, DoMatch\n");
   auto& gasParam = o2::tpc::ParameterGas::Instance();
   float vdrift = gasParam.DriftV;
 
@@ -1542,7 +1537,7 @@ void MatchTOF::selectBestMatches()
 bool MatchTOF::propagateToRefX(o2::track::TrackParCov& trc, float xRef, float stepInCm, o2::track::TrackLTIntegral& intLT)
 {
   // propagate track to matching reference X
-  o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrTGeo; // material correction method
+  o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT; // material correction method
   const float tanHalfSector = tan(o2::constants::math::SectorSpanRad / 2);
   bool refReached = false;
   float xStart = trc.getX();
@@ -1668,3 +1663,75 @@ void MatchTOF::fillTOFmatchTreeWithLabels(const char* trname, int cacheTOF, int
   }
   mTimerDBG.Stop();
 }
+
+//______________________________________________
+void MatchTOF::updateTimeDependentParams()
+{
+  ///< update parameters depending on time (once per TF)
+  auto& elParam = o2::tpc::ParameterElectronics::Instance();
+  auto& gasParam = o2::tpc::ParameterGas::Instance();
+  mTPCTBinMUS = elParam.ZbinWidth; // TPC bin in microseconds
+  mTPCTBinMUSInv = 1. / mTPCTBinMUS;
+  mTPCBin2Z = mTPCTBinMUS * gasParam.DriftV;
+
+  mBz = o2::base::Propagator::Instance()->getNominalBz();
+}
+
+//_________________________________________________________
+bool MatchTOF::makeConstrainedTPCTrack(int matchedID, o2::dataformats::TrackTPCTOF& trConstr)
+{
+  auto& match = mMatchedTracks[matchedID];
+  const auto& tpcTrOrig = mTPCTracksArrayInp[match.getTrackIndex()];
+  const auto& tofCl = mTOFClustersArrayInp[match.getTOFClIndex()];
+  const auto& intLT = match.getLTIntegralOut();
+  // correct the time of the track
+  auto timeTOFMUS = (tofCl.getTime() - intLT.getTOF(o2::track::PID::Pion)) * 1e-6; // tof time in \mus, FIXME: account for time of flight to R TOF
+  auto timeTOFTB = timeTOFMUS * mTPCTBinMUSInv;                                    // TOF time in TPC timebins
+  auto deltaTBins = timeTOFTB - tpcTrOrig.getTime0();                              // time shift in timeBins
+  float timeErr = 0.010;                                                           // assume 10 ns error FIXME
+  auto dzCorr = deltaTBins * mTPCBin2Z;
+
+  if (mTPCClusterIdxStruct) {                                              // refit was requested
+    trConstr.o2::track::TrackParCov::operator=(tpcTrOrig.getOuterParam()); // seed for inward refit of constrained track, made from the outer param
+  } else {
+    trConstr.o2::track::TrackParCov::operator=(tpcTrOrig); // inner param, we just correct is position, w/o refit
+  }
+
+  auto zTrack = trConstr.getZ();
+
+  if (tpcTrOrig.hasASideClustersOnly()) {
+    zTrack += dzCorr;
+  } else if (tpcTrOrig.hasCSideClustersOnly()) {
+    zTrack -= dzCorr;
+  } else {
+    // TODO : special treatment of tracks crossing the CE
+  }
+  trConstr.setZ(zTrack);
+  trConstr.setTimeMUS(timeTOFMUS, timeErr);
+  trConstr.setRefMatch(matchedID);
+  if (mTPCClusterIdxStruct) { // refit was requested
+    float chi2 = 0;
+    mTPCRefitter->setTrackReferenceX(o2::globaltracking::MatchTPCITS::XTPCInnerRef);
+    if (mTPCRefitter->RefitTrackAsTrackParCov(trConstr, tpcTrOrig.getClusterRef(), timeTOFTB, &chi2, false, true) < 0) { // outward refit after resetting cov.mat.
+      LOG(DEBUG) << "Refit failed";
+      return false;
+    }
+    trConstr.setChi2Refit(chi2);
+  }
+
+  return true;
+}
+
+//_________________________________________________________
+void MatchTOF::checkRefitter()
+{
+  if (mTPCClusterIdxStruct) {
+    if (!mTPCTransform) { // eventually, should be updated at every TF?
+      mTPCTransform = o2::tpc::TPCFastTransformHelperO2::instance()->create(0);
+    }
+
+    mTPCRefitter = std::make_unique<o2::gpu::GPUO2InterfaceRefit>(mTPCClusterIdxStruct, mTPCTransform.get(), mBz,
+                                                                  mTPCTrackClusIdx.data(), mTPCRefitterShMap.data(),
+                                                                  nullptr, o2::base::Propagator::Instance());
+  }
+}

Detectors/GlobalTrackingWorkflow/tofworkflow/include/TOFWorkflow/RecoWorkflowWithTPCSpec.h

@@ -19,7 +19,7 @@ namespace o2
 namespace tof
 {
 
-o2::framework::DataProcessorSpec getTOFRecoWorkflowWithTPCSpec(bool useMC, bool useFIT);
+o2::framework::DataProcessorSpec getTOFRecoWorkflowWithTPCSpec(bool useMC, bool useFIT, bool doTPCRefit);
 
 } // end namespace tof
 } // end namespace o2

Detectors/GlobalTrackingWorkflow/tofworkflow/include/TOFWorkflow/TOFMatchedReaderSpec.h

@@ -19,6 +19,7 @@
 #include "Framework/DataProcessorSpec.h"
 #include "Framework/Task.h"
 #include "ReconstructionDataFormats/MatchInfoTOF.h"
+#include "ReconstructionDataFormats/TrackTPCTOF.h"
 #include "SimulationDataFormat/MCCompLabel.h"
 
 namespace o2
@@ -29,7 +30,7 @@ namespace tof
 class TOFMatchedReader : public o2::framework::Task
 {
  public:
-  TOFMatchedReader(bool useMC) : mUseMC(useMC) {}
+  TOFMatchedReader(bool useMC, bool tpcmatch, bool readTracks) : mUseMC(useMC), mTPCMatch(tpcmatch), mReadTracks(readTracks) {}
   ~TOFMatchedReader() override = default;
   void init(o2::framework::InitContext& ic) final;
   void run(o2::framework::ProcessingContext& pc) final;
@@ -38,19 +39,23 @@ class TOFMatchedReader : public o2::framework::Task
   void connectTree(const std::string& filename);
 
   bool mUseMC = false;
+  bool mTPCMatch = false;
+  bool mReadTracks = false;
+
   std::string mInFileName{"o2match_tof.root"};
   std::string mInTreeName{"matchTOF"};
   std::unique_ptr<TFile> mFile = nullptr;
   std::unique_ptr<TTree> mTree = nullptr;
   std::vector<o2::dataformats::MatchInfoTOF> mMatches, *mMatchesPtr = &mMatches;
+  std::vector<o2::dataformats::TrackTPCTOF> mTracks, *mTracksPtr = &mTracks;
   std::vector<o2::MCCompLabel> mLabelTOF, *mLabelTOFPtr = &mLabelTOF;
   std::vector<o2::MCCompLabel> mLabelTPC, *mLabelTPCPtr = &mLabelTPC;
   std::vector<o2::MCCompLabel> mLabelITS, *mLabelITSPtr = &mLabelITS;
 };
 
 /// create a processor spec
 /// read matched TOF clusters from a ROOT file
-framework::DataProcessorSpec getTOFMatchedReaderSpec(bool useMC);
+framework::DataProcessorSpec getTOFMatchedReaderSpec(bool useMC, bool tpcmatch = false, bool readTracks = false);
 
 } // namespace tof
 } // namespace o2

Detectors/GlobalTrackingWorkflow/tofworkflow/include/TOFWorkflow/TOFMatchedWriterSpec.h

@@ -24,7 +24,7 @@ namespace tof
 
 /// create a processor spec
 /// write TOF matching info in a root file
-o2::framework::DataProcessorSpec getTOFMatchedWriterSpec(bool useMC, const char* outdef = "o2match_tof.root");
+o2::framework::DataProcessorSpec getTOFMatchedWriterSpec(bool useMC, const char* outdef = "o2match_tof.root", bool writeTracks = false);
 
 } // namespace tof
 } // namespace o2

Detectors/GlobalTrackingWorkflow/tofworkflow/src/RecoWorkflowSpec.cxx

@@ -9,6 +9,7 @@
 // or submit itself to any jurisdiction.
 
 #include "TOFWorkflow/RecoWorkflowSpec.h"
+#include "Framework/ConfigParamRegistry.h"
 #include "Framework/ControlService.h"
 #include "Framework/DataProcessorSpec.h"
 #include "Framework/DataRefUtils.h"
@@ -21,6 +22,7 @@
 #include "ReconstructionDataFormats/TrackTPCITS.h"
 #include "DetectorsBase/GeometryManager.h"
 #include "DetectorsBase/Propagator.h"
+#include "DetectorsCommonDataFormats/NameConf.h"
 #include <gsl/span>
 #include "TStopwatch.h"
 
@@ -55,6 +57,16 @@ class TOFDPLRecoWorkflowTask
     // nothing special to be set up
     o2::base::GeometryManager::loadGeometry();
     o2::base::Propagator::initFieldFromGRP("o2sim_grp.root");
+    std::string matLUTPath = ic.options().get<std::string>("material-lut-path");
+    std::string matLUTFile = o2::base::NameConf::getMatLUTFileName(matLUTPath);
+    if (o2::base::NameConf::pathExists(matLUTFile)) {
+      auto* lut = o2::base::MatLayerCylSet::loadFromFile(matLUTFile);
+      o2::base::Propagator::Instance()->setMatLUT(lut);
+      LOG(INFO) << "Loaded material LUT from " << matLUTFile;
+    } else {
+      LOG(INFO) << "Material LUT " << matLUTFile << " file is absent, only TGeo can be used";
+    }
+
     mTimer.Stop();
     mTimer.Reset();
   }
@@ -159,7 +171,8 @@ o2::framework::DataProcessorSpec getTOFRecoWorkflowSpec(bool useMC, bool useFIT)
     inputs,
     outputs,
     AlgorithmSpec{adaptFromTask<TOFDPLRecoWorkflowTask>(useMC, useFIT)},
-    Options{/* for the moment no options */}};
+    Options{
+      {"material-lut-path", VariantType::String, "", {"Path of the material LUT file"}}}};
 }
 
 } // end namespace tof