[PWGDQ] Add optional barrel association time margin cut to data and M…

PWGDQ/Core/VarManager.cxx

@@ -41,6 +41,7 @@
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <map>
 #include <numeric>
 #include <tuple>
@@ -1041,6 +1042,12 @@ void VarManager::SetDefaultVarNames()
   fgVariableUnits[kTrackTimeRes] = "ns";
   fgVariableNames[kTrackTimeResRelative] = "Relative resolution of the track time";
   fgVariableUnits[kTrackTimeResRelative] = "";
+  fgVariableNames[kTrackAssocDeltaTime] = "Track-collision association #Deltat";
+  fgVariableUnits[kTrackAssocDeltaTime] = "ns";
+  fgVariableNames[kTrackAssocTimeThreshold] = "Track-collision association time threshold";
+  fgVariableUnits[kTrackAssocTimeThreshold] = "ns";
+  fgVariableNames[kTrackAssocDeltaTimeNorm] = "|#Deltat| / threshold";
+  fgVariableUnits[kTrackAssocDeltaTimeNorm] = "";
   fgVariableNames[kDetectorMap] = "DetectorMap";
   fgVariableUnits[kDetectorMap] = "";
   fgVariableNames[kHasITS] = "HasITS";
@@ -2369,6 +2376,9 @@ void VarManager::SetDefaultVarNames()
   fgVarNamesMap["kTrackTime"] = kTrackTime;
   fgVarNamesMap["kTrackTimeRes"] = kTrackTimeRes;
   fgVarNamesMap["kTrackTimeResRelative"] = kTrackTimeResRelative;
+  fgVarNamesMap["kTrackAssocDeltaTime"] = kTrackAssocDeltaTime;
+  fgVarNamesMap["kTrackAssocTimeThreshold"] = kTrackAssocTimeThreshold;
+  fgVarNamesMap["kTrackAssocDeltaTimeNorm"] = kTrackAssocDeltaTimeNorm;
   fgVarNamesMap["kDetectorMap"] = kDetectorMap;
   fgVarNamesMap["kHasITS"] = kHasITS;
   fgVarNamesMap["kHasTRD"] = kHasTRD;
@@ -2918,3 +2928,64 @@ void VarManager::SetDefaultVarNames()
   fgVarNamesMap["kInnerTOFnSigmaHe3"] = kInnerTOFnSigmaHe3;
   fgVarNamesMap["kInnerTOFnSigmaAl"] = kInnerTOFnSigmaAl;
 }
