Integrate TRD extra data support in AOD production workflow

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -19,6 +19,10 @@
 #include "DataFormatsGlobalTracking/RecoContainer.h"
 #include "DataFormatsPHOS/Cell.h"
 #include "DataFormatsTRD/TrackTRD.h"
+#include "TRDBase/PadCalibrationsAliases.h"
+#include "DataFormatsTRD/NoiseCalibration.h"
+#include "DataFormatsTRD/CalGain.h"
+#include "DataFormatsTRD/Constants.h"
 #include "DetectorsBase/GRPGeomHelper.h"
 #include "DetectorsBase/Propagator.h"
 #include "Framework/DataProcessorSpec.h"
@@ -215,7 +219,7 @@ enum struct AODProducerStreamerFlags : uint8_t {
 class AODProducerWorkflowDPL : public Task
 {
  public:
-  AODProducerWorkflowDPL(GID::mask_t src, std::shared_ptr<DataRequest> dataRequest, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool enableSV, bool useMC = true, bool enableFITextra = false) : mUseMC(useMC), mEnableSV(enableSV), mEnableFITextra(enableFITextra), mInputSources(src), mDataRequest(dataRequest), mGGCCDBRequest(gr) {}
+  AODProducerWorkflowDPL(GID::mask_t src, std::shared_ptr<DataRequest> dataRequest, std::shared_ptr<o2::base::GRPGeomRequest> gr, bool enableSV, bool useMC = true, bool enableFITextra = false, bool enableTRDextra = false) : mUseMC(useMC), mEnableSV(enableSV), mEnableFITextra(enableFITextra), mEnableTRDextra(enableTRDextra), mInputSources(src), mDataRequest(dataRequest), mGGCCDBRequest(gr) {}
   ~AODProducerWorkflowDPL() override = default;
   void init(InitContext& ic) final;
   void run(ProcessingContext& pc) final;
@@ -250,6 +254,10 @@ class AODProducerWorkflowDPL : public Task
   o2::dataformats::MeanVertexObject mVtx;
   float mMaxPropXiu{5.0f}; // max X_IU for which track is to be propagated if mPropTracks is true. (other option: o2::constants::geom::XTPCInnerRef + 0.1f)
 
+  const o2::trd::LocalGainFactor* mTRDLocalGain; // TRD local gain factors from krypton calibration
+  const o2::trd::CalGain* mTRDGainCalib;         // TRD time-dependent gain calib at chamber level
+  const o2::trd::NoiseStatusMCM* mTRDNoiseMap;   // TRD noise map
+
   std::unordered_set<GIndex> mGIDUsedBySVtx;
   std::unordered_set<GIndex> mGIDUsedByStr;
 
@@ -261,6 +269,7 @@ class AODProducerWorkflowDPL : public Task
   bool mUseSigFiltMC = false; // enable signal filtering for MC with embedding
   bool mEnableSV = true; // enable secondary vertices
   bool mEnableFITextra = false;
+  bool mEnableTRDextra = false;
   bool mFieldON = false;
   const float cSpeed = 0.029979246f; // speed of light in TOF units
 
@@ -278,6 +287,7 @@ class AODProducerWorkflowDPL : public Task
   TStopwatch mTimer;
   bool mEMCselectLeading{false};
   uint64_t mEMCALTrgClassMask = 0;
+  size_t mCurrentTRDTrigID = 0; // current index of the TRD trigger record, to speed up search
 
   // unordered map connects global indices and table indices of barrel tracks
   std::unordered_map<GIndex, int> mGIDToTableID;
@@ -525,6 +535,9 @@ class AODProducerWorkflowDPL : public Task
   template <typename TracksQACursorType>
   void addToTracksQATable(TracksQACursorType& tracksQACursor, TrackQA& trackQAInfoHolder);
 
+  template <typename TRDsExtraCursorType>
+  void addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, int trkTableIdx);
+
   template <typename mftTracksCursorType, typename AmbigMFTTracksCursorType>
   void addToMFTTracksTable(mftTracksCursorType& mftTracksCursor, AmbigMFTTracksCursorType& ambigMFTTracksCursor,
                            GIndex trackID, const o2::globaltracking::RecoContainer& data, int collisionID,
@@ -544,7 +557,7 @@ class AODProducerWorkflowDPL : public Task
   // helper for track tables
   // * fills tables collision by collision
   // * interaction time is for TOF information
-  template <typename TracksCursorType, typename TracksCovCursorType, typename TracksExtraCursorType, typename TracksQACursorType, typename AmbigTracksCursorType,
+  template <typename TracksCursorType, typename TracksCovCursorType, typename TracksExtraCursorType, typename TracksQACursorType, typename TRDsExtraCursorType, typename AmbigTracksCursorType,
             typename MFTTracksCursorType, typename MFTTracksCovCursorType, typename AmbigMFTTracksCursorType,
             typename FwdTracksCursorType, typename FwdTracksCovCursorType, typename AmbigFwdTracksCursorType, typename FwdTrkClsCursorType>
   void fillTrackTablesPerCollision(int collisionID,
@@ -556,6 +569,7 @@ class AODProducerWorkflowDPL : public Task
                                    TracksCovCursorType& tracksCovCursor,
                                    TracksExtraCursorType& tracksExtraCursor,
                                    TracksQACursorType& tracksQACursor,
+                                   TRDsExtraCursorType& trdsExtraCursor,
                                    AmbigTracksCursorType& ambigTracksCursor,
                                    MFTTracksCursorType& mftTracksCursor,
                                    MFTTracksCovCursorType& mftTracksCovCursor,
@@ -680,7 +694,7 @@ class AODProducerWorkflowDPL : public Task
 };
 
 /// create a processor spec
-framework::DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, bool enableST, bool useMC, bool CTPConfigPerRun, bool enableFITextra);
+framework::DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, bool enableST, bool useMC, bool CTPConfigPerRun, bool enableFITextra, bool enableTRDextra);
 
 // helper interface for calo cells to "befriend" emcal and phos cells
 class CellHelper

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -33,7 +33,6 @@
 #include "DataFormatsPHOS/TriggerRecord.h"
 #include "DataFormatsPHOS/EventHandler.h"
 #include "DataFormatsTPC/TrackTPC.h"
