TRD: correction for pile-up

Detectors/TRD/qc/include/TRDQC/Tracking.h

@@ -23,6 +23,7 @@
 #include "DataFormatsTRD/Constants.h"
 #include "ReconstructionDataFormats/TrackTPCITS.h"
 #include "ReconstructionDataFormats/GlobalTrackID.h"
+#include "DataFormatsTRD/TrackTriggerRecord.h"
 #include "DataFormatsTPC/TrackTPC.h"
 #include "DetectorsBase/Propagator.h"
 #include "GPUTRDRecoParam.h"
@@ -103,6 +104,10 @@ class Tracking
     mLocalGain = localGain;
   }
 
+  // quantities necessary for pile-up correction
+  void setTriggeredBCFT0(std::vector<int> t) { mTriggeredBCFT0 = t; }
+  void setFirstOrbit(uint32_t o) { mFirstOrbit = o; }
+
  private:
   float mMaxSnp{o2::base::Propagator::MAX_SIN_PHI};  ///< max snp when propagating tracks
   float mMaxStep{o2::base::Propagator::MAX_STEP};    ///< maximum step for propagation
@@ -115,12 +120,20 @@ class Tracking
   std::vector<TrackQC> mTrackQC;
 
   // input from DPL
-  gsl::span<const o2::dataformats::TrackTPCITS> mTracksITSTPC; ///< ITS-TPC seeding tracks
-  gsl::span<const o2::tpc::TrackTPC> mTracksTPC;               ///< TPC seeding tracks
-  gsl::span<const TrackTRD> mTracksITSTPCTRD;                  ///< TRD tracks reconstructed from TPC or ITS-TPC seeds
-  gsl::span<const TrackTRD> mTracksTPCTRD;                     ///< TRD tracks reconstructed from TPC or TPC seeds
-  gsl::span<const Tracklet64> mTrackletsRaw;                   ///< array of raw tracklets needed for TRD refit
-  gsl::span<const CalibratedTracklet> mTrackletsCalib;         ///< array of calibrated tracklets needed for TRD refit
+  gsl::span<const o2::dataformats::TrackTPCITS> mTracksITSTPC;       ///< ITS-TPC seeding tracks
+  gsl::span<const o2::tpc::TrackTPC> mTracksTPC;                     ///< TPC seeding tracks
+  gsl::span<const TrackTRD> mTracksITSTPCTRD;                        ///< TRD tracks reconstructed from TPC or ITS-TPC seeds
+  gsl::span<const TrackTRD> mTracksTPCTRD;                           ///< TRD tracks reconstructed from TPC or TPC seeds
+  gsl::span<const TrackTriggerRecord> mTrackTriggerRecordsITSTPCTRD; ///< TRD tracks reconstructed from TPC or ITS-TPC seeds
+  gsl::span<const TrackTriggerRecord> mTrackTriggerRecordsTPCTRD;    ///< TRD tracks reconstructed from TPC or TPC seeds
+  gsl::span<const Tracklet64> mTrackletsRaw;                         ///< array of raw tracklets needed for TRD refit
+  gsl::span<const CalibratedTracklet> mTrackletsCalib;               ///< array of calibrated tracklets needed for TRD refit
+
+  // quantities necessary for pile-up correction
+  std::vector<int> mTriggeredBCFT0; ///< array with the FT0 trigger times
+  int mCurrentTriggerRecord;
+  uint32_t mFirstOrbit;
+  int mCurrentTrackId;
 
   // corrections from ccdb, some need to be loaded only once hence an init flag
   o2::trd::LocalGainFactor mLocalGain; ///< local gain factors from krypton calibration

Detectors/TRD/qc/src/Tracking.cxx

@@ -37,15 +37,23 @@ void Tracking::setInput(const o2::globaltracking::RecoContainer& input)
   mTracksTPCTRD = input.getTPCTRDTracks<TrackTRD>();
   mTrackletsRaw = input.getTRDTracklets();
   mTrackletsCalib = input.getTRDCalibratedTracklets();