+
+//__________________________________________________________________
+bool VarManager::computeBarrelAssocTimeCompat(float trackTime, float trackTimeRes, bool timeResIsRange, bool isPVContributor,
+                                              int64_t origBC, float origCollTime, int origNumContrib,
+                                              int64_t collBC, float collTime, float collTimeRes,
+                                              float nSigma, float timeMargin, int bcWindowForOneSigma,
+                                              int usePVAssociation, int maxPvContribLowMult, float* values)
+{
+  //
+  // Mirrors the time-based association of Common/Core/CollisionAssociation.h::runAssocWithTime for central barrel tracks.
+  // Every arithmetic step is kept identical to the associator (including the double->int64 truncations),
+  //   so that re-applying it on skimmed associations with a smaller timeMargin reproduces a skim produced directly with that margin.
+  // Only tracks assigned to a collision in the AO2D are handled (this is always the case for tracks written by the DQ table makers).
+  //
+  if (!values) {
+    values = fgValues;
+  }
+  values[kTrackAssocDeltaTime] = -9999.f;
+  values[kTrackAssocTimeThreshold] = 0.f;
+  values[kTrackAssocDeltaTimeNorm] = 9999.f;
+
+  constexpr double BunchSpacingNS = o2::constants::lhc::LHCBunchSpacingNS;
+
+  // (1) BC pre-window (CollisionAssociation.h: trackBCCache, bcOffsetMax, bcOffsetWindow)
+  //     NOTE: the associator uses the raw trackTime here also for PV contributors
+  const int64_t bcOffsetMax = static_cast<int64_t>(bcWindowForOneSigma * nSigma + timeMargin / BunchSpacingNS);
+  const int64_t trackBCCache = static_cast<int64_t>(static_cast<double>(origBC) + static_cast<double>(trackTime) / BunchSpacingNS);
+  if (std::abs(trackBCCache - collBC) > bcOffsetMax) {
+    return false;
+  }
+
+  // (2) PV-contributor handling (usePVAssociation: 0 off, 1 OnlySameBc, 2 SameBcAndLowMult)
+  constexpr int OnlySameBc = 1;
+  constexpr int SameBcAndLowMult = 2;
+  const bool pvMode = (usePVAssociation == OnlySameBc && isPVContributor) ||
+                      (usePVAssociation == SameBcAndLowMult && isPVContributor && origNumContrib > maxPvContribLowMult);
+  float tTrack = trackTime;
+  float tTrackRes = trackTimeRes;
+  if (pvMode) {
+    tTrack = origCollTime;                          // time of the ORIGINAL collision
+    tTrackRes = static_cast<float>(BunchSpacingNS); // 1 BC
+  }
+
+  // (3) time difference and threshold
+  const int64_t bcOffset = origBC - collBC;
+  const float collTimeRes2 = collTimeRes * collTimeRes;
+  const float deltaTime = tTrack - collTime + bcOffset * static_cast<float>(BunchSpacingNS);
+  float threshold = 0.f;
+  if (pvMode) {
+    threshold = tTrackRes; // margin NOT applied
+  } else if (timeResIsRange) {
+    threshold = tTrackRes + nSigma * std::sqrt(collTimeRes2) + timeMargin;
+  } else {
+    threshold = nSigma * std::sqrt(collTimeRes2 + tTrackRes * tTrackRes) + timeMargin;
+  }
+
+  values[kTrackAssocDeltaTime] = deltaTime;
+  values[kTrackAssocTimeThreshold] = threshold;
+  values[kTrackAssocDeltaTimeNorm] = (threshold > 0.f) ? std::abs(deltaTime) / threshold : 9999.f;
+  return std::abs(deltaTime) < threshold;
+}

PWGDQ/Core/VarManager.h

@@ -578,6 +578,9 @@ class VarManager : public TObject
     kTrackTime,
     kTrackTimeRes,
     kTrackTimeResRelative,
+    kTrackAssocDeltaTime,     // track time - time of the associated collision (BC offset included), ns
+    kTrackAssocTimeThreshold, // time-compatibility threshold of the track-to-collision associator (margin included), ns
+    kTrackAssocDeltaTimeNorm, // |kTrackAssocDeltaTime| / kTrackAssocTimeThreshold; < 1 means the association is time compatible
     kDetectorMap,
     kHasITS,
     kHasTRD,
@@ -1474,6 +1477,22 @@ class VarManager : public TObject
   static void FillTrackEMCal(T const& cluster, float trackP = -1.0f, float deltaEta = -999.0f, float deltaPhi = -999.0f, float* values = nullptr);
   template <uint32_t fillMap, typename T, typename C>
   static void FillTrackCollision(T const& track, C const& collision, float* values = nullptr);
+  // Re-evaluation of the time compatibility of a (barrel track, collision) association, mirroring
+  //   Common/Core/CollisionAssociation.h::runAssocWithTime. Fills kTrackAssocDeltaTime, kTrackAssocTimeThreshold, kTrackAssocDeltaTimeNorm.
+  //   orig* : collision originally assigned to the track in the AO2D (reference of trackTime)
+  //   coll* : collision of the association under test
+  //   The last 5 parameters must match the configuration of track-to-collision-associator used at skimming time
+  //   (timeMargin and nSigma may be smaller in order to tighten the association at analysis level).
+  static bool computeBarrelAssocTimeCompat(float trackTime, float trackTimeRes, bool timeResIsRange, bool isPVContributor,
+                                           int64_t origBC, float origCollTime, int origNumContrib,
+                                           int64_t collBC, float collTime, float collTimeRes,
+                                           float nSigma, float timeMargin, int bcWindowForOneSigma,
+                                           int usePVAssociation, int maxPvContribLowMult, float* values = nullptr);
+  // Same as above for DQ skimmed tables: track = ReducedTracks+ReducedTracksBarrel, collision/origCollision = ReducedEvents+ReducedEventsExtended
+  template <typename T, typename C>
+  static bool isBarrelAssocTimeCompatible(T const& track, C const& collision, C const& origCollision,
+                                          float nSigma, float timeMargin, int bcWindowForOneSigma,
+                                          int usePVAssociation, int maxPvContribLowMult, float* values = nullptr);
   template <int candidateType, uint32_t fillMap, typename T1, typename T2, typename C>
   static void FillTrackCollisionMC(T1 const& track, T2 const& MotherTrack, C const& collision, float* values = nullptr);
   template <int candidateType, typename T1>
