[PWGDQ] Fix event mixing bug in MixingHandler and tableReader

PWGDQ/Core/MixingHandler.cxx

@@ -62,6 +62,8 @@
 {
   fVariables[var] = fVariableLimits.size();
   fVariableLimits.push_back(binLims);
+  // FillEvent() only fills variables marked as used
+  VarManager::SetUseVariable(var);
 }
 
 /*
@@ -100,7 +102,7 @@
 {
   // 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)
@@ -125,13 +127,16 @@
 
   // loop over the variables and find out in which bin the value of the variable for the event is located
   std::vector<int> bin;
+  // number of bins per variable in the iteration order of fVariables (fVariableLimits is in insertion order)
+  std::vector<int> nBins;
   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);
+    nBins.push_back(fVariableLimits[pos].size() - 1);
   }
 
   // Hash the bin values to define a unique category
@@ -149,7 +154,7 @@
       if (iv2 == iv1) {
         tempCategory *= bin[iv2];
       } else {
-        tempCategory *= (fVariableLimits[iv2].size() - 1);
+        tempCategory *= nBins[iv2];
       }
     }
     category += tempCategory;
@@ -167,15 +172,40 @@
     return -1;
   }
 
-  // Search for the position of the variable "var" in the internal variable list of the handler
-  int ivar = fVariables.at(var);
+  // number of bins and position of var in the iteration order of fVariables, as used by FindEventCategory()
+  std::vector<int> nBins;
+  int ivar = -1;
+  for (auto const& [v, pos] : fVariables) {
+    if (v == var) {
+      ivar = static_cast<int>(nBins.size());
+    }
+    nBins.push_back(fVariableLimits[pos].size() - 1);
+  }
+  if (ivar < 0) {
+    return -1;
+  }
 
   // extract the bin position in variable "var" from the category
   int norm = 1;
-  for (int i = fVariables.size() - 1; i > ivar; --i) {
-    norm *= (fVariableLimits[i].size() - 1);
+  for (size_t i = nBins.size() - 1; i > static_cast<size_t>(ivar); --i) {
+    norm *= nBins[i];
   }
   int truncatedCategory = category - (category % norm);
   truncatedCategory /= norm;
-  return truncatedCategory % (fVariableLimits[ivar].size() - 1);
+  return truncatedCategory % nBins[ivar];
+}
+
+//_________________________________________________________________________
+void MixingHandler::SetCategoryBinCenters(int category, float* values) const
+{
+  //
+  // set the mixing variables to the bin centers of this category (used for the leftover mixing)
+  //
+  for (auto const& [var, pos] : fVariables) {
+    int bin = GetBinFromCategory(static_cast<VarManager::Variables>(var), category);
+    if (bin < 0) {
+      continue;
+    }
+    values[var] = 0.5 * (fVariableLimits[pos][bin] + fVariableLimits[pos][bin + 1]);
+  }
 }

PWGDQ/Core/MixingHandler.h

@@ -23,9 +23,10 @@
 
 #include <Rtypes.h>
 
+#include <algorithm>
 #include <array>
 #include <cstdint>
 #include <iostream>
 #include <map>
 #include <vector>
 
@@ -33,17 +34,30 @@
 {
 
  public:
+  // number of track cuts which fit in the 32-bit filtering masks
+  static constexpr int NMaxCuts = 32;
+  // smallest pool depth for which a mixed pair can be built
+  static constexpr int16_t MinPoolDepth = 2;
+
   // Struct to define track properties relevant for mixing and few utility functions
   struct MixingTrack {
     float pt;
     float eta;
     float phi;
     uint32_t filteringFlags;
+    // globalIndex is unique only within a dataframe, so the dataframe sequence is part of the track identity
+    uint64_t dataFrameSequence = 0;
+    uint64_t trackGlobalIndex = 0;
+    bool IsSamePhysicalTrack(const MixingTrack& other) const
+    {
+      return dataFrameSequence == other.dataFrameSequence && trackGlobalIndex == other.trackGlobalIndex;
+    }
     // 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;
+      std::cout << "pt: " << pt << ", eta: " << eta << ", phi: " << phi << ", filteringFlags: " << filteringFlags
+                << ", dataframe: " << dataFrameSequence << ", track: " << trackGlobalIndex << std::endl;
     }
   };
 
@@ -70,6 +84,19 @@
     }
     // Clear bits in the filtering mask.
     void ClearFilteringMask(uint32_t mask) { filteringMask &= ~mask; }
+    // clear the cut bits from all tracks and the filtering mask; remove tracks with no active bits left
+    void ClearBits(uint32_t mask)
+    {
+      for (auto& track : tracks1) {
+        track.ClearBit(mask);
+      }
+      tracks1.erase(std::remove_if(tracks1.begin(), tracks1.end(), [](auto const& track) { return track.filteringFlags == 0; }), tracks1.end());
+      for (auto& track : tracks2) {
+        track.ClearBit(mask);
+      }
+      tracks2.erase(std::remove_if(tracks2.begin(), tracks2.end(), [](auto const& track) { return track.filteringFlags == 0; }), tracks2.end());
+      ClearFilteringMask(mask);
+    }
     // 1) increment the counters for a given track cut bit mask and if the counters reached the pool depth,
     // 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
@@ -108,21 +135,21 @@
     }
     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();
       }
@@ -157,6 +184,35 @@
       CleanPool();
       events.push_back(event);
     }
+    // fixed-block mixing: AddEvent() until GetMixingMask() reports full cuts, mix, then ClearBits()
+    void AddEvent(const MixingEvent& event) { events.push_back(event); }
+    // bit mask of the cuts for which at least poolDepth events are in the pool
+    uint32_t GetMixingMask(int16_t poolDepth) const
+    {
+      std::array<int16_t, NMaxCuts> counts = {0};
+      for (auto const& event : events) {
+        for (int icut = 0; icut < NMaxCuts; ++icut) {
+          if (event.filteringMask & (static_cast<uint32_t>(1) << icut)) {
+            counts[icut]++;
+          }
+        }
+      }
+      uint32_t fullMask = 0;
+      for (int icut = 0; icut < NMaxCuts; ++icut) {
+        if (counts[icut] >= poolDepth) {
+          fullMask |= static_cast<uint32_t>(1) << icut;
+        }
+      }
+      return fullMask;
+    }
+    // clear the given cut bits from all events in the pool and remove the events with no tracks left
+    void ClearBits(uint32_t mask)
+    {
+      for (auto& event : events) {
+        event.ClearBits(mask);
+      }
+      CleanPool();
+    }
     // getter for the events in the pool
     const std::vector<MixingEvent>& GetEvents() const { return events; }
 
@@ -176,17 +232,22 @@
   // setters
   void AddMixingVariable(int var, const std::vector<float>& binLims);
   void SetPoolDepth(int16_t depth) { fPoolDepth = depth; }
+  // remove all pools (e.g. at a run change)
+  void ClearPools() { fPools.clear(); }
 
   // 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);
   MixingPool& GetPool(int category) { return fPools[category]; }
+  std::map<int, MixingPool>& GetPools() { return fPools; }
   int16_t GetPoolDepth() const { return fPoolDepth; }
 
   void Init();
   int FindEventCategory(float* values);
   int GetBinFromCategory(VarManager::Variables var, int category) const;
+  // set the mixing variables to the bin centers of the given category
+  void SetCategoryBinCenters(int category, float* values) const;
 
  private:
   MixingHandler(const MixingHandler& handler);

PWGDQ/Tasks/tableReader.cxx

@@ -273,6 +273,10 @@ struct AnalysisEventSelection {
 
     if (fMixHandler != nullptr) {
       int hh = fMixHandler->FindEventCategory(VarManager::fgValues);
+      // events outside the mixing limits (-1) get a distinct negative hash so that they are not mixed with each other
+      if (hh < 0) {
+        hh = -1 - static_cast<int>(event.globalIndex());
+      }
       hash(hh);
     }
   }