[PWGDQ] Adding functionality for event mixing across TFs

PWGDQ/Core/MixingHandler.cxx

@@ -27,9 +27,11 @@
 
 //_________________________________________________________________________
 MixingHandler::MixingHandler() : TNamed(),
-                                 fIsInitialized(kFALSE),
+                                 fIsInitialized(false),
                                  fVariableLimits(),
-                                 fVariables()
+                                 fVariables(),
+                                 fPoolDepth(0),
+                                 fPools()
 {
   //
   // default constructor
@@ -38,9 +40,11 @@
 
 //_________________________________________________________________________
 MixingHandler::MixingHandler(const char* name, const char* title) : TNamed(name, title),
-                                                                    fIsInitialized(kFALSE),
+                                                                    fIsInitialized(false),
                                                                     fVariableLimits(),
-                                                                    fVariables()
+                                                                    fVariables(),
+                                                                    fPoolDepth(0),
+                                                                    fPools()
 {
   //
   // Named constructor
@@ -56,29 +60,13 @@
 }
 
 //_________________________________________________________________________
-void MixingHandler::AddMixingVariable(int var, int nBins, float* binLims)
+void MixingHandler::AddMixingVariable(int var, std::vector<float> binLims)
 {
-  //
-  // add a mixing variable
-  //
-  fVariables.push_back(var);
-  TArrayF varBins;
-  varBins.Set(nBins, binLims);
-  fVariableLimits.push_back(varBins);
-  VarManager::SetUseVariable(var);
-}
-
-//_________________________________________________________________________
-void MixingHandler::AddMixingVariable(int var, int nBins, std::vector<float> binLims)
-{
-
-  float* bins = new float[nBins];
-  for (int i = 0; i < nBins; ++i) {
-    bins[i] = binLims[i];
-  }
-  AddMixingVariable(var, nBins, bins);
+  fVariables[var] = fVariableLimits.size();
+  fVariableLimits.push_back(binLims);
 }
 
+/*
 //_________________________________________________________________________
 int MixingHandler::GetMixingVariable(VarManager::Variables var)
 {
@@ -90,7 +78,9 @@
   }
   return -1;
 }
+*/
 
+/*
 //_________________________________________________________________________
 std::vector<float> MixingHandler::GetMixingVariableLimits(VarManager::Variables var)
 {
@@ -105,21 +95,21 @@
     }
   }
   return binLimits;
-}
+}*/
 
 //_________________________________________________________________________
 void MixingHandler::Init()
 {
-  //
-  // Initialization of pools
-  //       The correct event category will be retrieved using the function FindEventCategory()
-  //
-  int size = 1;
-  for (auto v : fVariableLimits) {
-    size *= (v.GetSize() - 1);
+  // loop over all variables and create a mixing pool for each category defined by the binning of the variables
+  int nCategories = 1;
+  for (auto& var : fVariables) {
+    nCategories *= (fVariableLimits[var.second].size() - 1);
+  }
+  // add elements in the map for each category (the key is the category and the value is an empty pool)
+  for (int i = 0; i < nCategories; i++) {
+    fPools[i] = MixingPool();
   }
-  (void)size;
-  fIsInitialized = kTRUE;
+  fIsInitialized = true;
 }
 
 //_________________________________________________________________________
@@ -135,16 +125,21 @@
     Init();
   }
 
+  // loop over the variables and find out in which bin the value of the variable for the event is located
   std::vector<int> bin;
-  int iVar = 0;
-  for (auto v = fVariableLimits.begin(); v != fVariableLimits.end(); v++, iVar++) {
-    int binValue = TMath::BinarySearch((*v).GetSize(), (*v).GetArray(), values[fVariables[iVar]]);
-    bin.push_back(binValue);
-    if (bin[iVar] == -1 || bin[iVar] == (*v).GetSize() - 1) {
+  for (auto [var, pos] : fVariables) {
+    // check that the value is within limits, if not return -1 to exclude the event from mixing
+    size_t binValue = std::distance(fVariableLimits[pos].begin(), std::upper_bound(fVariableLimits[pos].begin(), fVariableLimits[pos].end(), values[var]));
+    if (binValue == 0 || binValue == fVariableLimits[pos].size()) {
       return -1; // all variables must be inside limits
     }
+    bin.push_back(binValue - 1);
   }
 
+  // Hash the bin values to define a unique category
+  // The hashing is done such that the original bin values can be retrieved from the category
+  // For example, for 3 variables with n1, n2, n3 bins respectively, the category for bin values (b1, b2, b3) would be:
+  // category = b1*(n2*n3) + b2*(n3) + b3
   int category = 0;
   int tempCategory = 1;
   int iv1 = 0;
@@ -156,7 +151,7 @@
       if (iv2 == iv1) {
         tempCategory *= bin[iv2];
       } else {
-        tempCategory *= (fVariableLimits[iv2].GetSize() - 1);
+        tempCategory *= (fVariableLimits[iv2].size() - 1);
       }
     }
     category += tempCategory;
@@ -175,19 +170,14 @@
   }
 
   // Search for the position of the variable "var" in the internal variable list of the handler
-  int tempVar = 0;
-  for (auto v = fVariables.begin(); v != fVariables.end(); v++, tempVar++) {
-    if (*v == var) {
-      break;
-    }
-  }
+  int ivar = fVariables.at(var);
 
   // extract the bin position in variable "var" from the category
   int norm = 1;
-  for (int i = fVariables.size() - 1; i > tempVar; --i) {
-    norm *= (fVariableLimits[i].GetSize() - 1);
+  for (int i = fVariables.size() - 1; i > ivar; --i) {
+    norm *= (fVariableLimits[i].size() - 1);
   }
   int truncatedCategory = category - (category % norm);
   truncatedCategory /= norm;
-  return truncatedCategory % (fVariableLimits[tempVar].GetSize() - 1);
+  return truncatedCategory % (fVariableLimits[ivar].size() - 1);
 }

PWGDQ/Core/MixingHandler.h

@@ -24,24 +24,153 @@
 
 #include <Rtypes.h>
 
+#include <array>
+#include <iostream>
+#include <map>
 #include <vector>
 
 class MixingHandler : public TNamed
 {
 
  public:
+  // Struct to define track properties relevant for mixing and few utility functions
+  struct MixingTrack {
+    float pt;
+    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; }
+    void Print() const
+    {
+      std::cout << "pt: " << pt << ", eta: " << eta << ", phi: " << phi << ", filteringFlags: " << filteringFlags << std::endl;
+    }
+  };
+
+  // Struct to define events used in mixing and few utility functions.
+  // An event is defined as two vectors of tracks (typically the legs of a two-body
+  // decay or the two-particles in a correlation analysis)
+  struct MixingEvent {
+    std::vector<MixingTrack> tracks1;
+    std::vector<MixingTrack> tracks2;
+    // bit map for active filtering bits of all the tracks
+    uint32_t filteringMask = 0;
+    // counters to keep track of how many times the event was used for mixing (for each track cut separately)
+    std::array<short, 64> counters = {0};
+    // add a track to the event and update the filtering mask accordingly
+    void AddTrack1(const MixingTrack& track)
+    {
+      tracks1.push_back(track);
+      filteringMask |= track.filteringFlags;
+    }
+    void AddTrack2(const MixingTrack& track)
+    {
+      tracks2.push_back(track);
+      filteringMask |= track.filteringFlags;
+    }
+    // flip bits in the filtering mask
+    void FlipFilteringMask(int64_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
+    // 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)) {
+          counters[i]++;
+          if (counters[i] >= poolDepth) {
+            for (auto& track : tracks1) {
+              track.FlipBit(1ULL << i);
+              if (track.filteringFlags == 0) {
+                track = tracks1.back();
+                tracks1.pop_back();
+              }
+            }
+            for (auto& track : tracks2) {
+              track.FlipBit(1ULL << i);
+              if (track.filteringFlags == 0) {
+                track = tracks2.back();
+                tracks2.pop_back();
+              }
+            }
+            FlipFilteringMask(1ULL << i);
+          }
+        }
+      }
+    }
+    void Print() const
+    {
+      std::cout << "Event filtering mask: ";
+      for (int i = 0; i < 32; i++) {
+        if (filteringMask & (1ULL << i)) {
+          std::cout << "1";
+        } else {
+          std::cout << "0";
+        }
+      }
+      std::cout << std::endl;
+      for (int i = 0; i < 32; i++) {
+        if (filteringMask & (1ULL << i)) {
+          std::cout << "Counter " << i << ": " << counters[i] << std::endl;
+        }
+      }
+      std::cout << "Tracks 1: " << std::endl;
+      for (const auto& track : tracks1) {
+        track.Print();
+      }
+      std::cout << "Tracks 2: " << std::endl;
+      for (const auto& track : tracks2) {
+        track.Print();
+      }
+    }
+  };
+
+  struct MixingPool {
+    std::vector<MixingEvent> events;
+
+    // check which events in the pool are empty (i.e. no active tracks for mixing) and remove them from the pool
+    void CleanPool()
+    {
+      events.erase(std::remove_if(events.begin(), events.end(),
+                                  [](const MixingEvent& event) { return event.tracks1.empty() && event.tracks2.empty(); }),
+                   events.end());
+    }
+    // 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
+    // (e.g. incrementing the counters and removing the tracks that reached the pool depth for a given cut)
+    void UpdatePool(const MixingEvent& event, short poolDepth)
+    {
+      for (auto& event : events) {
+        event.IncrementCounters(event.filteringMask, poolDepth);
+      }
+      CleanPool();
+      events.push_back(event);
+    }
+    // getter for the events in the pool
+    const std::vector<MixingEvent>& GetEvents() const { return events; }
+
+    void Print() const
+    {
+      std::cout << "Mixing pool with " << events.size() << " events:" << std::endl;
+      for (const auto& event : events) {
+        event.Print();
+      }
+    }
+  };
+
   MixingHandler();
   MixingHandler(const char* name, const char* title);
   virtual ~MixingHandler();
 
   // setters