-#include "DataFormatsTRD/TriggerRecord.h"
 #include "DataFormatsZDC/BCRecData.h"
 #include "DataFormatsZDC/ZDCEnergy.h"
 #include "DataFormatsZDC/ZDCTDCData.h"
@@ -45,6 +44,9 @@
 #include "CommonDataFormat/InteractionRecord.h"
 #include "DataFormatsTRD/TrackTRD.h"
 #include "DataFormatsTRD/TrackTriggerRecord.h"
+#include "DataFormatsTRD/CalibratedTracklet.h"
+#include "DataFormatsTRD/TriggerRecord.h"
+#include "DataFormatsTRD/Tracklet64.h"
 #include "DataFormatsGlobalTracking/RecoContainer.h"
 #include "Framework/AnalysisDataModel.h"
 #include "Framework/ConfigParamRegistry.h"
@@ -390,6 +392,121 @@ void AODProducerWorkflowDPL::addToTracksQATable(TracksQACursorType& tracksQACurs
     mTrackQCRetainOnlydEdx ? std::numeric_limits<int8_t>::min() : trackQAInfoHolder.dTofdZ);
 }
 
+template <typename TRDsExtraCursorType>
+void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, int trkTableIdx)
+{
+  int q0s[6] = {-1}, q1s[6] = {-1}, q2s[6] = {-1};
+  float q0sCor[6] = {-1}, q1sCor[6] = {-1}, q2sCor[6] = {-1};
+  float ttgls[6] = {-999}, tphis[6] = {-999};
+
+  auto contributorsGID = recoData.getSingleDetectorRefs(trkIdx);
+  if (!contributorsGID[GIndex::Source::TRD].isIndexSet()) { // should be redunant
+    return;
+  }
+  const auto& trk = recoData.getTrack<o2::trd::TrackTRD>(contributorsGID[GIndex::Source::TRD]);
+  o2::track::TrackPar trkC{contributorsGID[GIndex::Source::ITSTPC].isIndexSet() ? recoData.getTPCITSTrack(contributorsGID[GIndex::Source::ITSTPC]).getParamOut() : recoData.getTPCTrack(contributorsGID[GIndex::Source::TPC]).getParamOut()};
+  const auto& trklets = recoData.getTRDTracklets();
+  const auto& ctrklets = recoData.getTRDCalibratedTracklets();
+  for (int iLay{0}; iLay < 6; ++iLay) {
+    q0s[iLay] = q1s[iLay] = q2s[iLay] = -1;
+    q0sCor[iLay] = q1sCor[iLay] = q2sCor[iLay] = -1;
+    tphis[iLay] = ttgls[iLay] = -999;
+    auto trkltId = trk.getTrackletIndex(iLay);
+    if (trkltId < 0) {
+      continue;
+    }
+    const auto& tracklet = trklets[trkltId];
+    if (mTRDNoiseMap->isTrackletFromNoisyMCM(tracklet)) {
+      continue;
+    }
+    // we need to propagate into TRD local system
+    int trkltDet = tracklet.getDetector();
+    int trkltSec = trkltDet / 30;
+    if (trkltSec != o2::math_utils::angle2Sector(trkC.getAlpha())) {
+      if (!trkC.rotate(o2::math_utils::sector2Angle(trkltSec))) {
+        break;
+      }
+    }
+    if (!o2::base::Propagator::Instance()->PropagateToXBxByBz(trkC, ctrklets[trkltId].getX(), o2::base::Propagator::MAX_SIN_PHI, o2::base::Propagator::MAX_STEP, mMatCorr)) {
+      break;
+    }
+
+    auto tphi = trkC.getSnp() / std::sqrt((1.f - trkC.getSnp()) * (1.f + trkC.getSnp()));
+    auto trackletLength = std::sqrt(1.f + tphi * tphi + trkC.getTgl() * trkC.getTgl());
+    float cor = mTRDLocalGain->getValue(tracklet.getHCID() / 2, tracklet.getPadCol(), tracklet.getPadRow()) * mTRDGainCalib->getMPVdEdx(tracklet.getDetector()) / o2::trd::constants::MPVDEDXDEFAULT * trackletLength;
+    q0s[iLay] = tracklet.getQ0();
+    q1s[iLay] = tracklet.getQ1();
+    q2s[iLay] = tracklet.getQ2();
+    q0sCor[iLay] = (float)tracklet.getQ0() / cor;
+    q1sCor[iLay] = (float)tracklet.getQ1() / cor;
+    q2sCor[iLay] = (float)tracklet.getQ2() / cor;
+    ttgls[iLay] = trkC.getTgl();
+    tphis[iLay] = tphi;
+
+    // z-row merging, we want to merge only with tracklets from the same trigger record
+    if (trk.getIsCrossingNeighbor(iLay) && trk.getHasNeighbor()) {
+      // find the trigger the tracklet belongs to
+      auto trigsTRD = recoData.getTRDTriggerRecords();
+      size_t trdSelID = -1;
+
+      const auto& trig = trigsTRD[mCurrentTRDTrigID];
+      bool foundTRDTrigger = false;
+      // first check current trigger
+      if (trkltId >= trig.getFirstTracklet() && trkltId < trig.getFirstTracklet() + trig.getNumberOfTracklets()) {
+        trdSelID = mCurrentTRDTrigID;
+        foundTRDTrigger = true;
+      } else {
+        // then check next trigger
+        if (mCurrentTRDTrigID < trigsTRD.size() - 1) {
+          const auto& trig = trigsTRD[mCurrentTRDTrigID + 1];
+          if (trkltId >= trig.getFirstTracklet() && trkltId < trig.getFirstTracklet() + trig.getNumberOfTracklets()) {
+            trdSelID = mCurrentTRDTrigID + 1;
+            foundTRDTrigger = true;
+          }
+        }
+      }
+
+      size_t low = 0, up = trigsTRD.size() - 1;
+
+      // otherwise binary search
+      while (low <= up && !foundTRDTrigger) {
+        trdSelID = low + std::floor((up - low) / 2);
+        const auto& trig = trigsTRD[trdSelID];
+        if (trig.getFirstTracklet() > trkltId) {
+          up = trdSelID - 1;
+        } else {
+          if (trig.getFirstTracklet() + trig.getNumberOfTracklets() <= trkltId) {
+            low = trdSelID + 1;
+          } else {
+            foundTRDTrigger = true;
+          }
+        }
+      }
+      //-------------------
+      mCurrentTRDTrigID = trdSelID;
+      const auto& trigSel = trigsTRD[trdSelID];
+
+      // loop on other tracklets from the same trigger record
+      for (const auto& trklt : trklets.subspan(trigSel.getFirstTracklet(), trigSel.getNumberOfTracklets())) {
+        if (tracklet.getTrackletWord() == trklt.getTrackletWord() || tracklet.getDetector() != trklt.getDetector()) {
+          continue;
+        }
+        if (std::abs(tracklet.getPadCol() - trklt.getPadCol()) <= 1 && std::abs(tracklet.getPadRow() - trklt.getPadRow()) == 1) {
+          cor = mTRDLocalGain->getValue(trklt.getHCID() / 2, trklt.getPadCol(), trklt.getPadRow()) * mTRDGainCalib->getMPVdEdx(tracklet.getDetector()) / o2::trd::constants::MPVDEDXDEFAULT * trackletLength;
+          q0s[iLay] += trklt.getQ0();
+          q1s[iLay] += trklt.getQ1();
+          q2s[iLay] += trklt.getQ2();
+          q0sCor[iLay] += (float)trklt.getQ0() / cor;
+          q1sCor[iLay] += (float)trklt.getQ1() / cor;
+          q2sCor[iLay] += (float)trklt.getQ2() / cor;
+        }
+      }
+    }
+  }
+
+  trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
+}
+
 template <typename mftTracksCursorType, typename AmbigMFTTracksCursorType>
 void AODProducerWorkflowDPL::addToMFTTracksTable(mftTracksCursorType& mftTracksCursor, AmbigMFTTracksCursorType& ambigMFTTracksCursor,
                                                  GIndex trackID, const o2::globaltracking::RecoContainer& data, int collisionID,
@@ -430,8 +547,7 @@ void AODProducerWorkflowDPL::addToMFTTracksTable(mftTracksCursorType& mftTracksC
     ambigMFTTracksCursor(mTableTrMFTID, bcSlice);
   }
 }
