AOD: protoype TRD signal and eProb encoding

f3sch · f3sch · commit fa4a19930c7a · 2025-03-18T08:33:18.000+01:00
Note to self: think about inverse-log encoding pid
signal and give 8 more bits to dEdx.

Signed-off-by: Felix Schlepper &lt;felix.schlepper@cern.ch&gt;
diff --git a/Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h b/Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h
@@ -524,7 +524,7 @@ class AODProducerWorkflowDPL : public Task
   void addToTracksQATable(TracksQACursorType& tracksQACursor, TrackQA& trackQAInfoHolder);
 
   template <typename TRDsExtraCursorType>
-  void addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, int trkTableIdx);
+  void addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, TrackExtraInfo& extraInfo, int trkTableIdx);
 
   template <typename mftTracksCursorType, typename AmbigMFTTracksCursorType>
   void addToMFTTracksTable(mftTracksCursorType& mftTracksCursor, AmbigMFTTracksCursorType& ambigMFTTracksCursor,
diff --git a/Detectors/AOD/src/AODProducerWorkflowSpec.cxx b/Detectors/AOD/src/AODProducerWorkflowSpec.cxx
@@ -349,7 +349,7 @@ void AODProducerWorkflowDPL::addToTracksExtraTable(TracksExtraCursorType& tracks
                     truncateFloatFraction(extraInfoHolder.trdChi2, mTrackChi2),
                     truncateFloatFraction(extraInfoHolder.tofChi2, mTrackChi2),
                     truncateFloatFraction(extraInfoHolder.tpcSignal, mTrackSignal),
-                    truncateFloatFraction(extraInfoHolder.trdSignal, mTrackSignal),
+                    extraInfoHolder.trdSignal, // byte encoded value do not truncate
                     truncateFloatFraction(extraInfoHolder.length, mTrackSignal),
                     truncateFloatFraction(extraInfoHolder.tofExpMom, mTrack1Pt),
                     truncateFloatFraction(extraInfoHolder.trackEtaEMCAL, mTrackPosEMCAL),
@@ -390,24 +390,24 @@ void AODProducerWorkflowDPL::addToTracksQATable(TracksQACursorType& tracksQACurs
 }
 
 template <typename TRDsExtraCursorType>
-void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, int trkTableIdx)
+void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContainer& recoData, TRDsExtraCursorType& trdExtraCursor, const GIndex& trkIdx, TrackExtraInfo& extraInfo, int trkTableIdx)
 {
-  static int q0s[6] = {-1}, q1s[6] = {-1}, q2s[6] = {-1};
-  static float q0sCor[6] = {-1}, q1sCor[6] = {-1}, q2sCor[6] = {-1};
-  static float ttgls[6] = {-999}, tphis[6] = {-999};
+  const int NLAYERS = o2::trd::TrackTRD::EGPUTRDTrack::kNLayers;
+  int q0s[NLAYERS] = {-1}, q1s[NLAYERS] = {-1}, q2s[NLAYERS] = {-1};
+  float q0sCor[NLAYERS] = {-1}, q1sCor[NLAYERS] = {-1}, q2sCor[NLAYERS] = {-1};
+  float ttgls[NLAYERS] = {-999}, tphis[NLAYERS] = {-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]);
-  auto trkC = trk;
   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 trkC = trk; // local copy to propagate
+
+  float dEdx{0.};
+  for (int iLay{0}; iLay < NLAYERS; ++iLay) {
     auto trkltId = trk.getTrackletIndex(iLay);
     if (trkltId < 0) {
       continue;
@@ -421,11 +421,11 @@ void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContai
     int trkltSec = trkltDet / 30;
     if (trkltSec != o2::math_utils::angle2Sector(trkC.getAlpha())) {
       if (!trkC.rotate(o2::math_utils::sector2Angle(trkltSec))) {
-        break;
+        continue;
       }
     }
     if (!o2::base::Propagator::Instance()->PropagateToXBxByBz(trkC, ctrklets[trkltId].getX(), o2::base::Propagator::MAX_SIN_PHI, o2::base::Propagator::MAX_STEP, mMatCorr)) {
-      break;
+      continue;
     }
 
     auto tphi = trkC.getSnp() / std::sqrt((1.f - trkC.getSnp()) * (1.f + trkC.getSnp()));
@@ -440,26 +440,18 @@ void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContai
     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;
-        }
-      }
-    }
+    dEdx += (q0sCor[iLay] + q1sCor[iLay] + q2sCor[iLay]);
   }
