Skip to content

Commit 349180c

Browse files
authored
[PWGDQ] Bug fix for mixing and add BDT score table (#16665)
1 parent 65a75bb commit 349180c

4 files changed

Lines changed: 124 additions & 75 deletions

File tree

PWGDQ/Core/MixingHandler.h

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
#include <Rtypes.h>
2626

27+
#include <algorithm>
2728
#include <array>
2829
#include <iostream>
2930
#include <map>
@@ -39,8 +40,8 @@ class MixingHandler : public TNamed
3940
float eta;
4041
float phi;
4142
uint32_t filteringFlags;
42-
// flip a bit to zero (needed when a track was already used in mixing for that bit for the required pool depth)
43-
void FlipBit(int64_t mask) { filteringFlags ^= mask; }
43+
// Clear a bit once the track was used in mixing for that bit for the required pool depth.
44+
void ClearBit(uint32_t mask) { filteringFlags &= ~mask; }
4445
void Print() const
4546
{
4647
std::cout << "pt: " << pt << ", eta: " << eta << ", phi: " << phi << ", filteringFlags: " << filteringFlags << std::endl;
@@ -68,32 +69,40 @@ class MixingHandler : public TNamed
6869
tracks2.push_back(track);
6970
filteringMask |= track.filteringFlags;
7071
}
71-
// flip bits in the filtering mask
72-
void FlipFilteringMask(int64_t mask) { filteringMask ^= mask; }
72+
// Clear bits in the filtering mask.
73+
void ClearFilteringMask(uint32_t mask) { filteringMask &= ~mask; }
7374
// 1) increment the counters for a given track cut bit mask and if the counters reached the pool depth,
74-
// 2) flip the corresponding bit in the tracks filtering flags to exclude them from further mixing
75+
// 2) clear the corresponding bit in the tracks filtering flags to exclude them from further mixing
7576
// 3) for each track, if there are no more active bits in the filtering mask, then remove the track from the event
7677
void IncrementCounters(uint32_t mask, short poolDepth)
7778
{
7879
for (int i = 0; i < 32; i++) {
79-
if (mask & (1ULL << i)) {
80+
uint32_t bitMask = (static_cast<uint32_t>(1) << i);
81+
if (mask & bitMask) {
8082
counters[i]++;
8183
if (counters[i] >= poolDepth) {
8284
for (auto& track : tracks1) {
83-
track.FlipBit(1ULL << i);
84-
if (track.filteringFlags == 0) {
85-
track = tracks1.back();
86-
tracks1.pop_back();
85+
track.ClearBit(bitMask);
86+
}
87+
for (auto track = tracks1.begin(); track != tracks1.end();) {
88+
if (track->filteringFlags == 0) {
89+
track = tracks1.erase(track);
90+
} else {
91+
++track;
8792
}
8893
}
94+
8995
for (auto& track : tracks2) {
90-
track.FlipBit(1ULL << i);
91-
if (track.filteringFlags == 0) {
92-
track = tracks2.back();
93-
tracks2.pop_back();
96+
track.ClearBit(bitMask);
97+
}
98+
for (auto track = tracks2.begin(); track != tracks2.end();) {
99+
if (track->filteringFlags == 0) {
100+
track = tracks2.erase(track);
101+
} else {
102+
++track;
94103
}
95104
}
96-
FlipFilteringMask(1ULL << i);
105+
ClearFilteringMask(bitMask);
97106
}
98107
}
99108
}
@@ -131,9 +140,13 @@ class MixingHandler : public TNamed
131140
// check which events in the pool are empty (i.e. no active tracks for mixing) and remove them from the pool
132141
void CleanPool()
133142
{
134-
events.erase(std::remove_if(events.begin(), events.end(),
135-
[](const MixingEvent& event) { return event.tracks1.empty() && event.tracks2.empty(); }),
136-
events.end());
143+
for (auto event = events.begin(); event != events.end();) {
144+
if (event->tracks1.empty() && event->tracks2.empty()) {
145+
event = events.erase(event);
146+
} else {
147+
++event;
148+
}
149+
}
137150
}
138151
// The function that performs the mixing is called outside this class, but the pool provides the events and tracks to be mixed and takes care of updating the events after mixing
139152
// (e.g. incrementing the counters and removing the tracks that reached the pool depth for a given cut)

PWGDQ/DataModel/ReducedInfoTables.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -826,6 +826,9 @@ DECLARE_SOA_COLUMN(PairDCAxyz, pairDCAxyz, float);
826826
DECLARE_SOA_COLUMN(PairDCAxy, pairDCAxy, float); //! Pair DCAxy to PV from KFParticle
827827
DECLARE_SOA_COLUMN(DeviationPairKF, deviationPairKF, float); //! Pair chi2 deviation to PV from KFParticle
828828
DECLARE_SOA_COLUMN(DeviationxyPairKF, deviationxyPairKF, float); //! Pair chi2 deviation to PV in XY from KFParticle
829+
DECLARE_SOA_COLUMN(BdtBackground, bdtBackground, float); //! BDT output score for the background class
830+
DECLARE_SOA_COLUMN(BdtPrompt, bdtPrompt, float); //! BDT output score for the prompt class
831+
DECLARE_SOA_COLUMN(BdtNonprompt, bdtNonprompt, float); //! BDT output score for the nonprompt class
829832
// DECLARE_SOA_INDEX_COLUMN(ReducedMuon, reducedmuon2); //!
830833
DECLARE_SOA_COLUMN(CosThetaHE, costhetaHE, float); //! Cosine in the helicity frame
831834
DECLARE_SOA_COLUMN(PhiHE, phiHe, float); //! Phi in the helicity frame
@@ -936,6 +939,12 @@ DECLARE_SOA_TABLE_STAGED(DielectronsAll, "RTDIELECTRONALL", //!
936939
reducedpair::Lz,
937940
reducedpair::Lxy);
938941

942+
DECLARE_SOA_TABLE_STAGED(DielectronsMls, "RTDIELECTRONML", //!
943+
reducedpair::CentFT0C,
944+
reducedpair::BdtBackground,
945+
reducedpair::BdtPrompt,
946+
reducedpair::BdtNonprompt);
947+
939948
DECLARE_SOA_TABLE(DimuonsAll, "AOD", "RTDIMUONALL", //!
940949
collision::PosX, collision::PosY, collision::PosZ, collision::NumContrib,
941950
evsel::Selection, reducedpair::EventSelection,
@@ -1010,6 +1019,7 @@ using DimuonExtra = DimuonsExtra::iterator;
10101019
using DileptonFlow = DileptonsFlow::iterator;
10111020
using DileptonInfo = DileptonsInfo::iterator;
10121021
using DielectronAll = DielectronsAll::iterator;
1022+
using DielectronMl = DielectronsMls::iterator;
10131023
using DimuonAll = DimuonsAll::iterator;
10141024
using DileptonMiniTree = DileptonsMiniTree::iterator;
10151025
using DileptonMiniTreeGen = DileptonsMiniTreeGen::iterator;

PWGDQ/Tasks/tableReader.cxx

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1026,6 +1026,7 @@ struct AnalysisSameEventPairing {
10261026
Produces<aod::DielectronsInfo> dielectronInfoList;
10271027
Produces<aod::DimuonsExtra> dimuonExtraList;
10281028
Produces<aod::DielectronsAll> dielectronAllList;
1029+
Produces<aod::DielectronsMls> dielectronMlList;
10291030
Produces<aod::DimuonsAll> dimuonAllList;
10301031
Produces<aod::DileptonFlow> dileptonFlowList;
10311032
Produces<aod::DileptonsInfo> dileptonInfoList;
@@ -1380,6 +1381,7 @@ struct AnalysisSameEventPairing {
13801381
dileptonFlowList.reserve(1);
13811382
if (fConfigFlatTables.value) {
13821383
dielectronAllList.reserve(1);
1384+
dielectronMlList.reserve(1);
13831385
dimuonAllList.reserve(1);
13841386
}
13851387
if (useMiniTree.fConfigMiniTree) {
@@ -1466,6 +1468,8 @@ struct AnalysisSameEventPairing {
14661468
}
14671469
}
14681470
if constexpr ((TPairType == pairTypeEE) && (TTrackFillMap & VarManager::ObjTypes::ReducedTrackBarrelPID) > 0) {
1471+
isSelectedBDT = false;
1472+
fOutputMlPsi2ee.clear();
14691473
if (applyBDT) {
14701474
std::vector<float> dqInputFeatures = fDQMlResponse.getInputFeatures(t1, t2, VarManager::fgValues);
14711475

@@ -1497,7 +1501,7 @@ struct AnalysisSameEventPairing {
14971501

14981502
LOG(debug) << "Model index: " << modelIndex << ", pT: " << VarManager::fgValues[VarManager::kPt] << ", centrality (kCentFT0C): " << VarManager::fgValues[VarManager::kCentFT0C];
14991503
isSelectedBDT = fDQMlResponse.isSelectedMl(dqInputFeatures, modelIndex, fOutputMlPsi2ee);
1500-
VarManager::FillBdtScore(fOutputMlPsi2ee); // TODO: check if this is needed or not
1504+
VarManager::FillBdtScore(fOutputMlPsi2ee);
15011505
}
15021506

15031507
if (applyBDT && !isSelectedBDT)
@@ -1512,6 +1516,10 @@ struct AnalysisSameEventPairing {
15121516
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],
15131517
VarManager::fgValues[VarManager::kKFPairDeviationFromPV], VarManager::fgValues[VarManager::kKFPairDeviationxyFromPV],
15141518
VarManager::fgValues[VarManager::kKFMassGeoTop], VarManager::fgValues[VarManager::kKFChi2OverNDFGeoTop], VarManager::fgValues[VarManager::kVertexingTauzProjected], VarManager::fgValues[VarManager::kVertexingTauxyProjected], VarManager::fgValues[VarManager::kVertexingLzProjected], VarManager::fgValues[VarManager::kVertexingLxyProjected]);
1519+
dielectronMlList(VarManager::fgValues[VarManager::kCentFT0C],
1520+
(applyBDT && fOutputMlPsi2ee.size() > 0) ? VarManager::fgValues[VarManager::kBdtBackground] : -999.f,
1521+
(applyBDT && fOutputMlPsi2ee.size() > 1) ? VarManager::fgValues[VarManager::kBdtPrompt] : -999.f,
1522+
(applyBDT && fOutputMlPsi2ee.size() > 2) ? VarManager::fgValues[VarManager::kBdtNonprompt] : -999.f);
15151523
}
15161524
}
15171525
if constexpr (TPairType == pairTypeMuMu) {

0 commit comments

Comments
 (0)