[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,25 @@
     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];
 }

PWGDQ/Core/MixingHandler.h

@@ -25,7 +25,7 @@
 
 #include <array>
 #include <cstdint>
 #include <iostream>
 #include <map>
 #include <vector>
 
@@ -33,17 +33,31 @@
 {
 
  public:
+  // number of track cuts which fit in the 32-bit filtering masks
+  static constexpr int NMaxCuts = 32;
+
   // 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;
+    // electric charge of the track; 0 means "not set" and disables the charge dependent pair variables
+    int8_t sign = 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 << ", sign: " << static_cast<int>(sign)
+                << ", filteringFlags: " << filteringFlags
+                << ", dataframe: " << dataFrameSequence << ", track: " << trackGlobalIndex << std::endl;
     }
   };
 
@@ -108,21 +122,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 +171,17 @@
       CleanPool();
       events.push_back(event);
     }
+    // Same, but the stored events are aged only for the cuts in agingMask. Passing the filtering mask of the
+    // incoming event ages an event only for the cuts for which a mixed pair was actually produced, so that the
+    // pool depth is a number of mixed partners and not a number of arrivals.
+    void UpdatePool(const MixingEvent& event, int16_t poolDepth, uint32_t agingMask)
+    {
+      for (auto& poolEvent : events) {
+        poolEvent.IncrementCounters(agingMask, poolDepth);
+      }
+      CleanPool();
+      events.push_back(event);
+    }
     // getter for the events in the pool
     const std::vector<MixingEvent>& GetEvents() const { return events; }
 
@@ -176,6 +201,8 @@
   // 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(); }

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);
     }
   }

PWGDQ/Tasks/tableReader_withAssoc.h

@@ -494,6 +494,10 @@ struct AnalysisEventSelection {
       // create the mixing hash and publish it into the hash table
       if (fMixHandler != nullptr) {
         int hh = fMixHandler->FindEventCategory(dqtablereader_helpers::varValues());
+        // 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);
       }
     }
