Skip to content

Commit 1c73b3b

Browse files
iarseneIonut Cristian Arsenealibuild
authored
[PWGDQ] Adding functionality for event mixing across TFs (#16528)
Co-authored-by: Ionut Cristian Arsene <iarsene@cern.ch> Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
1 parent feb2781 commit 1c73b3b

5 files changed

Lines changed: 359 additions & 103 deletions

File tree

PWGDQ/Core/MixingHandler.cxx

Lines changed: 39 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,11 @@ ClassImp(MixingHandler);
2727

2828
//_________________________________________________________________________
2929
MixingHandler::MixingHandler() : TNamed(),
30-
fIsInitialized(kFALSE),
30+
fIsInitialized(false),
3131
fVariableLimits(),
32-
fVariables()
32+
fVariables(),
33+
fPoolDepth(0),
34+
fPools()
3335
{
3436
//
3537
// default constructor
@@ -38,9 +40,11 @@ MixingHandler::MixingHandler() : TNamed(),
3840

3941
//_________________________________________________________________________
4042
MixingHandler::MixingHandler(const char* name, const char* title) : TNamed(name, title),
41-
fIsInitialized(kFALSE),
43+
fIsInitialized(false),
4244
fVariableLimits(),
43-
fVariables()
45+
fVariables(),
46+
fPoolDepth(0),
47+
fPools()
4448
{
4549
//
4650
// Named constructor
@@ -56,29 +60,13 @@ MixingHandler::~MixingHandler()
5660
}
5761

5862
//_________________________________________________________________________
59-
void MixingHandler::AddMixingVariable(int var, int nBins, float* binLims)
63+
void MixingHandler::AddMixingVariable(int var, std::vector<float> binLims)
6064
{
61-
//
62-
// add a mixing variable
63-
//
64-
fVariables.push_back(var);
65-
TArrayF varBins;
66-
varBins.Set(nBins, binLims);
67-
fVariableLimits.push_back(varBins);
68-
VarManager::SetUseVariable(var);
69-
}
70-
71-
//_________________________________________________________________________
72-
void MixingHandler::AddMixingVariable(int var, int nBins, std::vector<float> binLims)
73-
{
74-
75-
float* bins = new float[nBins];
76-
for (int i = 0; i < nBins; ++i) {
77-
bins[i] = binLims[i];
78-
}
79-
AddMixingVariable(var, nBins, bins);
65+
fVariables[var] = fVariableLimits.size();
66+
fVariableLimits.push_back(binLims);
8067
}
8168

69+
/*
8270
//_________________________________________________________________________
8371
int MixingHandler::GetMixingVariable(VarManager::Variables var)
8472
{
@@ -90,7 +78,9 @@ int MixingHandler::GetMixingVariable(VarManager::Variables var)
9078
}
9179
return -1;
9280
}
81+
*/
9382

83+
/*
9484
//_________________________________________________________________________
9585
std::vector<float> MixingHandler::GetMixingVariableLimits(VarManager::Variables var)
9686
{
@@ -105,21 +95,21 @@ std::vector<float> MixingHandler::GetMixingVariableLimits(VarManager::Variables
10595
}
10696
}
10797
return binLimits;
108-
}
98+
}*/
10999

110100
//_________________________________________________________________________
111101
void MixingHandler::Init()
112102
{
113-
//
114-
// Initialization of pools
115-
// The correct event category will be retrieved using the function FindEventCategory()
116-
//
117-
int size = 1;
118-
for (auto v : fVariableLimits) {
119-
size *= (v.GetSize() - 1);
103+
// loop over all variables and create a mixing pool for each category defined by the binning of the variables
104+
int nCategories = 1;
105+
for (auto& var : fVariables) {
106+
nCategories *= (fVariableLimits[var.second].size() - 1);
107+
}
108+
// add elements in the map for each category (the key is the category and the value is an empty pool)
109+
for (int i = 0; i < nCategories; i++) {
110+
fPools[i] = MixingPool();
120111
}
121-
(void)size;
122-
fIsInitialized = kTRUE;
112+
fIsInitialized = true;
123113
}
124114

125115
//_________________________________________________________________________
@@ -135,16 +125,21 @@ int MixingHandler::FindEventCategory(float* values)
135125
Init();
136126
}
137127

128+
// loop over the variables and find out in which bin the value of the variable for the event is located
138129
std::vector<int> bin;
139-
int iVar = 0;
140-
for (auto v = fVariableLimits.begin(); v != fVariableLimits.end(); v++, iVar++) {
141-
int binValue = TMath::BinarySearch((*v).GetSize(), (*v).GetArray(), values[fVariables[iVar]]);
142-
bin.push_back(binValue);
143-
if (bin[iVar] == -1 || bin[iVar] == (*v).GetSize() - 1) {
130+
for (auto [var, pos] : fVariables) {
131+
// check that the value is within limits, if not return -1 to exclude the event from mixing
132+
size_t binValue = std::distance(fVariableLimits[pos].begin(), std::upper_bound(fVariableLimits[pos].begin(), fVariableLimits[pos].end(), values[var]));
133+
if (binValue == 0 || binValue == fVariableLimits[pos].size()) {
144134
return -1; // all variables must be inside limits
145135
}
136+
bin.push_back(binValue - 1);
146137
}
147138

139+
// Hash the bin values to define a unique category
140+
// The hashing is done such that the original bin values can be retrieved from the category
141+
// For example, for 3 variables with n1, n2, n3 bins respectively, the category for bin values (b1, b2, b3) would be:
142+
// category = b1*(n2*n3) + b2*(n3) + b3
148143
int category = 0;
149144
int tempCategory = 1;
150145
int iv1 = 0;
@@ -156,7 +151,7 @@ int MixingHandler::FindEventCategory(float* values)
156151
if (iv2 == iv1) {
157152
tempCategory *= bin[iv2];
158153
} else {
159-
tempCategory *= (fVariableLimits[iv2].GetSize() - 1);
154+
tempCategory *= (fVariableLimits[iv2].size() - 1);
160155
}
161156
}
162157
category += tempCategory;
@@ -175,19 +170,14 @@ int MixingHandler::GetBinFromCategory(VarManager::Variables var, int category) c
175170
}
176171

177172
// Search for the position of the variable "var" in the internal variable list of the handler
178-
int tempVar = 0;
179-
for (auto v = fVariables.begin(); v != fVariables.end(); v++, tempVar++) {
180-
if (*v == var) {
181-
break;
182-
}
183-
}
173+
int ivar = fVariables.at(var);
184174

185175
// extract the bin position in variable "var" from the category
186176
int norm = 1;
187-
for (int i = fVariables.size() - 1; i > tempVar; --i) {
188-
norm *= (fVariableLimits[i].GetSize() - 1);
177+
for (int i = fVariables.size() - 1; i > ivar; --i) {
178+
norm *= (fVariableLimits[i].size() - 1);
189179
}
190180
int truncatedCategory = category - (category % norm);
191181
truncatedCategory /= norm;
192-
return truncatedCategory % (fVariableLimits[tempVar].GetSize() - 1);
182+
return truncatedCategory % (fVariableLimits[ivar].size() - 1);
193183
}

PWGDQ/Core/MixingHandler.h

Lines changed: 141 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,24 +24,153 @@
2424

2525
#include <Rtypes.h>
2626

27+
#include <array>
28+
#include <iostream>
29+
#include <map>
2730
#include <vector>
2831

2932
class MixingHandler : public TNamed
3033
{
3134

3235
public:
36+
// Struct to define track properties relevant for mixing and few utility functions
37+
struct MixingTrack {
38+
float pt;
39+
float eta;
40+
float phi;
41+
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; }
44+
void Print() const
45+
{
46+
std::cout << "pt: " << pt << ", eta: " << eta << ", phi: " << phi << ", filteringFlags: " << filteringFlags << std::endl;
47+
}
48+
};
49+
50+
// Struct to define events used in mixing and few utility functions.
51+
// An event is defined as two vectors of tracks (typically the legs of a two-body
52+
// decay or the two-particles in a correlation analysis)
53+
struct MixingEvent {
54+
std::vector<MixingTrack> tracks1;
55+
std::vector<MixingTrack> tracks2;
56+
// bit map for active filtering bits of all the tracks
57+
uint32_t filteringMask = 0;
58+
// counters to keep track of how many times the event was used for mixing (for each track cut separately)
59+
std::array<short, 64> counters = {0};
60+
// add a track to the event and update the filtering mask accordingly
61+
void AddTrack1(const MixingTrack& track)
62+
{
63+
tracks1.push_back(track);
64+
filteringMask |= track.filteringFlags;
65+
}
66+
void AddTrack2(const MixingTrack& track)
67+
{
68+
tracks2.push_back(track);
69+
filteringMask |= track.filteringFlags;
70+
}
71+
// flip bits in the filtering mask
72+
void FlipFilteringMask(int64_t mask) { filteringMask ^= mask; }
73+
// 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+
// 3) for each track, if there are no more active bits in the filtering mask, then remove the track from the event
76+
void IncrementCounters(uint32_t mask, short poolDepth)
77+
{
78+
for (int i = 0; i < 32; i++) {
79+
if (mask & (1ULL << i)) {
80+
counters[i]++;
81+
if (counters[i] >= poolDepth) {
82+
for (auto& track : tracks1) {
83+
track.FlipBit(1ULL << i);
84+
if (track.filteringFlags == 0) {
85+
track = tracks1.back();
86+
tracks1.pop_back();
87+
}
88+
}
89+
for (auto& track : tracks2) {
90+
track.FlipBit(1ULL << i);
91+
if (track.filteringFlags == 0) {
92+
track = tracks2.back();
93+
tracks2.pop_back();
94+
}
95+
}
96+
FlipFilteringMask(1ULL << i);
97+
}
98+
}
99+
}
100+
}
101+
void Print() const
102+
{
103+
std::cout << "Event filtering mask: ";
104+
for (int i = 0; i < 32; i++) {
105+
if (filteringMask & (1ULL << i)) {
106+
std::cout << "1";
107+
} else {
108+
std::cout << "0";
109+
}
110+
}
111+
std::cout << std::endl;
112+
for (int i = 0; i < 32; i++) {
113+
if (filteringMask & (1ULL << i)) {
114+
std::cout << "Counter " << i << ": " << counters[i] << std::endl;
115+
}
116+
}
117+
std::cout << "Tracks 1: " << std::endl;
118+
for (const auto& track : tracks1) {
119+
track.Print();
120+
}
121+
std::cout << "Tracks 2: " << std::endl;
122+
for (const auto& track : tracks2) {
123+
track.Print();
124+
}
125+
}
126+
};
127+
128+
struct MixingPool {
129+
std::vector<MixingEvent> events;
130+
131+
// check which events in the pool are empty (i.e. no active tracks for mixing) and remove them from the pool
132+
void CleanPool()
133+
{
134+
events.erase(std::remove_if(events.begin(), events.end(),
135+
[](const MixingEvent& event) { return event.tracks1.empty() && event.tracks2.empty(); }),
136+
events.end());
137+
}
138+
// 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
139+
// (e.g. incrementing the counters and removing the tracks that reached the pool depth for a given cut)
140+
void UpdatePool(const MixingEvent& event, short poolDepth)
141+
{
142+
for (auto& event : events) {
143+
event.IncrementCounters(event.filteringMask, poolDepth);
144+
}
145+
CleanPool();
146+
events.push_back(event);
147+
}
148+
// getter for the events in the pool
149+
const std::vector<MixingEvent>& GetEvents() const { return events; }
150+
151+
void Print() const
152+
{
153+
std::cout << "Mixing pool with " << events.size() << " events:" << std::endl;
154+
for (const auto& event : events) {
155+
event.Print();
156+
}
157+
}
158+
};
159+
33160
MixingHandler();
34161
MixingHandler(const char* name, const char* title);
35162
virtual ~MixingHandler();
36163

37164
// setters
38-
void AddMixingVariable(int var, int nBins, float* binLims);
39-
void AddMixingVariable(int var, int nBins, std::vector<float> binLims);
165+
void AddMixingVariable(int var, std::vector<float> binLims);
166+
void SetPoolDepth(short depth) { fPoolDepth = depth; }
40167

41168
// getters
42-
int GetNMixingVariables() const { return fVariables.size(); }
43-
int GetMixingVariable(VarManager::Variables var); // returns the position in the internal varible list of the handler. Useful for checks, mostly
44-
std::vector<float> GetMixingVariableLimits(VarManager::Variables var);
169+
// int GetNMixingVariables() const { return fVariables.size(); }
170+
// int GetMixingVariable(VarManager::Variables var); // returns the position in the internal varible list of the handler. Useful for checks, mostly
171+
// std::vector<float> GetMixingVariableLimits(VarManager::Variables var);
172+
MixingPool& GetPool(int category) { return fPools[category]; }
173+
short GetPoolDepth() const { return fPoolDepth; }
45174

46175
void Init();
47176
int FindEventCategory(float* values);
@@ -54,10 +183,14 @@ class MixingHandler : public TNamed
54183
// User options
55184
bool fIsInitialized; // check if the mixing handler is initialized
56185

57-
std::vector<TArrayF> fVariableLimits;
58-
std::vector<int> fVariables;
186+
// bin limits for the variables used for mixing, the number of vectors corresponds to the number of variables and the content of each vector corresponds to the bin limits for that variable
187+
std::vector<std::vector<float>> fVariableLimits;
188+
std::map<int, int> fVariables; // key: variable, value: position in the internal variable list of the handler (used to map the variables to the values passed to FindEventCategory)
189+
190+
short fPoolDepth; // number of events to be kept in each pool
191+
std::map<int, MixingPool> fPools; // key: category, value: pool of events corresponding to that category
59192

60-
ClassDef(MixingHandler, 1);
193+
ClassDef(MixingHandler, 2);
61194
};
62195

63196
#endif // PWGDQ_CORE_MIXINGHANDLER_H_

0 commit comments

Comments
 (0)