-  void AddMixingVariable(int var, int nBins, float* binLims);
-  void AddMixingVariable(int var, int nBins, std::vector<float> binLims);
+  void AddMixingVariable(int var, std::vector<float> binLims);
+  void SetPoolDepth(short depth) { fPoolDepth = depth; }
 
   // getters
-  int GetNMixingVariables() const { return fVariables.size(); }
-  int GetMixingVariable(VarManager::Variables var); // returns the position in the internal varible list of the handler. Useful for checks, mostly
-  std::vector<float> GetMixingVariableLimits(VarManager::Variables var);
+  // int GetNMixingVariables() const { return fVariables.size(); }
+  // int GetMixingVariable(VarManager::Variables var); // returns the position in the internal varible list of the handler. Useful for checks, mostly
+  // std::vector<float> GetMixingVariableLimits(VarManager::Variables var);
+  MixingPool& GetPool(int category) { return fPools[category]; }
+  short GetPoolDepth() const { return fPoolDepth; }
 
   void Init();
   int FindEventCategory(float* values);
@@ -54,10 +183,14 @@
   // User options
   bool fIsInitialized; // check if the mixing handler is initialized
 
-  std::vector<TArrayF> fVariableLimits;
-  std::vector<int> fVariables;
+  // 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
+  std::vector<std::vector<float>> fVariableLimits;
+  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)
+
+  short fPoolDepth;                 // number of events to be kept in each pool
+  std::map<int, MixingPool> fPools; // key: category, value: pool of events corresponding to that category
 
-  ClassDef(MixingHandler, 1);
+  ClassDef(MixingHandler, 2);
 };
 
 #endif // PWGDQ_CORE_MIXINGHANDLER_H_