+  dEdx /= (float)trkC.getNtracklets();
+
+  aod::track::extensions::TRDSignalEncoding enc;
+  enc.setDEdx(dEdx);
+  enc.setEProb(trk.getSignal());
+  extraInfo.trdSignal = enc.getSignal();
 
-  trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
+  if (mEnableTRDextra) {
+    trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
+  }
 }
 
 template <typename mftTracksCursorType, typename AmbigMFTTracksCursorType>
@@ -612,12 +604,13 @@ void AODProducerWorkflowDPL::fillTrackTablesPerCollision(int collisionID,
           if (!isProp) {
             addToTracksTable(tracksCursor, tracksCovCursor, trOrig, collisionID, aod::track::TrackIU);
           }
-          addToTracksExtraTable(tracksExtraCursor, extraInfoHolder);
 
-          if (mEnableTRDextra && trackIndex.includesDet(GIndex::Source::TRD)) {
-            addToTRDsExtra(data, trdsExtraCursor, trackIndex, mTableTrID);
+          if (trackIndex.includesDet(GIndex::Source::TRD)) {
+            addToTRDsExtra(data, trdsExtraCursor, trackIndex, extraInfoHolder, mTableTrID);
           }
 
+          addToTracksExtraTable(tracksExtraCursor, extraInfoHolder);
+
           //  collecting table indices of barrel tracks for V0s table
           if (extraInfoHolder.bcSlice[0] >= 0 && collisionID < 0) {
             ambigTracksCursor(mTableTrID, extraInfoHolder.bcSlice);
@@ -2606,7 +2599,6 @@ AODProducerWorkflowDPL::TrackExtraInfo AODProducerWorkflowDPL::processBarrelTrac
   if (contributorsGID[GIndex::Source::TRD].isIndexSet()) {                                        // ITS-TPC-TRD-TOF, TPC-TRD-TOF, TPC-TRD, ITS-TPC-TRD
     const auto& trdOrig = data.getTrack<o2::trd::TrackTRD>(contributorsGID[GIndex::Source::TRD]); // refitted TRD trac
     extraInfoHolder.trdChi2 = trdOrig.getChi2();
-    extraInfoHolder.trdSignal = trdOrig.getSignal();
     extraInfoHolder.trdPattern = getTRDPattern(trdOrig);
     if (extraInfoHolder.trackTimeRes < 0.) { // time is not set yet, this is possible only for TPC-TRD and ITS-TPC-TRD tracks, since those with TOF are set upstream
       // TRD is triggered: time uncertainty is within a BC
diff --git a/Framework/Core/include/Framework/AnalysisDataModel.h b/Framework/Core/include/Framework/AnalysisDataModel.h
@@ -394,6 +394,23 @@ DECLARE_SOA_DYNAMIC_COLUMN(TPCDeltaTBwd, tpcDeltaTBwd, //! Delta Backward of tra
 
 } // namespace v001
 
+namespace v002::extensions
+{
+using TRDEncoding = o2::aod::track::extensions::TRDSignalEncoding;
+DECLARE_SOA_DYNAMIC_COLUMN(TRDDEdx, trdDEDx, //! TRD dEdx signal (a.u.), returns 0. if no valid signal is provided
+                           [](float signal) -> float {
+                             TRDEncoding enc;
+                             enc.setSignal(signal);
+                             return enc.getDEdx();
+                           });
+DECLARE_SOA_DYNAMIC_COLUMN(TRDEProb, trdEProb, //! TRD electron probability [0., 1.), returns 0. if no valid signal is provided
+                           [](float signal) -> float {
+                             TRDEncoding enc;
+                             enc.setSignal(signal);
+                             return enc.getDEdx();
+                           });
+} // namespace v002::extensions
+
 DECLARE_SOA_DYNAMIC_COLUMN(HasITS, hasITS, //! Flag to check if track has a ITS match
                            [](uint8_t detectorMap) -> bool { return detectorMap & o2::aod::track::ITS; });
 DECLARE_SOA_DYNAMIC_COLUMN(HasTPC, hasTPC, //! Flag to check if track has a TPC match
@@ -617,7 +634,9 @@ DECLARE_SOA_TABLE_FULL_VERSIONED(StoredTracksExtra_002, "TracksExtra", "AOD", "T
                                  track::TPCNClsFindable, track::TPCNClsFindableMinusFound, track::TPCNClsFindableMinusPID, track::TPCNClsFindableMinusCrossedRows,
                                  track::TPCNClsShared, track::v001::extensions::TPCDeltaTFwd<track::TrackTimeRes, track::Flags>, track::v001::extensions::TPCDeltaTBwd<track::TrackTimeRes, track::Flags>,
                                  track::TRDPattern, track::ITSChi2NCl, track::TPCChi2NCl, track::TRDChi2, track::TOFChi2,
-                                 track::TPCSignal, track::TRDSignal, track::Length, track::TOFExpMom,
+                                 track::TPCSignal,
+                                 track::TRDSignal, track::v002::extensions::TRDDEdx<track::TRDSignal>, track::v002::extensions::TRDEProb<track::TRDSignal>,
+                                 track::Length, track::TOFExpMom,
                                  track::PIDForTracking<track::Flags>,
                                  track::IsPVContributor<track::Flags>,
                                  track::HasITS<track::v001::DetectorMap>, track::HasTPC<track::v001::DetectorMap>,
diff --git a/Framework/Core/include/Framework/DataTypes.h b/Framework/Core/include/Framework/DataTypes.h
@@ -76,6 +76,7 @@ enum TRDTrackPattern : uint8_t {
   HasNeighbor = 0x40,
   HasCrossing = 0x80,
 };
+
 namespace extensions
 {
 struct TPCTimeErrEncoding {
@@ -121,6 +122,68 @@ struct TPCTimeErrEncoding {
     return static_cast<float>(encoding.deltas.timeBackward) / timeScaler * TPCBinNS;
   }
 };
+
+class TRDSignalEncoding
+{
+  union {
+    float signal{0};
+    struct {
+      uint16_t dEdx{0};  // range [0., 512.)
+      uint16_t eProb{0}; // range [0., 1.)
+    } __attribute((packed)) data;
+  } encoding{0};
+  static_assert(sizeof(float) == 2 * sizeof(uint16_t));
+
+ public:
+  static constexpr float dEdxMax{512.f};
+  static constexpr float dEdxScaler{(1 << 16) / dEdxMax};
+  static constexpr float eProbMax{1.f};
+  static constexpr float eProbScaler{(1 << 16) / eProbMax};
+
+  void setSignal(float signal)
+  {
+    if (signal == -1. || signal == -999.f) {
+      // these two signal produce valid looking values
+      // but for backward comp. and not to confuse we mask these two specifically
+      encoding.signal = 0.;
+    } else {
+      encoding.signal = signal;
+    }
+  }
+
+  void setDEdx(float v)
+  {
+    if (v >= 0. && v < dEdxMax) {
+      encoding.data.dEdx = static_cast<uint16_t>(v * dEdxScaler);
+    } else {
+      encoding.data.dEdx = 0.f;
+    }
+  }
+  void setEProb(float v)
+  {
+    if (v >= 0. && v < eProbMax) {
+      encoding.data.eProb = static_cast<uint16_t>(v * eProbScaler);
+    } else {
+      encoding.data.eProb = 0.f;
+    }
+  }
+
+  float getSignal() const noexcept
+  {
+    return encoding.signal;
+  }
+
+  float getDEdx() const noexcept
+  {
+    return static_cast<float>(encoding.data.dEdx) / dEdxScaler;
+  }
+
+  float getEProb() const noexcept
+  {
+    return static_cast<float>(encoding.data.eProb) / eProbScaler;
+  }
+};
+
 } // namespace extensions
 
 // Reference radius for extrapolated tracks
