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 <[email protected]>

Author: f3sch

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -404,7 +404,7 @@ class AODProducerWorkflowDPL : public Task
     float trdChi2 = -999.f;
     float tofChi2 = -999.f;
     float tpcSignal = -999.f;
-    float trdSignal = -999.f;
+    float trdSignal = aod::track::extensions::TRDSignalEncoding::invalidSignal;
     float length = -999.f;
     float tofExpMom = -999.f;
     float trackEtaEMCAL = -999.f;
@@ -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,
                     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,6 +460,23 @@ void AODProducerWorkflowDPL::addToTRDsExtra(const o2::globaltracking::RecoContai
     }
   }
 
+  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
+  // TODO
+
+  extraInfo.trdSignal = enc.getSignal();
+
   trdExtraCursor(trkTableIdx, q0s, q1s, q2s, q0sCor, q1sCor, q2sCor, ttgls, tphis);
 }
 
@@ -612,12 +630,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);
+            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);

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,66 @@ 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{(1 << 8)};
+  static constexpr float dEdxScaler{(1 << 16) / dEdxMax};
+  static constexpr float eProbMax{1.};
+  static constexpr float eProbScaler{(1 << 16) / eProbMax};
+  // results in valid values, so we mark this and return 0.
+  // also luckily -999.f is accurately represented even if
+  // truncated by mantissa bits so all on-disk signals are accurate.
+  static constexpr float invalidSignal{-999.f};
+
+  void setSignal(float signal)
+  {
+    if (signal == invalidSignal) {
+      encoding.signal = 0.f;
+    } else {
+      encoding.signal = signal;
+    }
+  }
+
+  void setDEdx(float v)
+  {
+    if (v >= 0. && v < dEdxMax) {
+      encoding.data.dEdx = static_cast<uint16_t>(v * dEdxScaler);
+    }
+  }
+  void setEProb(float v)
+  {
+    if (v >= 0. && v < eProbMax) {
+      encoding.data.eProb = static_cast<uint16_t>(v * eProbScaler);
+    }
+  }
+
+  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