[PWGCF] Add histograms for purity and fraction extraction to track track task in femto framework (#18512)

Author: ariedel-cern

PWGCF/Femto/Core/femtoUtils.h

@@ -356,6 +356,59 @@ float unBinLogUnsigned(T binned, float magMin, float magMax)
   return mag;
 }
 
+// ---------------------------------------------------------------------------
+// purity/fraction packing: up to 4 species' log-binned nSigma values (see
+// binLogSigned<uint16_t>) packed bitwise into a single int64_t column. Slot
+// assignment (which species occupies which of the 4 slots) is decided at
+// producer init time from which requirePidXxx flags are active, and exposed
+// via a self-documenting histogram -- these functions only do the fixed-slot
+// bit shifting, nothing species-specific.
+// ---------------------------------------------------------------------------
+
+constexpr int PurityNSlots = 4;
+
+/// pack 4 already-binned uint16_t codes (one per slot) into a single int64_t
+inline int64_t packPurity4(uint16_t slot0, uint16_t slot1, uint16_t slot2, uint16_t slot3)
+{
+  return (static_cast<int64_t>(slot0) << 48) |
+         (static_cast<int64_t>(slot1) << 32) |
+         (static_cast<int64_t>(slot2) << 16) |
+         static_cast<int64_t>(slot3);
+}
+
+/// extract the raw uint16_t code of one of the 4 slots (0..3) from a packed int64_t
+inline uint16_t unpackPurity4(int64_t packed, int slot)
+{
+  switch (slot) {
+    case 0:
+      return static_cast<uint16_t>((packed >> 48) & 0xFFFF);
+    case 1:
+      return static_cast<uint16_t>((packed >> 32) & 0xFFFF);
+    case 2:
+      return static_cast<uint16_t>((packed >> 16) & 0xFFFF);
+    case 3:
+      return static_cast<uint16_t>(packed & 0xFFFF);
+    default:
+      LOG(fatal) << "unpackPurity4: invalid slot " << slot << " (must be 0-3). Breaking...";
+      return 0;
+  }
+}
+
+/// pack 4 raw signed nSigma values (one per slot, unused slots pass 0) directly into a packed int64_t
+inline int64_t packPurityNSigma4(float nSigma0, float nSigma1, float nSigma2, float nSigma3, float magMin, float magMax)
+{
+  return packPurity4(binLogSigned<uint16_t>(nSigma0, magMin, magMax),
+                     binLogSigned<uint16_t>(nSigma1, magMin, magMax),
+                     binLogSigned<uint16_t>(nSigma2, magMin, magMax),
+                     binLogSigned<uint16_t>(nSigma3, magMin, magMax));
+}
+
+/// unpack the nSigma value of one of the 4 slots (0..3) from a packed int64_t
+inline float unpackPurityNSigma4(int64_t packed, int slot, float magMin, float magMax)
+{
+  return unBinLogSigned<uint16_t>(unpackPurity4(packed, slot), magMin, magMax);
+}
+
 }; // namespace utils
 }; // namespace o2::analysis::femto
 //

PWGCF/Femto/Core/modes.h