-
-template <typename TracksCursorType, typename TracksCovCursorType, typename TracksExtraCursorType, typename TracksQACursorType, typename AmbigTracksCursorType,
+template <typename TracksCursorType, typename TracksCovCursorType, typename TracksExtraCursorType, typename TracksQACursorType, typename TRDsExtraCursor, typename AmbigTracksCursorType,
           typename MFTTracksCursorType, typename MFTTracksCovCursorType, typename AmbigMFTTracksCursorType,
           typename FwdTracksCursorType, typename FwdTracksCovCursorType, typename AmbigFwdTracksCursorType, typename FwdTrkClsCursorType>
 void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
@@ -443,6 +559,7 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
                                                          TracksCovCursorType& tracksCovCursor,
                                                          TracksExtraCursorType& tracksExtraCursor,
                                                          TracksQACursorType& tracksQACursor,
+                                                         TRDsExtraCursor& trdsExtraCursor,
                                                          AmbigTracksCursorType& ambigTracksCursor,
                                                          MFTTracksCursorType& mftTracksCursor,
                                                          MFTTracksCovCursorType& mftTracksCovCursor,
@@ -553,7 +670,9 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
             addToTracksTable(tracksCursor, tracksCovCursor, trOrig, collisionID, aod::track::TrackIU);
           }
           addToTracksExtraTable(tracksExtraCursor, extraInfoHolder);