+  mTrackTriggerRecordsITSTPCTRD = input.getITSTPCTRDTriggers();
+  mTrackTriggerRecordsTPCTRD = input.getTPCTRDTriggers();
 }
 
 void Tracking::run()
 {
+  mCurrentTriggerRecord = 0;
+  mCurrentTrackId = 0;
   for (const auto& trkTrd : mTracksTPCTRD) {
     checkTrack(trkTrd, true);
+    mCurrentTrackId++;
   }
+  mCurrentTriggerRecord = 0;
+  mCurrentTrackId = 0;
   for (const auto& trkTrd : mTracksITSTPCTRD) {
     checkTrack(trkTrd, false);
+    mCurrentTrackId++;
   }
 }
 
@@ -65,6 +73,59 @@ void Tracking::checkTrack(const TrackTRD& trkTrd, bool isTPCTRD)
     qcStruct.dEdxTotTPC = isTPCTRD ? mTracksTPC[id].getdEdx().dEdxTotTPC : mTracksTPC[mTracksITSTPC[id].getRefTPC()].getdEdx().dEdxTotTPC;
   }
 
+  // find corresponding track trigger record to get track timing
+  int triggeredBC = 0;
+  for (; mCurrentTriggerRecord < (isTPCTRD ? mTrackTriggerRecordsTPCTRD.size() : mTrackTriggerRecordsITSTPCTRD.size()); mCurrentTriggerRecord++) {
+    auto& tRecord = (isTPCTRD ? mTrackTriggerRecordsTPCTRD[mCurrentTriggerRecord] : mTrackTriggerRecordsITSTPCTRD[mCurrentTriggerRecord]);
+    if (mCurrentTrackId >= tRecord.getFirstTrack() && mCurrentTrackId < tRecord.getFirstTrack() + tRecord.getNumberOfTracks()) {
+      uint32_t currentOrbit = tRecord.getBCData().orbit;
+      triggeredBC = tRecord.getBCData().bc + (currentOrbit - mFirstOrbit) * o2::constants::lhc::LHCMaxBunches;
+      break;
+    }
+  }
+
+  // Find most probable BCs and RMS for pile-up correction and error. Same BC is assumed for all tracklets
+  float tCorrPileUp = 0.;
+  float tErrPileUp2 = 0;
+  float maxProb = 0.f;
+  // The uncertainty is the RMS wrt the default correction of all possible corrections weighted by their probability
+  float sumCorr = 0.f;
+  float sumCorr2 = 0.f;
+  float sumProb = 0.f;
+  for (int iBC = 0; iBC < mTriggeredBCFT0.size(); iBC++) {
+    int deltaBC = roundf(mTriggeredBCFT0[iBC] - triggeredBC);
+    if (deltaBC <= mRecoParam.getPileUpRangeBefore()) {
+      continue;
+    }
+    if (deltaBC >= mRecoParam.getPileUpRangeAfter()) {
+      break;
+    }
+    // collect the charges
+    std::array<int, 6> q0;
+    std::array<int, 6> q1;
+    for (int iLy = 0; iLy < NLAYER; iLy++) {
+      int trkltId = trkTrd.getTrackletIndex(iLy);
+      if (trkltId < 0) {
+        q0[iLy] = -1;
+        q1[iLy] = -1;
+      } else {
+        q0[iLy] = mTrackletsRaw[trkltId].getQ0();
+        q1[iLy] = mTrackletsRaw[trkltId].getQ1();
+      }
+    }
+    // get pile-up probability
+    float probBC = mRecoParam.getPileUpProbTrack(deltaBC, q0, q1);
+    sumCorr += probBC * deltaBC;
+    sumCorr2 += probBC * deltaBC * deltaBC;
+    sumProb += probBC;
+    if (probBC > maxProb) {
+      maxProb = probBC;
+      tCorrPileUp = -deltaBC;
+    }
+  }
+  if (sumProb > 1e-6)
+    tErrPileUp2 = sumCorr2 / sumProb - 2 * tCorrPileUp * sumCorr / sumProb + tCorrPileUp * tCorrPileUp;
+
   for (int iLayer = 0; iLayer < NLAYER; ++iLayer) {
     int trkltId = trkTrd.getTrackletIndex(iLayer);
     if (trkltId < 0) {
@@ -93,9 +154,16 @@ void Tracking::checkTrack(const TrackTRD& trkTrd, bool isTPCTRD)
     if (!((trk.getSigmaZ2() < (padLength * padLength / 12.f)) && (std::fabs(mTrackletsCalib[trkltId].getZ() - trk.getZ()) < padLength))) {
       tiltCorrUp = 0.f;
     }
-    std::array<float, 2> trkltPosUp{mTrackletsCalib[trkltId].getY() - tiltCorrUp, zPosCorrUp};
+
+    // conversion from slope in pad per time bin to slope in cm per BC = tracklets[trkltIdx].getSlopeFloat() * padWidth / BCperTimeBin
+    float slopeFactor = mTrackletsRaw[trkltId].getSlopeFloat() * pad->getWidthIPad() / 4.f;
+    float yCorrPileUp = tCorrPileUp * slopeFactor;
+    float yAddErrPileUp2 = tErrPileUp2 * slopeFactor * slopeFactor;
+
+    std::array<float, 2> trkltPosUp{mTrackletsCalib[trkltId].getY() - tiltCorrUp + yCorrPileUp, zPosCorrUp};
     std::array<float, 3> trkltCovUp;
     mRecoParam.recalcTrkltCov(tilt, trk.getSnp(), pad->getRowSize(tracklet.getPadRow()), trkltCovUp);
+    trkltCovUp[0] += yAddErrPileUp2;
     auto chi2trklt = trk.getPredictedChi2(trkltPosUp, trkltCovUp);
 
     qcStruct.trackProp[iLayer] = trk;

Detectors/TRD/workflow/include/TRDWorkflow/TRDGlobalTrackingQCSpec.h

@@ -30,6 +30,8 @@
 #include "DataFormatsParameters/GRPObject.h"
 #include "ReconstructionDataFormats/GlobalTrackID.h"
 #include "DataFormatsGlobalTracking/RecoContainer.h"
+#include "DataFormatsFT0/RecPoints.h"
+#include "FT0Reconstruction/InteractionTag.h"
 #include "TRDQC/Tracking.h"
 #include <cstring>
 
@@ -47,7 +49,7 @@ namespace trd
 class TRDGlobalTrackingQC : public Task
 {
  public:
-  TRDGlobalTrackingQC(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool tpcAvailable) : mDataRequest(dr), mGGCCDBRequest(gr), mTPCavailable(tpcAvailable) {}
+  TRDGlobalTrackingQC(std::shared_ptr<DataRequest> dr, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool tpcAvailable, o2::dataformats::GlobalTrackID::mask_t src) : mDataRequest(dr), mGGCCDBRequest(gr), mTPCavailable(tpcAvailable), mTrkMask(src) {}
   ~TRDGlobalTrackingQC() override = default;
   void init(InitContext& ic) final
   {
@@ -67,6 +69,23 @@ class TRDGlobalTrackingQC : public Task
     updateTimeDependentParams(pc); // Make sure this is called after recoData.collectData, which may load some conditions
     mQC.reset();
     mQC.setInput(recoData);
+    std::vector<int> triggeredBCFT0;
+    if (mTrkMask[GTrackID::FT0]) { // pile-up tagging was requested
+      auto ft0recPoints = recoData.getFT0RecPoints();
+      uint32_t firstOrbit = 0;
+      for (size_t ft0id = 0; ft0id < ft0recPoints.size(); ft0id++) {
+        const auto& f0rec = ft0recPoints[ft0id];
+        if (ft0id == 0) {
+          firstOrbit = f0rec.getInteractionRecord().orbit;
+          mQC.setFirstOrbit(firstOrbit);
+        }
+        if (o2::ft0::InteractionTag::Instance().isSelected(f0rec)) {
+          uint32_t currentOrbit = f0rec.getInteractionRecord().orbit;
+          triggeredBCFT0.push_back(f0rec.getInteractionRecord().bc + (currentOrbit - firstOrbit) * o2::constants::lhc::LHCMaxBunches);
+        }
+      }
+    }
+    mQC.setTriggeredBCFT0(triggeredBCFT0);
     mQC.run();
     pc.outputs().snapshot(Output{"TRD", "TRACKINGQC", 0}, mQC.getTrackQC());
   }
@@ -94,6 +113,7 @@ class TRDGlobalTrackingQC : public Task
     }
   }
 
+  o2::dataformats::GlobalTrackID::mask_t mTrkMask; ///< seeding track sources (TPC, ITS-TPC)
   std::shared_ptr<DataRequest> mDataRequest;
   std::shared_ptr<o2::base::GRPGeomRequest> mGGCCDBRequest;
   bool mTPCavailable{false};
@@ -133,7 +153,7 @@ DataProcessorSpec getTRDGlobalTrackingQCSpec(o2::dataformats::GlobalTrackID::mas
     "trd-tracking-qc",
     dataRequest->inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<TRDGlobalTrackingQC>(dataRequest, ggRequest, isTPCavailable)},
+    AlgorithmSpec{adaptFromTask<TRDGlobalTrackingQC>(dataRequest, ggRequest, isTPCavailable, src)},
     Options{}};
 }
 