@@ -53,6 +53,7 @@ enum class Mode : uint32_t {
   kMc = bit<Mode>(2),
   kSe = bit<Mode>(3),
   kMe = bit<Mode>(4),
+  kFitVars = bit<Mode>(5), // purity/fraction histograms (tracks: packed nSigma tables; V0s: cpa) -- gates booking at init time
   kReco_Qa = kReco | kQa,
   kReco_Mc = kReco | kMc,
   kReco_Qa_Mc = kReco | kQa | kMc,
@@ -62,6 +63,10 @@ enum class Mode : uint32_t {
   kMe_Reco_Mc = kMe | kReco | kMc,
   kSe_Mc = kSe | kMc,
   kMe_Mc = kMe | kMc,
+  kSe_Reco_FitVars = kSe | kReco | kFitVars,
+  kMe_Reco_FitVars = kMe | kReco | kFitVars,
+  kSe_Reco_Mc_FitVars = kSe | kReco | kMc | kFitVars,
+  kMe_Reco_Mc_FitVars = kMe | kReco | kMc | kFitVars,
 };
 
 enum class System : uint32_t {

PWGCF/Femto/Core/pairBuilder.h

@@ -91,7 +91,9 @@ class PairTrackTrackBuilder
             std::map<T12, std::vector<o2::framework::AxisSpec>> const& trackHistSpec2,
             std::map<T13, std::vector<o2::framework::AxisSpec>> const& pairHistSpec,
             std::map<T14, std::vector<o2::framework::AxisSpec>> const& cprHistSpec,
-            std::map<T15, std::vector<o2::framework::AxisSpec>> const& pairCleanerHistSpec)
+            std::map<T15, std::vector<o2::framework::AxisSpec>> const& pairCleanerHistSpec,
+            int puritySlotMask1 = 1,
+            int puritySlotMask2 = 1)
   {
 
     // check if correlate the same tracks or not
@@ -106,8 +108,8 @@ class PairTrackTrackBuilder
     if (mSameSpecies) {
       mTrackCleaner1.init(confTrackCleaner1);
 
-      mTrackHistManagerSe1.template init<modeSe>(registry, trackHistSpec1, confTrackSelection1);
-      mTrackHistManagerMe1.template init<modeMe>(registry, trackHistSpec1, confTrackSelection1);
+      mTrackHistManagerSe1.template init<modeSe>(registry, trackHistSpec1, confTrackSelection1, puritySlotMask1);
+      mTrackHistManagerMe1.template init<modeMe>(registry, trackHistSpec1, confTrackSelection1, puritySlotMask1);
 
       mPairHistManagerSe.setMass(confTrackSelection1.pdgCodeAbs.value, confTrackSelection1.pdgCodeAbs.value);
       mPairHistManagerSe.setCharge(confTrackSelection1.chargeAbs.value, confTrackSelection1.chargeAbs.value);
@@ -119,10 +121,10 @@ class PairTrackTrackBuilder
     } else {
       mTrackCleaner1.init(confTrackCleaner1);
       mTrackCleaner2.init(confTrackCleaner2);
-      mTrackHistManagerSe1.template init<modeSe>(registry, trackHistSpec1, confTrackSelection1);
-      mTrackHistManagerSe2.template init<modeSe>(registry, trackHistSpec2, confTrackSelection2);
-      mTrackHistManagerMe1.template init<modeMe>(registry, trackHistSpec1, confTrackSelection1);
-      mTrackHistManagerMe2.template init<modeMe>(registry, trackHistSpec2, confTrackSelection2);
+      mTrackHistManagerSe1.template init<modeSe>(registry, trackHistSpec1, confTrackSelection1, puritySlotMask1);
+      mTrackHistManagerSe2.template init<modeSe>(registry, trackHistSpec2, confTrackSelection2, puritySlotMask2);
+      mTrackHistManagerMe1.template init<modeMe>(registry, trackHistSpec1, confTrackSelection1, puritySlotMask1);
+      mTrackHistManagerMe2.template init<modeMe>(registry, trackHistSpec2, confTrackSelection2, puritySlotMask2);
 
       mPairHistManagerSe.setMass(confTrackSelection1.pdgCodeAbs.value, confTrackSelection2.pdgCodeAbs.value);
       mPairHistManagerSe.setCharge(confTrackSelection1.chargeAbs.value, confTrackSelection2.chargeAbs.value);

PWGCF/Femto/Core/trackBuilder.h

@@ -598,6 +598,9 @@ struct TrackBuilderProducts : o2::framework::ProducesGroup {
   o2::framework::Produces<o2::aod::FDeuteronPids> producedDeuteronPids;
   o2::framework::Produces<o2::aod::FTritonPids> producedTritonPids;
   o2::framework::Produces<o2::aod::FHeliumPids> producedHeliumPids;
+  o2::framework::Produces<o2::aod::FItsPurity> producedItsPurity;
+  o2::framework::Produces<o2::aod::FTpcPurity> producedTpcPurity;
+  o2::framework::Produces<o2::aod::FTofPurity> producedTofPurity;
 };
 
 struct ConfTrackTables : o2::framework::ConfigurableGroup {
@@ -615,6 +618,9 @@ struct ConfTrackTables : o2::framework::ConfigurableGroup {
   o2::framework::Configurable<int> produceDeuteronPids{"produceDeuteronPids", -1, "Produce DeuteronPids (-1: auto; 0 off; 1 on)"};
   o2::framework::Configurable<int> produceTritonPids{"produceTritonPids", -1, "Produce TritonPids (-1: auto; 0 off; 1 on)"};
   o2::framework::Configurable<int> produceHeliumPids{"produceHeliumPids", -1, "Produce HeliumPids (-1: auto; 0 off; 1 on)"};
+  o2::framework::Configurable<int> produceItsPurity{"produceItsPurity", -1, "Produce ItsPurity (-1: auto; 0 off; 1 on)"};
+  o2::framework::Configurable<int> produceTpcPurity{"produceTpcPurity", -1, "Produce TpcPurity (-1: auto; 0 off; 1 on)"};
+  o2::framework::Configurable<int> produceTofPurity{"produceTofPurity", -1, "Produce TofPurity (-1: auto; 0 off; 1 on)"};
 };
 
 template <auto& SelectionHistName, auto& FilterHistName>
@@ -642,6 +648,52 @@ class TrackBuilder
     mProduceDeuteronPids = utils::enableTable("FDeuteronPids_001", table.produceDeuteronPids.value, initContext);
     mProduceTritonPids = utils::enableTable("FTritonPids_001", table.produceTritonPids.value, initContext);
     mProduceHeliumPids = utils::enableTable("FHeliumPids_001", table.produceHeliumPids.value, initContext);
+    mProduceItsPurity = utils::enableTable("FItsPurity_001", table.produceItsPurity.value, initContext);
+    mProduceTpcPurity = utils::enableTable("FTpcPurity_001", table.produceTpcPurity.value, initContext);
+    mProduceTofPurity = utils::enableTable("FTofPurity_001", table.produceTofPurity.value, initContext);
+
+    if (mProduceItsPurity || mProduceTpcPurity || mProduceTofPurity) {
+      // assign each active species (requirePidXxx == true) to the next free slot (0-3), in
+      // canonical order (PuritySpeciesName); species with requirePidXxx == false get no slot.
+      const std::array<bool, PuritySpeciesLast> requirePid = {
+        config.requirePidElectron.value, config.requirePidPion.value, config.requirePidKaon.value,
+        config.requirePidProton.value, config.requirePidDeuteron.value, config.requirePidTriton.value, config.requirePidHelium.value};
+
+      int nextSlot = 0;
+      for (int i = 0; i < PuritySpeciesLast; i++) {
+        if (!requirePid[i]) {
+          continue;
+        }
+        if (nextSlot >= utils::PurityNSlots) {
+          LOG(fatal) << "More than " << utils::PurityNSlots << " species have requirePidXxx enabled, "
+                     << "but the purity tables only have " << utils::PurityNSlots << " slots. "
+                     << "Disable requirePidXxx for species you don't need purity info for. Breaking...";
+        }
+        mPuritySlot[i] = nextSlot;
+        nextSlot++;
+      }
+
+      // one bin per species (own SelectionName -> own cutculator group), plus the usual
+      // trailing "All analyzed"/"All passed" pair that cutculator always skips by position.
+      // BitPosition___X (unassigned species) is deliberately excluded from cutculator's
+      // "neutral" category, so nothing is selectable for a species that isn't configured here.
+      // Picking the one assigned bit for a species gives exactly 1<<slot -- the one-hot
+      // puritySlotMask value to paste into the downstream TrackSelection, with no changes
+      // needed in cutculator.py itself.
+      registry->add("hPuritySlotAssignment", "Purity slot assignment", o2::framework::HistType::kTH1F, {{PuritySpeciesLast + 2, -0.5, PuritySpeciesLast + 2 - 0.5}});
+      auto hSlotAssignment = registry->get<TH1>(HIST("hPuritySlotAssignment"));
+      for (int i = 0; i < PuritySpeciesLast; i++) {
+        std::string bitPosition = mPuritySlot[i] >= 0 ? std::to_string(mPuritySlot[i]) : "X";
+        std::string value = mPuritySlot[i] >= 0 ? ("Slot " + std::to_string(mPuritySlot[i])) : "Not configured";
+        std::string label = std::string("SelectionName___") + PuritySpeciesName[i] +
+                            ":::MinimalCut___0:::OptionalCut___0:::BitPosition___" + bitPosition +
+                            ":::Value___" + value;
+        hSlotAssignment->GetXaxis()->SetBinLabel(i + 1, label.c_str());
+        hSlotAssignment->Fill(i, 1);
+      }
+      hSlotAssignment->GetXaxis()->SetBinLabel(PuritySpeciesLast + 1, "All analyzed");
+      hSlotAssignment->GetXaxis()->SetBinLabel(PuritySpeciesLast + 2, "All passed");
+    }
 
     if (mProduceTracks && mProduceLiteTracks) {
       LOG(fatal) << "FTracks and FLiteTracks are mutually exclusive -- enable only one. "
@@ -660,7 +712,7 @@ class TrackBuilder
                  << "or switch to FTracks if CollisionTables.produceCollisions is enabled instead.";
     }
 
-    if (mProduceTracks || mProduceLiteTracks || mProduceTrackMasks || mProduceTrackMass || mProduceTrackDcas || mProduceTrackExtras || mProduceElectronPids || mProducePionPids || mProduceKaonPids || mProduceProtonPids || mProduceDeuteronPids || mProduceTritonPids || mProduceHeliumPids) {
+    if (mProduceTracks || mProduceLiteTracks || mProduceTrackMasks || mProduceTrackMass || mProduceTrackDcas || mProduceTrackExtras || mProduceElectronPids || mProducePionPids || mProduceKaonPids || mProduceProtonPids || mProduceDeuteronPids || mProduceTritonPids || mProduceHeliumPids || mProduceItsPurity || mProduceTpcPurity || mProduceTofPurity) {
       mFillAnyTable = true;
     } else {
       LOG(info) << "No tables configured, Selection object will not be configured...";
@@ -797,6 +849,41 @@ class TrackBuilder
       }
       trackProducts.producedHeliumPids(itsHe, track.tpcNSigmaHe(), track.tofNSigmaHe());
     }
+    if (mProduceItsPurity || mProduceTpcPurity || mProduceTofPurity) {
+      // raw nSigma per species, same values already used for the FXxxPids blocks above --
+      // unused slots (species without a requirePidXxx flag, or not covered by this track
+      // type) stay at the 0 placeholder, same contract as the missing-Tof case
+      std::array<float, PuritySpeciesLast> itsNSigma = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
+      std::array<float, PuritySpeciesLast> tpcNSigma = {track.tpcNSigmaEl(), track.tpcNSigmaPi(), track.tpcNSigmaKa(), track.tpcNSigmaPr(), track.tpcNSigmaDe(), track.tpcNSigmaTr(), track.tpcNSigmaHe()};
+      std::array<float, PuritySpeciesLast> tofNSigma = {track.tofNSigmaEl(), track.tofNSigmaPi(), track.tofNSigmaKa(), track.tofNSigmaPr(), track.tofNSigmaDe(), track.tofNSigmaTr(), track.tofNSigmaHe()};
+
+      if constexpr (requires(T1 t) { t.itsNSigmaEl(); }) {
+        itsNSigma = {track.itsNSigmaEl(), track.itsNSigmaPi(), track.itsNSigmaKa(), track.itsNSigmaPr(), track.itsNSigmaDe(), track.itsNSigmaTr(), track.itsNSigmaHe()};
+      }
+
+      std::array<float, utils::PurityNSlots> itsSlots = {0.f, 0.f, 0.f, 0.f};
+      std::array<float, utils::PurityNSlots> tpcSlots = {0.f, 0.f, 0.f, 0.f};
+      std::array<float, utils::PurityNSlots> tofSlots = {0.f, 0.f, 0.f, 0.f};
+      bool hasTof = track.hasTOF();
+      for (int i = 0; i < PuritySpeciesLast; i++) {
+        if (mPuritySlot[i] < 0) {
+          continue;
+        }
+        itsSlots[mPuritySlot[i]] = itsNSigma[i];
+        tpcSlots[mPuritySlot[i]] = tpcNSigma[i];
+        tofSlots[mPuritySlot[i]] = hasTof ? tofNSigma[i] : 0.f; // PurityHasTof is the gate; 0 here is just a placeholder
+      }
+
+      if (mProduceItsPurity) {
+        trackProducts.producedItsPurity(utils::packPurityNSigma4(itsSlots[0], itsSlots[1], itsSlots[2], itsSlots[3], o2::aod::femtotracks::NSigmaPurityMagMin, o2::aod::femtotracks::NSigmaPurityMagMax));
+      }
+      if (mProduceTpcPurity) {
+        trackProducts.producedTpcPurity(utils::packPurityNSigma4(tpcSlots[0], tpcSlots[1], tpcSlots[2], tpcSlots[3], o2::aod::femtotracks::NSigmaPurityMagMin, o2::aod::femtotracks::NSigmaPurityMagMax));
+      }
+      if (mProduceTofPurity) {
+        trackProducts.producedTofPurity(utils::packPurityNSigma4(tofSlots[0], tofSlots[1], tofSlots[2], tofSlots[3], o2::aod::femtotracks::NSigmaPurityMagMin, o2::aod::femtotracks::NSigmaPurityMagMax), hasTof);
+      }
+    }
     return true;
   }
 
@@ -892,6 +979,16 @@ class TrackBuilder
   bool mProduceDeuteronPids = false;
   bool mProduceTritonPids = false;
   bool mProduceHeliumPids = false;
+  bool mProduceItsPurity = false;
+  bool mProduceTpcPurity = false;
+  bool mProduceTofPurity = false;
+
+  // canonical species order for purity slot assignment -- fixed, not configurable (see
+  // init() for how requirePidXxx flags map onto slots 0-3)
+  static constexpr int PuritySpeciesLast = 7;
+  static constexpr std::array<const char*, PuritySpeciesLast> PuritySpeciesName = {"Electron", "Pion", "Kaon", "Proton", "Deuteron", "Triton", "Helium"};
+  // mPuritySlot[i] = assigned slot (0-3) for species i (PuritySpeciesName order), or -1 if not assigned
+  std::array<int, PuritySpeciesLast> mPuritySlot = {-1, -1, -1, -1, -1, -1, -1};
 
   std::unordered_map<int64_t, int64_t> indexMap; // for mapping tracks to daughters of lambdas, cascades and resonances ...
 };
@@ -900,6 +997,12 @@ struct TrackBuilderDerivedToDerivedProducts : o2::framework::ProducesGroup {
   o2::framework::Produces<o2::aod::StoredFTracks> producedTracks;
   o2::framework::Produces<o2::aod::StoredFTrackMasks> producedTrackMasks;
   o2::framework::Produces<o2::aod::StoredFTrackMass> producedTrackMass;
+  // purity/fraction, copied through if the input track type has them joined (same
+  // compile-time-detected, opt-in-by-type pattern as producedTrackMass above)
+  o2::framework::Produces<o2::aod::StoredFItsPurity> producedItsPurity;
+  o2::framework::Produces<o2::aod::StoredFTpcPurity> producedTpcPurity;
+  o2::framework::Produces<o2::aod::StoredFTofPurity> producedTofPurity;
+  o2::framework::Produces<o2::aod::StoredFTrackDcas> producedTrackDcas;
 };
 
 struct ConfTrackTablesDerivedToDerived : o2::framework::ConfigurableGroup {
@@ -977,6 +1080,14 @@ class TrackBuilderDerivedToDerived
     if constexpr (utils::HasMass<T1>) {
       trackProducts.producedTrackMass(track.mass());
     }
+    if constexpr (requires(T1 t) { t.itsPackedNSigma(); }) {
+      trackProducts.producedItsPurity(track.itsPackedNSigma());
+      trackProducts.producedTpcPurity(track.tpcPackedNSigma());
+      trackProducts.producedTofPurity(track.tofPackedNSigma(), track.purityHasTof());
+    }
+    if constexpr (requires(T1 t) { t.dcaXY(); }) {
+      trackProducts.producedTrackDcas(track.dcaXY(), track.dcaZ());
+    }
 
     const int64_t idx = trackProducts.producedTracks.lastIndex();
     indexMap.emplace(track.globalIndex(), idx);