Bug fix for mixing and add BDT score table

JinjooSeo · JinjooSeo · commit 43acec20b60f · 2026-06-15T13:05:52.000+02:00
diff --git a/PWGDQ/Core/MixingHandler.h b/PWGDQ/Core/MixingHandler.h
@@ -24,6 +24,7 @@
 
 #include <Rtypes.h>
 
+#include <algorithm>
 #include <array>
 #include <iostream>
 #include <map>
@@ -39,8 +40,8 @@ class MixingHandler : public TNamed
     float eta;
     float phi;
     uint32_t filteringFlags;
-    // flip a bit to zero (needed when a track was already used in mixing for that bit for the required pool depth)
-    void FlipBit(int64_t mask) { filteringFlags ^= mask; }
+    // Clear a bit once the track was used in mixing for that bit for the required pool depth.
+    void ClearBit(uint32_t mask) { filteringFlags &= ~mask; }
     void Print() const
     {
       std::cout << "pt: " << pt << ", eta: " << eta << ", phi: " << phi << ", filteringFlags: " << filteringFlags << std::endl;
@@ -68,32 +69,29 @@ class MixingHandler : public TNamed
       tracks2.push_back(track);
       filteringMask |= track.filteringFlags;
     }
-    // flip bits in the filtering mask
-    void FlipFilteringMask(int64_t mask) { filteringMask ^= mask; }
+    // Clear bits in the filtering mask.
+    void ClearFilteringMask(uint32_t mask) { filteringMask &= ~mask; }
     // 1) increment the counters for a given track cut bit mask and if the counters reached the pool depth,