-
+          if (mEnableTRDextra && trackIndex.includesDet(GIndex::Source::TRD)) {
+            addToTRDsExtra(data, trdsExtraCursor, trackIndex, mTableTrID);
+          }
           //  collecting table indices of barrel tracks for V0s table
           if (extraInfoHolder.bcSlice[0] >= 0 && collisionID < 0) {
             ambigTracksCursor(mTableTrID, extraInfoHolder.bcSlice);
@@ -1934,6 +2053,12 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   auto cpvClustersCursor = createTableCursor<o2::aod::CPVClusters>(pc);
   auto originCursor = createTableCursor<o2::aod::Origins>(pc);
 
+  /// Extra tables
+  o2::framework::Produces<o2::aod::TRDsExtra> trdExtraCursor;
+  if (mEnableTRDextra) {
+    trdExtraCursor = createTableCursor<o2::aod::TRDsExtra>(pc);
+  }
+
   // Declare MC cursors type without adding the output for a table
   o2::framework::Produces<o2::aod::McCollisionLabels> mcColLabelsCursor;
   o2::framework::Produces<o2::aod::McCollisions> mcCollisionsCursor;
@@ -2294,14 +2419,16 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
     }
   }
 
+  mCurrentTRDTrigID = 0; // reinitialize index for TRD trigger record search
   // 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
   // fixme: interaction time is undefined for unassigned tracks (?)