Detectors/TRD/workflow/include/TRDWorkflow/TRDGlobalTrackingSpec.h

@@ -109,6 +109,7 @@ class TRDGlobalTracking : public o2::framework::Task
 #endif
   std::array<float, 5> mCovDiagInner{}; ///< total cov.matrix extra diagonal error from TrackTuneParams
   std::array<float, 5> mCovDiagOuter{}; ///< total cov.matrix extra diagonal error from TrackTuneParams
+  std::vector<int> mTriggeredBCFT0;     ///< array with the FT0 trigger times
   // PID
   PIDPolicy mPolicy{PIDPolicy::DEFAULT}; ///< Model to load an evaluate
   std::unique_ptr<PIDBase> mBase;        ///< PID engine

Detectors/TRD/workflow/src/TRDGlobalTrackingSpec.cxx

@@ -415,6 +415,24 @@ void TRDGlobalTracking::run(ProcessingContext& pc)
   }
   LOGF(info, "%i tracks are loaded into the TRD tracker. Out of those %i ITS-TPC tracks and %i TPC tracks", nTracksLoadedITSTPC + nTracksLoadedTPC, nTracksLoadedITSTPC, nTracksLoadedTPC);
 
+  // Load the FT0 triggered BCs if this is requested
+
+  if (mTrkMask[GTrackID::FT0]) { // pile-up tagging was requested
+    auto ft0recPoints = inputTracks.getFT0RecPoints();
+    uint32_t firstOrbit = 0;
+    for (size_t ft0id = 0; ft0id < ft0recPoints.size(); ft0id++) {
+      const auto& f0rec = ft0recPoints[ft0id];
+      if (ft0id == 0)
+        firstOrbit = f0rec.getInteractionRecord().orbit;
+      if (o2::ft0::InteractionTag::Instance().isSelected(f0rec)) {
+        uint32_t currentOrbit = f0rec.getInteractionRecord().orbit;
+        mTriggeredBCFT0.push_back(f0rec.getInteractionRecord().bc + (currentOrbit - firstOrbit) * o2::constants::lhc::LHCMaxBunches);
+      }
+    }
+  }
+
+  mTracker->SetFT0TriggeredBC(mTriggeredBCFT0.data(), mTriggeredBCFT0.size());
+
   // start the tracking
   // mTracker->DumpTracks();
   mChainTracking->DoTRDGPUTracking<GPUTRDTrackerKernels::o2Version>(mTracker);
