Add PV-related rejections to cosmic matcher for interleaved mode

Author: shahor02

DataFormats/Reconstruction/include/ReconstructionDataFormats/DCA.h

@@ -54,13 +54,16 @@ class DCA
     mZ = z;
   }
 
+  GPUd() void setY(float y) { mY = y; }
+  GPUd() void setZ(float z) { mZ = z; }
   GPUd() auto getY() const { return mY; }
   GPUd() auto getZ() const { return mZ; }
   GPUd() auto getR2() const { return mY * mY + mZ * mZ; }
   GPUd() auto getSigmaY2() const { return mCov[0]; }
   GPUd() auto getSigmaYZ() const { return mCov[1]; }
   GPUd() auto getSigmaZ2() const { return mCov[2]; }
   GPUd() const auto& getCovariance() const { return mCov; }
+  GPUd() float calcChi2() const;
 
   void print() const;
 

DataFormats/Reconstruction/src/DCA.cxx

@@ -9,6 +9,8 @@
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 
+#include "GPUCommonMath.h"
+#include "CommonConstants/MathConstants.h"
 #include "ReconstructionDataFormats/DCA.h"
 #include <iostream>
 
@@ -20,6 +22,16 @@ namespace o2
 namespace dataformats
 {
 
+float DCA::calcChi2() const
+{
+  // Estimate the chi2 for DCA
+  const auto sdd = mCov[0], sdz = mCov[1], szz = mCov[2], det = sdd * szz - sdz * sdz;
+  if (o2::gpu::CAMath::Abs(det) < o2::constants::math::Almost0) {
+    return constants::math::VeryBig;
+  }
+  return (mY * (szz * mY - sdz * mZ) + mZ * (sdd * mZ - mY * sdz)) / det;
+}
+
 std::ostream& operator<<(std::ostream& os, const o2::dataformats::DCA& d)
 {
   // stream itself

Detectors/GlobalTracking/include/GlobalTracking/MatchCosmics.h

@@ -83,13 +83,16 @@ class MatchCosmics
     TBracket tBracket;      ///< bracketing time-bins
     GTrackID origID;        ///< track origin id
     int matchID = MinusOne; ///< entry (none if MinusOne) of its match in the vector of matches
+    short vtIDMin = -1;     ///< id of the 1st compatible vertex
+    short vtIDMax = -1;     ///< id of the last compatible vertex
   };
   void setTPCCorrMaps(const o2::gpu::TPCFastTransformPOD* maph);
   void setTPCVDrift(const o2::tpc::VDriftCorrFact& v);
   void setITSROFrameLengthMUS(float fums) { mITSROFrameLengthMUS = fums; }
   void setITSDict(const o2::itsmft::TopologyDictionary* dict) { mITSDict = dict; }
   void process(const o2::globaltracking::RecoContainer& data);
   void setUseMC(bool mc) { mUseMC = mc; }
+  void setUsePVInfo(bool v) { mUsePVInfo = v; }
   void init();
   void end();
 
@@ -146,6 +149,7 @@ class MatchCosmics
   float mTPCTBinMUS = 0.;         ///< TPC time bin duration in microseconds
   float mBz = 0;                  ///< nominal Bz
   bool mFieldON = true;
+  bool mUsePVInfo = false;
   bool mUseMC = true;
   float mITSROFrameLengthMUS = 0.;
   float mQ2PtCutoff = 1e9;

Detectors/GlobalTracking/include/GlobalTracking/MatchCosmicsParams.h

@@ -17,22 +17,28 @@
 #include "CommonUtils/ConfigurableParam.h"
 #include "CommonUtils/ConfigurableParamHelper.h"
 #include "DetectorsBase/Propagator.h"
+#include "ReconstructionDataFormats/GlobalTrackID.h"
 
 namespace o2
 {
 namespace globaltracking
 {
 
 struct MatchCosmicsParams : public o2::conf::ConfigurableParamHelper<MatchCosmicsParams> {
+  float dcaCutChi2[o2::dataformats::GlobalTrackID::NSources] = {};           // optional (>0) chi2 cut on track DCA to any of its compatible vertices
   float systSigma2[o2::track::kNParams] = {0.01f, 0.01f, 1e-4f, 1e-4f, 0.f}; // extra error to be added at legs comparison
   float crudeNSigma2Cut[o2::track::kNParams] = {49.f, 49.f, 49.f, 49.f, 49.f};
-  float crudeChi2Cut = 999.;
-  float timeToleranceMUS = 0.;
-  float maxStep = 10.;
-  float maxSnp = 0.99;
-  float minSeedPt = 0.10;  // use only tracks above this pT (scaled with field)
-  float nSigmaTError = 4.; // number of sigmas on track time error for matching (except for TPC which provides an interval)
+  float crudeChi2Cut = 999.f;
+  float timeToleranceMUS = 0.f;
+  float maxStep = 10.f;
+  float maxSnp = 0.99f;
+  float minSeedPt = 0.10f;     // use only tracks above this pT (scaled with field)
+  float nSigmaTError = 4.f;    // number of sigmas on track time error for matching (except for TPC which provides an interval)
+  float tpcExtraZError2 = 1.f; // extra error^2 on the TPC-only track Z coordinate
+  float fiducialRIP = 1.0f;    // consider track having |Y@x=0|< this as passing DCA cut (if requested)
+  float fiducialZIP = 20.f;    // consider track having |Z@x=0|< this as passing DCA cut (if requested)
   bool allowTPCOnly = true;
+  bool discardPVContributors = true; // used only if the global option --use-pv-info is requested
   o2::base::Propagator::MatCorrType matCorr = o2::base::Propagator::MatCorrType::USEMatCorrLUT;
 
   O2ParamDef(MatchCosmicsParams, "cosmicsMatch");

Detectors/GlobalTracking/src/MatchCosmics.cxx

@@ -23,6 +23,9 @@
 #include "ReconstructionDataFormats/TrackTPCITS.h"
 #include "ReconstructionDataFormats/TrackTPCTOF.h"
 #include "ReconstructionDataFormats/MatchInfoTOF.h"
+#include "ReconstructionDataFormats/PrimaryVertex.h"
+#include "ReconstructionDataFormats/VtxTrackRef.h"
+#include "ReconstructionDataFormats/DCA.h"
 #include "ITStracking/IOUtils.h"
 #include "ITSBase/GeometryTGeo.h"
 #include "TPCBase/ParameterElectronics.h"
@@ -35,6 +38,7 @@
 #include "TPCFastTransformPOD.h"
 #include <algorithm>
 #include <numeric>
+#include <unordered_map>
 
 using namespace o2::globaltracking;
 
@@ -52,21 +56,53 @@ void MatchCosmics::process(const o2::globaltracking::RecoContainer& data)
 
   createSeeds(data);
   int ntr = mSeeds.size();
-
+  const auto prop = o2::base::Propagator::Instance();
   // propagate to DCA to origin
   const o2::math_utils::Point3D<float> v{0., 0., 0};
   for (int i = 0; i < ntr; i++) {
     auto& trc = mSeeds[i];
-    if (!o2::base::Propagator::Instance()->propagateToDCABxByBz(v, trc, mMatchParams->maxStep, mMatchParams->matCorr)) {
-      trc.matchID = Reject; // reject track
+    if (!trc.matchID == Reject) {
+      if (!prop->propagateToDCABxByBz(v, trc, mMatchParams->maxStep, mMatchParams->matCorr)) {
+        trc.matchID = Reject; // reject track
+        continue;
+      }
+      if (mMatchParams->dcaCutChi2[trc.origID.getSource()] > 0.f && mUsePVInfo && (std::abs(trc.getY()) < mMatchParams->fiducialRIP && std::abs(trc.getZ()) < mMatchParams->fiducialZIP)) {
+        // do the propagation only if we are in the fiducial IP range.
+        for (int iv = trc.vtIDMin; iv < trc.vtIDMax; iv++) {
+          const auto& pv = data.getPrimaryVertex(iv);
+          o2::track::TrackParCov trcatPV(trc);
+          o2::dataformats::DCA dca;
+          if (!trcatPV.propagateToDCA(pv, mBz, &dca)) {
+            trc.matchID = Reject;
+            break;
+          }
+          if (trc.origID.getSource() == GTrackID::TPC) { // correct the track Z position for the vertex time
+            const auto& trcTPC = data.getTPCTrack(trc.origID.getSource());
+            float deltaZ = trcTPC.hasBothSidesClusters() ? 0.f : (pv.getTimeStamp().getTimeStamp() - trcTPC.getTime0() * 8 * o2::constants::lhc::LHCBunchSpacingMUS) * mTPCVDrift;
+            dca.setZ(dca.getZ() + (trcTPC.hasASideClustersOnly() ? deltaZ : -deltaZ));
+          }
+          if (dca.calcChi2() < mMatchParams->dcaCutChi2[trc.origID.getSource()]) {
+            trc.matchID = Reject;
+            break;
+          }
+        }
+      }
     }
   }
 
-  // sort in time bracket lower edge
+  // sort in time bracket lower edge, putting rejected tracks in the end
   std::vector<int> sortID(ntr);
   std::iota(sortID.begin(), sortID.end(), 0);
-  std::sort(sortID.begin(), sortID.end(), [this](int a, int b) { return mSeeds[a].tBracket.getMin() < mSeeds[b].tBracket.getMin(); });
-
+  std::sort(sortID.begin(), sortID.end(), [this](int a, int b) { return (mSeeds[a].matchID != Reject) && (mSeeds[a].tBracket.getMin() < mSeeds[b].tBracket.getMin()); });
+  int lastValid = ntr - 1;
+  for (; lastValid >= 0; lastValid--) {
+    if (mSeeds[sortID[lastValid]].matchID != Reject) {
+      break;
+    }
+  }
+  ntr = lastValid >= 0 ? lastValid + 1 : 0;
+  LOGP(info, "Collected {} seeds, validated: {}", mSeeds.size(), ntr);
+  sortID.resize(ntr);
   for (int i = 0; i < ntr; i++) {
     for (int j = i + 1; j < ntr; j++) {
       if (checkPair(sortID[i], sortID[j]) == RejTime) {
@@ -424,6 +460,8 @@ MatchCosmics::RejFlag MatchCosmics::checkPair(int i, int j)
       break;
     }
     bool ignoreZ = seed0.origID.getSource() == o2d::GlobalTrackID::TPC || seed1.origID.getSource() == o2d::GlobalTrackID::TPC;
+    // RSTODO this is simplification: one should constraint the TPC-only track Z by the time of other candidate (at least their difference of both tracks are TPC only).
+    // If the shift is large, eventually the tracks need to be refitted.
     if (!ignoreZ) { // cut on Z makes no sense for TPC only tracks
       auto dZ = seed0.getZ() - seed1Inv.getZ();
       if (dZ * dZ > (seed0.getSigmaZ2() + seed1Inv.getSigmaZ2()) * mMatchParams->crudeNSigma2Cut[o2::track::kZ]) {
@@ -492,8 +530,9 @@ void MatchCosmics::createSeeds(const o2::globaltracking::RecoContainer& data)
   // Scan all inputs and create seeding tracks
 
   mSeeds.clear();
+  std::unordered_map<GTrackID, int> trackEntry;
 
-  auto creator = [this](auto& _tr, GTrackID _origID, float t0, float terr) {
+  auto creator = [this, &trackEntry](auto& _tr, GTrackID _origID, float t0, float terr) {
     if constexpr (std::is_base_of_v<o2::track::TrackParCov, std::decay_t<decltype(_tr)>>) {
       if (std::abs(_tr.getQ2Pt()) > this->mQ2PtCutoff) {
         return true;
@@ -512,7 +551,11 @@ void MatchCosmics::createSeeds(const o2::globaltracking::RecoContainer& data)
         terr *= this->mMatchParams->nSigmaTError;
       }
       terr += this->mMatchParams->timeToleranceMUS;
+      trackEntry[_origID] = mSeeds.size();
       mSeeds.emplace_back(TrackSeed{_tr, {t0 - terr, t0 + terr}, _origID, MinusOne});
+      if constexpr (isTPCTrack<decltype(_tr)>()) {
+        mSeeds.back().setCov(this->mMatchParams->tpcExtraZError2 + _tr.getCov()[o2::track::kSigZ2], o2::track::kSigZ2);
+      }
       return true;
     } else {
       return false;
@@ -521,7 +564,32 @@ void MatchCosmics::createSeeds(const o2::globaltracking::RecoContainer& data)
 
   data.createTracksVariadic(creator);
 
-  LOG(info) << "collected " << mSeeds.size() << " seeds";
+  if (mUsePVInfo) {                                         // if needed, veto with the primary vertex info
+    auto trackIndex = data.getPrimaryVertexMatchedTracks(); // Global ID's for associated tracks
+    auto vtxRefs = data.getPrimaryVertexMatchedTrackRefs(); // references from vertex to these track IDs
+    int nv = vtxRefs.size() - 1;                            // The last entry is for unassigned tracks, no need to check them
+    const auto propagator = o2::base::Propagator::Instance();
+    for (int iv = 0; iv < nv; iv++) {
+      const auto& vtref = vtxRefs[iv];
+      int it = vtref.getFirstEntry(), itLim = it + vtref.getEntries();
+      for (; it < itLim; it++) {
+        auto tvid = trackIndex[it];
+        auto entry = trackEntry.find(tvid);
+        if (entry == trackEntry.end()) {
+          continue;
+        }
+        auto& seed = mSeeds[entry->second];
+        if (seed.matchID == Reject || (mMatchParams->discardPVContributors && tvid.isPVContributor())) {
+          seed.matchID = Reject;
+          continue;
+        }
+        if (seed.vtIDMin < 0) {
+          seed.vtIDMin = iv;
+        }
+        seed.vtIDMax = std::max(short(iv), seed.vtIDMax);
+      }
+    }
+  }
 }
 
 //________________________________________________________

Detectors/GlobalTrackingWorkflow/include/GlobalTrackingWorkflow/CosmicsMatchingSpec.h

@@ -24,7 +24,7 @@ namespace globaltracking
 {
 
 /// create a processor spec
-framework::DataProcessorSpec getCosmicsMatchingSpec(o2::dataformats::GlobalTrackID::mask_t src, bool useMC);
+framework::DataProcessorSpec getCosmicsMatchingSpec(o2::dataformats::GlobalTrackID::mask_t src, bool usePV, bool useMC);
 
 } // namespace globaltracking
 } // namespace o2

Detectors/GlobalTrackingWorkflow/src/CosmicsMatchingSpec.cxx

@@ -62,7 +62,7 @@ namespace globaltracking
 class CosmicsMatchingSpec : public Task
 {
  public:
-  CosmicsMatchingSpec(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool useMC) : mDataRequest(dr), mGGCCDBRequest(gr), mUseMC(useMC) {}
+  CosmicsMatchingSpec(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool usePV, bool useMC) : mDataRequest(dr), mGGCCDBRequest(gr), mUsePVInfo(usePV), mUseMC(useMC) {}
   ~CosmicsMatchingSpec() override = default;
   void init(InitContext& ic) final;
   void run(ProcessingContext& pc) final;
@@ -77,6 +77,7 @@ class CosmicsMatchingSpec : public Task
   const o2::gpu::TPCFastTransformPOD* mCorrMap{nullptr};
   o2::globaltracking::MatchCosmics mMatching; // matching engine
   bool mUseMC = true;
+  bool mUsePVInfo = false;
   TStopwatch mTimer;
 };
 
@@ -87,6 +88,7 @@ void CosmicsMatchingSpec::init(InitContext& ic)
   o2::base::GRPGeomHelper::instance().setRequest(mGGCCDBRequest);
   mMatching.setDebugFlag(ic.options().get<int>("debug-tree-flags"));
   mMatching.setUseMC(mUseMC);
+  mMatching.setUsePVInfo(mUsePVInfo);
   //
 }
 
@@ -160,7 +162,7 @@ void CosmicsMatchingSpec::endOfStream(EndOfStreamContext& ec)
        mTimer.CpuTime(), mTimer.RealTime(), mTimer.Counter() - 1);
 }
 
-DataProcessorSpec getCosmicsMatchingSpec(GTrackID::mask_t src, bool useMC)
+DataProcessorSpec getCosmicsMatchingSpec(GTrackID::mask_t src, bool usePV, bool useMC)
 {
   std::vector<OutputSpec> outputs;
   Options opts{
@@ -171,6 +173,9 @@ DataProcessorSpec getCosmicsMatchingSpec(GTrackID::mask_t src, bool useMC)
 
   dataRequest->requestTracks(src, useMC);
   dataRequest->requestClusters(src, false); // no MC labels for clusters needed for refit only
+  if (usePV) {
+    dataRequest->requestPrimaryVertices(useMC);
+  }
 
   outputs.emplace_back("GLO", "COSMICTRC", 0, Lifetime::Timeframe);
   if (useMC) {
@@ -192,7 +197,7 @@ DataProcessorSpec getCosmicsMatchingSpec(GTrackID::mask_t src, bool useMC)
     "cosmics-matcher",
     dataRequest->inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<CosmicsMatchingSpec>(dataRequest, ggRequest, useMC)},
+    AlgorithmSpec{adaptFromTask<CosmicsMatchingSpec>(dataRequest, ggRequest, usePV, useMC)},
     opts};
 }
 

Detectors/GlobalTrackingWorkflow/src/cosmics-match-workflow.cxx

@@ -51,6 +51,7 @@ void customize(std::vector<o2::framework::ConfigParamSpec>& workflowOptions)
     {"disable-mc", o2::framework::VariantType::Bool, false, {"disable MC propagation even if available"}},
     {"disable-root-input", o2::framework::VariantType::Bool, false, {"disable root-files input reader"}},
     {"disable-root-output", o2::framework::VariantType::Bool, false, {"disable root-files output writer"}},
+    {"use-pv-info", o2::framework::VariantType::Bool, false, {"request primary vertex for relevant cuts in the collision/cosmics interleaved data"}},
     {"track-sources", VariantType::String, std::string{GID::ALL}, {"comma-separated list of sources to use"}},
     {"configKeyValues", VariantType::String, "", {"Semicolon separated key=value strings ..."}}};
   o2::itsmft::DPLAlpideParamInitializer::addITSConfigOption(options);
@@ -102,14 +103,19 @@ WorkflowSpec defineDataProcessing(ConfigContext const& configcontext)
   if (sclOpt.requestCTPLumi) {
     src = src | GID::getSourcesMask("CTP");
   }
+
   GID::mask_t srcCl = src;
   GID::mask_t dummy;
   if (!configcontext.options().get<bool>("disable-root-input")) {
     specs.emplace_back(o2::tpc::getTPCScalerSpec(sclOpt.lumiType == o2::tpc::LumiScaleType::TPCScaler, sclOpt.enableMShapeCorrection, sclOpt));
   }
-  specs.emplace_back(o2::globaltracking::getCosmicsMatchingSpec(src, useMC));
+  bool usePV = configcontext.options().get<bool>("use-pv-info");
+  specs.emplace_back(o2::globaltracking::getCosmicsMatchingSpec(src, usePV, useMC));
 
   o2::globaltracking::InputHelper::addInputSpecs(configcontext, specs, src, src, src, useMC, dummy); // clusters MC is not needed
+  if (usePV) {
+    o2::globaltracking::InputHelper::addInputSpecsPVertex(configcontext, specs, useMC); // P-vertex is always needed
+  }
 
   if (!disableRootOut) {
     specs.emplace_back(o2::globaltracking::getTrackCosmicsWriterSpec(useMC));