@@ -1427,6 +1431,8 @@ struct AnalysisSameEventPairing {
 
   HistogramManager* fHistMan = nullptr;
   MixingHandler fMixingHandler;
+  // dataframe counter, part of the track identity in the mixing pools
+  uint64_t fMixingDataFrameSequence = 0;
 
   o2::analysis::DQMlResponse<float> fDQMlResponse;
   std::vector<float> fOutputMlPsi2ee; // TODO: check this is needed or not
@@ -1714,6 +1720,9 @@ struct AnalysisSameEventPairing {
     }
 
     if (fConfigRunMixingAcrossTFs) {
+      if (fNCutsBarrel > MixingHandler::NMaxCuts) {
+        LOGF(fatal, "Across-TF mixing supports at most %d barrel track-cut bits, got %d", MixingHandler::NMaxCuts, fNCutsBarrel);
+      }
       TString mixVarsString = fConfigMixingVariables.value;
       TString mixVarsJsonString = fConfigMixingVariablesJson.value;
       std::unique_ptr<TObjArray> objArray(mixVarsString.Tokenize(","));
@@ -1899,6 +1908,10 @@ struct AnalysisSameEventPairing {
   {
     if (events.size() > 0) { // Additional protection to avoid crashing of events.begin().runNumber()
       if (fCurrentRun != events.begin().runNumber()) {
+        if (fConfigRunMixingAcrossTFs) {
+          // do not mix events from different runs
+          fMixingHandler.ClearPools();
+        }
         initParamsFromCCDB(events.begin().timestamp(), events.begin().runNumber(), TTwoProngFitter);
         fCurrentRun = events.begin().runNumber();
       }
@@ -1981,6 +1994,10 @@ struct AnalysisSameEventPairing {
     // constexpr bool fillFlowReso = eventHasQvector || eventHasQvectorCentr;
     bool isSelectedBDT = false;
     fNPairPerEvent = 0;
+    uint64_t currentMixingDataFrameSequence = 0;
+    if (fConfigRunMixingAcrossTFs) {
+      currentMixingDataFrameSequence = ++fMixingDataFrameSequence;
+    }
 
     for (auto const& event : events) {
       if (!event.isEventSelected_bit(0)) {
@@ -2006,6 +2023,10 @@ struct AnalysisSameEventPairing {
         }
         VarManager::FillEventFlowResoFactor(ResoFlowSP, ResoFlowEP);
       }
+      int mixingCategory = -1;
+      if (fConfigRunMixingAcrossTFs) {
+        mixingCategory = fMixingHandler.FindEventCategory(dqtablereader_helpers::varValues());
+      }
 
       bool isFirst = true;
       for (auto const& [a1, a2] : o2::soa::combinations(groupedAssocs, groupedAssocs)) {
@@ -2478,6 +2499,9 @@ struct AnalysisSameEventPairing {
 
       if (fConfigRunMixingAcrossTFs) {
         // run event mixing across TFs
+        if (mixingCategory < 0) {
+          continue;
+        }
         // 1) create a MixingEvent and fill it with the relevant tracks
         MixingHandler::MixingEvent mixingEvent;
         uint32_t trackFilterForMixing = 0;
@@ -2488,16 +2512,19 @@ struct AnalysisSameEventPairing {
               continue;
             }
             auto t1 = assoc.template reducedtrack_as<TTracks>();
-            MixingHandler::MixingTrack mixingTrack(t1.pt(), t1.eta(), t1.phi(), trackFilterForMixing);
+            MixingHandler::MixingTrack mixingTrack(t1.pt(), t1.eta(), t1.phi(), trackFilterForMixing, currentMixingDataFrameSequence, static_cast<uint64_t>(assoc.reducedtrackId()), static_cast<int8_t>(t1.sign()));
             if (t1.sign() > 0) {
               mixingEvent.AddTrack1(mixingTrack);
             } else {
               mixingEvent.AddTrack2(mixingTrack);
             }
           }
         }
+        if (mixingEvent.tracks1.empty() && mixingEvent.tracks2.empty()) {
+          continue;
+        }
         // 2) run the mixing with the events in the pool corresponding to this event
-        auto& pool = fMixingHandler.GetPool(fMixingHandler.FindEventCategory(dqtablereader_helpers::varValues()));
+        auto& pool = fMixingHandler.GetPool(mixingCategory);
         for (auto const& poolEvent : pool.GetEvents()) {
           for (auto const& t1 : mixingEvent.tracks1) {
             // run +- pairing
@@ -2516,9 +2543,9 @@ struct AnalysisSameEventPairing {
             }
             // run ++ pairing
             for (auto const& t2 : poolEvent.tracks1) {
-              // check the two-track filter for the mixed pair
+              // check the two-track filter for the mixed pair and skip the same track associated to both collisions
               uint32_t mixedTwoTrackFilter = t1.filteringFlags & t2.filteringFlags;
-              if (!mixedTwoTrackFilter) {
+              if (!mixedTwoTrackFilter || t1.IsSamePhysicalTrack(t2)) {
                 continue;
               }
               VarManager::FillPairMEAcrossTFs(t1, t2);
@@ -2546,9 +2573,9 @@ struct AnalysisSameEventPairing {
             }
             // run -- pairing
             for (auto const& t2 : poolEvent.tracks2) {
-              // check the two-track filter for the mixed pair
+              // check the two-track filter for the mixed pair and skip the same track associated to both collisions
               uint32_t mixedTwoTrackFilter = t1.filteringFlags & t2.filteringFlags;
-              if (!mixedTwoTrackFilter) {
+              if (!mixedTwoTrackFilter || t1.IsSamePhysicalTrack(t2)) {
                 continue;
               }
               VarManager::FillPairMEAcrossTFs(t1, t2);
@@ -2561,7 +2588,7 @@ struct AnalysisSameEventPairing {
           }
         }
         // 3) add the current event to the pool
-        pool.UpdatePool(mixingEvent, fMixingHandler.GetPoolDepth());
+        pool.UpdatePool(mixingEvent, fMixingHandler.GetPoolDepth(), mixingEvent.filteringMask);
         // pool.Print();
       }
     } // end loop over events

PWGDQ/Tasks/tableReader_withAssoc_direct.cxx

@@ -494,6 +494,16 @@ struct AnalysisEventSelection {
       VarManager::FillBC(bc);
       VarManager::FillEvent<TEventFillMap>(event);
 
+      // the hash table is joined to the events by row order, so publish one row per event before any event selection
+      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);
+      }
+
       bool decision = false;
       if (fConfigQA) {
         fHistMan->FillHistClass("Event_BeforeCuts", VarManager::fgValues);
@@ -533,10 +543,6 @@ struct AnalysisEventSelection {
         auto& evIndices = fBCCollMap[bc.globalBC()];
         evIndices.push_back(event.globalIndex());
       }
-      if (fMixHandler != nullptr) {
-        int hh = fMixHandler->FindEventCategory(VarManager::fgValues);
-        hash(hh);
-      }
     }
   }