@@ -3471,6 +3490,18 @@ void VarManager::FillTrackCollision(T const& track, C const& collision, float* v
   }
 }
 
+template <typename T, typename C>
+bool VarManager::isBarrelAssocTimeCompatible(T const& track, C const& collision, C const& origCollision,
+                                             float nSigma, float timeMargin, int bcWindowForOneSigma,
+                                             int usePVAssociation, int maxPvContribLowMult, float* values)
+{
+  return computeBarrelAssocTimeCompat(track.trackTime(), track.trackTimeRes(),
+                                      (track.flags() & o2::aod::track::TrackTimeResIsRange) > 0, track.isPVContributor(),
+                                      static_cast<int64_t>(origCollision.globalBC()), origCollision.collisionTime(), origCollision.numContrib(),
+                                      static_cast<int64_t>(collision.globalBC()), collision.collisionTime(), collision.collisionTimeRes(),
+                                      nSigma, timeMargin, bcWindowForOneSigma, usePVAssociation, maxPvContribLowMult, values);
+}
+
 template <uint32_t fillMap, typename T, typename C, typename M, typename P>
 void VarManager::FillTrackCollisionMatCorr(T const& track, C const& collision, M const& materialCorr, P const& propagator, float* values)
 {

PWGDQ/Tasks/dqEfficiency_withAssoc.cxx

@@ -554,6 +554,17 @@ struct AnalysisTrackSelection {
   Configurable<std::string> fConfigAddTrackHistogram{"cfgAddTrackHistogram", "", "Comma separated list of histograms"};
   Configurable<std::string> fConfigAddJSONHistograms{"cfgAddJSONHistograms", "", "Histograms in JSON format"};
   Configurable<bool> fConfigPublishAmbiguity{"cfgPublishAmbiguity", true, "If true, publish ambiguity table and fill QA histograms"};
+  // Re-application of the track-to-collision time compatibility on the skimmed associations.
+  //   timeMargin and nSigma must be <= the values used by track-to-collision-associator at skimming time;
+  //   bcWindow, usePVAssociation and maxPvContributors must be identical to those values.
+  struct : ConfigurableGroup {
+    Configurable<bool> cfgAssocTimeCut{"cfgAssocTimeCut", false, "If true, reject associations failing the time compatibility computed with the parameters below"};
+    Configurable<float> cfgAssocTimeMargin{"cfgAssocTimeMargin", 0.f, "time margin (ns); must be <= skimming timeMargin"};
+    Configurable<float> cfgAssocNSigma{"cfgAssocNSigma", 4.f, "nSigmaForTimeCompat; must be <= skimming value"};
+    Configurable<int> cfgAssocBcWindowForOneSigma{"cfgAssocBcWindowForOneSigma", 60, "bcWindowForOneSigma; must be equal to the skimming value"};
+    Configurable<int> cfgAssocUsePVAssociation{"cfgAssocUsePVAssociation", 1, "usePVAssociation; must be equal to the skimming value"};
+    Configurable<int> cfgAssocMaxPvContributorsForLowMultReassoc{"cfgAssocMaxPvContributorsForLowMultReassoc", 10, "maxPvContributorsForLowMultReassoc; must be equal to the skimming value"};
+  } fConfigAssocTime;
   Configurable<std::string> fConfigCcdbUrl{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
   Configurable<std::string> fConfigCcdbPathTPC{"ccdb-path-tpc", "Users/z/zhxiong/TPCPID/PostCalib", "base path to the ccdb object"};
   Configurable<int64_t> fConfigNoLaterThan{"ccdb-no-later-than", std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count(), "latest acceptable timestamp of creation for the object"};
@@ -711,6 +722,20 @@ struct AnalysisTrackSelection {
       }
 
       auto track = assoc.template reducedtrack_as<TTracks>();
+
+      // Reject incompatible associations while preserving the row alignment of trackSel.
+      // The track time is referenced to its original collision, not the associated one.
+      if (fConfigAssocTime.cfgAssocTimeCut) {
+        auto origEvent = track.template reducedevent_as<TEvents>();
+        if (!VarManager::isBarrelAssocTimeCompatible(track, event, origEvent,
+                                                     fConfigAssocTime.cfgAssocNSigma, fConfigAssocTime.cfgAssocTimeMargin,
+                                                     fConfigAssocTime.cfgAssocBcWindowForOneSigma, fConfigAssocTime.cfgAssocUsePVAssociation,
+                                                     fConfigAssocTime.cfgAssocMaxPvContributorsForLowMultReassoc)) {
+          trackSel(0);
+          continue;
+        }
+      }
+
       VarManager::FillTrack<TTrackFillMap>(track);
       // compute quantities which depend on the associated collision, such as DCA
       VarManager::FillTrackCollision<TTrackFillMap>(track, event);

PWGDQ/Tasks/tableReader_withAssoc.h

@@ -673,6 +673,17 @@ struct AnalysisTrackSelection {
   o2::framework::Configurable<std::string> fConfigAddJSONHistograms{"cfgAddJSONHistograms", "", "Histograms in JSON format"};
   o2::framework::Configurable<bool> fConfigQA{"cfgQA", false, "If true, fill QA histograms"};
   o2::framework::Configurable<bool> fConfigPublishAmbiguity{"cfgPublishAmbiguity", true, "If true, publish ambiguity table and fill QA histograms"};
+  // Re-application of the track-to-collision time compatibility on the skimmed associations.
+  //   timeMargin and nSigma must be <= the values used by track-to-collision-associator at skimming time;
+  //   bcWindow, usePVAssociation and maxPvContributors must be identical to those values.
+  struct : o2::framework::ConfigurableGroup {
+    o2::framework::Configurable<bool> cfgAssocTimeCut{"cfgAssocTimeCut", false, "If true, reject associations failing the time compatibility computed with the parameters below"};
+    o2::framework::Configurable<float> cfgAssocTimeMargin{"cfgAssocTimeMargin", 0.f, "time margin (ns); must be <= skimming timeMargin"};
+    o2::framework::Configurable<float> cfgAssocNSigma{"cfgAssocNSigma", 4.f, "nSigmaForTimeCompat; must be <= skimming value"};
+    o2::framework::Configurable<int> cfgAssocBcWindowForOneSigma{"cfgAssocBcWindowForOneSigma", 60, "bcWindowForOneSigma; must be equal to the skimming value"};
+    o2::framework::Configurable<int> cfgAssocUsePVAssociation{"cfgAssocUsePVAssociation", 1, "usePVAssociation; must be equal to the skimming value"};
+    o2::framework::Configurable<int> cfgAssocMaxPvContributorsForLowMultReassoc{"cfgAssocMaxPvContributorsForLowMultReassoc", 10, "maxPvContributorsForLowMultReassoc; must be equal to the skimming value"};
+  } fConfigAssocTime;
 
   o2::framework::Configurable<std::string> fConfigCcdbUrl{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
   o2::framework::Configurable<std::string> fConfigCcdbPathTPC{"ccdb-path-tpc", "Users/z/zhxiong/TPCPID/PostCalib", "base path to the ccdb object"};
@@ -806,6 +817,20 @@ struct AnalysisTrackSelection {
       VarManager::FillEvent<TEventFillMap>(event);
 
       auto track = assoc.template reducedtrack_as<TTracks>();
+
+      // Reject incompatible associations while preserving the row alignment of trackSel.
+      // The track time is referenced to its original collision, not the associated one.
+      if (fConfigAssocTime.cfgAssocTimeCut) {
+        auto origEvent = track.template reducedevent_as<TEvents>();
+        if (!VarManager::isBarrelAssocTimeCompatible(track, event, origEvent,
+                                                     fConfigAssocTime.cfgAssocNSigma, fConfigAssocTime.cfgAssocTimeMargin,
+                                                     fConfigAssocTime.cfgAssocBcWindowForOneSigma, fConfigAssocTime.cfgAssocUsePVAssociation,
+                                                     fConfigAssocTime.cfgAssocMaxPvContributorsForLowMultReassoc)) {
+          trackSel(0);
+          continue;
+        }
+      }
+
       filterMap = static_cast<uint32_t>(0);
       VarManager::FillTrack<TTrackFillMap>(track);
       // compute quantities which depend on the associated collision, such as DCA