-  fillTrackTablesPerCollision(-1, std::uint64_t(-1), trackRef, primVerGIs, recoData, tracksCursor, tracksCovCursor, tracksExtraCursor, tracksQACursor,
+  fillTrackTablesPerCollision(-1, std::uint64_t(-1), trackRef, primVerGIs, recoData, tracksCursor, tracksCovCursor, tracksExtraCursor, tracksQACursor, trdExtraCursor,
                               ambigTracksCursor, mftTracksCursor, mftTracksCovCursor, ambigMFTTracksCursor,
                               fwdTracksCursor, fwdTracksCovCursor, ambigFwdTracksCursor, fwdTrkClsCursor, bcsMap);
 
+  mCurrentTRDTrigID = 0; // reinitialize index for TRD trigger record search
   // filling collisions and tracks into tables
   collisionID = 0;
   collisionsCursor.reserve(primVertices.size());
@@ -2340,7 +2467,7 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
 
     auto& trackRef = primVer2TRefs[collisionID];
     // passing interaction time in [ps]
-    fillTrackTablesPerCollision(collisionID, globalBC, trackRef, primVerGIs, recoData, tracksCursor, tracksCovCursor, tracksExtraCursor, tracksQACursor, ambigTracksCursor,
+    fillTrackTablesPerCollision(collisionID, globalBC, trackRef, primVerGIs, recoData, tracksCursor, tracksCovCursor, tracksExtraCursor, tracksQACursor, trdExtraCursor, ambigTracksCursor,
                                 mftTracksCursor, mftTracksCovCursor, ambigMFTTracksCursor,
                                 fwdTracksCursor, fwdTracksCovCursor, ambigFwdTracksCursor, fwdTrkClsCursor, bcsMap);
     collisionID++;
@@ -3011,6 +3138,11 @@ void AODProducerWorkflowDPL::updateTimeDependentParams(ProcessingContext& pc)
     mFieldON = std::abs(o2::base::Propagator::Instance()->getNominalBz()) > 0.01;
 
     pc.inputs().get<o2::ctp::CTPConfiguration*>("ctpconfig");
+    if (mEnableTRDextra) {
+      mTRDLocalGain = pc.inputs().get<o2::trd::LocalGainFactor*>("trdlocalgainfactors").get();
+      mTRDNoiseMap = pc.inputs().get<o2::trd::NoiseStatusMCM*>("trdnoisemap").get();
+      mTRDGainCalib = pc.inputs().get<o2::trd::CalGain*>("trdgaincalib").get(); // time dependent gain
+    }
   }
   if (mPropTracks) {
     pc.inputs().get<o2::dataformats::MeanVertexObject*>("meanvtx");
@@ -3222,7 +3354,7 @@ void AODProducerWorkflowDPL::endOfStream(EndOfStreamContext& /*ec*/)
   mStreamer.reset();
 }
 
-DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, bool enableStrangenessTracking, bool useMC, bool CTPConfigPerRun, bool enableFITextra)
+DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, bool enableStrangenessTracking, bool useMC, bool CTPConfigPerRun, bool enableFITextra, bool enableTRDextra)
 {
   auto dataRequest = std::make_shared<DataRequest>();
   dataRequest->inputs.emplace_back("ctpconfig", "CTP", "CTPCONFIG", 0, Lifetime::Condition, ccdbParamSpec("CTP/Config/Config", CTPConfigPerRun));
@@ -3313,6 +3445,13 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, boo
     OutputSpec{"TFF", "TFFilename"},
     OutputSpec{"AMD", "AODMetadataKeys"},
     OutputSpec{"AMD", "AODMetadataVals"}};
+  /// Extra tables
+  if (enableTRDextra) {
+    outputs.push_back(OutputForTable<TRDsExtra>::spec());
+    dataRequest->inputs.emplace_back("trdlocalgainfactors", "TRD", "LOCALGAINFACTORS", 0, Lifetime::Condition, ccdbParamSpec("TRD/Calib/LocalGainFactor"));
+    dataRequest->inputs.emplace_back("trdnoisemap", "TRD", "NOISEMAP", 0, Lifetime::Condition, ccdbParamSpec("TRD/Calib/NoiseMapMCM"));
+    dataRequest->inputs.emplace_back("trdgaincalib", "TRD", "CALGAIN", 0, Lifetime::Condition, ccdbParamSpec("TRD/Calib/CalGain"));
+  }
 
   if (useMC) {
     outputs.insert(outputs.end(),
@@ -3336,7 +3475,7 @@ DataProcessorSpec getAODProducerWorkflowSpec(GID::mask_t src, bool enableSV, boo
     "aod-producer-workflow",
     dataRequest->inputs,
     outputs,
-    AlgorithmSpec{adaptFromTask<AODProducerWorkflowDPL>(src, dataRequest, ggRequest, enableSV, useMC, enableFITextra)},
+    AlgorithmSpec{adaptFromTask<AODProducerWorkflowDPL>(src, dataRequest, ggRequest, enableSV, useMC, enableFITextra, enableTRDextra)},
     Options{
       ConfigParamSpec{"run-number", VariantType::Int64, -1L, {"The run-number. If left default we try to get it from DPL header."}},
       ConfigParamSpec{"aod-timeframe-id", VariantType::Int64, -1L, {"Set timeframe number"}},