Implement TRD extra data processing in AODProducerWorkflowSpec.cxx

This commit enables the workflow to generate TRD-specific extra tables in the AOD, applying calibration and noise corrections to improve reconstruction fidelity for TRD tracks.

Author: KangkanGoswami

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,79 @@ 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()) * 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
+    if (trk.getIsCrossingNeighbor(iLay) && trk.getHasNeighbor()) {
+      for (const auto& trklt : trklets) {
+        if (tracklet.getTrackletWord() == trklt.getTrackletWord()) {
+          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()) * 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 +505,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 +517,7 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
                                                          TracksCovCursorType& tracksCovCursor,
                                                          TracksExtraCursorType& tracksExtraCursor,
                                                          TracksQACursorType& tracksQACursor,
+                                                         TRDsExtraCursor& trdsExtraCursor,
                                                          AmbigTracksCursorType& ambigTracksCursor,
                                                          MFTTracksCursorType& mftTracksCursor,
                                                          MFTTracksCovCursorType& mftTracksCovCursor,
@@ -553,7 +628,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 +2011,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;
@@ -2298,7 +2381,7 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   // 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);
 
@@ -2340,7 +2423,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 +3094,10 @@ 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();
+    }
   }
   if (mPropTracks) {
     pc.inputs().get<o2::dataformats::MeanVertexObject*>("meanvtx");
@@ -3222,7 +3309,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 +3400,12 @@ 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"));
+  }
 
   if (useMC) {
     outputs.insert(outputs.end(),
@@ -3336,7 +3429,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"}},