@@ -797,6 +815,45 @@ bool TRDGlobalTracking::refitTRDTrack(TrackTRD& trk, float& chi2, bool inwards,
     }
   }
 
+  // Find most probable BCs and RMS for pile-up correction and error. Same BC is assumed for all tracklets
+  float tCorrPileUp = 0.;
+  float tErrPileUp2 = 0;
+  float maxProb = 0.f;
+  // The uncertainty is the RMS wrt the default correction of all possible corrections weighted by their probability
+  float sumCorr = 0.f;
+  float sumCorr2 = 0.f;
+  float sumProb = 0.f;
+  for (int iBC = 0; iBC < mTriggeredBCFT0.size(); iBC++) {
+    int deltaBC = roundf(mTriggeredBCFT0[iBC] - mChainTracking->mIOPtrs.trdTriggerTimes[trk.getCollisionId()] / o2::constants::lhc::LHCBunchSpacingMUS);
+    if (deltaBC <= mRecoParam.getPileUpRangeBefore() || deltaBC >= mRecoParam.getPileUpRangeAfter()) {
+      continue;
+    }
+    // collect the charges
+    std::array<int, 6> q0;
+    std::array<int, 6> q1;
+    for (int iLy = 0; iLy < NLAYER; iLy++) {
+      int trkltId = trk.getTrackletIndex(iLy);
+      if (trkltId < 0) {
+        q0[iLy] = -1;
+        q1[iLy] = -1;
+      } else {
+        q0[iLy] = mTrackletsRaw[trkltId].getQ0();
+        q1[iLy] = mTrackletsRaw[trkltId].getQ1();
+      }
+    }
+    // get pile-up probability
+    float probBC = mRecoParam.getPileUpProbTrack(deltaBC, q0, q1);
+    sumCorr += probBC * deltaBC;
+    sumCorr2 += probBC * deltaBC * deltaBC;
+    sumProb += probBC;
+    if (probBC > maxProb) {
+      maxProb = probBC;
+      tCorrPileUp = -deltaBC;
+    }
+  }
+  if (sumProb > 1e-6)
+    tErrPileUp2 = sumCorr2 / sumProb - 2 * tCorrPileUp * sumCorr / sumProb + tCorrPileUp * tCorrPileUp;
+
   if (inwards) {
     // reset covariance to something big for inwards refit
     trkParam->resetCovariance(100);
@@ -827,9 +884,15 @@ bool TRDGlobalTracking::refitTRDTrack(TrackTRD& trk, float& chi2, bool inwards,
       tiltCorrUp = 0.f;
     }
 
-    std::array<float, 2> trkltPosUp{mTrackletsCalib[trkltId].getY() - tiltCorrUp, zPosCorrUp};
+    // conversion from slope in pad per time bin to slope in cm per BC = tracklets[trkltIdx].getSlopeFloat() * padWidth / BCperTimeBin
+    float slopeFactor = mTrackletsRaw[trkltId].getSlopeFloat() * pad->getWidthIPad() / 4.f;
+    float yCorrPileUp = tCorrPileUp * slopeFactor;
+    float yAddErrPileUp2 = tErrPileUp2 * slopeFactor * slopeFactor;
+
+    std::array<float, 2> trkltPosUp{mTrackletsCalib[trkltId].getY() - tiltCorrUp + yCorrPileUp, zPosCorrUp};
     std::array<float, 3> trkltCovUp;
     mRecoParam.recalcTrkltCov(tilt, trkParam->getSnp(), pad->getRowSize(mTrackletsRaw[trkltId].getPadRow()), trkltCovUp);
+    trkltCovUp[0] += yAddErrPileUp2;
 
     chi2 += trkParam->getPredictedChi2(trkltPosUp, trkltCovUp);
     if (!trkParam->update(trkltPosUp, trkltCovUp)) {

GPU/GPUTracking/DataTypes/GPUTRDRecoParam.cxx

@@ -77,3 +77,71 @@ void GPUTRDRecoParam::recalcTrkltCov(const float tilt, const float snp, const fl
   cov[1] = c2 * tilt * (sz2 - sy2);
   cov[2] = c2 * (t2 * sy2 + sz2);
 }
+
+float GPUTRDRecoParam::getPileUpProbTracklet(int nBC, int Q0, int Q1) const
+{
+  // get the probability that the tracklet with charges Q0 and Q1 belongs to a given BC, with a (signed) distance nBC from the TRD-triggered BC
+  // parametrization depends on whether charges are 0 or not
+  int maxBC = mPileUpRangeAfter;
+  int minBC = mPileUpRangeBefore;
+  int maxProbBC = mPileUpMaxProb;
+  if (Q0 > 0 && Q1 > 0) {
+    maxBC = mPileUpRangeAfter11;
+    minBC = mPileUpRangeBefore11;
+    maxProbBC = mPileUpMaxProb11;
+  }
+  if (Q0 == 0 && Q1 > 0) {
+    maxBC = mPileUpRangeAfter01;
+    minBC = mPileUpRangeBefore01;
+    maxProbBC = mPileUpMaxProb01;
+  }
+  if (Q0 > 0 && Q1 == 0) {
+    maxBC = mPileUpRangeAfter10;
+    minBC = mPileUpRangeBefore10;
+    maxProbBC = mPileUpMaxProb10;
+  }
+  if (Q0 == 0 && Q1 == 0) {
+    maxBC = mPileUpRangeAfter00;
+    minBC = mPileUpRangeBefore00;
+    maxProbBC = mPileUpMaxProb00;
+  }
+
+  // prob is 0 if the BC is too far, maximal for a given nBC, and with two linear functions in between. The maximum is chosen so that the integral is 1.
+  if (nBC <= minBC || nBC >= maxBC)
+    return 0.;
+  float maxProb = 2. / (maxBC - minBC);
+  if (nBC > minBC && nBC <= maxProbBC)
+    return maxProb / (maxProbBC - minBC) * (nBC - minBC);
+  else
+    return maxProb / (maxProbBC - maxBC) * (nBC - maxBC);
+}
+
+float GPUTRDRecoParam::getPileUpProbTrack(int nBC, std::array<int, 6> Q0, std::array<int, 6> Q1) const
+{
+  // get the probability that the track belongs to a given BC, with a (signed) distance nBC from the TRD-triggered BC
+  // it depends on the individual probabilities for every of its tracklets.
+  //
+  // If P(BC|L0,L1,...) is the probability that the track belongs to a given BC, given the information on the tracklet charges in L0,L1, ...
+  // P(BC|L0,L1,...) proportional to P(BC)*P(L0,L1,...|BC), prop to P(BC)*P(L0|BC)*P(L1|BC)*... since for a given track and BC, charge in different layers are independent
+  // prop to P(BC) * P(BC|L0)/P(BC) * P(BC|L1)/P(BC) * ...
+  //
+  // P(BC) is the probability with no charge information: we start from this probability, and each tracklet adds new information on pileup probability
+
+  // basic probability, if we had no info on the charges
+  float probNoInfo = GPUTRDRecoParam::getPileUpProbTracklet(nBC, -1, -1);
+
+  float probTrack = probNoInfo;
+  if (probNoInfo < 1e-6f)
+    return 0.;
+
+  // For each tracklet, we add the info on its charge
+  for (int i = 0; i < 6; i++) {
+    // negative charge values if the tracklet is not present
+    if (Q0[i] < 0 || Q1[i] < 0)
+      continue;
+    float probTracklet = GPUTRDRecoParam::getPileUpProbTracklet(nBC, Q0[i], Q1[i]);
+    probTrack *= probTracklet / probNoInfo;
+  }
+
+  return probTrack;
+}