-    // 2) flip the corresponding bit in the tracks filtering flags to exclude them from further mixing
+    // 2) clear the corresponding bit in the tracks filtering flags to exclude them from further mixing
     // 3) for each track, if there are no more active bits in the filtering mask, then remove the track from the event
     void IncrementCounters(uint32_t mask, short poolDepth)
     {
       for (int i = 0; i < 32; i++) {
-        if (mask & (1ULL << i)) {
+        uint32_t bitMask = (static_cast<uint32_t>(1) << i);
+        if (mask & bitMask) {
           counters[i]++;
           if (counters[i] >= poolDepth) {
-            for (auto& track : tracks1) {
-              track.FlipBit(1ULL << i);
-              if (track.filteringFlags == 0) {
-                track = tracks1.back();
-                tracks1.pop_back();
+            auto clearExpiredBit = [bitMask](auto& tracks) {
+              for (auto& track : tracks) {
+                track.ClearBit(bitMask);
               }
-            }
-            for (auto& track : tracks2) {
-              track.FlipBit(1ULL << i);
-              if (track.filteringFlags == 0) {
-                track = tracks2.back();
-                tracks2.pop_back();
-              }
-            }
-            FlipFilteringMask(1ULL << i);
+              tracks.erase(std::remove_if(tracks.begin(), tracks.end(),
+                                          [](const MixingTrack& track) { return track.filteringFlags == 0; }),
+                           tracks.end());
+            };
+            clearExpiredBit(tracks1);
+            clearExpiredBit(tracks2);
+            ClearFilteringMask(bitMask);
           }
         }
       }
diff --git a/PWGDQ/DataModel/ReducedInfoTables.h b/PWGDQ/DataModel/ReducedInfoTables.h
@@ -826,6 +826,9 @@ DECLARE_SOA_COLUMN(PairDCAxyz, pairDCAxyz, float);
 DECLARE_SOA_COLUMN(PairDCAxy, pairDCAxy, float);                                         //! Pair DCAxy to PV from KFParticle
 DECLARE_SOA_COLUMN(DeviationPairKF, deviationPairKF, float);                             //! Pair chi2 deviation to PV from KFParticle
 DECLARE_SOA_COLUMN(DeviationxyPairKF, deviationxyPairKF, float);                         //! Pair chi2 deviation to PV in XY from KFParticle
+DECLARE_SOA_COLUMN(BdtBackground, bdtBackground, float);                                 //! BDT output score for the background class
+DECLARE_SOA_COLUMN(BdtPrompt, bdtPrompt, float);                                         //! BDT output score for the prompt class
+DECLARE_SOA_COLUMN(BdtNonprompt, bdtNonprompt, float);                                   //! BDT output score for the nonprompt class
 // DECLARE_SOA_INDEX_COLUMN(ReducedMuon, reducedmuon2); //!
 DECLARE_SOA_COLUMN(CosThetaHE, costhetaHE, float);             //! Cosine in the helicity frame
 DECLARE_SOA_COLUMN(PhiHE, phiHe, float);                       //! Phi in the helicity frame
@@ -936,6 +939,12 @@ DECLARE_SOA_TABLE_STAGED(DielectronsAll, "RTDIELECTRONALL", //!
                          reducedpair::Lz,
                          reducedpair::Lxy);
 
+DECLARE_SOA_TABLE_STAGED(DielectronsMls, "RTDIELECTRONML", //!
+                         reducedpair::CentFT0C,
+                         reducedpair::BdtBackground,
+                         reducedpair::BdtPrompt,
+                         reducedpair::BdtNonprompt);
+
 DECLARE_SOA_TABLE(DimuonsAll, "AOD", "RTDIMUONALL", //!
                   collision::PosX, collision::PosY, collision::PosZ, collision::NumContrib,
                   evsel::Selection, reducedpair::EventSelection,
@@ -1010,6 +1019,7 @@ using DimuonExtra = DimuonsExtra::iterator;
 using DileptonFlow = DileptonsFlow::iterator;
 using DileptonInfo = DileptonsInfo::iterator;
 using DielectronAll = DielectronsAll::iterator;
+using DielectronMl = DielectronsMls::iterator;
 using DimuonAll = DimuonsAll::iterator;
 using DileptonMiniTree = DileptonsMiniTree::iterator;
 using DileptonMiniTreeGen = DileptonsMiniTreeGen::iterator;
diff --git a/PWGDQ/Tasks/tableReader.cxx b/PWGDQ/Tasks/tableReader.cxx
@@ -1026,6 +1026,7 @@ struct AnalysisSameEventPairing {
   Produces<aod::DielectronsInfo> dielectronInfoList;
   Produces<aod::DimuonsExtra> dimuonExtraList;
   Produces<aod::DielectronsAll> dielectronAllList;
+  Produces<aod::DielectronsMls> dielectronMlList;
   Produces<aod::DimuonsAll> dimuonAllList;
   Produces<aod::DileptonFlow> dileptonFlowList;
   Produces<aod::DileptonsInfo> dileptonInfoList;
@@ -1380,6 +1381,7 @@ struct AnalysisSameEventPairing {
     dileptonFlowList.reserve(1);
     if (fConfigFlatTables.value) {
       dielectronAllList.reserve(1);
+      dielectronMlList.reserve(1);
       dimuonAllList.reserve(1);
     }
     if (useMiniTree.fConfigMiniTree) {
@@ -1466,6 +1468,8 @@ struct AnalysisSameEventPairing {
         }
       }
       if constexpr ((TPairType == pairTypeEE) && (TTrackFillMap & VarManager::ObjTypes::ReducedTrackBarrelPID) > 0) {
+        isSelectedBDT = false;
+        fOutputMlPsi2ee.clear();
         if (applyBDT) {
           std::vector<float> dqInputFeatures = fDQMlResponse.getInputFeatures(t1, t2, VarManager::fgValues);
 
@@ -1497,7 +1501,7 @@ struct AnalysisSameEventPairing {
 
           LOG(debug) << "Model index: " << modelIndex << ", pT: " << VarManager::fgValues[VarManager::kPt] << ", centrality (kCentFT0C): " << VarManager::fgValues[VarManager::kCentFT0C];
           isSelectedBDT = fDQMlResponse.isSelectedMl(dqInputFeatures, modelIndex, fOutputMlPsi2ee);
-          VarManager::FillBdtScore(fOutputMlPsi2ee); // TODO: check if this is needed or not
+          VarManager::FillBdtScore(fOutputMlPsi2ee);
         }
 
         if (applyBDT && !isSelectedBDT)
@@ -1512,6 +1516,10 @@ struct AnalysisSameEventPairing {
                             VarManager::fgValues[VarManager::kKFMass], VarManager::fgValues[VarManager::kKFChi2OverNDFGeo], VarManager::fgValues[VarManager::kVertexingLxyz], VarManager::fgValues[VarManager::kVertexingLxyzOverErr], VarManager::fgValues[VarManager::kVertexingLxy], VarManager::fgValues[VarManager::kVertexingLxyOverErr], VarManager::fgValues[VarManager::kVertexingTauxy], VarManager::fgValues[VarManager::kVertexingTauxyErr], VarManager::fgValues[VarManager::kKFCosPA], VarManager::fgValues[VarManager::kKFJpsiDCAxyz], VarManager::fgValues[VarManager::kKFJpsiDCAxy],
                             VarManager::fgValues[VarManager::kKFPairDeviationFromPV], VarManager::fgValues[VarManager::kKFPairDeviationxyFromPV],
                             VarManager::fgValues[VarManager::kKFMassGeoTop], VarManager::fgValues[VarManager::kKFChi2OverNDFGeoTop], VarManager::fgValues[VarManager::kVertexingTauzProjected], VarManager::fgValues[VarManager::kVertexingTauxyProjected], VarManager::fgValues[VarManager::kVertexingLzProjected], VarManager::fgValues[VarManager::kVertexingLxyProjected]);
+          dielectronMlList(VarManager::fgValues[VarManager::kCentFT0C],
+                           (applyBDT && fOutputMlPsi2ee.size() > 0) ? VarManager::fgValues[VarManager::kBdtBackground] : -999.f,
+                           (applyBDT && fOutputMlPsi2ee.size() > 1) ? VarManager::fgValues[VarManager::kBdtPrompt] : -999.f,
+                           (applyBDT && fOutputMlPsi2ee.size() > 2) ? VarManager::fgValues[VarManager::kBdtNonprompt] : -999.f);
         }
       }
       if constexpr (TPairType == pairTypeMuMu) {
diff --git a/PWGDQ/Tasks/tableReader_withAssoc.cxx b/PWGDQ/Tasks/tableReader_withAssoc.cxx
