AOD: protoype TRD signal and eProb encoding

Note to self: think about inverse-log encoding pid
signal and give 8 more bits to dEdx.

Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>

Author: f3sch

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,

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,7 +390,7 @@ 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};
@@ -404,6 +404,7 @@ void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContai
   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;
@@ -459,7 +460,26 @@ void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContai
     }
   }
 
-  trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
+  aod::track::extensions::TRDSignalEncoding enc;
+
+  // dEdX signal
+  float integral{0.}, ntracklets{0.};
+  for (int i{0}; i < 6; ++i) {
+    if (q0sCor[i] > 0) {
+      integral += (q0sCor[i] + q1sCor[i] + q2sCor[i]);
+      ++ntracklets;
+    }
+  }
+  enc.setDEdx(integral / ntracklets);
+
+  // pid signal
+  enc.setEProb(trk.getSignal());
+
+  extraInfo.trdSignal = enc.getSignal();
+
+  if (mEnableTRDextra) {
+    trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
+  }
 }
 
 template <typename mftTracksCursorType, typename AmbigMFTTracksCursorType>
@@ -612,12 +632,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 +2627,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

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>,

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,62 @@ struct TPCTimeErrEncoding {
     return static_cast<float>(encoding.deltas.timeBackward) / timeScaler * TPCBinNS;
   }
 };
+
+class TRDSignalEncoding
+{
+  union {
+    float signal{0};
+    struct {
+      uint16_t dEdx{0};  // range [0., 256.)
+      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{256.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)
+  {
+    encoding.signal = (signal < 0.f) ? 0.f : 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